minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Kamil
2026-09-07 18:01:25 -07:00
committed by GitHub
parent 4dcecc7891
commit 792543f98c
7 changed files with 139 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
+63 -45
View File
@@ -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()
python/sglang/srt/managers/scheduler_components/request_receiver.py
+34 -9
View File
@@ -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:
test/registered/unit/managers/test_mm_shm_error_consensus.py
+1
View File
@@ -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
+1
View File
@@ -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
-2
View File
@@ -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
-1
View File
@@ -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
+40 -42
View File
@@ -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__":