From febb360519875d95dfc25997a7e7d73ba1dc8377 Mon Sep 17 00:00:00 2001 From: Mick Date: Sun, 6 Sep 2026 10:15:33 +0800 Subject: [PATCH] [VLM] retire aborted disaggregated prefill results (#36988) --- python/sglang/srt/disaggregation/prefill.py | 86 ++++++-- .../test_disaggregation_basic.py | 104 --------- ...st_disaggregation_chunked_prefill_abort.py | 117 ++++++++++ .../test_prefill_abort_result_cleanup.py | 208 ++++++++++++++++++ .../test_generation_auxiliary_output.py | 2 + 5 files changed, 396 insertions(+), 121 deletions(-) create mode 100644 test/registered/disaggregation/test_disaggregation_chunked_prefill_abort.py create mode 100644 test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 22c265640..67248b3e2 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -729,6 +729,7 @@ class SchedulerDisaggregationPrefillMixin: result.indexer_topk_output = None logprob_pt = 0 + aborted_reqs: List[Req] = [] assert batch.spec_info is result.next_draft_input draft_input = result.next_draft_input draft_hidden_states_cpu = None @@ -763,6 +764,13 @@ class SchedulerDisaggregationPrefillMixin: if req.inflight_middle_chunks <= 0: req.time_stats.set_prefill_finished_time() + if is_aborted(req): + if self._retire_aborted_prefill_result(req): + req.time_stats.set_completion_time() + aborted_reqs.append(req) + advance_logprob_pt(i, req) + continue + # Test hook: exercise the release/requeue retry path. if req.pending_bootstrap and should_force_retry(req): self.optimistic_release_and_requeue(req) @@ -770,6 +778,24 @@ class SchedulerDisaggregationPrefillMixin: continue req.output_ids.append(next_token_id) + if req.grammar is not None: + try: + req.grammar.accept_token(next_token_id) + except ValueError as e: + error_message = f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}" + prepare_abort( + req, + error_message, + status_code=HTTPStatus.INTERNAL_SERVER_ERROR, + ) + req.grammar.finished = req.finished() + if is_aborted(req): + if self._retire_aborted_prefill_result(req): + req.time_stats.set_completion_time() + aborted_reqs.append(req) + advance_logprob_pt(i, req) + continue + maybe_cache_unfinished_req(req, self.tree_cache) self.disagg_prefill_inflight_queue.append(req) if self.spec_algorithm.is_eagle() and draft_input is not None: @@ -808,17 +834,6 @@ class SchedulerDisaggregationPrefillMixin: self.send_kv_chunk(req, last_chunk=True) req.time_stats.set_prefill_transfer_queue_entry_time() - if req.grammar is not None: - try: - req.grammar.accept_token(next_token_id) - except ValueError as e: - error_message = f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}" - prepare_abort( - req, - error_message, - status_code=HTTPStatus.INTERNAL_SERVER_ERROR, - ) - req.grammar.finished = req.finished() else: # being chunked reqs' prefill is not finished req.inflight_middle_chunks -= 1 @@ -831,16 +846,18 @@ class SchedulerDisaggregationPrefillMixin: req.extend_range is not None and req.extend_range.end >= len(req.origin_input_ids) ) - if req.pending_bootstrap and not still_chunking: - self.optimistic_release_and_requeue(req) + # Abort is terminal. Do not requeue an aborted optimistic + # request merely because bootstrap is still pending. + if is_aborted(req): + if not still_chunking and self._retire_aborted_prefill_result(req): + req.time_stats.set_completion_time() + aborted_reqs.append(req) advance_logprob_pt(i, req) req.time_stats.set_last_chunked_prefill_finish_time() continue - # Optimistic bootstrap can fail while this overlapped chunk is - # already running. Drop aborted chunks instead of sending KV. - if is_aborted(req): - self.clear_pending_chunk_send(req) + if req.pending_bootstrap and not still_chunking: + self.optimistic_release_and_requeue(req) advance_logprob_pt(i, req) req.time_stats.set_last_chunked_prefill_finish_time() continue @@ -876,6 +893,12 @@ class SchedulerDisaggregationPrefillMixin: auxiliary_output_starts, ) + if aborted_reqs: + self.output_streamer.stream_output( + aborted_reqs, + any(req.return_logprob for req in aborted_reqs), + ) + can_run_cuda_graph = result.can_run_cuda_graph self.metrics_reporter.report_prefill_stats( batch=batch, @@ -1030,6 +1053,35 @@ class SchedulerDisaggregationPrefillMixin: """ self.disagg_prefill_pending_chunk_rids.discard(req.rid) + def _retire_aborted_prefill_result(self: Scheduler, req: Req) -> bool: + """Release an aborted request when its last prefill result is safe.""" + self.clear_pending_chunk_send(req) + owns_resources = ( + req.kv.holds_kv or req.kv.holds_mamba or req.metadata_buffer_index >= 0 + ) + if not owns_resources: + # A bootstrap failure or earlier abort already retired it. + return False + + sender = req.disagg_kv_sender + if sender is not None: + try: + sender.abort() + except Exception: + # Transport notification is best effort; local ownership must + # still be released or the next idle invariant check will fail. + logger.exception("Failed to notify KV sender of abort for %s", req.rid) + + if req.to_finish is not None and not req.finished(): + req.update_finish_state() + maybe_release_metadata_buffer(req, self.req_to_metadata_buffer_idx_allocator) + req.pending_bootstrap = False + if self.enable_hicache_storage: + self.tree_cache.release_aborted_request(req.rid) + if req.kv.holds_kv or req.kv.holds_mamba: + release_kv_cache(req, self.tree_cache, is_insert=False) + return True + def handle_bootstrap_failure(self: Scheduler, req: Req) -> None: self.clear_pending_chunk_send(req) error_message = ( diff --git a/test/registered/disaggregation/test_disaggregation_basic.py b/test/registered/disaggregation/test_disaggregation_basic.py index b994a144e..c06476970 100644 --- a/test/registered/disaggregation/test_disaggregation_basic.py +++ b/test/registered/disaggregation/test_disaggregation_basic.py @@ -1,13 +1,9 @@ import asyncio import json import os -import threading -import time import unittest -import uuid from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace -from typing import Any import aiohttp import openai @@ -28,7 +24,6 @@ from sglang.test.server_fixtures.disaggregation_fixture import ( from sglang.test.test_utils import ( DEFAULT_DRAFT_MODEL_EAGLE3, DEFAULT_MODEL_NAME_FOR_TEST, - DEFAULT_SMALL_MODEL_NAME_FOR_TEST, DEFAULT_TARGET_MODEL_EAGLE3, ) @@ -754,104 +749,5 @@ class TestDisaggregationPauseResumePrefillLeak(PDDisaggregationServerBase): ) -PD_CHUNKED_ABORT_EXTRA_ARGS = [ - "--max-running-requests", - "4", - "--chunked-prefill-size", - "64", -] -_CHUNKED_ABORT_LONG_PROMPT = ( - "The quick brown fox jumps over the lazy dog. " - "Pack my box with five dozen liquor jugs. " - "Sphinx of black quartz, judge my vow. " -) * 900 - - -def _decode_response(response: requests.Response) -> Any: - try: - return response.json() - except ValueError: - return response.text - - -def _is_abort_result(status_code: int, body: Any) -> bool: - if status_code == 200: - reason = ( - body.get("meta_info", {}).get("finish_reason", {}) - if isinstance(body, dict) - else {} - ) - return isinstance(reason, dict) and reason.get("type") == "abort" - - if status_code not in (500, 503): - return False - - text = body if isinstance(body, str) else str(body) - return "abort" in text.lower() - - -class TestDisaggChunkedPrefillAbort(PDDisaggregationServerBase): - @classmethod - def setUpClass(cls): - super().setUpClass() - cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST - cls.extra_prefill_args = PD_CHUNKED_ABORT_EXTRA_ARGS - cls.extra_decode_args = PD_CHUNKED_ABORT_EXTRA_ARGS - cls.launch_all() - - def _post_abort(self, rid: str): - for url in (self.prefill_url, self.decode_url): - requests.post( - url + "/abort_request", - json={"rid": rid, "abort_all": False}, - timeout=10, - ) - - def test_abort_mid_chunked_prefill_by_rid(self): - rid = f"pd-chunked-prefill-abort-{uuid.uuid4().hex}" - result: dict[str, Any] = {} - - def run_generate(): - try: - response = requests.post( - self.lb_url + "/generate", - json={ - "rid": rid, - "text": f"{rid}\n{_CHUNKED_ABORT_LONG_PROMPT}", - "sampling_params": { - "temperature": 0, - "max_new_tokens": 4096, - "ignore_eos": True, - }, - }, - timeout=180, - ) - result["status_code"] = response.status_code - result["body"] = _decode_response(response) - except requests.RequestException as exc: - result["exception"] = repr(exc) - - thread = threading.Thread(target=run_generate) - thread.start() - - time.sleep(1.0) - abort_deadline = time.monotonic() + 8 - while thread.is_alive() and time.monotonic() < abort_deadline: - self._post_abort(rid) - time.sleep(0.2) - - thread.join(timeout=60) - self.assertFalse(thread.is_alive(), "Chunked-prefill abort request hung") - self.assertNotIn("exception", result, result.get("exception")) - self.assertTrue( - _is_abort_result(result["status_code"], result["body"]), - f"Expected chunked-prefill request to abort, got {result}", - ) - - for url in (self.lb_url, self.prefill_url, self.decode_url): - health = requests.get(url + "/health", timeout=10) - self.assertEqual(health.status_code, 200, health.text) - - if __name__ == "__main__": unittest.main() diff --git a/test/registered/disaggregation/test_disaggregation_chunked_prefill_abort.py b/test/registered/disaggregation/test_disaggregation_chunked_prefill_abort.py new file mode 100644 index 000000000..3550ca062 --- /dev/null +++ b/test/registered/disaggregation/test_disaggregation_chunked_prefill_abort.py @@ -0,0 +1,117 @@ +import threading +import time +import unittest +import uuid +from typing import Any + +import requests + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.server_fixtures.disaggregation_fixture import ( + PDDisaggregationServerBase, +) +from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST + +register_cuda_ci(est_time=120, stage="base-b", runner_config="2-gpu-large") + +PD_CHUNKED_ABORT_EXTRA_ARGS = [ + "--max-running-requests", + "4", + "--chunked-prefill-size", + "64", +] +_CHUNKED_ABORT_LONG_PROMPT = ( + "The quick brown fox jumps over the lazy dog. " + "Pack my box with five dozen liquor jugs. " + "Sphinx of black quartz, judge my vow. " +) * 900 + + +def _decode_response(response: requests.Response) -> Any: + try: + return response.json() + except ValueError: + return response.text + + +def _is_abort_result(status_code: int, body: Any) -> bool: + if status_code == 200: + reason = ( + body.get("meta_info", {}).get("finish_reason", {}) + if isinstance(body, dict) + else {} + ) + return isinstance(reason, dict) and reason.get("type") == "abort" + + if status_code not in (500, 503): + return False + + text = body if isinstance(body, str) else str(body) + return "abort" in text.lower() + + +class TestDisaggChunkedPrefillAbort(PDDisaggregationServerBase): + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST + cls.extra_prefill_args = PD_CHUNKED_ABORT_EXTRA_ARGS + cls.extra_decode_args = PD_CHUNKED_ABORT_EXTRA_ARGS + cls.launch_all() + + def _post_abort(self, rid: str): + for url in (self.prefill_url, self.decode_url): + requests.post( + url + "/abort_request", + json={"rid": rid, "abort_all": False}, + timeout=10, + ) + + def test_abort_mid_chunked_prefill_by_rid(self): + rid = f"pd-chunked-prefill-abort-{uuid.uuid4().hex}" + result: dict[str, Any] = {} + + def run_generate(): + try: + response = requests.post( + self.lb_url + "/generate", + json={ + "rid": rid, + "text": f"{rid}\n{_CHUNKED_ABORT_LONG_PROMPT}", + "sampling_params": { + "temperature": 0, + "max_new_tokens": 4096, + "ignore_eos": True, + }, + }, + timeout=180, + ) + result["status_code"] = response.status_code + result["body"] = _decode_response(response) + except requests.RequestException as exc: + result["exception"] = repr(exc) + + thread = threading.Thread(target=run_generate) + thread.start() + + time.sleep(1.0) + abort_deadline = time.monotonic() + 8 + while thread.is_alive() and time.monotonic() < abort_deadline: + self._post_abort(rid) + time.sleep(0.2) + + thread.join(timeout=60) + self.assertFalse(thread.is_alive(), "Chunked-prefill abort request hung") + self.assertNotIn("exception", result, result.get("exception")) + self.assertTrue( + _is_abort_result(result["status_code"], result["body"]), + f"Expected chunked-prefill request to abort, got {result}", + ) + + for url in (self.lb_url, self.prefill_url, self.decode_url): + health = requests.get(url + "/health", timeout=10) + self.assertEqual(health.status_code, 200, health.text) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py b/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py new file mode 100644 index 000000000..2d9bd5778 --- /dev/null +++ b/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py @@ -0,0 +1,208 @@ +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest +import torch + +from sglang.srt.disaggregation.prefill import SchedulerDisaggregationPrefillMixin +from sglang.srt.managers.schedule_batch import FINISH_ABORT, ReqKvInfo +from sglang.srt.managers.utils import GenerationBatchResult +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class _Req: + def __init__(self, *, inflight_middle_chunks: int, allocated: bool = True): + self.rid = "aborted-prefill" + self.inflight_middle_chunks = inflight_middle_chunks + self.kv = ReqKvInfo( + req_pool_idx=1 if allocated else None, + kv_allocated_len=1 if allocated else 0, + mamba_pool_idx=object() if allocated else None, + ) + self.metadata_buffer_index = 7 if allocated else -1 + self.pending_bootstrap = allocated + self.disagg_kv_sender = Mock() + self.to_finish = FINISH_ABORT() if allocated else None + self.finished_reason = None if allocated else FINISH_ABORT() + self.return_logprob = False + self.return_sampling_mask = False + self.grammar = None + self.output_ids = [] + self.origin_input_ids = list(range(100)) + self.extend_range = None + self.time_stats = SimpleNamespace( + set_prefill_finished_time=Mock(), + set_last_chunked_prefill_finish_time=Mock(), + set_completion_time=Mock(), + ) + + def finished(self): + return self.finished_reason is not None + + def update_finish_state(self): + self.finished_reason = self.to_finish + self.to_finish = None + + +class _Scheduler(SchedulerDisaggregationPrefillMixin): + def __init__(self): + self.batch_result_processor = SimpleNamespace( + snapshot_auxiliary_output_starts=Mock(return_value=[]), + move_logprobs_to_cpu=Mock(), + consume_auxiliary_output=Mock(), + ) + self.spec_algorithm = SimpleNamespace(is_eagle=lambda: False) + self.tree_cache = Mock() + self.disagg_prefill_inflight_queue = [] + self.disagg_prefill_pending_chunk_rids = {"aborted-prefill"} + self.send_kv_chunk = Mock() + self.output_streamer = Mock() + self.metrics_reporter = SimpleNamespace(report_prefill_stats=Mock()) + self.req_to_metadata_buffer_idx_allocator = Mock() + self.enable_hicache_storage = True + self.chunked_req = None + + +def _batch(req): + return SimpleNamespace( + reqs=[req], + spec_info=None, + prefill_stats=None, + dp_cooperation_info=None, + ) + + +def _result(): + return GenerationBatchResult(next_token_ids=torch.tensor([11])) + + +def _free_req(req, _tree_cache, *, is_insert): + assert is_insert is False + req.kv.req_pool_idx = None + req.kv.mark_kv_released() + req.kv.mamba_pool_idx = None + + +@patch("sglang.srt.disaggregation.prefill.release_kv_cache", side_effect=_free_req) +@patch("sglang.srt.disaggregation.prefill.maybe_cache_unfinished_req") +def test_aborted_final_result_releases_hybrid_cache( + maybe_cache_unfinished_req, release_kv_cache +): + scheduler = _Scheduler() + req = _Req(inflight_middle_chunks=0) + + scheduler.process_batch_result_disagg_prefill(_batch(req), _result()) + + release_kv_cache.assert_called_once_with(req, scheduler.tree_cache, is_insert=False) + maybe_cache_unfinished_req.assert_not_called() + req.disagg_kv_sender.abort.assert_called_once_with() + scheduler.req_to_metadata_buffer_idx_allocator.free.assert_called_once_with(7) + scheduler.tree_cache.release_aborted_request.assert_called_once_with(req.rid) + scheduler.output_streamer.stream_output.assert_called_once_with([req], False) + scheduler.send_kv_chunk.assert_not_called() + assert req.output_ids == [] + assert req.finished() + assert req.metadata_buffer_index == -1 + assert req.rid not in scheduler.disagg_prefill_pending_chunk_rids + + +@patch("sglang.srt.disaggregation.prefill.release_kv_cache", side_effect=_free_req) +def test_aborted_middle_result_releases_after_last_chunk(release_kv_cache): + scheduler = _Scheduler() + req = _Req(inflight_middle_chunks=1) + req.extend_range = SimpleNamespace(end=50) + + scheduler.process_batch_result_disagg_prefill(_batch(req), _result()) + + assert req.inflight_middle_chunks == 0 + release_kv_cache.assert_called_once_with(req, scheduler.tree_cache, is_insert=False) + scheduler.output_streamer.stream_output.assert_called_once_with([req], False) + + +@patch("sglang.srt.disaggregation.prefill.release_kv_cache", side_effect=_free_req) +def test_aborted_middle_result_waits_for_inflight_chunk(release_kv_cache): + scheduler = _Scheduler() + req = _Req(inflight_middle_chunks=1) + req.extend_range = SimpleNamespace(end=len(req.origin_input_ids)) + + scheduler.process_batch_result_disagg_prefill(_batch(req), _result()) + + release_kv_cache.assert_not_called() + scheduler.output_streamer.stream_output.assert_not_called() + + scheduler.process_batch_result_disagg_prefill(_batch(req), _result()) + + release_kv_cache.assert_called_once_with(req, scheduler.tree_cache, is_insert=False) + scheduler.output_streamer.stream_output.assert_called_once_with([req], False) + + +@patch("sglang.srt.disaggregation.prefill.release_kv_cache") +def test_delayed_result_ignores_already_retired_request(release_kv_cache): + scheduler = _Scheduler() + req = _Req(inflight_middle_chunks=0, allocated=False) + + scheduler.process_batch_result_disagg_prefill(_batch(req), _result()) + + release_kv_cache.assert_not_called() + req.disagg_kv_sender.abort.assert_not_called() + scheduler.output_streamer.stream_output.assert_not_called() + + +@patch("sglang.srt.disaggregation.prefill.release_kv_cache", side_effect=_free_req) +def test_sender_abort_failure_does_not_skip_local_cleanup(release_kv_cache): + scheduler = _Scheduler() + req = _Req(inflight_middle_chunks=0) + req.disagg_kv_sender.abort.side_effect = RuntimeError("transport is down") + + scheduler.process_batch_result_disagg_prefill(_batch(req), _result()) + + release_kv_cache.assert_called_once_with(req, scheduler.tree_cache, is_insert=False) + scheduler.req_to_metadata_buffer_idx_allocator.free.assert_called_once_with(7) + scheduler.output_streamer.stream_output.assert_called_once_with([req], False) + assert req.finished() + + +@patch("sglang.srt.disaggregation.prefill.release_kv_cache", side_effect=_free_req) +@patch("sglang.srt.disaggregation.prefill.maybe_cache_unfinished_req") +def test_grammar_rejection_retires_prefill_before_transfer( + maybe_cache_unfinished_req, release_kv_cache +): + scheduler = _Scheduler() + req = _Req(inflight_middle_chunks=0) + req.to_finish = None + req.grammar = Mock() + req.grammar.accept_token.side_effect = ValueError("invalid token") + + scheduler.process_batch_result_disagg_prefill(_batch(req), _result()) + + req.grammar.accept_token.assert_called_once_with(11) + assert req.grammar.finished + assert req.finished() + release_kv_cache.assert_called_once_with(req, scheduler.tree_cache, is_insert=False) + maybe_cache_unfinished_req.assert_not_called() + scheduler.send_kv_chunk.assert_not_called() + scheduler.output_streamer.stream_output.assert_called_once_with([req], False) + assert scheduler.disagg_prefill_inflight_queue == [] + + +def test_aborted_result_releases_mamba_allocated_before_kv(): + scheduler = _Scheduler() + scheduler.enable_hicache_storage = False + scheduler.tree_cache.supports_mamba.return_value = True + scheduler.tree_cache.req_to_token_pool.mamba_allocator.free = Mock() + req = _Req(inflight_middle_chunks=0, allocated=False) + req.kv.mamba_pool_idx = torch.tensor([3]) + req.to_finish = FINISH_ABORT() + req.finished_reason = None + + scheduler.process_batch_result_disagg_prefill(_batch(req), _result()) + + scheduler.tree_cache.req_to_token_pool.mamba_allocator.free.assert_called_once() + assert req.kv.mamba_pool_idx is None + scheduler.output_streamer.stream_output.assert_called_once_with([req], False) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/managers/test_generation_auxiliary_output.py b/test/registered/unit/managers/test_generation_auxiliary_output.py index f7462155e..cf3876daa 100644 --- a/test/registered/unit/managers/test_generation_auxiliary_output.py +++ b/test/registered/unit/managers/test_generation_auxiliary_output.py @@ -443,6 +443,8 @@ def test_disaggregated_prefill_consumes_auxiliary_output_after_commit(): req = SimpleNamespace( output_ids=[], finished_len=None, + to_finish=None, + finished_reason=None, inflight_middle_chunks=0, pending_bootstrap=False, return_logprob=False,