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

[PD] Do not admit intake-rejected requests to a PD handoff (#38935)

Browse Source 
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Shangming Cai
2026-09-16 12:44:49 +08:00
committed by GitHub
co-authored by Claude Opus 5
parent 100e1cd0d9
commit 3f8eb35ead
5 changed files with 395 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+8
View File
@@ -60,6 +60,7 @@ from sglang.srt.disaggregation.utils import (
get_qsa_pending_state_indices,
is_dsv4_c128_online_enabled,
is_mla_backend,
is_unadmitted_reject,
poll_and_all_reduce,
poll_and_all_reduce_pp,
poll_and_all_reduce_with_staging,
@@ -649,6 +650,13 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
dispatch happens later, after preallocation and ``send_metadata`` (see
``pop_preallocated``).
"""
# See `PrefillBootstrapQueue.add`. A retracted or rebootstrapping
# request owns a host KV backup that `retracted_queue` releases, and by
# this point carries none of the markers `is_unadmitted_reject` reads,
# so take the caller's word for it rather than sniffing.
if not is_retracted and not is_rebootstrap and is_unadmitted_reject(req):
self.scheduler.retire_unadmitted_request(req)
return
if self._check_if_req_exceed_kv_capacity(req):
return
python/sglang/srt/disaggregation/prefill.py
+8
View File
@@ -52,6 +52,7 @@ from sglang.srt.disaggregation.utils import (
is_aborted,
is_dsv4_c128_online_enabled,
is_mla_backend,
is_unadmitted_reject,
poll_and_all_reduce_attn_cp_tp_group,
poll_and_all_reduce_pp,
prepare_abort,
@@ -397,6 +398,13 @@ class PrefillBootstrapQueue:
return True
def add(self, req: Req, num_kv_heads: int) -> None:
# Rejected at intake: `set_finish_with_abort` left the verdict in
# `to_finish`, which `finished()` does not read, and swapped the prompt
# for a one-token stub. Bootstrapping it costs a handshake, a metadata
# buffer and a forward pass before anything unwinds it.
if is_unadmitted_reject(req):
self.scheduler.retire_unadmitted_request(req)
return
if not self.create_sender(req, num_kv_heads):
return
self.queue.append(req)
python/sglang/srt/disaggregation/utils.py
+29
View File
@@ -1704,6 +1704,35 @@ def prepare_abort(req: Req, error_message: str, status_code=None):
req.logprob.input_token_ids_logprobs_idx = []
def is_unadmitted_reject(req: Req) -> bool:
"""A request rejected at intake, before it acquired anything.
A preempted or resumed request can also carry a pending abort -- "Abort
method 3" marks a *running* request and `filter_batch` does not drop it,
since `finished()` is still False -- and its queue owns the release of
whatever it still holds.
`req.is_retracted` catches the two re-entries that declare nothing:
priority preemption and the pause/retract-all path both requeue through a
bare `_add_request_to_queue`. `release_req` always calls
`reset_for_retract`, which sets it, and its clear sites all run downstream
of these doors. The resource markers stay as a second line of defence --
on their own they miss a `seqlen <= 1` preemption, whose KV is already
freed and whose `retraction_backup` was never taken.
`DecodePreallocQueue.add` still gates on its own `is_retracted` /
`is_rebootstrap` parameters as well, since they state the caller's intent
rather than inferring it.
"""
return is_aborted(req) and not (
req.is_retracted
or req.kv.holds_kv
or req.kv.holds_mamba
or req.metadata_buffer_index >= 0
or req.kv.retraction_backup is not None
)
def is_aborted(req: Req) -> bool:
from sglang.srt.managers.schedule_batch import FINISH_ABORT
python/sglang/srt/managers/scheduler.py
+24
View File
@@ -3198,6 +3198,30 @@ class Scheduler(
self._retry_storage_prefetch(req)
return True
def retire_unadmitted_request(self, req: Req) -> None:
"""Finish a request the disaggregation queues rejected at their door."""
# `create_req` marks a streaming session in-flight, and the pre-abort
# detach lives in `StreamingSession.find_active_slot`, which only runs
# while scheduling; a session left in-flight rejects every later request.
if req.session is not None and req.session.streaming:
req.session.abort_req()
req.session = None
# `beam_coordinator.validate_and_init` counts the group in ahead of the
# checks that reject; no-op when the request has no group.
self.beam_coordinator.retire_group(req)
# PREFILL runs `_prefetch_kvcache` before its door, so even the
# one-token stub is registered with the cache by now:
# `prefetch_from_storage` arms the paced-retry set for this attempt's
# cache handle. Only a `finish`/ABORT and a `waiting_queue` sweep clear
# that, and a retired request reaches neither.
self._release_aborted_request(req)
# `update_finish_state` returns early once `finished()`, so an already
# set `finished_reason` is what the client receives; report the same.
reason = req.finished_reason or req.to_finish
req.time_stats.trace_ctx.abort(abort_info={"reason": reason.message})
req.update_finish_state()
self.output_streamer.stream_output([req], req.return_logprob)
def _add_request_to_queue(self, req: Req, is_retracted: bool = False):
if not self._set_or_validate_priority(req):
return
test/registered/unit/disaggregation/test_admission_abort_not_enqueued.py
+326
View File
@@ -0,0 +1,326 @@
"""A request rejected at intake must not be admitted to a PD handoff.
`Req.set_finish_with_abort()` records the rejection in `to_finish` (not
`finished_reason`) and replaces the prompt with a one-token stub, so
`req.finished()` stays False and the disaggregation admission queues let it
through: it completes a bootstrap handshake, reserves a metadata buffer,
initialises an RDMA sender and runs a forward pass on the stub before anything
unwinds it, and the decode worker sits on it until the transfer timeout.
Only disaggregation is affected. In NULL mode the stub costs one cheap forward
pass and the batch boundary returns the 400, so that path is deliberately left
alone -- `test_null_mode_is_deliberately_untouched` pins that.
The door must also not retire a *re-entering* request. "Abort method 3" sets
`to_finish` on a running request and `filter_batch` does not drop it, so a
retracted, preempted or resumed request can arrive carrying one while its queue
still owes a release. The prefill door detects that from the resources the
request still holds; the decode door cannot -- a retracted request holds none of
them by then -- so it gates on its own `is_retracted` / `is_rebootstrap` flags.
"""
import unittest
from array import array
from http import HTTPStatus
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from sglang.srt.disaggregation.decode import DecodePreallocQueue
from sglang.srt.disaggregation.prefill import PrefillBootstrapQueue
from sglang.srt.disaggregation.utils import DisaggregationMode, is_unadmitted_reject
from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req
from sglang.srt.managers.scheduler import Scheduler
from sglang.srt.runtime_context import publish, reset_context
from sglang.srt.sampling.sampling_params import SamplingParams
from sglang.srt.server_args import ServerArgs
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=3, suite="base-a-test-cpu")
ERROR_MSG = (
"Input length (1500 tokens) exceeds the maximum allowed length (1018 tokens)."
)
def _make_req(prompt_len=1500, session=None):
req = Req(
rid="admission-abort",
origin_input_text="",
origin_input_ids=array("q", list(range(prompt_len))),
sampling_params=SamplingParams(max_new_tokens=8),
session=session,
)
req.time_stats.trace_ctx = MagicMock()
return req
def _make_scheduler():
"""One stub for both halves.
`retire_unadmitted_request` is a spy wrapping the real method, so a door
test still asserts the call *and* runs the body -- otherwise the doors and
the retirement are only ever tested apart, and a step missing from the
retirement passes every door test.
"""
sched = SimpleNamespace(
output_streamer=MagicMock(),
beam_coordinator=MagicMock(),
_release_aborted_request=MagicMock(),
)
sched.retire_unadmitted_request = MagicMock(
side_effect=lambda req: Scheduler.retire_unadmitted_request(sched, req)
)
return sched
def _prefill_queue(sched):
q = SimpleNamespace(
scheduler=sched,
queue=[],
create_sender=MagicMock(return_value=True),
)
return q
def _decode_queue(sched):
q = SimpleNamespace(
scheduler=sched,
retracted_queue=[],
pending_reqs=[],
_check_if_req_exceed_kv_capacity=MagicMock(return_value=False),
_create_receiver_and_enqueue=MagicMock(
return_value=SimpleNamespace(kv_receiver=MagicMock())
),
_resolve_prefill_dp_rank=MagicMock(return_value=0),
)
return q
def _admit_prefill(q, req):
PrefillBootstrapQueue.add(q, req, 8)
def _admit_decode(q, req, **kw):
with patch(
"sglang.srt.disaggregation.decode._is_fake_transfer", return_value=False
):
DecodePreallocQueue.add(q, req, **kw)
class TestAdmissionAbortNotEnqueued(CustomTestCase):
def setUp(self):
reset_context()
self.addCleanup(reset_context)
publish(ServerArgs(model_path="dummy"), role="tokenizer")
# set_finish_with_abort() logs on TP rank 0, and ParallelContext.tp_rank
# reads through to the live process group -- none exists here.
p = patch(
"sglang.srt.managers.schedule_batch.get_parallel",
return_value=SimpleNamespace(tp_rank=0),
)
p.start()
self.addCleanup(p.stop)
def test_rejected_request_is_retired_at_each_pd_door(self):
for door in ("prefill", "decode"):
with self.subTest(door=door):
req = _make_req()
req.set_finish_with_abort(ERROR_MSG)
# Precondition: the rejection is pending, so `finished()` --
# what every admission gate keys off -- is still False.
self.assertIsInstance(req.to_finish, FINISH_ABORT)
self.assertFalse(req.finished())
self.assertEqual(len(req.origin_input_ids), 1)
sched = _make_scheduler()
if door == "prefill":
q = _prefill_queue(sched)
_admit_prefill(q, req)
q.create_sender.assert_not_called()
self.assertEqual(q.queue, [])
else:
q = _decode_queue(sched)
_admit_decode(q, req)
q._create_receiver_and_enqueue.assert_not_called()
self.assertEqual(q.pending_reqs, [])
sched.retire_unadmitted_request.assert_called_once_with(req)
def test_reentering_prefill_request_is_not_retired(self):
"""A preempted prefill requeue must reach its queue, which owns release.
Preemption runs `release_req`, so KV is gone by then; what survives is
the metadata buffer (`finalize_bootstrap` allocated it before the
request ever entered the running batch) and the host retraction backup.
"""
for marker in ("metadata_buffer_index", "retraction_backup"):
with self.subTest(marker=marker):
req = _make_req(prompt_len=16)
req.to_finish = FINISH_ABORT(
"Aborted by AbortReq.", HTTPStatus.SERVICE_UNAVAILABLE
)
# release_req already freed the KV row.
self.assertFalse(req.kv.holds_kv)
if marker == "metadata_buffer_index":
req.metadata_buffer_index = 7
else:
req.kv.retraction_backup = object()
self.assertFalse(is_unadmitted_reject(req))
sched = _make_scheduler()
q = _prefill_queue(sched)
_admit_prefill(q, req)
self.assertEqual(q.queue, [req])
sched.retire_unadmitted_request.assert_not_called()
self.assertIsInstance(req.to_finish, FINISH_ABORT)
def test_retracted_decode_request_is_not_retired(self):
"""The decode door cannot sniff this one, so it must trust its flags.
By the time `retract_decode` requeues, the request carries none of the
markers `is_unadmitted_reject` reads: `release_req` nulled
`kv.req_pool_idx`, `reset_for_retract` nulled `kv.mamba_pool_idx`, the
decode metadata buffer lives on `DecodeRequest` rather than `Req`, and
`add()` itself clears `retraction_mb_id`. Retiring it here would strand
the host pages `release_req` allocated, which only `retraction_restore`
or `retraction_discard` free.
`backup=None` is the short-sequence case: `retraction_backup()` returns
early for `seqlen <= 1` without setting it, so the resource predicate
alone says "unadmitted" and only the flag keeps the request safe.
"""
for flag in ("is_retracted", "is_rebootstrap"):
for backup in (object(), None):
with self.subTest(flag=flag, has_backup=backup is not None):
req = _make_req(prompt_len=16)
req.to_finish = FINISH_ABORT(
"Aborted by AbortReq.", HTTPStatus.SERVICE_UNAVAILABLE
)
# Exactly what retract_decode leaves behind.
self.assertFalse(req.kv.holds_kv)
self.assertFalse(req.kv.holds_mamba)
self.assertEqual(req.metadata_buffer_index, -1)
req.kv.retraction_backup = backup
sched = _make_scheduler()
q = _decode_queue(sched)
_admit_decode(q, req, **{flag: True})
sched.retire_unadmitted_request.assert_not_called()
self.assertIsInstance(req.to_finish, FINISH_ABORT)
if flag == "is_retracted":
self.assertEqual(q.retracted_queue, [req])
else:
q._create_receiver_and_enqueue.assert_called_once()
def test_preempted_reentry_is_not_retired(self):
"""Preemption requeues through a bare `_add_request_to_queue`.
The door sees no flag at all, and `release_req` -> `reset_for_retract`
is the only thing that marks the request. The resource markers miss
this shape on their own: the KV row is already freed and
`retraction_backup` is never taken for `seqlen <= 1`.
"""
for door in ("prefill", "decode"):
with self.subTest(door=door):
req = _make_req(prompt_len=16)
req.to_finish = FINISH_ABORT(
"Aborted by AbortReq.", HTTPStatus.SERVICE_UNAVAILABLE
)
req.is_retracted = True
# Exactly what a short preemption leaves behind.
self.assertFalse(req.kv.holds_kv)
self.assertIsNone(req.kv.retraction_backup)
self.assertEqual(req.metadata_buffer_index, -1)
self.assertFalse(is_unadmitted_reject(req))
sched = _make_scheduler()
if door == "prefill":
q = _prefill_queue(sched)
_admit_prefill(q, req)
self.assertEqual(q.queue, [req])
else:
q = _decode_queue(sched)
_admit_decode(q, req)
q._create_receiver_and_enqueue.assert_called_once()
sched.retire_unadmitted_request.assert_not_called()
def test_valid_request_still_admitted(self):
for door in ("prefill", "decode"):
with self.subTest(door=door):
req = _make_req(prompt_len=16)
sched = _make_scheduler()
if door == "prefill":
q = _prefill_queue(sched)
_admit_prefill(q, req)
self.assertEqual(q.queue, [req])
else:
q = _decode_queue(sched)
_admit_decode(q, req)
q._create_receiver_and_enqueue.assert_called_once()
sched.retire_unadmitted_request.assert_not_called()
def test_null_mode_is_deliberately_untouched(self):
"""The bug is disaggregation-only; NULL mode must keep enqueuing.
There the one-token stub costs a cheap forward pass and the batch
boundary promotes `to_finish` into the 400. Adding a guard here would
skip `StreamingSession.find_active_slot`'s pre-abort detach, which only
runs while scheduling.
"""
req = _make_req()
req.set_finish_with_abort(ERROR_MSG)
sched = SimpleNamespace(
disaggregation_mode=DisaggregationMode.NULL,
waiting_queue=[],
processed_tokens_counter=0,
_set_or_validate_priority=MagicMock(return_value=True),
_abort_on_queued_limit=MagicMock(return_value=False),
_prefetch_kvcache=MagicMock(),
)
Scheduler._add_request_to_queue(sched, req)
self.assertEqual(sched.waiting_queue, [req])
self.assertIsInstance(req.to_finish, FINISH_ABORT)
def test_retire_detaches_a_streaming_session(self):
"""Otherwise the session stays in-flight forever.
`create_req` marks it in-flight and the pre-abort detach lives in
`StreamingSession.find_active_slot`, which a retired request never
reaches; a stuck flag fails every later request on that session.
"""
session = MagicMock()
session.streaming = True
req = _make_req(session=session)
req.set_finish_with_abort(ERROR_MSG)
sched = _make_scheduler()
Scheduler.retire_unadmitted_request(sched, req)
# PREFILL arms the cache's paced-retry set via `_prefetch_kvcache`
# before the door, and only this call clears it again.
sched._release_aborted_request.assert_called_once_with(req)
session.abort_req.assert_called_once()
self.assertIsNone(req.session)
sched.beam_coordinator.retire_group.assert_called_once_with(req)
req.time_stats.trace_ctx.abort.assert_called_once()
# The original 400 survives, rather than being replaced downstream.
self.assertIsInstance(req.finished_reason, FINISH_ABORT)
self.assertEqual(req.finished_reason.status_code, HTTPStatus.BAD_REQUEST)
self.assertIsNone(req.to_finish)
sched.output_streamer.stream_output.assert_called_once()
def test_retire_leaves_a_non_streaming_session_alone(self):
session = MagicMock()
session.streaming = False
req = _make_req(session=session)
req.set_finish_with_abort(ERROR_MSG)
sched = _make_scheduler()
Scheduler.retire_unadmitted_request(sched, req)
session.abort_req.assert_not_called()
self.assertIs(req.session, session)
if __name__ == "__main__":
unittest.main()