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

fix: pause_generation should not populate running_batch on prefill nodes (#20273)

Browse Source
This commit is contained in:
Lawrence Wu
2026-04-03 16:16:06 -07:00
committed by GitHub
parent 5118295f7b
commit 9593d434c4
2 changed files with 214 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
+9 -2
View File
@@ -3278,7 +3278,14 @@ class Scheduler(
self.last_batch.filter_batch(
chunked_req_to_exclude=list(chunked_req_to_exclude)
)
if not self.last_batch.is_empty():
# Skip merge for disagg prefill: completed prefill requests are
# already in disagg_prefill_inflight_queue. Merging them into
# running_batch leaks them, since the prefill event loop never
# calls update_running_batch to clean them up.
if (
not self.last_batch.is_empty()
and self.disaggregation_mode != DisaggregationMode.PREFILL
):
if self.running_batch.is_empty():
self.running_batch = self.last_batch
else:
@@ -3287,7 +3294,7 @@ class Scheduler(
self.last_batch = None
self.cur_batch = None
if recv_req.mode == "retract":
if recv_req.mode == "retract" and not self.running_batch.is_empty():
self.running_batch.filter_batch(v1_spec_info_filtered=True)
if len(self.running_batch.reqs) != 0:
retracted_reqs = self.running_batch.retract_all(self.server_args)
test/registered/disaggregation/test_disaggregation_basic.py
+205
View File
@@ -1,8 +1,10 @@
import asyncio
import json
import os
import unittest
from types import SimpleNamespace
import aiohttp
import openai
import requests
from transformers import AutoTokenizer
@@ -450,5 +452,208 @@ class TestDisaggregationSimulatedRetract(PDDisaggregationServerBase):
self.assertGreater(metrics["score"], 0.62)
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
scheduling budget and blocks all subsequent prefills."""
MAX_RUNNING = 4
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
cls.start_prefill()
cls.start_decode()
cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill)
cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode)
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
"1",
"--max-running-requests",
str(cls.MAX_RUNNING),
"--enable-metrics",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
"1",
"--base-gpu-id",
"1",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_retract_pause_no_leak_on_prefill(self):
"""Retract-mode pause on a disagg prefill node must not leak prefill
requests into running_batch. Without the fix, each retract pause merges
last_batch into running_batch, but the prefill event loop never cleans
them up via update_running_batch. After enough cycles the
max-running-requests budget is exhausted and all new prefills hang."""
asyncio.run(self._run_pause_resume_leak_test("retract"))
def test_retract_pause_empty_running_batch(self):
"""Retract-mode pause must not crash when running_batch is empty.
Regression test for issue #20272."""
asyncio.run(self._run_pause_on_idle("retract"))
async def _run_pause_on_idle(self, mode):
"""Pause/resume on an idle prefill node (no in-flight requests)."""
async with aiohttp.ClientSession() as session:
async with session.post(
self.prefill_url + "/pause_generation",
json={"mode": mode},
timeout=aiohttp.ClientTimeout(total=10),
) as resp:
resp.raise_for_status()
async with session.post(
self.prefill_url + "/continue_generation",
json={},
timeout=aiohttp.ClientTimeout(total=10),
) as resp:
resp.raise_for_status()
# Verify the engine still works after pause/resume
async with session.post(
self.lb_url + "/generate",
json={
"text": "What is 1+1?",
"sampling_params": {"temperature": 0, "max_new_tokens": 1},
},
timeout=aiohttp.ClientTimeout(total=10),
) as resp:
resp.raise_for_status()
body = await resp.json()
self.assertIn("text", body)
self.assertGreater(len(body["text"]), 0)
async def _get_num_running_reqs(self, session):
"""Query sglang:num_running_reqs from prefill node's /metrics."""
async with session.get(
self.prefill_url + "/metrics",
timeout=aiohttp.ClientTimeout(total=5),
) as resp:
resp.raise_for_status()
text = await resp.text()
for line in text.splitlines():
# Match the gauge line, skip HELP/TYPE comments and
# per-priority breakdowns (which have priority="<int>")
if (
line.startswith("sglang:num_running_reqs{")
and "priority=" not in line
):
return int(float(line.split()[-1]))
return 0
async def _run_pause_resume_leak_test(self, mode):
NUM_WORKERS = 64
NUM_PAUSE_RESUME_CYCLES = self.MAX_RUNNING * 4
MAX_NEW_TOKENS = 1
LONG_PROMPT = "Tell me a story. " * 200
async def _background_worker(session, worker_id, cancel_event):
"""Send requests sequentially until cancelled."""
seq = 0
while not cancel_event.is_set():
try:
async with session.post(
self.lb_url + "/generate",
json={
"text": f"[w{worker_id}-{seq}] {LONG_PROMPT}",
"sampling_params": {
"temperature": 0,
"max_new_tokens": MAX_NEW_TOKENS,
},
},
timeout=aiohttp.ClientTimeout(total=30),
) as resp:
await resp.read()
except Exception:
pass
seq += 1
async def _post(session, url, json_data):
async with session.post(
url,
json=json_data,
timeout=aiohttp.ClientTimeout(total=30),
) as resp:
resp.raise_for_status()
cancel_event = asyncio.Event()
async with aiohttp.ClientSession() as session:
workers = [
asyncio.create_task(_background_worker(session, i, cancel_event))
for i in range(NUM_WORKERS)
]
for _ in range(NUM_PAUSE_RESUME_CYCLES):
await _post(
session,
self.prefill_url + "/pause_generation",
{"mode": mode},
)
await _post(
session,
self.prefill_url + "/continue_generation",
{},
)
await asyncio.sleep(0.1)
# Stop workers and abort all in-flight requests
cancel_event.set()
await _post(
session, self.prefill_url + "/abort_request", {"abort_all": True}
)
await _post(
session, self.decode_url + "/abort_request", {"abort_all": True}
)
await asyncio.gather(*workers, return_exceptions=True)
# Wait for abort cleanup, then check for leaked phantom requests.
# With the bug, running_batch accumulates phantom prefill requests
# that are never cleaned up.
await asyncio.sleep(2)
num_running = await self._get_num_running_reqs(session)
self.assertEqual(
num_running,
0,
f"Prefill node has {num_running} phantom running requests "
f"after abort — pause_generation is leaking into running_batch",
)
if __name__ == "__main__":
unittest.main()