diff --git a/python/sglang/srt/mem_cache/cp_shared_l2_pool.py b/python/sglang/srt/mem_cache/cp_shared_l2_pool.py index 058a6e699..592907990 100644 --- a/python/sglang/srt/mem_cache/cp_shared_l2_pool.py +++ b/python/sglang/srt/mem_cache/cp_shared_l2_pool.py @@ -804,6 +804,30 @@ class CpSharedL2PageAllocator: self._stats["cp_shared_l2_objects_evicted"] += 1 return dropped + def mark_object_committed(self, object_key: str) -> bool: + """B1 commit: mark an object committed directly, bypassing the per-(rank, + layer, payload) commit_layer quorum. + + Under the B1 collective-free allocator the all-ranks-done consensus is the + writing_check ReduceOp.MIN frontier (the same barrier that releases the + node's write lock), NOT a per-layer all-gather. When that frontier commits a + node, every CP rank calls this with the same object_key, so the committed set + transitions rank-uniformly (and feeds placement_digest). Does not populate + _commits_by_object -- the per-layer quorum (_has_full_commit) is unused under + B1; is_committed / release / split_committed_object read _committed_objects, + which this maintains. Returns True on the first commit, False if already + committed. Raises on an unknown object (fail-loud: the node's reservation + must outlive its write lock). + """ + self._validate_object_key(object_key) + if object_key not in self._ranges_by_object: + raise ValueError(f"cannot commit unknown object {object_key!r}") + if object_key in self._committed_objects: + return False + self._committed_objects.add(object_key) + self._stats["cp_shared_l2_objects_committed"] += 1 + return True + def split_committed_object( self, object_key: str, diff --git a/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py b/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py index 1a26c5153..e0abce0a1 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py +++ b/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py @@ -2784,5 +2784,83 @@ class TestCpSharedL2PlacementDigest(unittest.TestCase): self.assertNotEqual(a.placement_digest(), b.placement_digest()) +class TestCpSharedL2MarkObjectCommitted(unittest.TestCase): + """B1 commit path (2b.1b): mark_object_committed replaces the per-(rank,layer, + payload) quorum -- the writing_check MIN frontier is the consensus. Locks in + MF4 (split stays coherent for a B1-committed object with no per-layer commits).""" + + def make_allocator(self, pages=16): + return _cp_shared_l2_pool.CpSharedL2PageAllocator( + pages_per_payload={PAYLOAD_TARGET_KV: pages}, + slab_ids_by_payload={PAYLOAD_TARGET_KV: 10}, + expected_ranks=(0, 1, 2, 3), + expected_layers=(0, 1), + required_payloads=(PAYLOAD_TARGET_KV,), + ) + + def test_mark_commits_without_per_layer_quorum(self): + a = self.make_allocator() + a.reserve("n", PAYLOAD_TARGET_KV, 3) + self.assertFalse(a.is_committed("n")) + self.assertTrue(a.mark_object_committed("n")) + self.assertTrue(a.is_committed("n")) + + def test_mark_is_idempotent(self): + a = self.make_allocator() + a.reserve("n", PAYLOAD_TARGET_KV, 2) + self.assertTrue(a.mark_object_committed("n")) + self.assertFalse(a.mark_object_committed("n")) + + def test_mark_unknown_object_raises(self): + a = self.make_allocator() + with self.assertRaises(ValueError): + a.mark_object_committed("never-reserved") + + def test_mark_bumps_committed_stat(self): + a = self.make_allocator() + a.reserve("n", PAYLOAD_TARGET_KV, 2) + before = a.stats()["cp_shared_l2_objects_committed"] + a.mark_object_committed("n") + self.assertEqual(a.stats()["cp_shared_l2_objects_committed"], before + 1) + + def test_mark_changes_placement_digest(self): + # committed set is part of the digest, so commit transitions are visible to + # the cross-rank assert; and two ranks committing in lockstep stay equal. + a, b = self.make_allocator(), self.make_allocator() + a.reserve("n", PAYLOAD_TARGET_KV, 3) + b.reserve("n", PAYLOAD_TARGET_KV, 3) + self.assertEqual(a.placement_digest(), b.placement_digest()) + a.mark_object_committed("n") + self.assertNotEqual(a.placement_digest(), b.placement_digest()) + b.mark_object_committed("n") + self.assertEqual(a.placement_digest(), b.placement_digest()) + + def test_split_of_marked_committed_object_is_coherent(self): + # MF4: split a B1-committed object (no per-layer commit facts) -> children + # are committed, pages preserved, no error from the empty _commits_by_object. + a = self.make_allocator() + a.reserve("child", PAYLOAD_TARGET_KV, 6) + a.mark_object_committed("child") + parent_ranges, child_ranges = a.split_committed_object( + "child", + split_pages_by_payload={PAYLOAD_TARGET_KV: 2}, + parent_object_key="parent", + child_object_key="child", + ) + self.assertTrue(a.is_committed("parent")) + self.assertTrue(a.is_committed("child")) + self.assertEqual(parent_ranges[PAYLOAD_TARGET_KV].num_pages, 2) + self.assertEqual(child_ranges[PAYLOAD_TARGET_KV].num_pages, 4) + + def test_release_after_mark_frees_pages(self): + a = self.make_allocator(pages=8) + free0 = a.free_pages(PAYLOAD_TARGET_KV) + a.reserve("n", PAYLOAD_TARGET_KV, 4) + a.mark_object_committed("n") + self.assertTrue(a.release("n")) + self.assertEqual(a.free_pages(PAYLOAD_TARGET_KV), free0) + self.assertFalse(a.is_committed("n")) + + if __name__ == "__main__": unittest.main()