Abort during chunked prefill + PD peer-liveness abort (sgl-project#28086)

Author: cctry

python/sglang/srt/disaggregation/prefill.py

@@ -406,16 +406,55 @@ def maybe_prefetch_staging_for_batch(self: Scheduler, batch: ScheduleBatch) -> N
             if room is not None and room in kv_mgr.transfer_infos:
                 prefetch(room)
 
+    def resolve_waiting_queue_bootstrap(self: Scheduler) -> None:
+        """Resolve bootstrap status for waiting prefill requests before admission.
+
+        Covers the window between leaving the bootstrap queue and being admitted
+        into a running batch: aborts requests whose decode peer died, and
+        finalizes optimistic requests whose bootstrap completed so they skip
+        the post-forward bootstrap check.
+        """
+        candidates = [req for req in self.waiting_queue if not is_aborted(req)]
+        if not candidates:
+            return
+        polls = poll_and_all_reduce_attn_cp_tp_group(
+            [req.disagg_kv_sender for req in candidates],
+            self.attn_cp_cpu_group,
+            self.attn_tp_cpu_group,
+        )
+        failed = set()
+        for req, poll in zip(candidates, polls):
+            if poll == KVPoll.Failed:
+                self.handle_bootstrap_failure(req)
+                failed.add(req)
+            elif (
+                poll == KVPoll.WaitingForInput
+                and req.pending_bootstrap
+                and not should_force_retry(req)
+            ):
+                # Optimistic requests reserved a metadata buffer when popped, so
+                # finalize cannot fail here; if it ever does, the request stays
+                # pending and the post-forward check resolves it.
+                self.disagg_prefill_bootstrap_queue.finalize_bootstrap(req)
+        if failed:
+            self.waiting_queue = [
+                req for req in self.waiting_queue if req not in failed
+            ]
+
     @scheduler_nvtx_method("scheduler.get_next_batch_to_run")
     def get_next_disagg_prefill_batch_to_run(
         self: Scheduler,
     ) -> Optional[ScheduleBatch]:
+        self.process_pending_chunked_abort()
+
         # HACK (byronhsu): reset the batch_is_full flag because we never enter update_running_batch which resets it
         # Otherwise, it hangs under high concurrency
         self.running_batch.batch_is_full = False
 
         self.process_prefill_chunk()
 
+        self.resolve_waiting_queue_bootstrap()
+
         batch = self.get_new_batch_prefill()
         batch = self.dp_attn_adapter.maybe_prepare_mlp_sync_batch(batch)
 

python/sglang/srt/managers/scheduler.py

@@ -985,6 +985,7 @@ def init_chunked_prefill(self):
         elif self.chunked_prefill_size is not None and self.chunked_prefill_size <= 0:
             self.chunked_prefill_size = None
         self.chunked_req = None
+        self._pending_chunked_abort_req = None
         self.is_mixed_chunk = (
             self.chunked_prefill_size is not None
             and self.server_args.enable_mixed_chunk
@@ -2424,6 +2425,52 @@ def handle_batch_embedding_request(
     def stash_chunked_request(self, req: Req):
         maybe_cache_unfinished_req(req, self.tree_cache, chunked=True)
 
+    def process_pending_chunked_abort(self) -> None:
+        """Abort an in-flight chunked-prefill request once it is safe to do so.
+
+        ``abort_request`` only records the target in ``_pending_chunked_abort_req``
+        (tearing it down mid-iteration is unsafe). Clearing ``chunked_req`` here at
+        the top of the scheduling step stops the next chunk from launching; the
+        chunk already launched is drained when its result is resolved. Under overlap
+        the result lands a step later, so the batch-result processors keep
+        ``inflight_middle_chunks`` accounting intact and skip the aborted chunk:
+        ``process_batch_result_disagg_prefill`` via its ``is_aborted`` drop, and
+        ``process_batch_result_prefill`` via its chunked branch (the finished req
+        is excluded from streaming and its logprob offset is still accounted).
+        Mirrors ``handle_bootstrap_failure``.
+        """
+        req = self._pending_chunked_abort_req
+        if req is None:
+            return
+        if self.chunked_req is not req:
+            # Already past chunked prefill; the running-batch abort path handles
+            # it. Drop the marker once the request is actually gone.
+            if req.finished() or req.req_pool_idx is None:
+                self._pending_chunked_abort_req = None
+            return
+
+        prepare_abort(req, "Aborted")
+        req.time_stats.trace_ctx.abort(abort_info={"reason": "Aborted"})
+        req.to_finish = None
+        if self.disaggregation_mode == DisaggregationMode.PREFILL:
+            req.disagg_kv_sender.abort()
+            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.req_pool_idx is not None or self.tree_cache.supports_mamba()
+        ) and not req.kv_committed_freed:
+            release_kv_cache(req, self.tree_cache, is_insert=False)
+
+        self.chunked_req = None
+        self._chunked_req_scheduled_last_iter = False
+        self._pending_chunked_abort_req = None
+        self.ipc_channels.send_to_tokenizer.send_output(AbortReq(rid=req.rid), req)
+        logger.debug(f"Abort chunked prefill request. {req.rid=}")
+
     def _build_hisparse_decode_batch(self, reqs):
         """Build a ScheduleBatch for hisparse requests transitioning from staging to decode."""
         device = self.device
@@ -2467,6 +2514,8 @@ def _build_hisparse_decode_batch(self, reqs):
 
     @scheduler_nvtx_method("scheduler.get_next_batch_to_run")
     def get_next_batch_to_run(self) -> Optional[ScheduleBatch]:
+        self.process_pending_chunked_abort()
+
         if self.enable_fpm:
             self._fpm_batch_t0 = time.monotonic()
         self._abort_on_waiting_timeout()
@@ -3684,6 +3733,10 @@ def handle_rpc_request(self, recv_req: RpcReqInput):
         return RpcReqOutput(success, "" if not exec else str(exec))
 
     def abort_request(self, recv_req: AbortReq):
+        if (chunked_req := self.chunked_req) is not None:
+            if recv_req.abort_all or chunked_req.rid.startswith(recv_req.rid):
+                self._pending_chunked_abort_req = chunked_req
+
         # todo hisparse, release resources for abort requests in hisparse coordinator
         # Delete requests in the waiting queue
         to_del = []

python/sglang/srt/managers/scheduler_components/batch_result_processor.py

@@ -216,8 +216,12 @@ def process_batch_result_prefill(
             logprob_pt = 0
 
             for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
-                if req.finished() or req.is_retracted:
-                    # decode req in mixed batch or retracted req
+                if (
+                    req.finished() and req.inflight_middle_chunks <= 0
+                ) or req.is_retracted:
+                    # Decode req in a mixed batch, or a retracted req. Keep an
+                    # aborted middle chunk in the chunked branch long enough to
+                    # drain its accounting without streaming it.
                     continue
 
                 if req.inflight_middle_chunks <= 0:

python/sglang/test/manual/disaggregation/test_chunked_prefill_abort.py

@@ -0,0 +1,134 @@
+import threading
+import time
+import unittest
+import uuid
+from typing import Any
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+CHUNKED_PREFILL_SIZE = 64
+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 TestChunkedPrefillAbortE2E(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--disable-cuda-graph",
+                "--chunked-prefill-size",
+                str(CHUNKED_PREFILL_SIZE),
+                "--max-running-requests",
+                "4",
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        if cls.process:
+            kill_process_tree(cls.process.pid)
+
+    def test_abort_mid_chunked_prefill_by_rid(self):
+        rid = f"chunked-prefill-abort-{uuid.uuid4().hex}"
+        result: dict[str, Any] = {}
+
+        def run_generate():
+            try:
+                response = requests.post(
+                    self.base_url + "/generate",
+                    json={
+                        "rid": rid,
+                        "text": f"{rid}\n{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(0.5)
+        abort_deadline = time.monotonic() + 8
+        while thread.is_alive() and time.monotonic() < abort_deadline:
+            requests.post(
+                self.base_url + "/abort_request",
+                json={"rid": rid, "abort_all": False},
+                timeout=10,
+            )
+            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}",
+        )
+
+        health = requests.get(self.base_url + "/health", timeout=10)
+        self.assertEqual(health.status_code, 200, health.text)
+
+        follow_up = requests.post(
+            self.base_url + "/generate",
+            json={
+                "rid": f"chunked-prefill-after-abort-{uuid.uuid4().hex}",
+                "text": "The capital of France is",
+                "sampling_params": {"temperature": 0, "max_new_tokens": 4},
+            },
+            timeout=60,
+        )
+        follow_up_body = _decode_response(follow_up)
+        self.assertEqual(follow_up.status_code, 200, follow_up_body)
+        self.assertFalse(_is_abort_result(follow_up.status_code, follow_up_body))
+
+
+if __name__ == "__main__":
+    unittest.main()

python/sglang/test/manual/disaggregation/test_disaggregation_chunked_prefill_abort.py

@@ -0,0 +1,109 @@
+import threading
+import time
+import unittest
+import uuid
+from typing import Any
+
+import requests
+
+from sglang.test.server_fixtures.disaggregation_fixture import (
+    PDDisaggregationServerBase,
+)
+from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+
+PD_EXTRA_ARGS = ["--max-running-requests", "4", "--chunked-prefill-size", "64"]
+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_EXTRA_ARGS
+        cls.extra_decode_args = PD_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{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()