diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index 2b7976506..b0566e05f 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -158,6 +158,18 @@ class BaseKVSender(ABC): """ ... + def clear(self): + """ + Clear any internal states. + """ + pass + + def abort(self): + """ + Abort the current transfer. + """ + pass + class BaseKVReceiver(ABC): @abstractmethod diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 1da2a091d..38b8fc541 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -798,52 +798,32 @@ class SchedulerDisaggregationPrefillMixin: undone_reqs.append(req) continue - if req.pending_bootstrap and poll != KVPoll.Failed: - # prefill finished before bootstrap - if poll == KVPoll.WaitingForInput: - assert self.disagg_prefill_bootstrap_queue.finalize_bootstrap(req) + if req.pending_bootstrap: + # Parked: prefill finished before bootstrap completed. + if self.handle_pending_bootstrap(req, poll): self.send_kv_chunk(req, last_chunk=True) - undone_reqs.append(req) - elif poll in [KVPoll.WaitingForInput, KVPoll.Transferring]: + undone_reqs.append(req) + elif poll != KVPoll.Failed: + undone_reqs.append(req) + continue + + if poll in [KVPoll.WaitingForInput, KVPoll.Transferring]: + # todo: set Transferring correctly in backend undone_reqs.append(req) elif poll == KVPoll.Success: # transfer done release_kv_cache(req, self.tree_cache) # unlock the tree req.finished_reason = FINISH_LENGTH(length=0) # FIXME: clean up req's data in transfer engine - if hasattr(req.disagg_kv_sender, "clear"): - req.disagg_kv_sender.clear() + req.disagg_kv_sender.clear() done_reqs.append(req) req.time_stats.set_prefill_kv_transfer_finish_time() elif poll == KVPoll.Failed: - error_message = f"Prefill transfer failed for request rank={self.ps.tp_rank} {req.rid=} {req.bootstrap_room=}" - is_propagated = False - try: - req.disagg_kv_sender.failure_exception() - except Exception as e: - error_message += f" with exception {e}" - is_propagated = getattr(e, "is_from_another_rank", False) - # Mute error message for propagated exceptions to avoid duplicate logging - if is_propagated: - logger.debug(error_message) - else: - logger.warning(error_message) - req.time_stats.trace_ctx.abort(abort_info={"reason": error_message}) - release_kv_cache(req, self.tree_cache) # unlock the tree - prepare_abort( - req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR - ) + self.handle_inflight_transfer_failure(req) done_reqs.append(req) - if self.metrics_reporter.enable_metrics: - if req.pending_bootstrap: - self.metrics_collector.increment_bootstrap_failed_reqs() - else: - self.metrics_collector.increment_transfer_failed_reqs() else: - logger.warning_once( - f"Unexpected polling state {poll} for rid {req.rid} in inflight queue; " - f"treating as undone", + raise RuntimeError( + f"Unexpected poll state {poll} for req {req.rid} in inflight queue" ) - undone_reqs.append(req) for req in done_reqs: req.time_stats.set_completion_time() @@ -883,6 +863,32 @@ class SchedulerDisaggregationPrefillMixin: return done_reqs + def handle_inflight_transfer_failure( + self: Scheduler, req: Req + ) -> Optional[Exception]: + """Conclude an inflight request whose KV transfer failed.""" + error_message = ( + f"Prefill transfer failed for request rank={self.ps.tp_rank} " + f"{req.rid=} {req.bootstrap_room=}" + ) + exc: Optional[Exception] = None + try: + req.disagg_kv_sender.failure_exception() + except Exception as e: + exc = e + error_message += f" with exception {e}" + # Mute error message for propagated exceptions to avoid duplicate logging + if getattr(exc, "is_from_another_rank", False): + logger.debug(error_message) + else: + logger.warning(error_message) + req.time_stats.trace_ctx.abort(abort_info={"reason": error_message}) + release_kv_cache(req, self.tree_cache) # unlock the tree + prepare_abort(req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR) + if self.metrics_reporter.enable_metrics: + self.metrics_collector.increment_transfer_failed_reqs() + return exc + def get_transferred_rids(self: Scheduler) -> List[str]: """ Used by PP, get the transferred rids but **do not pop**