diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 51bf710cd..aa99d651f 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -964,6 +964,37 @@ class HiCacheController: device_indices, self.page_size, "physical_device_indices" ) + def node_has_undrained_write_ack(self, node_id: int) -> bool: + """Whether a final write ack for this node is still in ack_write_queue. + + The radix side drains acks via writing_check; until then the node's + previous backup owns radix/host state. Creating a second layer-write + state for the same node while its ack is undrained would put two acks + for one radix registration into the queue (one inc_lock_ref, two + completions) — the invariant is at most one in-flight ack per node. + """ + ack_write_queue = getattr(self, "ack_write_queue", None) + if not ack_write_queue: + return False + return any(node_id in ack.node_ids for ack in ack_write_queue) + + def cancel_layer_write_state(self, node_id: int) -> bool: + """Drop a pending per-layer write state during backup rollback. + + Without this, an exception between reservation and final ack leaves a + half-submitted state in pending_layer_writes: the next forward's layer + hooks keep driving it, eventually writing D2H into host slots the + rollback already evicted and appending an ack for a node the radix + side no longer has registered. Synchronize the write stream so any + already-issued per-layer copies finish before the caller frees the + host slots. + """ + state = self.pending_layer_writes.pop(node_id, None) + if state is None: + return False + self.write_stream.synchronize() + return True + def reserve_write_cp( self, device_indices: torch.Tensor, @@ -972,6 +1003,17 @@ class HiCacheController: ) -> HiCacheWriteReservation | HiCacheWriteFailure: from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata + if node_id >= 0 and self.node_has_undrained_write_ack(node_id): + # Refusing here (instead of crashing later in writing_check) keeps + # the failure on the existing HiCacheWriteFailure path: the caller + # skips this backup round and may retry after the ack drains. + logger.warning( + "[CacheCtrl-write] reserve_write_cp refused: node_id=%d still " + "has an undrained write ack; skipping duplicate backup", + node_id, + ) + return HiCacheWriteFailure(required_host_slots=0) + page_size = self.page_size layout = self.cp_shared_kv_layout logical_len = len(device_indices) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 7c68fcd27..a9ffabc4a 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -2090,9 +2090,39 @@ class HiRadixCache(RadixCache): self.dec_node_lock_ref(node) return node + def _remove_undrained_write_acks(self, node_id: int) -> bool: + """Scrub this node's final write acks out of the controller queue. + + Every rollback that clears the radix-side registration + (ongoing_write_through / pending_host_backups) MUST also scrub the + node's acks: an orphaned ack re-opens the _node_host_write_pending + guard for a fresh registration while the stale ack is still queued, + producing two acks for one registration (one inc_lock_ref, two + completions) and crashing writing_check on the second. + """ + retained_acks = [] + removed = False + for ack in self.cache_controller.ack_write_queue: + if node_id not in ack.node_ids: + retained_acks.append(ack) + continue + ack.finish_event.synchronize() + remaining_node_ids = [nid for nid in ack.node_ids if nid != node_id] + if remaining_node_ids: + retained_acks.append(ack._replace(node_ids=remaining_node_ids)) + removed = True + if removed: + self.cache_controller.ack_write_queue = retained_acks + return removed + def _rollback_pending_backup(self, node_id: int) -> TreeNode: pending = self.pending_host_backups.pop(node_id) node = pending.node + # Cancel the half-submitted per-layer state (so later layer hooks + # cannot write into the host slots evicted below) and scrub any + # already-appended final ack before releasing the registration. + self.cache_controller.cancel_layer_write_state(node_id) + self._remove_undrained_write_acks(node_id) self.cache_controller.evict_cp_host(pending.metadata) if node.cp_hicache is pending.metadata: node.cp_hicache = None @@ -2271,21 +2301,7 @@ class HiRadixCache(RadixCache): f"while per-layer writes are still pending; reason={reason}" ) - retained_acks = [] - removed_ack = False - for ack in self.cache_controller.ack_write_queue: - if prepared.node_id not in ack.node_ids: - retained_acks.append(ack) - continue - ack.finish_event.synchronize() - remaining_node_ids = [ - node_id for node_id in ack.node_ids if node_id != prepared.node_id - ] - if remaining_node_ids: - retained_acks.append(ack._replace(node_ids=remaining_node_ids)) - removed_ack = True - if removed_ack: - self.cache_controller.ack_write_queue = retained_acks + self._remove_undrained_write_acks(prepared.node_id) logger.debug( "[HiCache-write] rollback unattached prepared CP backup: node_id=%d logical_len=%d reason=%s", prepared.node_id, diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index 328a0286e..f8d21f684 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -3,7 +3,7 @@ import re import sys import types import unittest -from unittest.mock import patch +from unittest.mock import MagicMock, patch import torch @@ -1320,7 +1320,9 @@ class TestHiRadixCacheCPBackup(CustomTestCase): cache._uses_cp_hicache = True evicted = [] cache.cache_controller = types.SimpleNamespace( - evict_cp_host=lambda metadata: evicted.append(metadata) or 2 + evict_cp_host=lambda metadata: evicted.append(metadata) or 2, + cancel_layer_write_state=lambda node_id: False, + ack_write_queue=[], ) dec_locked = [] cache.dec_node_lock_ref = lambda node: dec_locked.append(node) @@ -2174,6 +2176,104 @@ class TestHiRadixCacheCPBackup(CustomTestCase): self.assertEqual(cache.ongoing_write_through, {53: attached_node}) self.assertEqual(dec_locked, []) + def test_reserve_write_cp_refuses_node_with_undrained_ack(self): + """Producer-side invariant: at most one in-flight ack per node_id. + A reservation for a node whose previous final ack is still queued is + refused on the normal HiCacheWriteFailure path instead of creating a + second layer-write state (which would enqueue a duplicate ack).""" + from sglang.srt.managers.cache_controller import ( + HiCacheController, + HiCacheWriteFailure, + ) + + controller = types.SimpleNamespace( + ack_write_queue=[ + HiCacheAck(MagicMock(), MagicMock(), [4614]), + ], + ) + controller.node_has_undrained_write_ack = types.MethodType( + HiCacheController.node_has_undrained_write_ack, controller + ) + result = HiCacheController.reserve_write_cp( + controller, + device_indices=torch.arange(4, dtype=torch.int64), + node_id=4614, + ) + self.assertIsInstance(result, HiCacheWriteFailure) + self.assertEqual(result.required_host_slots, 0) + # Different node id is not gated by the queue scan. + self.assertFalse( + HiCacheController.node_has_undrained_write_ack(controller, 4615) + ) + + def test_rollback_pending_backup_scrubs_acks_and_cancels_state(self): + """The write_backup exception rollback must not orphan an ack or a + half-submitted layer-write state: an orphaned ack re-opens the + registration guard while still queued, producing the duplicate-ack + crash; a live state would keep writing into evicted host slots.""" + cache = HiRadixCache.__new__(HiRadixCache) + node = TreeNode() + node.id = 4614 + metadata = object() + dec_locked = [] + cache.pending_host_backups = { + 4614: PendingHiCacheBackup( + node=node, + metadata=metadata, + logical_len=8, + submitted=True, + locked=True, + ) + } + cache.dec_node_lock_ref = lambda n: dec_locked.append(n.id) + + finish_event = MagicMock() + cancelled = [] + evicted = [] + cache.cache_controller = types.SimpleNamespace( + ack_write_queue=[ + HiCacheAck(MagicMock(), finish_event, [4614]), + HiCacheAck(MagicMock(), MagicMock(), [4615]), + ], + cancel_layer_write_state=lambda node_id: cancelled.append(node_id), + evict_cp_host=lambda meta: evicted.append(meta), + ) + node.cp_hicache = metadata + + returned = cache._rollback_pending_backup(4614) + + self.assertIs(returned, node) + self.assertEqual(cache.pending_host_backups, {}) + self.assertEqual(cancelled, [4614]) + self.assertEqual(evicted, [metadata]) + # The 4614 ack is scrubbed (after syncing its finish event); the + # unrelated 4615 ack is retained. + finish_event.synchronize.assert_called_once() + self.assertEqual( + [ack.node_ids for ack in cache.cache_controller.ack_write_queue], + [[4615]], + ) + self.assertEqual(dec_locked, [4614]) + self.assertIsNone(node.cp_hicache) + self.assertEqual(node.host_len, 0) + + def test_remove_undrained_write_acks_splits_shared_ack(self): + """A grouped ack covering several nodes only loses the rolled-back + node id; other nodes' completion is preserved.""" + cache = HiRadixCache.__new__(HiRadixCache) + cache.cache_controller = types.SimpleNamespace( + ack_write_queue=[ + HiCacheAck(MagicMock(), MagicMock(), [4614, 4615]), + ], + ) + removed = cache._remove_undrained_write_acks(4614) + self.assertTrue(removed) + self.assertEqual( + [ack.node_ids for ack in cache.cache_controller.ack_write_queue], + [[4615]], + ) + self.assertFalse(cache._remove_undrained_write_acks(9999)) + def test_writing_check_ignores_duplicate_ready_ack_for_same_node(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True