From 144ae7acb1fa00780fbb45faae1a6a6a79b5e35d Mon Sep 17 00:00:00 2001 From: leavelet Date: Wed, 24 Jun 2026 14:52:30 +0000 Subject: [PATCH] Regression test: CP HiCache L3 reload same-tick admit->reserve race Drives _cp_l3_reserve_reloads directly against a real CpSharedL2PageAllocator: - same-tick state (l3rl object committed-and-live + node attached -> attach_ok False) no longer raises 'duplicate live reservation'; the candidate is dropped and the request left un-held (re-matches the inserted node). - WITHOUT the recheck (attach_ok forced True = pre-fix path) the identical state raises the duplicate error -> proves the recheck is load-bearing. - an in-flight reload still piggybacks (recheck sits after the dedup, not before). - a genuine clean miss still reserves (no false-positive starvation). CPU unit test (register_cpu_ci, stage-a-test-cpu); runs in the dev-cu13 container. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../test_cp_l3_reload_reserve_race.py | 202 ++++++++++++++++++ 1 file changed, 202 insertions(+) create mode 100644 test/registered/unit/mem_cache/test_cp_l3_reload_reserve_race.py diff --git a/test/registered/unit/mem_cache/test_cp_l3_reload_reserve_race.py b/test/registered/unit/mem_cache/test_cp_l3_reload_reserve_race.py new file mode 100644 index 000000000..358520a49 --- /dev/null +++ b/test/registered/unit/mem_cache/test_cp_l3_reload_reserve_race.py @@ -0,0 +1,202 @@ +"""Regression tests for the CP HiCache L3 reload same-tick admit->reserve race. + +Bug (crash): _drain_l3_control_queues runs validate-candidates -> admit -> reserve +in one tick. _cp_l3_admit_reload pops the in-flight key and mark_object_committed()s +the l3rl: object (which stays LIVE in the allocator) while inserting the +reloaded node. A second same-suffix candidate, validated against the PRE-admit tree, +then reaches _cp_l3_reserve_reloads with its dedup guard missing (in-flight key gone) +and calls allocator.reserve() on the still-live object -> + ValueError: duplicate live reservation for object 'l3rl:' payload 'target_kv' +-> rank-uniform scheduler crash. + +Fix: re-validate _cp_l3_reload_attach_ok inside _cp_l3_reserve_reloads (after the +piggyback dedup, before reserve). If a same-tick admit inserted the suffix, the +candidate is dropped; the un-held request re-matches and hits the inserted L2 node. + +These tests exercise _cp_l3_reserve_reloads directly with a real CpSharedL2PageAllocator +and assert the four behaviors that pin the fix: (1) the racing candidate is dropped with +no duplicate reserve; (2) WITHOUT the recheck the same state raises the duplicate error +(proves the recheck is load-bearing); (3) an in-flight reload still piggybacks (the recheck +sits after the dedup, not before it); (4) a genuine clean miss still reserves (no +false-positive starvation). +""" + +import functools +import json +import sys +import types +import unittest + +import torch + +# Import-time stubs so hiradix_cache imports on a CPU-only host (mirrors the sibling +# mem_cache CPU unit tests). +if "pybase64" not in sys.modules: + pybase64_stub = types.ModuleType("pybase64") + pybase64_stub.b64encode = lambda *args, **kwargs: b"" + pybase64_stub.b64decode = lambda *args, **kwargs: b"" + sys.modules["pybase64"] = pybase64_stub +if "orjson" not in sys.modules: + orjson_stub = types.ModuleType("orjson") + orjson_stub.loads = lambda data, *args, **kwargs: json.loads( + data.decode() if isinstance(data, (bytes, bytearray)) else data + ) + orjson_stub.dumps = lambda obj, *args, **kwargs: json.dumps(obj).encode() + sys.modules["orjson"] = orjson_stub +sgl_kernel_stub = sys.modules.setdefault("sgl_kernel", types.ModuleType("sgl_kernel")) +sgl_kernel_stub.__file__ = getattr(sgl_kernel_stub, "__file__", "sgl_kernel_stub.py") +sgl_kernel_stub.__path__ = getattr(sgl_kernel_stub, "__path__", []) +if not hasattr(sgl_kernel_stub, "__getattr__"): + + def _sgl_kernel_getattr(name): + if name.startswith("__"): + raise AttributeError(name) + fn = lambda *args, **kwargs: None + setattr(sgl_kernel_stub, name, fn) + return fn + + sgl_kernel_stub.__getattr__ = _sgl_kernel_getattr + +from sglang.srt.mem_cache.cp_shared_l2_pool import ( + PAYLOAD_TARGET_KV, + CpSharedL2PageAllocator, +) +from sglang.srt.mem_cache.hiradix_cache import HiRadixCache +from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode, get_child_key +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="stage-a-test-cpu") + +PAGE_SIZE = 4 +RELOAD_HASH = "deadbeef" # shashes[c-1] +RELOAD_KEY = f"l3rl:{RELOAD_HASH}" + + +def _make_store(): + # Only the fields _cp_l3_reserve_reloads / _cp_l3_build_reload_owned_pages read. + return types.SimpleNamespace( + payloads=(PAYLOAD_TARGET_KV,), + cp_size=4, + cp_rank=0, + submit_reload=lambda reload_key, owned_pages: 12345, + ) + + +def _make_l3_cache(allocator, store): + cache = HiRadixCache.__new__(HiRadixCache) + cache.page_size = PAGE_SIZE + cache.load_back_threshold = 0 # -> min_pages = 1 + cache.cp_l3_reload_max_inflight = 1000 + cache.cp_l3_reload_ttl_ticks = 100 + cache.ongoing_l3_reload = {} + cache.cp_l3_reload_inflight_keys = {} + cache.cp_l3_reload_done = {} + cache.evictable_host_leaves = set() + cache.cp_l3_store = store + cache.cache_controller = types.SimpleNamespace( + cp_shared_l2_page_allocator=allocator + ) + cache.root_node = TreeNode(id=0) + cache.root_node.key = RadixKey(token_ids=[]) + cache.get_child_key_fn = functools.partial(get_child_key, page_size=PAGE_SIZE) + return cache + + +def _make_allocator(): + return CpSharedL2PageAllocator(pages_per_payload={PAYLOAD_TARGET_KV: 16}) + + +def _commit_live_object(allocator, num_pages=1): + """Simulate the same-tick admit: reserve + commit the l3rl object so it is LIVE.""" + allocator.reserve(RELOAD_KEY, PAYLOAD_TARGET_KV, num_pages) + allocator.mark_object_committed(RELOAD_KEY) + + +def _suffix_key(): + return RadixKey(token_ids=list(range(PAGE_SIZE))) # one full page + + +def _attach_existing_child(cache, suffix_key): + """Simulate the node the same-tick admit inserted under the (root) last_host_node, + so _cp_l3_reload_attach_ok now sees the suffix's first-page child present.""" + child = TreeNode(id=99) + child.parent = cache.root_node + child.key = suffix_key + cache.root_node.children[cache.get_child_key_fn(suffix_key)] = child + return child + + +class TestCpL3ReloadReserveRace(CustomTestCase): + def test_same_tick_admit_then_reserve_drops_candidate_no_duplicate(self): + # The bug's exact state: l3rl: committed-and-live (admit just committed it) + # AND its node attached under the last_host_node -> attach_ok is False. + allocator = _make_allocator() + _commit_live_object(allocator) + cache = _make_l3_cache(allocator, _make_store()) + skey = _suffix_key() + _attach_existing_child(cache, skey) + valid_cands = [("r2", cache.root_node, skey, [RELOAD_HASH])] + + # Must NOT raise "duplicate live reservation". + cache._cp_l3_reserve_reloads(valid_cands, [1]) + + # Candidate dropped: no new op, request left un-held (re-matches the inserted node). + self.assertEqual(cache.ongoing_l3_reload, {}) + self.assertNotIn("r2", cache.cp_l3_reload_done) + self.assertNotIn(RELOAD_KEY, cache.cp_l3_reload_inflight_keys) + # The pre-existing committed object is untouched (not re-reserved, not aborted). + self.assertTrue(allocator.is_committed(RELOAD_KEY)) + self.assertIn(PAYLOAD_TARGET_KV, allocator.object_ranges(RELOAD_KEY)) + + def test_without_recheck_same_state_raises_duplicate(self): + # Proves the recheck is load-bearing: with attach_ok forced True (the pre-fix code + # path, which never re-validated at reserve) the identical state crashes. + allocator = _make_allocator() + _commit_live_object(allocator) + cache = _make_l3_cache(allocator, _make_store()) + cache._cp_l3_reload_attach_ok = lambda *args, **kwargs: True # pre-fix behavior + skey = _suffix_key() + valid_cands = [("r2", cache.root_node, skey, [RELOAD_HASH])] + + with self.assertRaisesRegex(ValueError, "duplicate live reservation"): + cache._cp_l3_reserve_reloads(valid_cands, [1]) + + def test_inflight_piggyback_takes_precedence_over_recheck(self): + # The recheck sits AFTER the in-flight dedup: a same-suffix candidate must still + # piggyback on an in-flight op (and never touch the allocator), even with no node + # attached. allocator stays empty -> a stray reserve would be observable. + allocator = _make_allocator() + cache = _make_l3_cache(allocator, _make_store()) + op_x = 777 + cache.ongoing_l3_reload[op_x] = {"waiters": ["r1"]} + cache.cp_l3_reload_inflight_keys[RELOAD_KEY] = op_x + skey = _suffix_key() + valid_cands = [("r2", cache.root_node, skey, [RELOAD_HASH])] + + cache._cp_l3_reserve_reloads(valid_cands, [1]) + + self.assertIn("r2", cache.ongoing_l3_reload[op_x]["waiters"]) + self.assertIs(cache.cp_l3_reload_done["r2"], False) + self.assertEqual(len(cache.ongoing_l3_reload), 1) # no new op reserved + self.assertFalse(allocator.object_ranges(RELOAD_KEY)) # allocator untouched + + def test_clean_miss_still_reserves(self): + # No false-positive: a genuine clean miss (attach_ok True, object not live) still + # reserves a new reload. The owned-pages builder is stubbed (not under test here). + allocator = _make_allocator() + cache = _make_l3_cache(allocator, _make_store()) + cache._cp_l3_build_reload_owned_pages = lambda ranges, hashes, store: {} + skey = _suffix_key() + valid_cands = [("r3", cache.root_node, skey, [RELOAD_HASH])] + + cache._cp_l3_reserve_reloads(valid_cands, [1]) + + self.assertIn(12345, cache.ongoing_l3_reload) + self.assertEqual(cache.cp_l3_reload_inflight_keys[RELOAD_KEY], 12345) + self.assertIs(cache.cp_l3_reload_done["r3"], False) + self.assertTrue(allocator.object_ranges(RELOAD_KEY)) # reserved + + +if __name__ == "__main__": + unittest.main()