[test] Report token tps in fwd occupancy kit and force ignore_eos (#29332)
This commit is contained in:
@@ -82,27 +82,32 @@ class FwdOccupancyMixin:
|
||||
resp = requests.get(
|
||||
self.base_url + "/metrics", timeout=_METRICS_REQUEST_TIMEOUT
|
||||
)
|
||||
assert resp.status_code == 200, (
|
||||
f"/metrics returned {resp.status_code}; the test class's server "
|
||||
"must be launched with --enable-metrics"
|
||||
)
|
||||
assert "sglang:fwd_occupancy" in resp.text, (
|
||||
"sglang:fwd_occupancy gauge not exposed; set "
|
||||
"SGLANG_ENABLE_METRICS_DEVICE_TIMER=1 in the server's env and "
|
||||
"pass --enable-metrics"
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
raise AssertionError(
|
||||
f"/metrics returned {resp.status_code}; the test class's "
|
||||
"server must be launched with --enable-metrics"
|
||||
)
|
||||
if "sglang:fwd_occupancy" not in resp.text:
|
||||
raise AssertionError(
|
||||
"sglang:fwd_occupancy gauge not exposed; set "
|
||||
"SGLANG_ENABLE_METRICS_DEVICE_TIMER=1 in the server's env "
|
||||
"and pass --enable-metrics"
|
||||
)
|
||||
|
||||
def _fwd_occupancy_fire(self, prompt: str, max_new_tokens: int):
|
||||
"""Fire one /generate. Must not be called concurrently -- that
|
||||
would break the single-batch invariant."""
|
||||
"""Fire one /generate, return (completion_tokens, wall_time).
|
||||
Must not be called concurrently -- that would break the
|
||||
single-batch invariant."""
|
||||
t0 = time.perf_counter()
|
||||
try:
|
||||
requests.post(
|
||||
resp = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": prompt,
|
||||
"sampling_params": {
|
||||
"temperature": 0.0,
|
||||
"max_new_tokens": max_new_tokens,
|
||||
"ignore_eos": True,
|
||||
},
|
||||
},
|
||||
timeout=_GENERATE_REQUEST_TIMEOUT,
|
||||
@@ -110,7 +115,13 @@ class FwdOccupancyMixin:
|
||||
except requests.RequestException:
|
||||
# Final stats-vs-threshold is the signal; individual fire
|
||||
# failure isn't.
|
||||
pass
|
||||
return 0, 0.0
|
||||
elapsed = time.perf_counter() - t0
|
||||
try:
|
||||
tokens = resp.json().get("meta_info", {}).get("completion_tokens", 0)
|
||||
except ValueError: # non-JSON body
|
||||
tokens = 0
|
||||
return tokens, elapsed
|
||||
|
||||
def _fwd_occupancy_warmup(self):
|
||||
"""Fill cuda graphs + step the device-timer past its first NaN
|
||||
@@ -123,15 +134,19 @@ class FwdOccupancyMixin:
|
||||
|
||||
def _fwd_occupancy_measure(self):
|
||||
"""Background-fire one long single-batch request, scrape
|
||||
/metrics on the foreground; return non-NaN samples."""
|
||||
/metrics on the foreground; return (non-NaN samples,
|
||||
token_tps)."""
|
||||
samples = []
|
||||
request_done = threading.Event()
|
||||
result = {"completion_tokens": 0, "elapsed": 0.0}
|
||||
|
||||
def fire_one():
|
||||
try:
|
||||
self._fwd_occupancy_fire(
|
||||
self.fwd_occupancy_prompt,
|
||||
self.fwd_occupancy_max_new_tokens,
|
||||
result["completion_tokens"], result["elapsed"] = (
|
||||
self._fwd_occupancy_fire(
|
||||
self.fwd_occupancy_prompt,
|
||||
self.fwd_occupancy_max_new_tokens,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
request_done.set()
|
||||
@@ -146,12 +161,17 @@ class FwdOccupancyMixin:
|
||||
time.sleep(self.fwd_occupancy_scrape_interval)
|
||||
|
||||
firer.join(timeout=_GENERATE_REQUEST_TIMEOUT)
|
||||
return samples
|
||||
token_tps = (
|
||||
result["completion_tokens"] / result["elapsed"]
|
||||
if result["elapsed"] > 0
|
||||
else 0.0
|
||||
)
|
||||
return samples, token_tps
|
||||
|
||||
def test_fwd_occupancy(self):
|
||||
self._assert_metrics_device_timer_enabled()
|
||||
self._fwd_occupancy_warmup()
|
||||
samples = self._fwd_occupancy_measure()
|
||||
samples, token_tps = self._fwd_occupancy_measure()
|
||||
|
||||
self.assertGreaterEqual(
|
||||
len(samples),
|
||||
@@ -166,7 +186,8 @@ class FwdOccupancyMixin:
|
||||
samples_sorted = sorted(samples)
|
||||
median = statistics.median(samples_sorted)
|
||||
peak = samples_sorted[-1]
|
||||
p10 = samples_sorted[max(0, len(samples_sorted) // 10 - 1)]
|
||||
p10_idx = min(len(samples_sorted) - 1, max(0, len(samples_sorted) // 10))
|
||||
p10 = samples_sorted[p10_idx]
|
||||
print(
|
||||
"\n"
|
||||
+ tabulate.tabulate(
|
||||
@@ -176,6 +197,7 @@ class FwdOccupancyMixin:
|
||||
["peak", f"{peak:.2f}"],
|
||||
["p10", f"{p10:.2f}"],
|
||||
["threshold", f"{self.fwd_occupancy_threshold:.2f}"],
|
||||
["token tps", f"{token_tps:.2f}"],
|
||||
],
|
||||
headers=["fwd_occupancy", "value"],
|
||||
tablefmt="github",
|
||||
|
||||
Reference in New Issue
Block a user