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

python/sglang/srt/managers/scheduler.py

@@ -1806,21 +1806,6 @@ class Scheduler(
             ]
         )
 
-    def _abort_on_running_timeout(self, running_batch: ScheduleBatch):
-        # NOTE: this should be called before a batch is launched.
-        timeout_s = envs.SGLANG_REQ_RUNNING_TIMEOUT.get()
-        if timeout_s <= 0:
-            return
-        if running_batch.is_empty():
-            return
-
-        deadline = time.perf_counter() - timeout_s
-        for req in running_batch.reqs:
-            if not req.finished() and 0 < req.time_stats.forward_entry_time < deadline:
-                req.to_finish = FINISH_ABORT(
-                    "Request running timeout reached.", HTTPStatus.SERVICE_UNAVAILABLE
-                )
-
     def get_init_info(self) -> Dict[str, Any]:
         """Return scheduler initialization info for handshake.
 
@@ -2293,6 +2278,7 @@ class Scheduler(
             get_last_batch=lambda: self.last_batch,
             scripted_scheduler_hook=self.scripted_scheduler_hook,
             scheduler_stage_metrics=self.scheduler_stage_metrics,
+            poll_timeout_aborts=self._poll_timeout_aborts,
         )
 
     def init_dp_attn_adapter(self) -> None:
@@ -3249,34 +3235,58 @@ class Scheduler(
         req_to_abort.time_stats.trace_ctx.abort(abort_info={"reason": message})
         return req_to_abort.rid == recv_req.rid
 
-    def _abort_on_waiting_timeout(self):
+    def _poll_timeout_aborts(self) -> List[AbortReq]:
-        if (timeout_s := envs.SGLANG_REQ_WAITING_TIMEOUT.get()) <= 0:
+        """Emit aborts only; every rank must drop the same requests in the
-            return
+        same iteration, or the extend-vs-decode decision splits and the
+        collectives hang.
+        """
+        aborts: List[AbortReq] = []
 
-        deleted_reqs = set()
+        if (timeout_s := envs.SGLANG_REQ_WAITING_TIMEOUT.get()) > 0:
             deadline = time.perf_counter() - timeout_s
             for req in self.waiting_queue:
                 entry_time = req.time_stats.wait_queue_entry_time
                 if 0 < entry_time < deadline:
-                self._release_aborted_request(req.rid)
+                    aborts.append(
-                self.ipc_channels.send_to_tokenizer.send_output(
+                        AbortReq(
-                    _make_abort_req(
+                            rid=req.rid,
-                        req,
+                            abort_message="Request waiting timeout reached.",
                             finished_reason={
                                 "type": "abort",
                                 "status_code": HTTPStatus.SERVICE_UNAVAILABLE,
                                 "message": "Request waiting timeout reached.",
                             },
-                    ),
-                    req,
                         )
-                deleted_reqs.add(req)
+                    )
-                self.beam_coordinator.retire_group(req)
 
-        if deleted_reqs:
+        if (timeout_s := envs.SGLANG_REQ_RUNNING_TIMEOUT.get()) > 0:
-            self.waiting_queue = [
+            deadline = time.perf_counter() - timeout_s
-                req for req in self.waiting_queue if req not in deleted_reqs
+            if self.ps.pp_size == 1:
-            ]
+                inflight_batches = [self.running_batch, self.last_batch]
+            else:
+                inflight_batches = [*self.running_mbs, *self.mbs]
+            seen_rids = set()
+            for batch in inflight_batches:
+                if batch is None:
+                    continue
+                for req in batch.reqs:
+                    if req.rid in seen_rids or req.finished():
+                        continue
+                    seen_rids.add(req.rid)
+                    if 0 < req.time_stats.forward_entry_time < deadline:
+                        aborts.append(
+                            AbortReq(
+                                rid=req.rid,
+                                abort_message="Request running timeout reached.",
+                                finished_reason={
+                                    "type": "abort",
+                                    "status_code": HTTPStatus.SERVICE_UNAVAILABLE,
+                                    "message": "Request running timeout reached.",
+                                },
+                            )
+                        )
+
+        return aborts
 
     def handle_embedding_request(
         self,
@@ -3478,8 +3488,6 @@ class Scheduler(
 
         if self.enable_fpm:
             self._fpm_batch_t0 = time.monotonic()
-        self._abort_on_waiting_timeout()
-        self._abort_on_running_timeout(running_batch)
         if self.dllm_config is not None:
             self.dllm_manager.filter_finished_reqs()
 
@@ -5161,7 +5169,11 @@ class Scheduler(
             req = self.waiting_queue.pop(i)
             self._release_aborted_request(req.rid)
             self.beam_coordinator.retire_group(req)
-            self.ipc_channels.send_to_tokenizer.send_output(_make_abort_req(req), req)
+            # Without the initiator's reason the tokenizer falls back to a
+            # generic abort message.
+            self.ipc_channels.send_to_tokenizer.send_output(
+                _make_abort_req(req, finished_reason=recv_req.finished_reason), req
+            )
             # For disaggregation decode mode, the request in the waiting queue has KV cache allocated.
             if self.disaggregation_mode == DisaggregationMode.DECODE:
                 release_kv_cache(req, self.tree_cache)
@@ -5280,6 +5292,12 @@ class Scheduler(
                 # The request will still run one decode forward pass.
                 # Then we reuse all existing code to clean up the KV cache allocation.
                 logger.debug(f"Abort running request. {req.rid=}")
+                if recv_req.abort_message:
+                    # Timeout aborts carry an SLA message + 503 for the client.
+                    req.to_finish = FINISH_ABORT(
+                        recv_req.abort_message, HTTPStatus.SERVICE_UNAVAILABLE
+                    )
+                else:
                     req.to_finish = FINISH_ABORT()
 
     def _pause_engine(self) -> Tuple[List[Req], int]:

python/sglang/srt/managers/scheduler_components/request_receiver.py

@@ -20,6 +20,7 @@ from sglang.srt.disaggregation.utils import prepare_abort
 from sglang.srt.distributed.communication_op import attn_cp_tp_broadcast_pyobj
 from sglang.srt.environ import envs
 from sglang.srt.managers.io_struct import (
+    AbortReq,
     BatchTokenizedEmbeddingReqInput,
     BatchTokenizedGenerateReqInput,
     MMInputsProcessError,
@@ -78,6 +79,9 @@ class SchedulerRequestReceiver:
     get_last_batch: Callable[[], Any]
     scripted_scheduler_hook: Optional[ScriptedSchedulerHook] = None
     scheduler_stage_metrics: Optional[SchedulerStageMetricsRecorder] = None
+    # Emits AbortReqs for SGLANG_REQ_WAITING_TIMEOUT / _RUNNING_TIMEOUT;
+    # runs on the rank that owns the waiting queue.
+    poll_timeout_aborts: Callable[[], List[AbortReq]]
 
     def recv_limit_reached(self, num_recv_reqs: int) -> bool:
         if self.max_recv_per_poll < 0:
@@ -102,7 +106,17 @@ class SchedulerRequestReceiver:
         if self.input_blocker is not None:
             recv_reqs = self.input_blocker.handle(recv_reqs)
 
-        recv_reqs = self._broadcast_reqs_across_ranks(recv_reqs)
+        # Decided once and broadcast, so every rank sharing this waiting queue
+        # drops the same requests in the same iteration.
+        local_reqs = []
+        if (
+            self.ps.pp_rank == 0
+            and self.ps.attn_tp_rank == 0
+            and self.ps.attn_cp_rank == 0
+        ):
+            local_reqs = self.poll_timeout_aborts()
+
+        recv_reqs = self._broadcast_reqs_across_ranks(recv_reqs, local_reqs)
 
         if self.ps.pp_rank == 0:
             self.unwrap_pickle_wrapper(recv_reqs)
@@ -162,10 +176,18 @@ class SchedulerRequestReceiver:
                 recv_reqs = None
         return recv_reqs
 
-    def _broadcast_reqs_across_ranks(self, recv_reqs: Optional[List]) -> List:
+    def _broadcast_reqs_across_ranks(
+        self, recv_reqs: Optional[List], local_reqs: Optional[List] = None
+    ) -> List:
+        """local_reqs ride the work channel, which is scoped to the ranks
+        sharing one waiting queue; the control channel fans out from global
+        rank 0 and would overwrite every DP group's aborts but the first.
+        """
+        local_reqs = local_reqs or []
         if get_parallel().enable_dp_attention:
             if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0:
                 work_reqs, control_reqs = self._split_work_and_control_reqs(recv_reqs)
+                work_reqs.extend(local_reqs)
             else:
                 work_reqs = None
                 control_reqs = None
@@ -191,7 +213,10 @@ class SchedulerRequestReceiver:
                     src=self.tp_group.ranks[0],
                 )
             recv_reqs = work_reqs + control_reqs
-        elif self.ps.tp_size != 1:
+        else:
+            if recv_reqs is not None:
+                recv_reqs = [*recv_reqs, *local_reqs]
+            if self.ps.tp_size != 1:
                 recv_reqs = broadcast_pyobj(
                     recv_reqs,
                     self.tp_group.rank,

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
+The poll is pure bookkeeping over timestamps -- no model, no GPU, no draft
-worker -- so they are driven here directly instead of through a server. The
+worker -- so it is driven here directly instead of through a server. The 503
-e2e side (503 reaching the client, server stays up) is covered by
+reaching the client is covered by scheduler/test_scheduler_control.py.
-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):
+    def test_emits_only_reqs_past_the_deadline_and_keeps_queue_intact(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):
         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([a.rid for a in aborts], ["stale"])
-        self.assertEqual(s.ipc_channels.send_to_tokenizer.send_output.call_count, 1)
+        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(s._poll_timeout_aborts(), [])
-        self.assertEqual(len(s.waiting_queue), 1)
-        s.ipc_channels.send_to_tokenizer.send_output.assert_not_called()
 
     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(s._poll_timeout_aborts(), [])
-        self.assertEqual(len(s.waiting_queue), 1)
 
 
 class TestRunningTimeout(CustomTestCase):
-    @staticmethod
+    def test_emits_only_stale_unfinished_reqs_without_marking(self):
-    def _batch(reqs):
-        return SimpleNamespace(reqs=reqs, is_empty=lambda: not reqs)
-
-    def test_marks_only_stale_unfinished_reqs(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):
+    def test_req_in_both_running_and_last_batch_is_emitted_once(self):
-        s = _scheduler([])
+        stale = _req("stale", forward_entry=time.perf_counter() - 10)
-        req = _req("unstamped", forward_entry=0.0)
+        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.assertEqual(s._poll_timeout_aborts(), [])
-        self.assertIsNone(req.to_finish)
 
     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([]))
+            self.assertEqual(s._poll_timeout_aborts(), [])
-        req = _req("stale", forward_entry=time.perf_counter() - 100)
+        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.assertEqual(s._poll_timeout_aborts(), [])
-        self.assertIsNone(req.to_finish)
 
 
 if __name__ == "__main__":