[Scheduler] Make request-timeout aborts rank-consistent to fix TP collective hangs (#37143)

test/registered/unit/managers/test_mm_shm_error_consensus.py

@@ -121,6 +121,7 @@ def _receiver(tp_size: int = 1) -> SchedulerRequestReceiver:
         max_recv_per_poll=-1,
         stream_output=lambda *args, **kwargs: None,
         get_last_batch=lambda: None,
+        poll_timeout_aborts=lambda: [],
     )
 
 

test/registered/unit/managers/test_pp_cp_rank_offsets.py

@@ -63,6 +63,7 @@ def _make_receiver(ps: ParallelState) -> SchedulerRequestReceiver:
         max_recv_per_poll=-1,
         stream_output=lambda *args, **kwargs: None,
         get_last_batch=lambda: None,
+        poll_timeout_aborts=lambda: [],
     )
 
 

test/registered/unit/managers/test_scheduler_chunked_req_gate.py

@@ -76,8 +76,6 @@ def _make_chunk_cache(req_to_token_pool) -> ChunkCache:
 def _scheduler_for_get_next_batch(*, tree_cache, chunked_req) -> Scheduler:
     s = Scheduler.__new__(Scheduler)
     s.scheduler_stage_metrics = None
-    s._abort_on_waiting_timeout = MagicMock()
-    s._abort_on_running_timeout = MagicMock()
     s.dllm_config = None
     s.dllm_manager = None
     s.enable_hisparse = False

test/registered/unit/managers/test_scheduler_decision_batch_params.py

