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

[VLM] retire aborted disaggregated prefill results (#36988)

Browse Source
This commit is contained in:
Mick
2026-09-06 10:15:33 +08:00
committed by GitHub
parent f5819b09bf
commit febb360519
5 changed files with 396 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/prefill.py
+69 -17
View File
@@ -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 = (
test/registered/disaggregation/test_disaggregation_basic.py
-104
View File
@@ -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()
test/registered/disaggregation/test_disaggregation_chunked_prefill_abort.py
+117
View File
@@ -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()
test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py
+208
View File
@@ -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"]))
test/registered/unit/managers/test_generation_auxiliary_output.py
+2
View File
@@ -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,