From dfa168abe94c581e56a688f9fc00d675972a8967 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Thu, 11 Jun 2026 17:46:22 +0800 Subject: [PATCH] Prevent duplicate CP HiCache write acks from crashing split drain CP per-layer write-through can leave duplicated ready ack ids in the ack queue while pending-split insertion drains write visibility. The pre-scan observed both duplicates as registered before the first completion removed radix state, so the second completion raised KeyError and killed the scheduler. The write-check path now records completed ack ids until matching queue entries are drained. Duplicate completed acks are removed with a warning, while genuinely unknown or unregistered acks still fail fast. Constraint: CP HiCache pending-split drain may call writing_check outside the normal idle event path. Rejected: Blind pop(..., None) for all unknown acks | would silently hide truly unattached or corrupt ack state. Confidence: high Scope-risk: moderate Directive: Do not remove the completed-ack short-term memory unless duplicate and stale ack queues are proven impossible across TP MIN synchronization. Tested: Local py_compile for hiradix_cache.py and test_cp_hicache_metadata.py. Tested: Remote cjy-glm5-new py_compile and full test_cp_hicache_metadata.py, 121 passed. Not-tested: Full ETE prefill workload after restarting service. Co-authored-by: OmX --- python/sglang/srt/mem_cache/hiradix_cache.py | 39 +++++++++-- .../mem_cache/test_cp_hicache_metadata.py | 66 +++++++++++++++++++ 2 files changed, 101 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index b90377c98..7c68fcd27 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -10,7 +10,7 @@ import threading import time from dataclasses import dataclass, field from queue import Empty -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple import torch @@ -2877,13 +2877,36 @@ class HiRadixCache(RadixCache): self.write_backup(node) def writing_check(self, write_back=False): - def complete_write_ack_node(ack_id: int) -> TreeNode: + completed_write_ack_ids = getattr(self, "_completed_write_ack_ids", None) + if completed_write_ack_ids is None: + completed_write_ack_ids = set() + self._completed_write_ack_ids = completed_write_ack_ids + + def complete_write_ack_node(ack_id: int) -> Optional[TreeNode]: if ack_id in self.pending_host_backups: backuped_node = self._commit_pending_backup(ack_id) self.ongoing_write_through.pop(ack_id, None) + completed_write_ack_ids.add(ack_id) return backuped_node - backuped_node = self.ongoing_write_through.pop(ack_id) + backuped_node = self.ongoing_write_through.pop(ack_id, None) + if backuped_node is None: + if ack_id not in completed_write_ack_ids: + raise KeyError(ack_id) + # A duplicated ready ack can pass the pre-scan while the first + # occurrence still owns radix state, then become already-consumed + # by the time this loop reaches the second occurrence. Keep + # genuinely unattached acks deferred in the pre-scan, but make + # duplicate completion idempotent instead of crashing scheduler. + logger.warning( + "[HiCache-write] writing_check skipped already-consumed write ack: " + "ack_id=%d pending=%s ongoing=%s", + ack_id, + list(self.pending_host_backups.keys()), + list(self.ongoing_write_through.keys()), + ) + return None self.dec_node_lock_ref(backuped_node) + completed_write_ack_ids.add(ack_id) return backuped_node if write_back: @@ -2893,7 +2916,7 @@ class HiRadixCache(RadixCache): finish_event.synchronize() for ack_id in ack_list: backuped_node = complete_write_ack_node(ack_id) - if self.enable_storage: + if backuped_node is not None and self.enable_storage: self.write_backup_storage(backuped_node) self.cache_controller.ack_write_queue.clear() assert len(self.ongoing_write_through) == 0 @@ -2918,6 +2941,7 @@ class HiRadixCache(RadixCache): return ( ack_id in self.pending_host_backups or ack_id in self.ongoing_write_through + or ack_id in completed_write_ack_ids ) finish_count = 0 @@ -2966,10 +2990,17 @@ class HiRadixCache(RadixCache): finish_event.synchronize() for ack_id in ack_list: backuped_node = complete_write_ack_node(ack_id) + if backuped_node is None: + continue released_nodes.append(ack_id) if self.enable_storage: self.write_backup_storage(backuped_node) finish_count -= 1 + if completed_write_ack_ids: + remaining_ack_ids = set() + for ack in self.cache_controller.ack_write_queue: + remaining_ack_ids.update(ack.node_ids) + completed_write_ack_ids.intersection_update(remaining_ack_ids) if released_nodes: logger.debug( "[HiCache-write] writing_check released %d write locks: node_ids=%s remaining_ongoing=%d", 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 8677c0e4a..328a0286e 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -2174,6 +2174,72 @@ class TestHiRadixCacheCPBackup(CustomTestCase): self.assertEqual(cache.ongoing_write_through, {53: attached_node}) self.assertEqual(dec_locked, []) + def test_writing_check_ignores_duplicate_ready_ack_for_same_node(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.tp_world_size = 1 + cache.enable_storage = False + cache.pending_host_backups = {} + + class ReadyEvent: + def query(self): + return True + + def synchronize(self): + pass + + node = TreeNode() + node.id = 4614 + dec_locked = [] + cache.ongoing_write_through = {4614: node} + cache.dec_node_lock_ref = lambda released_node: dec_locked.append(released_node.id) + cache.cache_controller = types.SimpleNamespace( + ack_write_queue=[ + HiCacheAck(ReadyEvent(), ReadyEvent(), [4614]), + HiCacheAck(ReadyEvent(), ReadyEvent(), [4614]), + ], + ) + + cache.writing_check() + + self.assertEqual(cache.ongoing_write_through, {}) + self.assertEqual(cache.cache_controller.ack_write_queue, []) + self.assertEqual(dec_locked, [4614]) + + def test_writing_check_drops_stale_duplicate_ack_before_valid_ack(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.tp_world_size = 1 + cache.enable_storage = False + cache.pending_host_backups = {} + cache._completed_write_ack_ids = {4614} + + class ReadyEvent: + def query(self): + return True + + def synchronize(self): + pass + + node = TreeNode() + node.id = 4615 + dec_locked = [] + cache.ongoing_write_through = {4615: node} + cache.dec_node_lock_ref = lambda released_node: dec_locked.append(released_node.id) + cache.cache_controller = types.SimpleNamespace( + ack_write_queue=[ + HiCacheAck(ReadyEvent(), ReadyEvent(), [4614]), + HiCacheAck(ReadyEvent(), ReadyEvent(), [4615]), + ], + ) + + cache.writing_check() + + self.assertEqual(cache.ongoing_write_through, {}) + self.assertEqual(cache.cache_controller.ack_write_queue, []) + self.assertEqual(dec_locked, [4615]) + self.assertEqual(cache._completed_write_ack_ids, set()) + def test_write_backup_cp_failfast_on_unplanned_reservation_failure(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True