@@ -20,7 +20,6 @@ DECISION_METHODS = (
     Scheduler.get_next_batch_to_run,
     Scheduler.get_new_batch_prefill,
     Scheduler._get_new_batch_prefill_raw,
-    Scheduler._abort_on_running_timeout,
     Scheduler.is_disable_overlap_for_batch,
     SchedulerDisaggregationPrefillMixin.get_next_disagg_prefill_batch_to_run,
     SchedulerDisaggregationPrefillMixin.process_prefill_chunk,

test/registered/unit/managers/test_scheduler_timeouts.py

@@ -1,9 +1,8 @@
 """Boundary tests for the scheduler's waiting / running request timeouts.
 
-Both paths are pure bookkeeping over timestamps -- no model, no GPU, no draft
-worker -- so they are driven here directly instead of through a server. The
-e2e side (503 reaching the client, server stays up) is covered by
-scheduler/test_scheduler_control.py.
+The poll is pure bookkeeping over timestamps -- no model, no GPU, no draft
+worker -- so it is driven here directly instead of through a server. The 503
+reaching the client is covered by scheduler/test_scheduler_control.py.
 """
 
 import time
@@ -23,8 +22,6 @@ register_cpu_ci(est_time=12, suite="base-a-test-cpu")
 
 
 class _FakeReq:
-    """Must stay hashable: the waiting-timeout path collects drops in a set."""
-
     def __init__(self, rid, wait_entry=0.0, forward_entry=0.0, is_finished=False):
         self.rid = rid
         self.to_finish = None
@@ -48,7 +45,11 @@ def _req(
     return _FakeReq(rid, wait_entry, forward_entry, finished)
 
 
-def _scheduler(waiting_queue):
+def _batch(reqs):
+    return SimpleNamespace(reqs=reqs, is_empty=lambda: not reqs)
+
+
+def _scheduler(waiting_queue, running_reqs=(), last_batch_reqs=()):
     s = Scheduler.__new__(Scheduler)
     s.waiting_queue = waiting_queue
     s.enable_hierarchical_cache = False
@@ -56,6 +57,9 @@ def _scheduler(waiting_queue):
     s.enable_unified_cache_external_linker = False
     s.ipc_channels = SimpleNamespace(send_to_tokenizer=MagicMock())
     s.beam_coordinator = MagicMock()
+    s.ps = SimpleNamespace(pp_size=1)
+    s.running_batch = _batch(list(running_reqs))
+    s.last_batch = _batch(list(last_batch_reqs)) if last_batch_reqs else None
     return s
 
 
@@ -88,75 +92,69 @@ class TestQueuedLimitAbort(CustomTestCase):
 
 
 class TestWaitingTimeout(CustomTestCase):
-    def setUp(self):
-        patcher = patch(
-            "sglang.srt.managers.scheduler.get_serving",
-            return_value=SimpleNamespace(weight_version="v0"),
-        )
-        patcher.start()
-        self.addCleanup(patcher.stop)
-
-    def test_drops_only_reqs_past_the_deadline(self):
+    def test_emits_only_reqs_past_the_deadline_and_keeps_queue_intact(self):
         now = time.perf_counter()
         stale = _req("stale", wait_entry=now - 10)
         fresh = _req("fresh", wait_entry=now)
         s = _scheduler([stale, fresh])
 
         with envs.SGLANG_REQ_WAITING_TIMEOUT.override(1.0):
-            s._abort_on_waiting_timeout()
+            aborts = s._poll_timeout_aborts()
 
-        self.assertEqual([r.rid for r in s.waiting_queue], ["fresh"])
-        self.assertEqual(s.ipc_channels.send_to_tokenizer.send_output.call_count, 1)
+        self.assertEqual([a.rid for a in aborts], ["stale"])
+        self.assertEqual(aborts[0].finished_reason["type"], "abort")
+        # The poll emits only; removal happens on every rank via the broadcast.
+        self.assertEqual([r.rid for r in s.waiting_queue], ["stale", "fresh"])
 
-    def test_unset_entry_time_is_never_dropped(self):
+    def test_unset_entry_time_is_never_emitted(self):
         # 0 is the "not yet stamped" sentinel; the guard is `0 < entry_time`.
         s = _scheduler([_req("unstamped", wait_entry=0.0)])
         with envs.SGLANG_REQ_WAITING_TIMEOUT.override(1e-9):
-            s._abort_on_waiting_timeout()
-        self.assertEqual(len(s.waiting_queue), 1)
-        s.ipc_channels.send_to_tokenizer.send_output.assert_not_called()
+            self.assertEqual(s._poll_timeout_aborts(), [])
 
     def test_disabled_timeout_is_a_no_op(self):
         s = _scheduler([_req("stale", wait_entry=time.perf_counter() - 100)])
         with envs.SGLANG_REQ_WAITING_TIMEOUT.override(0):
-            s._abort_on_waiting_timeout()
-        self.assertEqual(len(s.waiting_queue), 1)
+            self.assertEqual(s._poll_timeout_aborts(), [])
 
 
 class TestRunningTimeout(CustomTestCase):
-    @staticmethod
-    def _batch(reqs):
-        return SimpleNamespace(reqs=reqs, is_empty=lambda: not reqs)
-
-    def test_marks_only_stale_unfinished_reqs(self):
+    def test_emits_only_stale_unfinished_reqs_without_marking(self):
         now = time.perf_counter()
         stale = _req("stale", forward_entry=now - 10)
         fresh = _req("fresh", forward_entry=now)
         done = _req("done", forward_entry=now - 10, finished=True)
-        s = _scheduler([])
+        s = _scheduler([], running_reqs=[stale, fresh, done])
 
         with envs.SGLANG_REQ_RUNNING_TIMEOUT.override(1.0):
-            s._abort_on_running_timeout(self._batch([stale, fresh, done]))
+            aborts = s._poll_timeout_aborts()
 
-        self.assertIsNotNone(stale.to_finish)
+        self.assertEqual([a.rid for a in aborts], ["stale"])
+        # to_finish is set by abort_request() on every rank, not by the poll.
+        self.assertIsNone(stale.to_finish)
         self.assertIsNone(fresh.to_finish)
         self.assertIsNone(done.to_finish, "a finished req must not be aborted")
 
-    def test_unset_forward_entry_time_is_never_marked(self):
-        s = _scheduler([])
-        req = _req("unstamped", forward_entry=0.0)
+    def test_req_in_both_running_and_last_batch_is_emitted_once(self):
+        stale = _req("stale", forward_entry=time.perf_counter() - 10)
+        s = _scheduler([], running_reqs=[stale], last_batch_reqs=[stale])
+        with envs.SGLANG_REQ_RUNNING_TIMEOUT.override(1.0):
+            aborts = s._poll_timeout_aborts()
+        self.assertEqual([a.rid for a in aborts], ["stale"])
+
+    def test_unset_forward_entry_time_is_never_emitted(self):
+        s = _scheduler([], running_reqs=[_req("unstamped", forward_entry=0.0)])
         with envs.SGLANG_REQ_RUNNING_TIMEOUT.override(1e-9):
-            s._abort_on_running_timeout(self._batch([req]))
-        self.assertIsNone(req.to_finish)
+            self.assertEqual(s._poll_timeout_aborts(), [])
 
     def test_empty_batch_and_disabled_timeout_are_no_ops(self):
         s = _scheduler([])
         with envs.SGLANG_REQ_RUNNING_TIMEOUT.override(1.0):
-            s._abort_on_running_timeout(self._batch([]))
-        req = _req("stale", forward_entry=time.perf_counter() - 100)
+            self.assertEqual(s._poll_timeout_aborts(), [])
+        stale = _req("stale", forward_entry=time.perf_counter() - 100)
+        s = _scheduler([], running_reqs=[stale])
         with envs.SGLANG_REQ_RUNNING_TIMEOUT.override(0):
-            s._abort_on_running_timeout(self._batch([req]))
-        self.assertIsNone(req.to_finish)
+            self.assertEqual(s._poll_timeout_aborts(), [])
 
 
 if __name__ == "__main__":