diff --git a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py new file mode 100644 index 000000000..5b0cbdbc8 --- /dev/null +++ b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py @@ -0,0 +1,252 @@ +"""Ownership contract for per-request bookkeeping clocks. + +Per-request accounting state (`decode_batch_idx` / `extend_batch_idx` iter +clocks, `kv_committed_len` / `kv_allocated_len` KV watermarks, +`spec_verify_ct`, and the `maybe_evict_swa()` call) must only be advanced by +the reviewed owner sites in _OWNER_SITES; spec-v2 draft workers must not +repeat any of them (the scheduler-driven mixin / resolve path already does). +A clock that runs fast fires SWA eviction in the overlap race window and +releases the SWA prefix lock early; neither shows up in e2e CI or the idle +leak checker, hence this AST-level guard. +""" + +import ast +import unittest +import warnings +from collections import Counter +from pathlib import Path + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=8, suite="base-a-test-cpu") + +_REPO_ROOT = Path(__file__).resolve().parents[4] +_SRT_DIR = _REPO_ROOT / "python" / "sglang" / "srt" +_SPECULATIVE_DIR = _SRT_DIR / "speculative" +assert _SRT_DIR.is_dir(), f"srt dir not found: {_SRT_DIR}" + +_TRACKED_ATTRS = ( + "decode_batch_idx", + "extend_batch_idx", + "kv_committed_len", + "kv_allocated_len", + "spec_verify_ct", +) +_EVICT_METHOD = "maybe_evict_swa" + +# {(path relative to srt/, scope, kind): mutation count}. Kind is the mutated +# attribute (`= 0` resets exempt) or "evict" for a `maybe_evict_swa()` call. +# Any added/removed/recounted site fails until reviewed here. +_SB = "managers/schedule_batch.py" +_MIXIN = ("speculative/eagle_info_v2.py", "EagleDraftInputV2Mixin.prepare_for_decode") +_RESOLVE = ( + "managers/scheduler_components/batch_result_processor.py", + "SchedulerBatchResultProcessor._resolve_spec_v2_tokens", +) +_SS = "session/streaming_session.py" +_OWNER_SITES = { + # non-spec scheduler + (_SB, "ScheduleBatch.prepare_for_decode", "decode_batch_idx"): 1, + (_SB, "ScheduleBatch.prepare_for_decode", "kv_committed_len"): 1, + (_SB, "ScheduleBatch.prepare_for_decode", "kv_allocated_len"): 1, + (_SB, "ScheduleBatch.prepare_for_extend", "extend_batch_idx"): 1, + (_SB, "ScheduleBatch.prepare_for_extend", "kv_committed_len"): 1, + (_SB, "ScheduleBatch.prepare_for_extend", "kv_allocated_len"): 1, + ("mem_cache/common.py", "alloc_for_extend", "evict"): 1, + ("mem_cache/common.py", "alloc_for_decode", "evict"): 1, + # spec v2: pre-claim in the scheduler-driven mixin, settle in resolve + (*_MIXIN, "decode_batch_idx"): 1, + (*_MIXIN, "evict"): 1, + (*_MIXIN, "kv_committed_len"): 1, + (*_MIXIN, "kv_allocated_len"): 1, + (*_RESOLVE, "kv_committed_len"): 2, + (*_RESOLVE, "spec_verify_ct"): 1, + # spec v1: each verify path owns its own settlement + ("speculative/eagle_info.py", "EagleVerifyInput.verify", "kv_committed_len"): 1, + ("speculative/eagle_info.py", "EagleVerifyInput.verify", "kv_allocated_len"): 1, + ("speculative/eagle_info.py", "EagleVerifyInput.verify", "spec_verify_ct"): 1, + ( + "speculative/ngram_info.py", + "NgramVerifyInput._fill_requests", + "spec_verify_ct", + ): 1, + ( + "speculative/ngram_info.py", + "NgramVerifyInput._free_cache", + "kv_committed_len", + ): 1, + ( + "speculative/ngram_info.py", + "NgramVerifyInput._free_cache", + "kv_allocated_len", + ): 1, + ("speculative/dflash_info.py", "DFlashVerifyInput.verify", "kv_committed_len"): 1, + ("speculative/dflash_info.py", "DFlashVerifyInput.verify", "kv_allocated_len"): 1, + ("speculative/dflash_info.py", "DFlashVerifyInput.verify", "spec_verify_ct"): 1, + # disaggregation decode prealloc + ( + "disaggregation/decode.py", + "DecodePreallocQueue._pre_alloc", + "kv_committed_len", + ): 1, + ( + "disaggregation/decode.py", + "DecodePreallocQueue._pre_alloc", + "kv_allocated_len", + ): 1, + # streaming session slot save/restore and tail trimming + (_SS, "SessionSlot.save_from_req", "kv_committed_len"): 1, + (_SS, "SessionSlot.save_from_req", "kv_allocated_len"): 1, + (_SS, "SessionSlot.restore_to_req", "kv_committed_len"): 1, + (_SS, "SessionSlot.restore_to_req", "kv_allocated_len"): 1, + (_SS, "StreamingSession._free_tail", "kv_committed_len"): 2, + (_SS, "StreamingSession._free_tail", "kv_allocated_len"): 2, + (_SS, "StreamingSession._trim_overshoot", "kv_committed_len"): 1, + (_SS, "StreamingSession._trim_overshoot", "kv_allocated_len"): 1, + (_SS, "StreamingSession.try_cache_finished_req", "kv_allocated_len"): 1, +} + + +def _iter_scoped_nodes(tree): + """Yield (node, dotted Class.method scope) for every node.""" + scope_of = {} + + def visit(node, scope): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + scope = f"{scope}.{node.name}" if scope else node.name + scope_of[node] = scope + for child in ast.iter_child_nodes(node): + visit(child, scope) + + visit(tree, "") + return scope_of.items() + + +def _is_zero_reset(node): + return isinstance(node, ast.Assign) and ( + isinstance(node.value, ast.Constant) and node.value.value == 0 + ) + + +def _scan_tree(tree): + """Count bookkeeping sites in an AST as Counter[(scope, kind)].""" + sites = Counter() + for node, scope in _iter_scoped_nodes(tree): + if isinstance(node, (ast.AugAssign, ast.Assign)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + for target in targets: + if ( + isinstance(target, ast.Attribute) + and target.attr in _TRACKED_ATTRS + and not _is_zero_reset(node) + ): + sites[(scope, target.attr)] += 1 + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == _EVICT_METHOD + ): + sites[(scope, "evict")] += 1 + return sites + + +def _parse(path: Path): + # utf-8-sig: some srt files carry a BOM that breaks plain-utf-8 ast.parse. + with warnings.catch_warnings(): + warnings.simplefilter("ignore", SyntaxWarning) + return ast.parse(path.read_text(encoding="utf-8-sig")) + + +def _scan_srt(): + """Count all bookkeeping sites in srt/ as Counter[(rel, scope, kind)].""" + found = Counter() + for path in sorted(_SRT_DIR.rglob("*.py")): + rel = path.relative_to(_SRT_DIR).as_posix() + for (scope, kind), count in _scan_tree(_parse(path)).items(): + found[(rel, scope, kind)] += count + return found + + +def _draft_worker_classes(): + """All transitive BaseDraftWorker subclasses under speculative/.""" + by_name = {} + for path in sorted(_SPECULATIVE_DIR.glob("*.py")): + rel = path.relative_to(_SRT_DIR).as_posix() + for node in ast.walk(_parse(path)): + if isinstance(node, ast.ClassDef): + bases = { + b.id if isinstance(b, ast.Name) else getattr(b, "attr", None) + for b in node.bases + } + by_name[node.name] = (rel, node, bases) + + workers = {"BaseDraftWorker"} + changed = True + while changed: + changed = False + for name, (_, _, bases) in by_name.items(): + if name not in workers and bases & workers: + workers.add(name) + changed = True + return [ + (rel, node) + for name, (rel, node, _) in sorted(by_name.items()) + if name in workers and name != "BaseDraftWorker" + ] + + +def _scan_class_subtree(class_node): + """Scan one ClassDef subtree; returns (method_scope, kind) sites.""" + module = ast.Module(body=[class_node], type_ignores=[]) + sites = set() + for scope, kind in _scan_tree(module): + # Strip the leading class name; keep method-level scope. + sites.add((scope.split(".", 1)[1] if "." in scope else scope, kind)) + return sites + + +class TestDecodeBookkeepingOwnership(CustomTestCase): + def test_bookkeeping_sites_match_owner_allowlist(self): + found = _scan_srt() + allow = Counter(_OWNER_SITES) + unexpected = found - allow + missing = allow - found + msg = [] + if unexpected: + msg.append( + "New bookkeeping mutation(s) beyond the recorded counts:\n " + + "\n ".join(f"{site} x{n}" for site, n in sorted(unexpected.items())) + + "\nThese are owned by the sites in _OWNER_SITES -- do not " + "repeat them; a genuinely new owner must be recorded there." + ) + if missing: + msg.append( + "Recorded site(s) no longer exist (update _OWNER_SITES):\n " + + "\n ".join(f"{site} x{n}" for site, n in sorted(missing.items())) + ) + self.assertFalse(msg, "\n\n".join(msg)) + + def test_spec_v2_draft_workers_do_no_scheduler_bookkeeping(self): + classes = _draft_worker_classes() + names = {node.name for _, node in classes} + # Discovery sanity: fail loudly instead of silently guarding nothing. + self.assertIn("EagleDraftWorker", names) + self.assertIn("FrozenKVMTPDraftWorker", names) + + violations = [] + for rel, node in classes: + for scope, kind in _scan_class_subtree(node): + violations.append((rel, f"{node.name}.{scope}", kind)) + self.assertFalse( + violations, + "Spec-v2 draft worker(s) repeat scheduler-owned bookkeeping:\n " + + "\n ".join(map(str, sorted(violations))) + + "\nUnder spec v2 the iter-clock ticks, `maybe_evict_swa`, and " + "KV watermark settlement are owned by the scheduler-driven " + "mixin / resolve path. Remove these from the worker.", + ) + + +if __name__ == "__main__": + unittest.main(verbosity=3)