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 <omx@oh-my-codex.dev>
This commit is contained in:
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user