Revert "Fix CP HiCache writing_check deadlock at the source: rank-replicated symmetric attach"
This reverts commit 85585fcf4b.
This commit is contained in:
@@ -1,185 +0,0 @@
|
||||
"""Regression: CP-HiCache radix probes must read the rank-replicated cp_backup_pending
|
||||
marker, not the drain-timed ongoing_write_through / pending_host_backups dicts.
|
||||
|
||||
The deadlock root cause: the attach/defer/prune/split probes read
|
||||
`_node_host_write_pending` (membership in the dicts that are drained ASYNCHRONOUSLY at
|
||||
per-rank ack completion), so `logical_len != len(value)` could diverge across CP ranks ->
|
||||
asymmetric `write_backup` catch-up -> asymmetric ack-append -> the `ack_queue_len==0` skip
|
||||
in `writing_check` splits the collective -> NCCL hang. The fix makes those probes read
|
||||
`TreeNode.cp_backup_pending`, which is set at the rank-replicated prepared-backup attach
|
||||
and cleared only at the MIN-committed commit/rollback -> rank-uniform by construction, so
|
||||
the attach is symmetric and `ack_write_queue` stays symmetric.
|
||||
|
||||
See docs_internal/cp_hicache_symmetric_attach_design.md. Run in the dev-cu13 container.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
# Prefer the real sgl_kernel; fall back to a stub on CPU-only hosts.
|
||||
try:
|
||||
import sgl_kernel # noqa: F401
|
||||
except (ImportError, RuntimeError):
|
||||
if "sgl_kernel" not in sys.modules:
|
||||
stub = types.ModuleType("sgl_kernel")
|
||||
stub.__path__ = []
|
||||
|
||||
def _getattr(name):
|
||||
if name.startswith("__"):
|
||||
raise AttributeError(name)
|
||||
fn = lambda *a, **k: None
|
||||
setattr(stub, name, fn)
|
||||
return fn
|
||||
|
||||
stub.__getattr__ = _getattr
|
||||
sys.modules["sgl_kernel"] = stub
|
||||
|
||||
from sglang.srt.mem_cache.hiradix_cache import ( # noqa: E402
|
||||
HiCachePendingBackupSplit,
|
||||
HiRadixCache,
|
||||
PendingHiCacheBackup,
|
||||
PreparedCpHiCacheBackup,
|
||||
)
|
||||
from sglang.srt.mem_cache.radix_cache import TreeNode # noqa: E402
|
||||
|
||||
|
||||
def _cp_cache():
|
||||
c = HiRadixCache.__new__(HiRadixCache)
|
||||
c._uses_cp_hicache = True
|
||||
c.ongoing_write_through = {}
|
||||
c.pending_host_backups = {}
|
||||
return c
|
||||
|
||||
|
||||
class TestRadixProbesReadReplicatedMarker(unittest.TestCase):
|
||||
def test_split_still_pending_reads_marker_not_dicts(self):
|
||||
c = _cp_cache()
|
||||
node = TreeNode(id=42)
|
||||
# In the drain-timed dicts but marker False -> probe must IGNORE the dicts.
|
||||
c.ongoing_write_through[42] = node
|
||||
c.pending_host_backups[42] = object()
|
||||
self.assertFalse(
|
||||
c._cp_node_split_still_pending(node),
|
||||
"probe must read the replicated marker, not ongoing_write_through/pending",
|
||||
)
|
||||
# Marker True, NOT in any dict -> probe must return True (reads the marker).
|
||||
c.ongoing_write_through.clear()
|
||||
c.pending_host_backups.clear()
|
||||
node.cp_backup_pending = True
|
||||
self.assertTrue(c._cp_node_split_still_pending(node))
|
||||
|
||||
def test_split_node_guard_raises_on_marker(self):
|
||||
c = _cp_cache()
|
||||
child = TreeNode(id=7)
|
||||
child.cp_backup_pending = True
|
||||
# The guard is the first statement in _split_node; it raises before touching
|
||||
# the key / child structure, so a stub key is fine.
|
||||
with self.assertRaises(HiCachePendingBackupSplit):
|
||||
c._split_node(None, child, 1)
|
||||
|
||||
def test_unprunable_ignores_lock_ref_reads_marker_and_pin(self):
|
||||
c = _cp_cache()
|
||||
node = TreeNode(id=9)
|
||||
node.children = {}
|
||||
# lock_ref is rank-local (backup lock) -> must be IGNORED now.
|
||||
node.lock_ref = 5
|
||||
node.host_ref_counter = 0
|
||||
node.cp_backup_pending = False
|
||||
self.assertFalse(
|
||||
c._cp_subtree_has_unprunable_state(node),
|
||||
"must not read lock_ref (its only cross-rank skew is the backup lock)",
|
||||
)
|
||||
# host_ref_counter (replicated pin) still blocks.
|
||||
node.host_ref_counter = 1
|
||||
self.assertTrue(c._cp_subtree_has_unprunable_state(node))
|
||||
# The replicated marker blocks.
|
||||
node.host_ref_counter = 0
|
||||
node.cp_backup_pending = True
|
||||
self.assertTrue(c._cp_subtree_has_unprunable_state(node))
|
||||
|
||||
def test_inc_hit_count_short_circuits_on_marker_not_dicts(self):
|
||||
c = _cp_cache()
|
||||
c.cache_controller = SimpleNamespace(write_policy="write_through")
|
||||
c.write_through_threshold = 1
|
||||
reached_backuped_check = {"hit": False}
|
||||
c._node_backuped = lambda n: reached_backuped_check.__setitem__("hit", True) or True
|
||||
|
||||
node = TreeNode(id=3)
|
||||
node.hit_count = 0
|
||||
# Marker True -> _inc_hit_count returns right after the marker check, before
|
||||
# the _node_backuped/threshold path.
|
||||
node.cp_backup_pending = True
|
||||
c._inc_hit_count(node)
|
||||
self.assertFalse(
|
||||
reached_backuped_check["hit"], "marker True must short-circuit _inc_hit_count"
|
||||
)
|
||||
# Marker False but node in the drain-timed dicts -> must NOT short-circuit
|
||||
# (proves it reads the marker, not the dicts).
|
||||
node.cp_backup_pending = False
|
||||
c.ongoing_write_through[node.id] = node
|
||||
c.pending_host_backups[node.id] = object()
|
||||
c._inc_hit_count(node)
|
||||
self.assertTrue(
|
||||
reached_backuped_check["hit"], "dicts must be ignored; marker False -> proceed"
|
||||
)
|
||||
|
||||
|
||||
class TestMarkerLifecycle(unittest.TestCase):
|
||||
def test_set_at_attach_cleared_at_commit(self):
|
||||
c = _cp_cache()
|
||||
c.inc_node_lock_ref = lambda n: None
|
||||
c.dec_node_lock_ref = lambda n: None
|
||||
node = TreeNode(id=None)
|
||||
node.value = torch.arange(4)
|
||||
meta = SimpleNamespace(owned_positions=torch.arange(4))
|
||||
prepared = PreparedCpHiCacheBackup(
|
||||
node_id=node.id, reservation=SimpleNamespace(), metadata=meta, logical_len=4
|
||||
)
|
||||
self.assertFalse(node.cp_backup_pending)
|
||||
c._attach_prepared_cp_backup(node, prepared)
|
||||
self.assertTrue(node.cp_backup_pending, "attach must set the marker")
|
||||
|
||||
c.pending_host_backups[node.id] = PendingHiCacheBackup(
|
||||
node=node, metadata=meta, logical_len=4, locked=True
|
||||
)
|
||||
c._commit_pending_backup(node.id)
|
||||
self.assertFalse(node.cp_backup_pending, "commit must clear the marker")
|
||||
self.assertEqual(node.host_len, 4)
|
||||
|
||||
def test_double_attach_fails_loud(self):
|
||||
c = _cp_cache()
|
||||
c.inc_node_lock_ref = lambda n: None
|
||||
node = TreeNode(id=5)
|
||||
node.value = torch.arange(2)
|
||||
node.cp_backup_pending = True # already pending
|
||||
prepared = PreparedCpHiCacheBackup(
|
||||
node_id=5, reservation=SimpleNamespace(), metadata=SimpleNamespace(), logical_len=2
|
||||
)
|
||||
with self.assertRaises(RuntimeError):
|
||||
c._attach_prepared_cp_backup(node, prepared)
|
||||
|
||||
def test_cleared_on_rollback(self):
|
||||
c = _cp_cache()
|
||||
c.dec_node_lock_ref = lambda n: None
|
||||
c.cache_controller = SimpleNamespace(
|
||||
cancel_layer_write_state=lambda nid: None,
|
||||
evict_cp_host=lambda m: None,
|
||||
ack_write_queue=[],
|
||||
)
|
||||
node = TreeNode(id=11)
|
||||
meta = SimpleNamespace()
|
||||
node.cp_hicache = meta
|
||||
node.cp_backup_pending = True
|
||||
c.pending_host_backups[node.id] = PendingHiCacheBackup(
|
||||
node=node, metadata=meta, logical_len=3, locked=True
|
||||
)
|
||||
c._rollback_pending_backup(node.id)
|
||||
self.assertFalse(node.cp_backup_pending, "rollback must clear the marker")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user