[test] Report token tps in fwd occupancy kit and force ignore_eos (#29332)

This commit is contained in:
Liangsheng Yin
2026-06-25 15:08:33 -07:00
committed by GitHub
parent b7d3c3016d
commit e4696ed62d
+42 -20
View File
@@ -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",