From f11656d98e0a9857f2ed235ba2d3442e4387f562 Mon Sep 17 00:00:00 2001 From: Zhiqiang Xie Date: Sun, 1 Nov 2026 01:18:35 -0800 Subject: [PATCH] [Fix] hicache: vote the prefetch timeout once per tick, not per poll Under the default timeout prefetch policy the per-request deadline verdict was agreed by a MAX collective issued from inside the admission loop, once per waiting request per scheduling pass. The loop's exits are evaluated per rank, so one rank could enter a collective its peers never issued (#30760), and even matched counts cost one gloo all_reduce per waiting request. check_hicache_events now votes once per tick: one packed MAX over every live prefetch in ongoing_prefetch order, evaluated on PP rank 0 with the existing per-request budget and relayed along PP by _all_reduce. For every agreed 1 the tick marks the deadline expired and terminates the operation, so the transfer worker and the hit drain stop spending on it; retirement stays with the admission poll or the drain. _can_terminate_prefetch under timeout reads the recorded verdict and issues no collective, so check_prefetch_progress is collective-free for every policy. Because the polls retire per rank, the vote's shape would depend on the poll sequence. The ready-count MIN sync's piggybacked digest now folds in the live prefetch set, so a divergent set raises on every rank in the same tick instead of hanging in a vote of a different length; the row width is unchanged, so the shm fast path is unaffected. Tests: the poll-only agreement case in test_storage_prefetch_lifecycle.py is adapted to agree through the tick; new cases pin zero collectives from polls and one from the tick, uneven polls before a verdict, the same-tick raise on divergent live sets after a verdict, and the PP relay of the vote. --- .../srt/mem_cache/unified_radix_cache.py | 85 +++++- .../test_shm_tag_exchange.py | 3 + .../mem_cache/test_hicache_pp_sync_drain.py | 87 ++++++- .../test_storage_prefetch_lifecycle.py | 243 +++++++++++++++++- 4 files changed, 396 insertions(+), 22 deletions(-) diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 455b22a563ed..af8b4f0061ff 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -4,6 +4,7 @@ import logging import threading import time +import zlib from array import array from contextlib import nullcontext from dataclasses import replace @@ -190,6 +191,7 @@ class _OngoingLoadBack(NamedTuple): class _StoragePrefetchWait(msgspec.Struct): timer: PollTimer = msgspec.field(default_factory=PollTimer) timeout_s: float = 0.0 + # Rank-agreed verdict of the per-tick vote; admission polls only read it. expired: bool = False @@ -2643,18 +2645,65 @@ def _can_terminate_prefetch(self, operation: PrefetchOperation) -> bool: # progress, rank-min queue drain); a vote here only adds collectives. return False if self.prefetch_stop_policy == "timeout": - # Wall clocks differ across ranks: PP0 decides, TP takes MAX (any - # rank timed out) and _all_reduce broadcasts the verdict along PP. - should_terminate = False - if self.pp_rank == 0: - should_terminate = self._prefetch_deadline_expired(operation) - should_terminate_tensor = torch.tensor( - int(should_terminate), dtype=torch.int, device="cpu" - ) - self._all_reduce(should_terminate_tensor, torch.distributed.ReduceOp.MAX) - return should_terminate_tensor.item() == 1 + # The verdict is agreed once per tick by _vote_prefetch_deadlines, + # which also terminates the operation (the poll then takes the + # is_terminated path); a poll only reads it, so ranks may poll + # different requests (batch_is_full, allocatable-request exits) + # without a collective. + wait = self._storage_prefetch_wait.get(operation.handle) + return wait is not None and wait.expired return True + def _vote_prefetch_deadlines(self) -> None: + """One MAX vote per tick over every live prefetch's `timeout` budget. + + Wall clocks differ across ranks: PP0 decides per operation, TP takes + MAX (any rank timed out) and _all_reduce relays the verdict along PP. + The packed tensor has one slot per `ongoing_prefetch` entry in + insertion order; this tick's ready-count sync asserted that set is + rank-identical (_hicache_sync_digest), and the drain between them + only retires in lockstep. With no live prefetch no rank issues a + collective. A wait that has not started yet (no blocked poll on PP0) + contributes 0 until the next tick. An expired operation is terminated + here so the transfer worker and the hit drain stop spending on it; + its retirement stays with the admission poll or the drain. + """ + if self.prefetch_stop_policy != "timeout" or not self.ongoing_prefetch: + return + operations = [info.operation for info in self.ongoing_prefetch.values()] + expired = [0] * len(operations) + if self.pp_rank == 0: + expired = [ + int(self._prefetch_deadline_expired(operation)) + for operation in operations + ] + verdicts = torch.tensor(expired, dtype=torch.int, device="cpu") + self._all_reduce(verdicts, torch.distributed.ReduceOp.MAX) + for operation, verdict in zip(operations, verdicts.tolist()): + if verdict: + self._mark_prefetch_deadline_expired(operation) + self.cache_controller.terminate_prefetch(operation) + + def _hicache_sync_digest(self) -> int: + """Value piggybacked on the ready-count MIN sync: equal on every rank + iff the write_back reclaim victim order matched and, when the deadline + vote reads it, so did the live prefetch set (`ongoing_prefetch` in + insertion order). The admission polls retire per request, so a rank + that left the loop earlier than its peers (#30760) is caught here, in + the same tick on every rank, instead of hanging in a vote of a + different length. Fits int64: 42-bit digest * 1000003 + crc32.""" + digest = self.tree_core.write_back_duplicate_reclaim_digest + if ( + self.enable_storage + and self.prefetch_stop_policy == "timeout" + and self.ongoing_prefetch + ): + live = "\0".join( + f"{handle.rid}\0{handle.attempt_id}" for handle in self.ongoing_prefetch + ) + digest = digest * 1000003 + zlib.crc32(live.encode()) + 1 + return digest + def has_ongoing_prefetch(self, handle: CacheRequestHandle) -> bool: return handle in self.ongoing_prefetch @@ -3396,6 +3445,9 @@ def _try_alloc_storage_hit(operation) -> bool: if info is None: return True # aborted/cleaned; nothing to retry if operation.is_terminated(): + if operation.deadline_expired: + # Cut by the tick vote while parked: same reason as the poll. + self._finish_storage_prefetch(request, 0, reason="deadline") self.revoke_pending_prefetch(request) return True @@ -3791,8 +3843,8 @@ def _sync_hicache_ready_counts( ) # Piggybacked TP check: [digest, -digest] MIN-reduces to [min, -max], - # equal iff reclaim victim order matched on every rank. - digest = self.tree_core.write_back_duplicate_reclaim_digest + # equal iff reclaim victim order and live prefetch set matched on every rank. + digest = self._hicache_sync_digest() exchange = self._shm_count_exchange() if exchange is not None: values = [write_acks, load_acks, *storage_queue_sizes, *fold_u64(digest)] @@ -3801,7 +3853,8 @@ def _sync_hicache_ready_counts( if len(values) <= exchange.width: rows = exchange.allgather(values) assert all(row[-2:] == rows[0][-2:] for row in rows), ( - "write_back duplicate-reclaim victims diverged across TP ranks" + "write_back duplicate-reclaim victims or live storage " + "prefetches diverged across TP ranks" ) mins = [min(col) for col in zip(*rows)] return mins[0], mins[1], tuple(mins[2:-2]), extra_pool_names @@ -3823,7 +3876,8 @@ def _sync_hicache_ready_counts( count_values = list(map(int, ready_counts.tolist())) assert digest == count_values[-2] and digest == -count_values[-1], ( - "write_back duplicate-reclaim victims diverged across PP/TP ranks" + "write_back duplicate-reclaim victims or live storage prefetches " + "diverged across PP/TP ranks" ) return ( count_values[0], @@ -4105,6 +4159,9 @@ def check_hicache_events(self) -> None: extra_release_counts=extra_release_counts, log_metrics=True, ) + if self.enable_storage: + # After the drain: a prefetch it retired is not voted on. + self._vote_prefetch_deadlines() if self.buffer_pipeline is not None: self.buffer_pipeline.flush_pending_writes() if self.enable_storage_metrics and self.storage_metrics_collector is not None: diff --git a/test/registered/unit/managers/scheduler_components/test_shm_tag_exchange.py b/test/registered/unit/managers/scheduler_components/test_shm_tag_exchange.py index faba5fe25b6f..e595c41642a6 100644 --- a/test/registered/unit/managers/scheduler_components/test_shm_tag_exchange.py +++ b/test/registered/unit/managers/scheduler_components/test_shm_tag_exchange.py @@ -50,6 +50,9 @@ def _hicache_tp_rank(rank: int, extra_pools: int) -> UnifiedRadixCache: ) cache.pp_rank, cache.pp_size, cache.pp_group = 0, 1, None cache.host_memory_mode = "cache" + # The tick reads the stop policy and the live prefetch set (deadline vote). + cache.prefetch_stop_policy = "wait_complete" + cache.ongoing_prefetch = {} cache.attn_cp_group = cache.attn_tp_group = None cache.tp_group, cache.tp_world_size = torch.distributed.group.WORLD, 2 cache.cache_controller = SimpleNamespace( diff --git a/test/registered/unit/mem_cache/test_hicache_pp_sync_drain.py b/test/registered/unit/mem_cache/test_hicache_pp_sync_drain.py index 8d335ce6a28c..7e705b682361 100644 --- a/test/registered/unit/mem_cache/test_hicache_pp_sync_drain.py +++ b/test/registered/unit/mem_cache/test_hicache_pp_sync_drain.py @@ -3,11 +3,17 @@ import unittest from queue import Queue from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import torch -from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache +from sglang.srt.mem_cache.base_prefix_cache import CacheRequestHandle +from sglang.srt.mem_cache.storage_prefetch import StoragePrefetchRetries +from sglang.srt.mem_cache.unified_radix_cache import ( + UnifiedRadixCache, + _OngoingPrefetch, + _StoragePrefetchWait, +) from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=10, suite="base-a-test-cpu") @@ -23,6 +29,10 @@ def _make_cache(self, pp_rank, write_ready, load_ready): cache.pp_rank = pp_rank cache.pp_size = 2 cache.host_memory_mode = "cache" + # The per-tick deadline vote runs under storage; no live prefetch + # means it issues no collective, which the asserts below pin. + cache.prefetch_stop_policy = "timeout" + cache.ongoing_prefetch = {} cache.enable_storage_metrics = False cache.storage_metrics_collector = None cache.buffer_pipeline = None @@ -48,6 +58,79 @@ def _make_cache(self, pp_rank, write_ready, load_ready): ) return cache + def _arm_live_prefetch(self, cache): + """Cache-mode storage with one live prefetch on this PP stage.""" + cache.tree_core.enable_storage = True + cache._drain_storage_control_queues_impl = MagicMock() + cache.storage_prefetch_retries = StoragePrefetchRetries() + cache._storage_prefetch_wait = {} + cc = cache.cache_controller + for name in ( + "prefetch_hit_queue", + "ack_prefetch_queue", + "ack_backup_queue", + "host_mem_release_queue", + ): + setattr(cc, name, Queue()) + cc.extra_host_mem_release_queues = {} + cc.terminate_prefetch = MagicMock() + handle = CacheRequestHandle("r", 0) + operation = SimpleNamespace( + handle=handle, + deadline_expired=False, + host_indices=None, + hash_value=[], + completed_tokens=0, + ) + cache.ongoing_prefetch[handle] = _OngoingPrefetch( + 0, None, None, operation, None, {} + ) + return handle, operation + + def test_pp_relays_the_deadline_vote_after_the_counts(self): + """PP0 evaluates the clocks and votes; a follower contributes zeros, + receives the verdict through the same tag after the counts, and marks + and terminates the operation without reading its own clock.""" + leader = self._make_cache(0, [True], [True]) + handle, operation = self._arm_live_prefetch(leader) + wait = _StoragePrefetchWait(timeout_s=1.0) + with patch("time.perf_counter", return_value=100.0): + wait.timer.observe(ready=False) + leader._storage_prefetch_wait[handle] = wait + with patch("time.perf_counter", return_value=103.0): + leader.check_hicache_events() + + counts, vote = leader._all_reduce.call_args_list + self.assertIs(counts.args[1], torch.distributed.ReduceOp.MIN) + self.assertEqual(vote.args[0].tolist(), [1]) + self.assertIs(vote.args[1], torch.distributed.ReduceOp.MAX) + self.assertTrue(wait.expired) + self.assertTrue(operation.deadline_expired) + leader.cache_controller.terminate_prefetch.assert_called_once_with(operation) + + follower = self._make_cache(1, [True], [True]) + handle, operation = self._arm_live_prefetch(follower) + follower._prefetch_deadline_expired = MagicMock() + relayed = [] + + def relay_from_pp0(data, _): + relayed.append(data.tolist()) + if len(relayed) == 1: + data[0], data[1] = 1, 1 # counts; digest slots stay local + else: + data.fill_(1) # PP0's verdict + + follower._all_reduce.side_effect = relay_from_pp0 + follower.check_hicache_events() + + self.assertEqual(relayed[1], [0]) + self.assertEqual(follower._all_reduce.call_count, 2) + follower._prefetch_deadline_expired.assert_not_called() + self.assertTrue(follower._storage_prefetch_wait[handle].expired) + self.assertTrue(operation.deadline_expired) + follower.cache_controller.terminate_prefetch.assert_called_once_with(operation) + follower.writing_check.assert_called_once_with(finish_count=1) + def test_pp_batches_write_and_load_counts_once(self): leader = self._make_cache(0, [True, False], [True, True]) leader.check_hicache_events() diff --git a/test/registered/unit/mem_cache/test_storage_prefetch_lifecycle.py b/test/registered/unit/mem_cache/test_storage_prefetch_lifecycle.py index 950b2eddc6ab..a4bd5770b9cc 100644 --- a/test/registered/unit/mem_cache/test_storage_prefetch_lifecycle.py +++ b/test/registered/unit/mem_cache/test_storage_prefetch_lifecycle.py @@ -1,6 +1,7 @@ """Staged L3 prefetch lifecycle through the buffer pipeline; no GPU kernels.""" import tempfile +import threading import unittest from array import array from collections import defaultdict, deque @@ -69,6 +70,9 @@ def _staged_fixture(full_match=2): cache.storage_metrics_collector = None cache.ongoing_prefetch = {} cache._prefetch_outcome_stats = defaultdict(int) + # Backup-retry state reset() sets; read by the tick of #47626 (inert here). + cache._host_capacity_backup_retries = {} + cache._backup_write_specs = {} cache.tree_components = [] cache.components = {} index_slots(cache, cache.components) @@ -80,6 +84,10 @@ def _staged_fixture(full_match=2): full_available_size=Mock(return_value=100) ) cc = HybridCacheController.__new__(HybridCacheController) + # Worker-thread state __init__ sets; read by the tick of #47617 (inert here, no workers started). + cc.storage_stop_event = threading.Event() + cc.prefetch_thread = cc.prefetch_io_aux_thread = None + cc.prefetch_sync_thread = cc.backup_thread = None cc.pp_prefetch_command_group = None cc.page_size = 2 cc.get_hash_str = get_hash_str @@ -150,11 +158,15 @@ def _hit_drain_fixture(): cache.storage_prefetch_retries = StoragePrefetchRetries() cache._storage_prefetch_wait = {} cache._prefetch_outcome_stats = defaultdict(int) + # Backup-retry state reset() sets; read by the tick of #47626 (inert here). + cache._host_capacity_backup_retries = {} + cache._backup_write_specs = {} cache._storage_prefetch_hit_remaining_by_reqid = {} cache._record_storage_prefetch_hit = Mock() cache.revoke_pending_prefetch = Mock() cache.buffer_pipeline = SimpleNamespace(pending_hit_allocs=deque()) cache.cache_controller = SimpleNamespace( + storage_workers_alive=lambda: True, prefetch_hit_queue=Queue(), ack_prefetch_queue=Queue(), ack_backup_queue=Queue(), @@ -201,6 +213,52 @@ def _timeout_fixture(mode="buffer_only", policy="timeout"): return cache, cache.ongoing_prefetch[_REQ].operation +def _arm_tick(cache, sync_counts=False): + """Let check_hicache_events run on the CPU fixture: the drains are + stubbed, the deadline vote and the polls are real. With sync_counts the + ready-count sync is real too (gloo path, empty queues), so its piggybacked + live-prefetch digest check runs.""" + cache.pp_size = 1 + cache.enable_storage = True + cache._drain_async_work = Mock() + cache.flush_pending_backups = Mock() + if sync_counts: + cc = cache.cache_controller + cache.tree_core.write_back_duplicate_reclaim_digest = 0 + cache._shm_count_exchange = lambda: None + cc.ack_write_queue = [] + cc.ack_load_queue = [] + for name in ( + "prefetch_hit_queue", + "ack_prefetch_queue", + "ack_backup_queue", + "host_mem_release_queue", + ): + setattr(cc, name, Queue()) + cc.extra_host_mem_release_queues = {} + else: + cache._sync_hicache_ready_counts = Mock(return_value=(0, 0, (0, 0, 0, 0), ())) + cache.writing_check = Mock() + cache.loading_check = Mock() + cache._drain_storage_control_queues_impl = Mock() + + +def _counting_all_reduce(reduces): + """Records each collective's agreed (post-reduce) value and op.""" + + def all_reduce(tensor, op): + torch.distributed.all_reduce(tensor, op=op) + reduces.append((tensor.tolist(), op)) + + return all_reduce + + +def _assert_ranks_agree(snapshot): + snapshots = [None, None] + torch.distributed.all_gather_object(snapshots, snapshot) + assert snapshots[0] == snapshots[1], snapshots + + def _two_rank_prefetch_deadline(rank, rendezvous): torch.distributed.init_process_group( "gloo", @@ -211,15 +269,26 @@ def _two_rank_prefetch_deadline(rank, rendezvous): ) try: cache, operation = _timeout_fixture() - cache._all_reduce = torch.distributed.all_reduce - # Only one rank has reached admission; its expiry must close both budgets. + _arm_tick(cache) + reduces = [] + cache._all_reduce = _counting_all_reduce(reduces) + # Only rank 0 has reached admission, and only its clock has run out. if rank == 0: with patch("time.perf_counter", return_value=100.0): cache._start_storage_prefetch_wait(operation) with patch("time.perf_counter", return_value=103.0): - assert cache.check_prefetch_progress(_REQ) - assert operation.deadline_expired - assert cache._storage_prefetch_wait[_REQ].expired + if rank == 0: + # A poll reads the verdict; it cannot expire a request alone, + # and the peer that never polled issues nothing to match. + assert not cache.check_prefetch_progress(_REQ) + assert reduces == [] + cache.check_hicache_events() + assert reduces == [([1], torch.distributed.ReduceOp.MAX)], reduces + assert operation.deadline_expired + assert cache._storage_prefetch_wait[_REQ].expired + verdict = cache.check_prefetch_progress(_REQ) + _assert_ranks_agree((verdict, operation.deadline_expired, len(reduces))) + assert verdict cache.prefetch_from_storage(_REQ, 0, list(range(8))) assert not cache.ongoing_prefetch assert not cache.storage_prefetch_retries._pending @@ -227,6 +296,90 @@ def _two_rank_prefetch_deadline(rank, rendezvous): torch.distributed.destroy_process_group() +def _two_rank_uneven_polls(rank, rendezvous): + """The #30760 shape before any verdict: ranks leave the admission loop at + different points, so rank 0 polls two waiting requests and rank 1 polls + one; a False poll is rank-local state only, so nothing diverges.""" + torch.distributed.init_process_group( + "gloo", + init_method=f"file://{rendezvous}", + rank=rank, + world_size=2, + timeout=timedelta(seconds=30), + ) + try: + cache, _ = _timeout_fixture() + _arm_tick(cache) + reduces = [] + cache._all_reduce = _counting_all_reduce(reduces) + other = CacheRequestHandle("r1", 0) + cache.prefetch_from_storage(other, 0, list(range(8))) + handles = [_REQ, other] + with patch("time.perf_counter", return_value=100.0): + for handle in handles: + assert not cache.check_prefetch_progress(handle) + with patch("time.perf_counter", return_value=103.0): + polled = handles[: 2 - rank] + for handle in polled: + assert not cache.check_prefetch_progress(handle) + assert reduces == [] + cache.check_hicache_events() + assert reduces == [([1, 1], torch.distributed.ReduceOp.MAX)], reduces + verdicts = [cache.check_prefetch_progress(handle) for handle in handles] + _assert_ranks_agree((verdicts, len(reduces), list(cache.ongoing_prefetch))) + assert verdicts == [True, True] + assert not cache.ongoing_prefetch + finally: + torch.distributed.destroy_process_group() + + +def _two_rank_divergent_live_sets(rank, rendezvous): + """The #30760 shape after a verdict: a True poll retires per rank, so a + rank that left the admission loop early holds one more live prefetch than + its peer. The next tick must raise on both ranks at the ready-count sync + instead of voting on tensors of different length (a hang until the gloo + timeout).""" + torch.distributed.init_process_group( + "gloo", + init_method=f"file://{rendezvous}", + rank=rank, + world_size=2, + timeout=timedelta(seconds=30), + ) + try: + cache, _ = _timeout_fixture() + _arm_tick(cache, sync_counts=True) + reduces = [] + cache._all_reduce = _counting_all_reduce(reduces) + other = CacheRequestHandle("r1", 0) + cache.prefetch_from_storage(other, 0, list(range(8))) + handles = [_REQ, other] + min_op, max_op = torch.distributed.ReduceOp.MIN, torch.distributed.ReduceOp.MAX + with patch("time.perf_counter", return_value=100.0): + for handle in handles: + assert not cache.check_prefetch_progress(handle) + with patch("time.perf_counter", return_value=103.0): + cache.check_hicache_events() + assert [op for _, op in reduces] == [min_op, max_op], reduces + assert reduces[1][0] == [1, 1], reduces + # Rank 1 trips batch_is_full after one request; rank 0 polls both. + for handle in handles[: 2 - rank]: + assert cache.check_prefetch_progress(handle) + assert len(cache.ongoing_prefetch) == rank + try: + cache.check_hicache_events() + except AssertionError as error: + outcome = "diverged" in str(error) + else: + outcome = "no raise" + # Same tick, every rank: the sync raised before any vote was issued. + _assert_ranks_agree((outcome, [op for _, op in reduces])) + assert outcome is True + assert [op for _, op in reduces] == [min_op, max_op, min_op] + finally: + torch.distributed.destroy_process_group() + + class TestWaitingStorageSpan(unittest.TestCase): def test_base_cache_has_no_storage_span(self): self.assertIsNone(BasePrefixCache.waiting_storage_span(None, _REQ)) @@ -301,9 +454,64 @@ def test_resident_mode_keeps_existing_accounting(self): class TestStoragePrefetchTimeout(unittest.TestCase): + def test_polls_never_reduce_and_the_tick_votes_once(self): + """N waiting requests polled in one pass issue no collective; the tick + issues exactly one MAX vote packed over every live prefetch, and only + that vote can expire a request, however late its poll.""" + cache, _ = _timeout_fixture() + polled = [_REQ] + [CacheRequestHandle(f"r{i}", 0) for i in range(1, 4)] + never_polled = CacheRequestHandle("parked", 0) + for handle in polled[1:] + [never_polled]: + cache.prefetch_from_storage(handle, 0, list(range(8))) + _arm_tick(cache) + reduces = [] + cache._all_reduce = lambda tensor, op: reduces.append((tensor.tolist(), op)) + max_op = torch.distributed.ReduceOp.MAX + + with patch("time.perf_counter", return_value=100.0): + for handle in polled: + self.assertFalse(cache.check_prefetch_progress(handle)) + self.assertEqual(reduces, []) + with patch("time.perf_counter", return_value=100.5): + cache.check_hicache_events() + self.assertEqual(reduces, [([0, 0, 0, 0, 0], max_op)]) + + # Budget 1.0 + 4 pages * 0.25 is long spent, yet a poll cannot cut. + with patch("time.perf_counter", return_value=200.0): + for handle in polled: + self.assertFalse(cache.check_prefetch_progress(handle)) + self.assertEqual(len(reduces), 1) + cache.check_hicache_events() + # A wait that never started (no blocked poll) contributes 0. + self.assertEqual(reduces[1], ([1, 1, 1, 1, 0], max_op)) + # The vote terminates (the worker and the hit drain stop spending + # on the span); the poll retires, with the deadline reason. + self.assertFalse( + cache.ongoing_prefetch[never_polled].operation.is_terminated() + ) + for handle in polled: + self.assertTrue(cache._storage_prefetch_wait[handle].expired) + self.assertTrue( + cache.ongoing_prefetch[handle].operation.is_terminated() + ) + self.assertTrue(cache.check_prefetch_progress(handle)) + self.assertEqual(list(cache.ongoing_prefetch), [never_polled]) + self.assertEqual(cache.storage_prefetch_retries._pending, {}) + self.assertEqual(len(reduces), 2) + + # The other policies never vote, and no live prefetch means no collective. + for policy in ("best_effort", "wait_complete"): + cache.prefetch_stop_policy = policy + cache.check_hicache_events() + cache.prefetch_stop_policy = "timeout" + cache.ongoing_prefetch.clear() + cache.check_hicache_events() + self.assertEqual(len(reduces), 2) + def test_expiry_is_agreed_even_before_a_peer_starts_waiting(self): """A rank that has not reached its admission poll must still adopt - the agreed expiry, or it would re-issue the prefetch alone.""" + the agreed expiry in the same tick, or it would re-issue the prefetch + alone; the poll on both ranks then returns the same verdict.""" with tempfile.TemporaryDirectory() as tmpdir: torch.multiprocessing.spawn( _two_rank_prefetch_deadline, @@ -311,6 +519,29 @@ def test_expiry_is_agreed_even_before_a_peer_starts_waiting(self): nprocs=2, ) + def test_uneven_polls_before_a_verdict_issue_no_collective(self): + """sgl-project/sglang#30760's trigger: a rank that polls more waiting + requests than its peer must not enter more collectives.""" + with tempfile.TemporaryDirectory() as tmpdir: + torch.multiprocessing.spawn( + _two_rank_uneven_polls, + args=(f"{tmpdir}/gloo",), + nprocs=2, + ) + + def test_divergent_live_prefetch_sets_raise_on_every_rank_in_the_same_tick( + self, + ): + """Uneven polls after a verdict leave the ranks with different live + prefetch sets; the next tick's ready-count sync fails loud on both + ranks before the vote could be issued with mismatched lengths.""" + with tempfile.TemporaryDirectory() as tmpdir: + torch.multiprocessing.spawn( + _two_rank_divergent_live_sets, + args=(f"{tmpdir}/gloo",), + nprocs=2, + ) + def _two_rank_retry_trace(rank, rendezvous): torch.distributed.init_process_group(