minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[CI][RFC] Replace black-jupyter with ruff-format (#37210)

Browse Source 
Co-authored-by: Alison Shao <a.shao@wustl.edu>
This commit is contained in:
Alex Nails
2026-09-02 19:46:08 -07:00
committed by GitHub
co-authored by Alison Shao
parent 2641e427be
commit 28262c20df
1411 changed files with 7766 additions and 8176 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/chunked_prefill/test_scripted_abort.py
+14 -14
View File
@@ -53,15 +53,15 @@ class TestAbortBasic(ScriptedTestCase):
"finished",
"unknown",
), f"after abort r should be finished/unknown, got {r.status}"
assert (
r.kv_pages == 0
), f"abort must release KV; r.kv_pages={r.kv_pages} after abort"
assert (
r.req is None or r.req.kv.req_pool_idx is None
), f"abort must release row; r.req={r.req} after abort"
assert (
r.lock_refs == 0
), f"abort must release lock_refs; r.lock_refs={r.lock_refs}"
assert r.kv_pages == 0, (
f"abort must release KV; r.kv_pages={r.kv_pages} after abort"
)
assert r.req is None or r.req.kv.req_pool_idx is None, (
f"abort must release row; r.req={r.req} after abort"
)
assert r.lock_refs == 0, (
f"abort must release lock_refs; r.lock_refs={r.lock_refs}"
)
def test_abort_at_chunk_0(self):
self.server.execute_script(self._script_abort_at_chunk_0)
@@ -436,9 +436,9 @@ class TestAbortBasic(ScriptedTestCase):
assert not r.is_chunking, "aborted gap req must stay out of chunking"
if r.req is not None:
assert (
r.req.inflight_middle_chunks == 0
), f"inflight_middle_chunks not cleared; got {r.req.inflight_middle_chunks}"
assert r.req.inflight_middle_chunks == 0, (
f"inflight_middle_chunks not cleared; got {r.req.inflight_middle_chunks}"
)
def test_abort_when_chunked_only_then_idle(self):
self.server.execute_script(self._script_abort_when_chunked_only_then_idle)
@@ -509,10 +509,10 @@ class TestAbortBasic(ScriptedTestCase):
yield from _drain_until_released(t, r1)
assert r1.kv_pages == 0, (
f"force_retract + abort same yield must release KV; got " f"{r1.kv_pages}"
f"force_retract + abort same yield must release KV; got {r1.kv_pages}"
)
assert r1.req is None or r1.req.kv.req_pool_idx is None, (
f"force_retract + abort same yield must release row; got " f"{r1.req}"
f"force_retract + abort same yield must release row; got {r1.req}"
)
assert r1.lock_refs == 0, (
f"force_retract + abort same yield must release lock_refs; "
test/manual/chunked_prefill/test_scripted_chunk_size.py
+3 -3
View File
@@ -73,9 +73,9 @@ class TestChunkSizeDefault(ScriptedTestCase):
r = t.start_req(prompt_len=1, max_new_tokens=2)
yield from run_until_finished(r)
assert r.finished
assert (
r.chunks_done == 0
), f"single-token prompt should not chunk, got chunks_done={r.chunks_done}"
assert r.chunks_done == 0, (
f"single-token prompt should not chunk, got chunks_done={r.chunks_done}"
)
def test_chunk_size_256_prompt_100x(self):
self.server.execute_script(self._script_chunk_size_256_prompt_100x)
test/manual/chunked_prefill/test_scripted_http_smoke.py
+3 -3
View File
@@ -26,9 +26,9 @@ class TestScriptedHttpSmoke(ScriptedTestCase):
break
yield
assert r.finished
assert (
saw_chunking
), "expected the req to hold the chunked_req slot at least once"
assert saw_chunking, (
"expected the req to hold the chunked_req slot at least once"
)
def test_two_reqs_finish(self):
self.server.execute_script(self._script_two_reqs_finish)
test/manual/chunked_prefill/test_scripted_hybrid_swa.py
+6 -6
View File
@@ -123,9 +123,9 @@ class TestSWAHalfWindowChunk(ScriptedTestCase):
r = t.start_req(prompt_len=2 * _SWA_WINDOW, max_new_tokens=4)
yield from run_until_finished(r, max_steps=800)
assert r.finished
assert (
r.chunks_done >= 4
), f"expected >=4 chunks for 2*window / (window/2), got {r.chunks_done}"
assert r.chunks_done >= 4, (
f"expected >=4 chunks for 2*window / (window/2), got {r.chunks_done}"
)
assert len(r.req.output_ids) == 4
@@ -171,9 +171,9 @@ class TestSWARadix(ScriptedTestCase):
)
yield from run_until_finished(r2, max_steps=800)
assert r2.finished
assert (
r2.req.cached_tokens > 0
), f"r2 must hit the radix prefix, got cached_tokens={r2.req.cached_tokens}"
assert r2.req.cached_tokens > 0, (
f"r2 must hit the radix prefix, got cached_tokens={r2.req.cached_tokens}"
)
if __name__ == "__main__":
test/manual/chunked_prefill/test_scripted_invariants.py
+15 -15
View File
@@ -38,9 +38,9 @@ class TestInvariantsBasic(ScriptedTestCase):
for _ in range(DEFAULT_MAX_STEPS):
if r.is_chunking:
observed_chunking = True
assert (
r.kv_pages > 0
), f"kv_pages must be > 0 while is_chunking; got {r.kv_pages}"
assert r.kv_pages > 0, (
f"kv_pages must be > 0 while is_chunking; got {r.kv_pages}"
)
if r.finished:
break
yield
@@ -63,9 +63,9 @@ class TestInvariantsBasic(ScriptedTestCase):
+ comp.get("decode", [])
+ comp.get("chunked", [])
)
assert (
r.rid in all_rids
), f"running but not in batch_composition: {comp}"
assert r.rid in all_rids, (
f"running but not in batch_composition: {comp}"
)
if r.finished:
return
yield
@@ -157,9 +157,9 @@ class TestInvariantsBasic(ScriptedTestCase):
t.flush_cache()
yield
final = t.engine_stats()
assert (
final["kv_pool_free"] >= baseline["kv_pool_free"]
), f"KV leak: {baseline['kv_pool_free']} -> {final['kv_pool_free']}"
assert final["kv_pool_free"] >= baseline["kv_pool_free"], (
f"KV leak: {baseline['kv_pool_free']} -> {final['kv_pool_free']}"
)
assert final["req_pool_free"] >= baseline["req_pool_free"]
def test_long_lived_engine_reps_chunked(self):
@@ -212,9 +212,9 @@ class TestInvariantsBasic(ScriptedTestCase):
t.flush_cache()
yield
final_kv = t.engine_stats()["kv_pool_free"]
assert (
final_kv >= baseline_kv
), f"KV leak after sustained chunked load: {baseline_kv} -> {final_kv}"
assert final_kv >= baseline_kv, (
f"KV leak after sustained chunked load: {baseline_kv} -> {final_kv}"
)
def test_round_robin_short_and_chunked(self):
self.server.execute_script(self._script_round_robin_short_and_chunked)
@@ -356,9 +356,9 @@ class TestInvariantsBasic(ScriptedTestCase):
)
yield from run_until_finished(r)
assert r.finished
assert (
r.chunks_done >= 2
), f"VERY_LONG_PROMPT_LEN should chunk; got chunks_done={r.chunks_done}"
assert r.chunks_done >= 2, (
f"VERY_LONG_PROMPT_LEN should chunk; got chunks_done={r.chunks_done}"
)
assert len(r.req.output_ids) == n, (
f"ignore_eos=True + max_new_tokens={n} must produce exactly "
f"{n} output tokens; got len(output_tokens)={len(r.req.output_ids)}"
test/manual/chunked_prefill/test_scripted_kv_pressure.py
+3 -3
View File
@@ -62,9 +62,9 @@ class TestKVPressureBasic(ScriptedTestCase):
f"long req must really chunk under pinned cache; got chunks_done="
f"{r_long.chunks_done}"
)
assert (
r_long.lock_refs == 0
), f"req {r_long.rid} leaked {r_long.lock_refs} lock_refs after finish"
assert r_long.lock_refs == 0, (
f"req {r_long.rid} leaked {r_long.lock_refs} lock_refs after finish"
)
t._release_exhausted_pools()
final_lock_refs = t.get_all_node_lock_refs()
test/manual/chunked_prefill/test_scripted_lifecycle.py
+7 -8
View File
@@ -171,9 +171,9 @@ class TestLifecycleBasic(ScriptedTestCase):
assert "running" in seen, f"never observed running status; seen={seen}"
assert seen[-1] == "finished", f"final status must be finished; seen={seen}"
finished_idx = seen.index("finished")
assert all(
s in ("finished",) for s in seen[finished_idx:]
), f"status regressed after finish; seen={seen}"
assert all(s in ("finished",) for s in seen[finished_idx:]), (
f"status regressed after finish; seen={seen}"
)
def test_long_prompt_only_one_decode(self):
self.server.execute_script(self._script_long_prompt_only_one_decode)
@@ -279,8 +279,7 @@ class TestLifecycleBasic(ScriptedTestCase):
assert r1.finished and r2.finished
assert r2.chunks_done == 0
assert r2.req.cached_tokens > 0, (
f"r2 must hit r1's radix prefix; got cached_tokens="
f"{r2.req.cached_tokens}"
f"r2 must hit r1's radix prefix; got cached_tokens={r2.req.cached_tokens}"
)
assert len(r2.req.output_ids) == 2
@@ -408,9 +407,9 @@ class TestLifecycleBasic(ScriptedTestCase):
t.flush_cache()
yield
final = t.engine_stats()["kv_pool_free"]
assert (
final >= baseline - 1
), f"KV pool drift: baseline={baseline}, final={final}"
assert final >= baseline - 1, (
f"KV pool drift: baseline={baseline}, final={final}"
)
def test_abort_all_during_chunked(self):
self.server.execute_script(self._script_abort_all_during_chunked)
test/manual/chunked_prefill/test_scripted_max_new_tokens.py
+2 -3
View File
@@ -143,14 +143,13 @@ class TestMaxNewTokensFirstDecodeAdjacent(ScriptedTestCase):
decode_records = _decode_records(batch_log, r.rid)
assert len(decode_records) == max_new_tokens, (
f"expected {max_new_tokens} decode forwards, got " f"{len(decode_records)}"
f"expected {max_new_tokens} decode forwards, got {len(decode_records)}"
)
rid_modes = [rec.mode for rec in rid_records]
first_decode_pos = rid_modes.index("decode")
assert first_decode_pos >= 1, (
f"first decode must be preceded by an extend chunk; rid_modes="
f"{rid_modes}"
f"first decode must be preceded by an extend chunk; rid_modes={rid_modes}"
)
assert rid_modes[first_decode_pos - 1] == "extend", (
f"record immediately before the first decode (in this rid's "
test/manual/chunked_prefill/test_scripted_multi_req.py
+9 -9
View File
@@ -20,9 +20,9 @@ def _drain_flush_then_assert_no_kv_leak(t: ScriptedContext, baseline: dict):
t.flush_cache()
yield
final = t.engine_stats()
assert (
final["kv_pool_free"] >= baseline["kv_pool_free"]
), f"KV leak: {baseline['kv_pool_free']} -> {final['kv_pool_free']}"
assert final["kv_pool_free"] >= baseline["kv_pool_free"], (
f"KV leak: {baseline['kv_pool_free']} -> {final['kv_pool_free']}"
)
class TestMultiReqBasic(ScriptedTestCase):
@@ -39,9 +39,9 @@ class TestMultiReqBasic(ScriptedTestCase):
yield
assert r1.is_chunking, "r1 should still be chunking"
assert (
not r2.is_chunking
), "r2 must wait for r1's chunk loop to clear before chunking"
assert not r2.is_chunking, (
"r2 must wait for r1's chunk loop to clear before chunking"
)
yield from run_until_all_finished([r1, r2])
assert r1.finished and r2.finished
@@ -144,9 +144,9 @@ class TestMultiReqBasic(ScriptedTestCase):
r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN + 8, max_new_tokens=2)
yield from run_until_finished(r2)
assert r1.finished and r2.finished
assert (
r2.chunks_done < r1.chunks_done
), "r2 reuses r1's cached prefix, so it should chunk fewer times"
assert r2.chunks_done < r1.chunks_done, (
"r2 reuses r1's cached prefix, so it should chunk fewer times"
)
def test_trickle_per_yield_50(self):
self.server.execute_script(self._script_trickle_per_yield_50)
test/manual/chunked_prefill/test_scripted_pp.py
+3 -3
View File
@@ -175,9 +175,9 @@ class TestPPPdmux(ScriptedTestCase):
def _script_pp_split_prefill_chunked_no_merge_assert(t: ScriptedContext):
r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
yield from run_until_finished(r, max_steps=800)
assert (
r.finished
), "engine died before req finished — merge_batch assert may have tripped"
assert r.finished, (
"engine died before req finished — merge_batch assert may have tripped"
)
assert r.chunks_done >= 2, (
f"pdmux + chunked path must produce >=2 chunks to exercise "
f"split_prefill_batch filter; got chunks_done={r.chunks_done}"
test/manual/chunked_prefill/test_scripted_priority.py
+3 -3
View File
@@ -30,9 +30,9 @@ class TestPriorityBasic(ScriptedTestCase):
t.pause_generation(mode="retract")
yield
assert (
r.status == "waiting"
), f"force-retracted chunked req must be back in waiting; got {r.status}"
assert r.status == "waiting", (
f"force-retracted chunked req must be back in waiting; got {r.status}"
)
assert r.kv_pages == 0, f"retract must release KV; got {r.kv_pages}"
t.continue_generation()
test/manual/chunked_prefill/test_scripted_radix.py
+9 -9
View File
@@ -76,9 +76,9 @@ class TestRadixBasic(ScriptedTestCase):
f"after eviction r2 must re-chunk from scratch; "
f"chunks_done={r2.chunks_done} cached_tokens={r2.req.cached_tokens}"
)
assert (
r2.req.cached_tokens == 0
), f"eviction must clear r1's prefix; cached_tokens={r2.req.cached_tokens}"
assert r2.req.cached_tokens == 0, (
f"eviction must clear r1's prefix; cached_tokens={r2.req.cached_tokens}"
)
assert r2.kv_pages == 0
assert r2.lock_refs == 0
@@ -219,9 +219,9 @@ class TestRadixBasic(ScriptedTestCase):
r = t.start_req(prompt_len=prompt_len, max_new_tokens=2)
yield from run_until_finished(r, max_steps=400)
assert r.finished
assert (
r.chunks_done == 0
), f"full prefix hit must skip chunked path; got chunks_done={r.chunks_done}"
assert r.chunks_done == 0, (
f"full prefix hit must skip chunked path; got chunks_done={r.chunks_done}"
)
def test_radix_evict_race_concurrent_chunked_admit(self):
self.server.execute_script(
@@ -319,9 +319,9 @@ class TestRadixNoTailChunked(ScriptedTestCase):
"test must observe r as the in-flight chunked_req at least once; the "
"no-tail else branch was never exercised"
)
assert (
r.kv_pages == 0
), f"finished chunked req must release KV; got {r.kv_pages}"
assert r.kv_pages == 0, (
f"finished chunked req must release KV; got {r.kv_pages}"
)
class TestRadixHitCountInvariant(ScriptedTestCase):
test/manual/chunked_prefill/test_scripted_regression.py
+12 -12
View File
@@ -100,9 +100,9 @@ class TestRegressionBasic(ScriptedTestCase):
f"observed max={observed_max} (pre-fix bug would bump to 2 "
f"at the last-chunk admit boundary)"
)
assert (
cleared_inflight
), "inflight_middle_chunks should be 0 once the chunk loop clears"
assert cleared_inflight, (
"inflight_middle_chunks should be 0 once the chunk loop clears"
)
yield from run_until_finished(r)
assert r.finished
@@ -247,15 +247,15 @@ class TestRegressionBasic(ScriptedTestCase):
t.abort(r)
yield from _drain_until_released(t, r)
assert (
r.req.kv.req_pool_idx is None
), f"96d4749094: abort must release row; got row_idx={r.req.kv.req_pool_idx!r}"
assert (
r.kv_pages == 0
), f"96d4749094: abort must release KV; got kv_pages={r.kv_pages}"
assert (
r.lock_refs == 0
), f"96d4749094: abort must release lock_ref; got lock_refs={r.lock_refs}"
assert r.req.kv.req_pool_idx is None, (
f"96d4749094: abort must release row; got row_idx={r.req.kv.req_pool_idx!r}"
)
assert r.kv_pages == 0, (
f"96d4749094: abort must release KV; got kv_pages={r.kv_pages}"
)
assert r.lock_refs == 0, (
f"96d4749094: abort must release lock_ref; got lock_refs={r.lock_refs}"
)
assert not r.is_chunking
assert r.req.inflight_middle_chunks == 0
assert sum(t.get_all_node_lock_refs().values()) == baseline_refs
test/manual/chunked_prefill/test_scripted_sampling.py
+14 -16
View File
@@ -34,8 +34,7 @@ class TestSamplingBasic(ScriptedTestCase):
if r.rid in rec.rids and rec.mode == "decode"
]
assert len(decode_records) == 0, (
f"max_new_tokens=0 must run zero decode forwards; got "
f"{len(decode_records)}"
f"max_new_tokens=0 must run zero decode forwards; got {len(decode_records)}"
)
def test_max_new_tokens_one_long_chunked(self):
@@ -100,8 +99,7 @@ class TestSamplingBasic(ScriptedTestCase):
assert r.chunks_done >= 2
assert len(r.req.output_ids) == 16
assert isinstance(r.req.finished_reason, FINISH_LENGTH), (
f"ignore_eos=True must finish via length cap; got "
f"{r.req.finished_reason!r}"
f"ignore_eos=True must finish via length cap; got {r.req.finished_reason!r}"
)
def test_return_logprob_top_logprobs_chunked(self):
@@ -176,9 +174,9 @@ class TestSamplingBasic(ScriptedTestCase):
)
yield from run_until_finished(r)
assert r.finished
assert (
r.chunks_done >= 2
), f"prompt should span multiple chunks, got chunks_done={r.chunks_done}"
assert r.chunks_done >= 2, (
f"prompt should span multiple chunks, got chunks_done={r.chunks_done}"
)
assert r.req.logprob is not None
input_lp = r.req.logprob.input_token_logprobs_val
assert len(input_lp) == prompt_len, (
@@ -201,9 +199,9 @@ class TestSamplingBasic(ScriptedTestCase):
)
yield from run_until_finished(r)
assert r.finished
assert (
r.chunks_done >= 3
), f"prompt should span 3+ chunks, got chunks_done={r.chunks_done}"
assert r.chunks_done >= 3, (
f"prompt should span 3+ chunks, got chunks_done={r.chunks_done}"
)
assert r.req.logprob is not None
input_lp = r.req.logprob.input_token_logprobs_val
assert len(input_lp) == prompt_len - start_len, (
@@ -246,9 +244,9 @@ class TestSamplingBasic(ScriptedTestCase):
)
yield from run_until_finished(r_eos, max_steps=2000)
assert r_eos.finished
assert (
r_eos.chunks_done >= 2
), f"scenario 1 should chunk; got chunks_done={r_eos.chunks_done}"
assert r_eos.chunks_done >= 2, (
f"scenario 1 should chunk; got chunks_done={r_eos.chunks_done}"
)
assert isinstance(r_eos.req.finished_reason, FINISH_MATCHED_TOKEN), (
f"a stop token the model deterministically produces under greedy must "
f"finish via the matched-token path; got {r_eos.req.finished_reason!r}"
@@ -261,9 +259,9 @@ class TestSamplingBasic(ScriptedTestCase):
)
yield from run_until_finished(r_length)
assert r_length.finished
assert (
r_length.chunks_done >= 2
), f"scenario 2 should chunk; got chunks_done={r_length.chunks_done}"
assert r_length.chunks_done >= 2, (
f"scenario 2 should chunk; got chunks_done={r_length.chunks_done}"
)
assert isinstance(r_length.req.finished_reason, FINISH_LENGTH), (
f"ignore_eos=True + max_new_tokens=4 chunked must finish via "
f"length cap; got {r_length.req.finished_reason!r}"
test/manual/chunked_prefill/test_scripted_special_case.py
+67 -68
View File
@@ -45,9 +45,9 @@ class TestSpecialCaseBasic(ScriptedTestCase):
for _ in range(DEFAULT_MAX_STEPS):
if r.is_chunking:
saw_chunking = True
assert (
not t.is_idle
), "scheduler must not idle while chunked_req is in flight"
assert not t.is_idle, (
"scheduler must not idle while chunked_req is in flight"
)
if r.finished:
break
yield
@@ -102,9 +102,9 @@ class TestSpecialCaseBasic(ScriptedTestCase):
break
yield
assert (
t.scheduler.chunked_req is None
), f"abort must clear the chunked slot; got {t.scheduler.chunked_req!r}"
assert t.scheduler.chunked_req is None, (
f"abort must clear the chunked slot; got {t.scheduler.chunked_req!r}"
)
assert r.kv_pages == 0
assert r.lock_refs == 0
@@ -193,19 +193,19 @@ class TestSpecialCaseBasic(ScriptedTestCase):
if r1.is_chunking:
saw_r1_chunking = True
comp = t.batch_composition()
assert r1.rid in comp.get(
"chunked", []
), f"mid-chunk r1 must occupy the chunked role; got {comp!r}"
assert r1.rid not in comp.get(
"running", []
), f"chunked r1 must be excluded from the running role; got {comp!r}"
assert r1.rid in comp.get("chunked", []), (
f"mid-chunk r1 must occupy the chunked role; got {comp!r}"
)
assert r1.rid not in comp.get("running", []), (
f"chunked r1 must be excluded from the running role; got {comp!r}"
)
if r1.finished and r2.finished:
break
yield
assert r1.finished and r2.finished
assert (
saw_r1_chunking
), "r1 must have chunked at some point to exercise the exclude branch"
assert saw_r1_chunking, (
"r1 must have chunked at some point to exercise the exclude branch"
)
@unittest.skip(
"pdmux split_prefill_batch requires the pdmux topology — "
@@ -257,9 +257,9 @@ class TestSpecialCaseBasic(ScriptedTestCase):
t.pause_generation(mode="retract")
yield
assert (
t.scheduler.chunked_req is None
), f"pause(retract) must clear chunked_req; got {t.scheduler.chunked_req!r}"
assert t.scheduler.chunked_req is None, (
f"pause(retract) must clear chunked_req; got {t.scheduler.chunked_req!r}"
)
assert not r.finished, "retract must re-queue r, not finish or abort it"
assert r.status == "waiting", (
f"retracted chunked req must return to the waiting queue; "
@@ -268,9 +268,9 @@ class TestSpecialCaseBasic(ScriptedTestCase):
t.continue_generation()
yield from run_until_finished(r)
assert (
r.finished
), "continue_generation must drive the re-queued req to completion"
assert r.finished, (
"continue_generation must drive the re-queued req to completion"
)
def test_retract_during_gap_inflight_middle_chunks_positive(self):
self.server.execute_script(
@@ -304,9 +304,9 @@ class TestSpecialCaseBasic(ScriptedTestCase):
t.continue_generation()
yield from run_until_finished(r, max_steps=2000)
assert (
r.finished
), "continue_generation must drive the re-queued req to completion"
assert r.finished, (
"continue_generation must drive the re-queued req to completion"
)
assert r.kv_pages == 0
assert len(r.req.output_ids) == 2
@@ -344,9 +344,9 @@ class TestSpecialCaseBasic(ScriptedTestCase):
break
yield
assert r.finished
assert (
saw_chunking
), "test must observe the dual-queue chunked state at least once"
assert saw_chunking, (
"test must observe the dual-queue chunked state at least once"
)
assert saw_dedup, (
"test must observe the chunked req with a committed prefix so the "
"dedup subtraction is actually exercised"
@@ -491,9 +491,9 @@ class TestSpecialCaseBasic(ScriptedTestCase):
break
yield
assert r.finished
assert (
saw_mid_chunk
), "test must observe the fill_ids reset boundary at least once"
assert saw_mid_chunk, (
"test must observe the fill_ids reset boundary at least once"
)
assert r.finished
def test_chunked_req_slot_cleared_when_chunk_completes(self):
@@ -516,9 +516,9 @@ class TestSpecialCaseBasic(ScriptedTestCase):
yield
assert r.finished
assert saw_chunking, "req should have occupied the chunked_req slot mid-chunk"
assert (
s.chunked_req is None
), f"chunked_req slot must clear after last chunk; got {s.chunked_req!r}"
assert s.chunked_req is None, (
f"chunked_req slot must clear after last chunk; got {s.chunked_req!r}"
)
def test_second_chunked_admit_blocked_when_chunked_req_set(self):
self.server.execute_script(
@@ -571,9 +571,9 @@ class TestSpecialCaseBasic(ScriptedTestCase):
progressed = False
for _ in range(DEFAULT_MAX_STEPS):
if r.is_chunking:
assert (
not t.is_idle
), "scheduler must not go idle while a chunked req is in flight"
assert not t.is_idle, (
"scheduler must not go idle while a chunked req is in flight"
)
cur_chunks_done = r.chunks_done
if cur_chunks_done > prev_chunks_done:
progressed = True
@@ -681,9 +681,9 @@ class TestSpecialCaseMixedChunk(ScriptedTestCase):
prompt_token=310,
)
yield from run_until(r, lambda h: h.is_chunking)
assert (
t.last_batch_forward_mode != "MIXED"
), f"return_logprob must disable mixed-chunk path; got {t.last_batch_forward_mode!r}"
assert t.last_batch_forward_mode != "MIXED", (
f"return_logprob must disable mixed-chunk path; got {t.last_batch_forward_mode!r}"
)
yield from run_until_finished(r)
def test_mixed_chunk_with_running_batch(self):
@@ -700,9 +700,9 @@ class TestSpecialCaseMixedChunk(ScriptedTestCase):
yield
yield from run_until(r_chunk, lambda h: h.is_chunking)
assert (
t.last_batch_forward_mode == "MIXED"
), f"chunked admission with running batch must enter MIXED; got {t.last_batch_forward_mode!r}"
assert t.last_batch_forward_mode == "MIXED", (
f"chunked admission with running batch must enter MIXED; got {t.last_batch_forward_mode!r}"
)
for _ in range(DEFAULT_MAX_STEPS * 2):
if r_chunk.finished and r_dec.finished:
break
@@ -743,9 +743,9 @@ class TestSpecialCaseNoChunking(ScriptedTestCase):
def _script_chunk_size_negative_disables_chunking(t: ScriptedContext):
r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
for _ in range(DEFAULT_MAX_STEPS):
assert (
not r.is_chunking
), "chunked_prefill_size=-1 should disable chunked path"
assert not r.is_chunking, (
"chunked_prefill_size=-1 should disable chunked path"
)
if r.finished:
return
yield
@@ -825,9 +825,9 @@ class TestSpecialCaseHiCache(ScriptedTestCase):
yield
assert r.finished
assert saw_chunking, "test must observe r mid-chunk at least once"
assert (
first_chunk_snap is not None
), "test must snapshot cached_tokens at the first chunk boundary"
assert first_chunk_snap is not None, (
"test must snapshot cached_tokens at the first chunk boundary"
)
def test_hicache_cached_tokens_set_once_invariant(self):
self.server.execute_script(
@@ -862,9 +862,9 @@ class TestSpecialCaseHiCache(ScriptedTestCase):
break
yield
assert r.finished
assert (
saw_chunking
), "test must observe the req mid-chunk (chunks_done >= 1) at least once"
assert saw_chunking, (
"test must observe the req mid-chunk (chunks_done >= 1) at least once"
)
assert snap is not None, "test must snapshot the cached_tokens_* breakdown"
@@ -1011,24 +1011,23 @@ class TestSpecialCaseRetractMerge(ScriptedTestCase):
t.pause_generation(mode="retract")
yield
assert (
s.last_batch is None
), "retract must clear last_batch after merging the extend chunk batch"
assert s.last_batch is None, (
"retract must clear last_batch after merging the extend chunk batch"
)
assert len(s.running_batch.reqs) == 0, (
"the merged extend chunk batch must be retracted out of running_batch, "
f"not stranded; got {len(s.running_batch.reqs)} reqs"
)
assert (
r.status == "waiting"
), f"retracted chunked req must return to the waiting queue; got {r.status!r}"
assert r.status == "waiting", (
f"retracted chunked req must return to the waiting queue; got {r.status!r}"
)
assert r.kv_pages == 0
t.continue_generation()
yield from run_until_finished(r)
assert r.finished
assert len(r.req.output_ids) == 2, (
f"resumed req must emit exactly max_new_tokens; got "
f"{len(r.req.output_ids)}"
f"resumed req must emit exactly max_new_tokens; got {len(r.req.output_ids)}"
)
@@ -1183,9 +1182,9 @@ class TestSpecialCaseRetractedStain(ScriptedTestCase):
assert r.finished
req = r.req
assert (
req.retracted_stain is True
), "retract must set retracted_stain so the cached-token recount is suppressed"
assert req.retracted_stain is True, (
"retract must set retracted_stain so the cached-token recount is suppressed"
)
assert req.cached_tokens == cached_before, (
f"retracted_stain must suppress re-adding pre_len-already_computed on "
f"resume; cached_tokens grew from {cached_before} to {req.cached_tokens}"
@@ -1237,17 +1236,17 @@ class TestSpecialCaseMiddleChunkNoToken(ScriptedTestCase):
f"middle chunk must not append an output token; got "
f"output_ids len {len(r.req.output_ids)}"
)
assert (
r.status != "finished"
), "middle chunk must not finish the req (skip_stream_req)"
assert r.status != "finished", (
"middle chunk must not finish the req (skip_stream_req)"
)
if r.finished:
break
yield
assert r.finished
assert saw_middle_chunk, "test must observe r mid-chunk at least once"
assert (
len(r.req.output_ids) >= 1
), "output tokens must appear only after the chunked prefill completes"
assert len(r.req.output_ids) >= 1, (
"output tokens must appear only after the chunked prefill completes"
)
if __name__ == "__main__":