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: