From 792543f98c25bf1ab23c114513a20d44f1d0a6bc Mon Sep 17 00:00:00 2001 From: Kamil Date: Tue, 8 Sep 2026 04:01:25 +0300 Subject: [PATCH] [Scheduler] Make request-timeout aborts rank-consistent to fix TP collective hangs (#37143) --- python/sglang/srt/managers/scheduler.py | 108 ++++++++++-------- .../scheduler_components/request_receiver.py | 43 +++++-- .../managers/test_mm_shm_error_consensus.py | 1 + .../unit/managers/test_pp_cp_rank_offsets.py | 1 + .../test_scheduler_chunked_req_gate.py | 2 - .../test_scheduler_decision_batch_params.py | 1 - .../unit/managers/test_scheduler_timeouts.py | 82 +++++++------ 7 files changed, 139 insertions(+), 99 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 85b9c203c..f408205aa 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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): - if (timeout_s := envs.SGLANG_REQ_WAITING_TIMEOUT.get()) <= 0: - return + def _poll_timeout_aborts(self) -> List[AbortReq]: + """Emit aborts only; every rank must drop the same requests in the + same iteration, or the extend-vs-decode decision splits and the + collectives hang. + """ + aborts: List[AbortReq] = [] - deleted_reqs = set() - 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) - self.ipc_channels.send_to_tokenizer.send_output( - _make_abort_req( - req, - 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 (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: + aborts.append( + AbortReq( + rid=req.rid, + abort_message="Request waiting timeout reached.", + finished_reason={ + "type": "abort", + "status_code": HTTPStatus.SERVICE_UNAVAILABLE, + "message": "Request waiting timeout reached.", + }, + ) + ) - if deleted_reqs: - self.waiting_queue = [ - req for req in self.waiting_queue if req not in deleted_reqs - ] + if (timeout_s := envs.SGLANG_REQ_RUNNING_TIMEOUT.get()) > 0: + deadline = time.perf_counter() - timeout_s + 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,7 +5292,13 @@ 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=}") - req.to_finish = FINISH_ABORT() + 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]: raise NotImplementedError() diff --git a/python/sglang/srt/managers/scheduler_components/request_receiver.py b/python/sglang/srt/managers/scheduler_components/request_receiver.py index e99e86f35..49d2ba7a9 100644 --- a/python/sglang/srt/managers/scheduler_components/request_receiver.py +++ b/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,13 +213,16 @@ class SchedulerRequestReceiver: src=self.tp_group.ranks[0], ) recv_reqs = work_reqs + control_reqs - elif self.ps.tp_size != 1: - recv_reqs = broadcast_pyobj( - recv_reqs, - self.tp_group.rank, - self.tp_cpu_group, - src=self.tp_group.ranks[0], - ) + 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, + self.tp_cpu_group, + src=self.tp_group.ranks[0], + ) return recv_reqs def unwrap_pickle_wrapper(self, recv_reqs: Optional[List]) -> None: diff --git a/test/registered/unit/managers/test_mm_shm_error_consensus.py b/test/registered/unit/managers/test_mm_shm_error_consensus.py index 51d951709..110071037 100644 --- a/test/registered/unit/managers/test_mm_shm_error_consensus.py +++ b/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: [], ) diff --git a/test/registered/unit/managers/test_pp_cp_rank_offsets.py b/test/registered/unit/managers/test_pp_cp_rank_offsets.py index bea857174..21e61d723 100644 --- a/test/registered/unit/managers/test_pp_cp_rank_offsets.py +++ b/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: [], ) diff --git a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py index 3aac9df0f..e8c37eb94 100644 --- a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py +++ b/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 diff --git a/test/registered/unit/managers/test_scheduler_decision_batch_params.py b/test/registered/unit/managers/test_scheduler_decision_batch_params.py index 049e07e40..63b51aed2 100644 --- a/test/registered/unit/managers/test_scheduler_decision_batch_params.py +++ b/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, diff --git a/test/registered/unit/managers/test_scheduler_timeouts.py b/test/registered/unit/managers/test_scheduler_timeouts.py index fba263af5..ca67a8c8d 100644 --- a/test/registered/unit/managers/test_scheduler_timeouts.py +++ b/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__":