diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index ca4a7121d..fe3379c95 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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) diff --git a/test/registered/disaggregation/test_disaggregation_basic.py b/test/registered/disaggregation/test_disaggregation_basic.py index 6b03753ef..a8056f9ad 100644 --- a/test/registered/disaggregation/test_disaggregation_basic.py +++ b/test/registered/disaggregation/test_disaggregation_basic.py @@ -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="") + 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()