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

[PDD] Add true request retraction for PDD (#25372)

Browse Source 
Signed-off-by: Ata Fatahi <immrata@gmail.com>
This commit is contained in:
Ata Fatahi
2026-07-09 15:33:01 +08:00
committed by GitHub
parent 666a09fe2a
commit 866ae6848f
12 changed files with 820 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/disaggregation/test_disaggregation_basic.py
+159
View File
@@ -291,6 +291,165 @@ class TestDisaggregationSimulatedRetract(PDDisaggregationServerBase):
self.assertGreater(metrics["score"], 0.62)
class TestDisaggregationPauseResumeDecodeRetract(PDDisaggregationServerBase):
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
cls.launch_all()
def test_retract_pause_decode_running_batch(self):
"""Retract-mode pause on a disagg decode node must preserve in-flight
requests that are already in running_batch."""
asyncio.run(self._run_pause_on_decode_running_batch("retract"))
def test_retract_weight_update_decode_running_batch(self):
"""Retract pause + weight update on a disagg decode node.
This guards the core reason retract mode exists: while paused, the
running_batch AND the rebootstrap preallocation queue are empty, so the
scheduler is fully idle and the post-update cache flush succeeds (a
regression here trips ``assert ..., "Cache flush failed after updating
weights"`` and crashes the decode worker). On continue, the retracted
requests rebootstrap-recompute their prefix KV under the updated weights
and resume to completion.
"""
asyncio.run(
self._run_pause_on_decode_running_batch("retract", weight_update=True)
)
async def _get_decode_num_running_reqs(self, session):
"""Query current decode running_batch size from /v1/loads."""
async with session.get(
self.decode_url + "/v1/loads?include=core",
timeout=aiohttp.ClientTimeout(total=5),
) as resp:
resp.raise_for_status()
body = await resp.json()
return sum(load["num_running_reqs"] for load in body["loads"])
async def _wait_for_decode_running_batch(self, session, timeout):
deadline = asyncio.get_running_loop().time() + timeout
while asyncio.get_running_loop().time() < deadline:
if await self._get_decode_num_running_reqs(session) > 0:
return
await asyncio.sleep(0.2)
self.fail("Timed out waiting for decode running_batch to become non-empty")
async def _run_pause_on_decode_running_batch(self, mode, weight_update=False):
num_requests = 2
max_new_tokens = 512
prompt = "Write a detailed numbered explanation of distributed inference. " * 12
async def _post(session, url, json_data, timeout=30):
async with session.post(
url,
json=json_data,
timeout=aiohttp.ClientTimeout(total=timeout),
) as resp:
resp.raise_for_status()
return await resp.json()
async def _generate(session, request_id):
return await _post(
session,
self.lb_url + "/generate",
{
"text": f"Request {request_id}: {prompt}",
"background": True,
"sampling_params": {
"temperature": 0,
"ignore_eos": True,
"max_new_tokens": max_new_tokens,
},
},
timeout=180,
)
async with aiohttp.ClientSession() as session:
tasks = [
asyncio.create_task(_generate(session, i)) for i in range(num_requests)
]
decode_paused = False
try:
await self._wait_for_decode_running_batch(session, timeout=30)
await asyncio.sleep(0.1)
self.assertTrue(
any(not task.done() for task in tasks),
"All requests finished before decode retract pause was issued.",
)
await _post(
session,
self.decode_url + "/pause_generation",
{"mode": mode},
)
decode_paused = True
await asyncio.sleep(1)
if weight_update:
# Reload the same weights from disk while retract-paused. The
# update mechanism (disk/tensor/distributed/ipc) is irrelevant
# here: they all share flush_cache_after_weight_update, whose
# flush asserts the scheduler is fully idle. This must not
# crash, proving retracted reqs are not stuck in the prealloc
# queue.
wu = await _post(
session,
self.decode_url + "/update_weights_from_disk",
{"model_path": self.model},
timeout=180,
)
self.assertTrue(
wu.get("success", False),
f"update_weights_from_disk failed during retract pause: {wu}",
)
await _post(session, self.decode_url + "/continue_generation", {})
decode_paused = False
responses = await asyncio.wait_for(asyncio.gather(*tasks), timeout=180)
finally:
if decode_paused:
try:
await _post(
session, self.decode_url + "/continue_generation", {}
)
except Exception:
pass
unfinished = [task for task in tasks if not task.done()]
if unfinished:
for url in [self.prefill_url, self.decode_url]:
try:
await _post(
session,
url + "/abort_request",
{"abort_all": True},
)
except Exception:
pass
for task in unfinished:
task.cancel()
await asyncio.gather(*unfinished, return_exceptions=True)
for response in responses:
self.assertIn("text", response)
self.assertGreater(len(response["text"]), 0)
self.assertGreater(
sum(
response.get("meta_info", {}).get("num_retractions", 0)
for response in responses
),
0,
"Expected pause_generation(retract) to retract a running decode request.",
)
class TestDisaggregationPauseResumePrefillLeak(PDDisaggregationServerBase):
"""Regression test: pause_generation must not leak prefill requests into
running_batch. With a small --max-running-requests the leak fills the
test/registered/unit/disaggregation/test_register_to_bootstrap.py
+5
View File
@@ -165,9 +165,13 @@ class TestRegisterToBootstrap(CustomTestCase):
"rank_port",
"page_size",
"kv_cache_dtype",
# Self-registered HTTP API port used to derive the PD retract
# rebootstrap /generate URL on the decode side.
"prefill_http_port",
]
for field in required_fields:
self.assertIn(field, payload)
self.assertEqual(payload["prefill_http_port"], 30000)
@patch("sglang.srt.disaggregation.common.conn.time")
@patch("sglang.srt.disaggregation.common.conn.requests.put")
@@ -266,6 +270,7 @@ class TestRegisterToBootstrap(CustomTestCase):
mgr.server_args = MagicMock()
mgr.server_args.kv_cache_dtype = "auto"
mgr.server_args.load_balance_method = "follow_bootstrap_room"
mgr.server_args.port = 30000
return mgr
test/registered/unit/managers/test_priority_scheduling_disaggregation.py
+205 -3
View File
@@ -1,8 +1,11 @@
import json
import sys
import threading
import unittest
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import numpy as np
import torch
from sglang.srt.disaggregation.decode import ( # noqa: E402
@@ -10,7 +13,7 @@ from sglang.srt.disaggregation.decode import ( # noqa: E402
SchedulerDisaggregationDecodeMixin,
)
from sglang.srt.disaggregation.utils import DisaggregationMode # noqa: E402
from sglang.srt.managers.schedule_batch import FINISH_ABORT # noqa: E402
from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req # noqa: E402
from sglang.srt.managers.scheduler import Scheduler # noqa: E402
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
@@ -96,6 +99,7 @@ class TestDecodePreallocQueuePriority(unittest.TestCase):
waiting_for_input=True,
kv_receiver=MagicMock(),
metadata_buffer_index=-1,
is_rebootstrap=False,
)
def _new_queue(self, decode_reqs, *, low_priority_values_first: bool = False):
@@ -205,6 +209,204 @@ class TestDecodePreallocQueuePriority(unittest.TestCase):
)
class TestDecodePreallocQueueRebootstrapPayload(unittest.TestCase):
"""The decode scheduler builds the rebootstrap ``/generate`` payload; the
dispatch itself now lives on the kv manager (see
``TestCommonKVManagerPrefillRecompute``)."""
def _sampling_params(self):
return SimpleNamespace(
temperature=0.0,
top_p=1.0,
top_k=-1,
min_p=0.0,
frequency_penalty=0.0,
presence_penalty=0.0,
repetition_penalty=1.0,
ignore_eos=False,
skip_special_tokens=True,
spaces_between_special_tokens=True,
no_stop_trim=False,
)
def _new_req(self):
return SimpleNamespace(
rid="rid-0",
origin_input_ids=np.array([1, 2], dtype=np.int32),
output_ids=[np.int32(3), np.int32(4)],
sampling_params=self._sampling_params(),
bootstrap_host="127.0.0.1",
bootstrap_port=30000,
bootstrap_room=7,
priority=10,
extra_key=None,
routing_key=None,
disagg_prefill_dp_rank=None,
)
def test_build_rebootstrap_payload_converts_numpy_ids_to_json_lists(self):
req = self._new_req()
# build_rebootstrap_payload lives on Req; exercise it unbound with a
# namespace that carries the attributes it reads.
payload = Req.build_rebootstrap_payload(req)
# origin_input_ids + output_ids, coerced to plain python ints.
self.assertEqual(payload["input_ids"], [1, 2, 3, 4])
self.assertTrue(all(type(x) is int for x in payload["input_ids"]))
self.assertEqual(payload["sampling_params"]["max_new_tokens"], 1)
self.assertEqual(payload["bootstrap_room"], 7)
# The prefill /generate URL is derived from bootstrap info on the decode
# side, not sent in the payload; and the boundary token is replayed via
# the decode-side override, so neither belongs in the payload.
self.assertNotIn("pd_rebootstrap_prefill_url", payload)
self.assertNotIn("pd_rebootstrap_forced_output_id", payload)
# Must be JSON-serializable (numpy scalars would raise here).
json.dumps(payload)
class TestCommonKVManagerPrefillRecompute(unittest.TestCase):
"""The kv manager owns the shared executor + HTTP session and routes any
rebootstrap ``/generate`` failure through ``kv_receiver.abort()`` ->
``KVPoll.Failed`` so the scheduler's normal transfer-failure streaming runs.
"""
def _new_manager(self):
from sglang.srt.disaggregation.common.conn import CommonKVManager
mgr = CommonKVManager.__new__(CommonKVManager)
mgr._prefill_recompute_executor = None
mgr._prefill_recompute_executor_lock = threading.Lock()
mgr._prefill_recompute_sessions = threading.local()
mgr.waiting_timeout = 300
mgr.failure_records = {}
mgr.failure_lock = threading.Lock()
# Only the attn-tp/attn-cp group leader on the first PP stage issues the
# single rebootstrap /generate; default the mock manager to that leader.
mgr.attn_tp_rank = 0
mgr.attn_cp_rank = 0
mgr.pp_rank = 0
# Decode-side prefill info cache; the rebootstrap /generate URL is derived
# from here (bootstrap_addr host + self-registered prefill_http_port)
# instead of a router-injected pd_rebootstrap_prefill_url.
mgr.prefill_info_table = {}
return mgr
def _register_prefill_info(self, mgr, bootstrap_addr, http_port):
from sglang.srt.disaggregation.common.conn import PrefillServerInfo
mgr.prefill_info_table[bootstrap_addr] = PrefillServerInfo(
attn_tp_size=1,
attn_cp_size=1,
dp_size=1,
pp_size=1,
page_size=1,
kv_cache_dtype=None,
follow_bootstrap_room=True,
prefill_http_port=http_port,
)
def _payload(self):
return {
"input_ids": [1, 2, 3, 4],
"rid": "rid-0",
}
def test_submit_dispatches_run_to_shared_executor(self):
mgr = self._new_manager()
mgr._prefill_recompute_executor = MagicMock()
receiver = MagicMock(bootstrap_room=7, bootstrap_addr="127.0.0.1:8998")
self._register_prefill_info(mgr, "127.0.0.1:8998", 30000)
mgr.submit_prefill_recompute(receiver, self._payload())
mgr._prefill_recompute_executor.submit.assert_called_once()
args = mgr._prefill_recompute_executor.submit.call_args[0]
self.assertEqual(args[0], mgr._run_prefill_recompute)
self.assertIs(args[1], receiver)
# URL derived from bootstrap_addr host + registered prefill_http_port.
self.assertEqual(args[2], "http://127.0.0.1:30000")
receiver.abort.assert_not_called()
def test_submit_is_noop_on_non_leader_ranks(self):
# A retracted request is replicated across every rank in its attention
# TP/CP group and every PP stage; only the group/first-stage leader must
# POST the single /generate, or the prefill recomputes it once per rank.
for attn_tp_rank, attn_cp_rank, pp_rank in (
(1, 0, 0),
(0, 1, 0),
(0, 0, 1),
):
with self.subTest(
attn_tp_rank=attn_tp_rank,
attn_cp_rank=attn_cp_rank,
pp_rank=pp_rank,
):
mgr = self._new_manager()
mgr.attn_tp_rank = attn_tp_rank
mgr.attn_cp_rank = attn_cp_rank
mgr.pp_rank = pp_rank
mgr._prefill_recompute_executor = MagicMock()
receiver = MagicMock(bootstrap_room=7, bootstrap_addr="127.0.0.1:8998")
self._register_prefill_info(mgr, "127.0.0.1:8998", 30000)
mgr.submit_prefill_recompute(receiver, self._payload())
mgr._prefill_recompute_executor.submit.assert_not_called()
receiver.abort.assert_not_called()
self.assertEqual(mgr.failure_records, {})
def test_submit_unresolved_url_fails_via_abort(self):
mgr = self._new_manager()
mgr._prefill_recompute_executor = MagicMock()
# No prefill_info registered for this bootstrap_addr -> URL unresolved.
receiver = MagicMock(bootstrap_room=7, bootstrap_addr="127.0.0.1:8998")
mgr.submit_prefill_recompute(receiver, self._payload())
receiver.abort.assert_called_once()
mgr._prefill_recompute_executor.submit.assert_not_called()
self.assertIn(7, mgr.failure_records)
def test_run_aborts_on_http_error(self):
mgr = self._new_manager()
session = MagicMock()
session.post.return_value = SimpleNamespace(status_code=500, text="boom")
mgr._prefill_recompute_sessions.session = session
receiver = MagicMock(bootstrap_room=7)
mgr._run_prefill_recompute(receiver, "http://prefill", self._payload())
session.post.assert_called_once()
receiver.abort.assert_called_once()
self.assertIn(7, mgr.failure_records)
def test_run_aborts_on_exception(self):
mgr = self._new_manager()
session = MagicMock()
session.post.side_effect = RuntimeError("network down")
mgr._prefill_recompute_sessions.session = session
receiver = MagicMock(bootstrap_room=7)
mgr._run_prefill_recompute(receiver, "http://prefill", self._payload())
receiver.abort.assert_called_once()
self.assertIn(7, mgr.failure_records)
def test_run_success_does_not_abort(self):
mgr = self._new_manager()
session = MagicMock()
session.post.return_value = SimpleNamespace(status_code=200, text="")
mgr._prefill_recompute_sessions.session = session
receiver = MagicMock(bootstrap_room=7)
mgr._run_prefill_recompute(receiver, "http://prefill", self._payload())
session.post.assert_called_once()
receiver.abort.assert_not_called()
self.assertEqual(mgr.failure_records, {})
class TestDecodePrebuiltPriority(unittest.TestCase):
def test_waiting_queue_is_sorted_before_prebuilt_selection(self):
scheduler = Scheduler.__new__(Scheduler)
@@ -229,8 +431,8 @@ class TestDecodePrebuiltPriority(unittest.TestCase):
)
scheduler.future_map = MagicMock()
scheduler.policy = MagicMock()
scheduler.policy.calc_priority.side_effect = (
lambda waiting_queue, _: waiting_queue.sort(key=lambda req: -req.priority)
scheduler.policy.calc_priority.side_effect = lambda waiting_queue, _: (
waiting_queue.sort(key=lambda req: -req.priority)
)
new_batch = MagicMock()
test/registered/unit/managers/test_scheduler_pause_generation.py
+54 -1
View File
@@ -1,5 +1,6 @@
import unittest
from collections import deque
from types import SimpleNamespace
from unittest.mock import MagicMock
from sglang.test.ci.ci_register import register_cpu_ci
@@ -7,7 +8,11 @@ from sglang.test.test_utils import maybe_stub_sgl_kernel
maybe_stub_sgl_kernel()
from sglang.srt.managers.io_struct import PauseGenerationReqInput
from sglang.srt.disaggregation.utils import DisaggregationMode
from sglang.srt.managers.io_struct import (
ContinueGenerationReqInput,
PauseGenerationReqInput,
)
from sglang.srt.managers.scheduler import Scheduler
from sglang.srt.managers.scheduler_components.pool_stats_observer import PoolStats
@@ -31,6 +36,7 @@ class TestSchedulerPauseGeneration(unittest.TestCase):
scheduler.tree_cache.protected_size.return_value = 0
scheduler.req_to_token_pool = MagicMock()
scheduler.result_queue = deque()
scheduler.disaggregation_mode = DisaggregationMode.NULL
# Support _kv_snap diagnostic logging in patched schedulers
scheduler.token_to_kv_pool_allocator = MagicMock()
scheduler.token_to_kv_pool_allocator.available_size.return_value = 1000
@@ -127,6 +133,53 @@ class TestSchedulerPauseGeneration(unittest.TestCase):
self.assertEqual(scheduler._add_request_to_queue.call_count, 2)
self.assertIsNone(scheduler.chunked_req)
def test_pd_decode_retract_requeues_for_rebootstrap(self):
"""PD decode retract should rebootstrap instead of resuming stale CPU KV."""
scheduler = self._new_scheduler()
scheduler.disaggregation_mode = DisaggregationMode.DECODE
scheduler.last_batch = None
scheduler.running_batch.reqs = [MagicMock()]
scheduler.running_batch.is_empty.return_value = False
scheduler._add_request_to_queue = MagicMock()
scheduler.disagg_decode_prealloc_queue = MagicMock()
req = SimpleNamespace(
output_ids=[10, 11, 12],
time_stats=MagicMock(),
)
scheduler.running_batch.retract_all.return_value = [req]
scheduler.running_batch.filter_batch = MagicMock()
scheduler.server_args = MagicMock()
scheduler.pause_generation(PauseGenerationReqInput(mode="retract"))
scheduler._add_request_to_queue.assert_not_called()
scheduler.disagg_decode_prealloc_queue.hold_rebootstrap.assert_called_once_with(
req
)
self.assertEqual(req.output_ids, [10, 11])
self.assertEqual(req.pd_rebootstrap_forced_output_id, 12)
self.assertTrue(req.pd_rebootstrap_in_progress)
# Rebootstrap recomputes the KV from the prefill, so the retract must skip
# the device->host KV offload rather than offload-then-delete it.
scheduler.running_batch.retract_all.assert_called_once_with(
scheduler.server_args, offload_kv=False
)
def test_pd_decode_continue_releases_held_rebootstrap(self):
"""continue_generation must enqueue staged rebootstrap reqs on resume."""
scheduler = self._new_scheduler()
scheduler.disaggregation_mode = DisaggregationMode.DECODE
scheduler.disagg_decode_prealloc_queue = MagicMock()
scheduler._engine_paused = True
scheduler.continue_generation(
ContinueGenerationReqInput(torch_empty_cache=False)
)
scheduler.disagg_decode_prealloc_queue.enqueue_held_rebootstrap.assert_called_once_with()
self.assertFalse(scheduler._engine_paused)
def test_abort_drains_overlap_queue(self):
"""abort with overlap enabled should drain the result_queue."""
scheduler = self._new_scheduler()
test/registered/unit/mem_cache/test_decode_radix_lock_ref.py
+4
View File
@@ -310,6 +310,10 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
decode_req = MagicMock()
decode_req.req = req
decode_req.waiting_for_input = True
# Non-rebootstrap request: exercise the normal decode radix-cache path
# (a truthy MagicMock would disable use_decode_radix_cache via the
# `not decode_req.is_rebootstrap` gate in pop_preallocated).
decode_req.is_rebootstrap = False
queue.queue = [decode_req]
queue.pending_reqs = []