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) <noreply@anthropic.com>

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:<hash> 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:<hash>' 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:<hash> 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()