[Cherry-Pick] Fix PD interaction and error response

fastdeploy/engine/common_engine.py

@@ -959,14 +959,18 @@ def _fetch_request():
                             status, msg = self.split_connector.check_decode_allocated(task)
                             task.metrics.ask_decode_resource_finish_time = time.time()
                             if not status:
-                                self.llm_logger.error(f"{task.request_id} prefill failed with msg:{msg}.")
+                                error_msg = (
+                                    f"PD Error: prefill failed to apply for resource from decode, "
+                                    f"req: {task.request_id}, msg:{msg}."
+                                )
+                                self.llm_logger.error(error_msg)
                                 self.scheduler.put_results(
                                     [
                                         RequestOutput(
                                             request_id=task.request_id,
                                             finished=True,
                                             error_code=500,
-                                            error_msg=msg,
+                                            error_msg=error_msg,
                                         )
                                     ]
                                 )
@@ -1072,14 +1076,17 @@ def _fetch_request():
                     if self.cfg.scheduler_config.splitwise_role == "decode":
                         for task in tasks:
                             if task.task_type == RequestType.PREEMPTED:
-                                msg = f"{task.request_id} decode not enough blocks, need to be rescheduled."
+                                msg = (
+                                    f"PD Error: decode does not have enough blocks for "
+                                    f"preallocated request. req:{task.request_id} "
+                                )
                                 self.llm_logger.error(msg)
                                 self.scheduler.put_results(
                                     [
                                         RequestOutput(
                                             request_id=task.request_id,
                                             finished=True,
-                                            error_code=500,
+                                            error_code=502,
                                             error_msg=msg,
                                         )
                                     ]

fastdeploy/engine/request.py

@@ -205,6 +205,7 @@ def __init__(
             self.metrics = RequestMetrics()
         else:
             self.metrics = metrics
+        self.metrics.prompt_token_ids_len = self.prompt_token_ids_len
         # from ChatCompletionRequest or CompletionRequest
         self.user = user
         self.metadata = metadata
@@ -872,6 +873,7 @@ class RequestMetrics:
     speculate_metrics: Optional[SpeculateMetrics] = None
 
     # cache related
+    prompt_token_ids_len: Optional[int] = None
     gpu_cache_token_num: Optional[int] = 0
     cpu_cache_token_num: Optional[int] = 0
     storage_cache_token_num: Optional[int] = 0

fastdeploy/engine/sched/resource_manager_v1.py

@@ -1309,6 +1309,7 @@ def get_prefix_cached_blocks(self, request: Request):
                 request.metrics.storage_cache_token_num = metrics["storage_match_token_num"]
                 request.metrics.cpu_cache_prepare_time = metrics["cpu_cache_prepare_time"]
                 request.metrics.storage_cache_prepare_time = metrics["storage_cache_prepare_time"]
+                request.metrics.prompt_token_ids_len = request.prompt_token_ids_len
 
                 main_process_metrics.prefix_cache_token_num.inc(request.num_computed_tokens)
                 main_process_metrics.prefix_gpu_cache_token_num.inc(request.metrics.gpu_cache_token_num)
@@ -1439,7 +1440,6 @@ def preallocate_resource_in_d(self, request: Request):
             request.disaggregate_info["block_tables"] = request.block_tables
             allocated_position = self.get_available_position()
             request.idx = allocated_position
-            self.tasks_list[request.idx] = request
             self.stop_flags[request.idx] = False
             self.requests[request.request_id] = request
             self.req_dict[request.request_id] = allocated_position
@@ -1483,6 +1483,8 @@ def add_prefilled_request(self, request_output: RequestOutput):
             request.metrics = copy.deepcopy(request_output.metrics)
             request.metrics.decode_inference_start_time = time.time()
             request.metrics.update_decoder_start_time()
+
+            self.tasks_list[request.idx] = request
             self.running.append(request)
 
     def _free_blocks(self, request: Request):

fastdeploy/entrypoints/openai/protocol.py

@@ -268,7 +268,7 @@ class ChatCompletionResponseChoice(BaseModel):
     logprobs: Optional[LogProbs] = None
     draft_logprobs: Optional[LogProbs] = None
     prompt_logprobs: Optional[PromptLogprobs] = None
-    finish_reason: Optional[Literal["stop", "length", "tool_calls", "recover_stop", "abort"]]
+    finish_reason: Optional[Literal["stop", "length", "tool_calls", "recover_stop", "abort", "pd_reschedule"]]
     speculate_metrics: Optional[SpeculateMetrics] = None
 
 
@@ -333,7 +333,7 @@ class ChatCompletionResponseStreamChoice(BaseModel):
     logprobs: Optional[LogProbs] = None
     draft_logprobs: Optional[LogProbs] = None
     prompt_logprobs: Optional[PromptLogprobs] = None
-    finish_reason: Optional[Literal["stop", "length", "tool_calls", "recover_stop", "abort"]] = None
+    finish_reason: Optional[Literal["stop", "length", "tool_calls", "recover_stop", "abort", "pd_reschedule"]] = None
     arrival_time: Optional[float] = None
     speculate_metrics: Optional[SpeculateMetrics] = None
 
@@ -369,7 +369,7 @@ class CompletionResponseChoice(BaseModel):
     draft_logprobs: Optional[CompletionLogprobs] = None
     prompt_logprobs: Optional[PromptLogprobs] = None
     reasoning_content: Optional[str] = None
-    finish_reason: Optional[Literal["stop", "length", "tool_calls", "recover_stop", "abort"]] = None
+    finish_reason: Optional[Literal["stop", "length", "tool_calls", "recover_stop", "abort", "pd_reschedule"]] = None
     tool_calls: Optional[List[DeltaToolCall | ToolCall]] = None
     speculate_metrics: Optional[SpeculateMetrics] = None
 
@@ -415,7 +415,7 @@ class CompletionResponseStreamChoice(BaseModel):
     prompt_tokens: Optional[str] = None
     completion_tokens: Optional[str] = None
     reasoning_content: Optional[str] = None
-    finish_reason: Optional[Literal["stop", "length", "tool_calls", "recover_stop", "abort"]] = None
+    finish_reason: Optional[Literal["stop", "length", "tool_calls", "recover_stop", "abort", "pd_reschedule"]] = None
     tool_calls: Optional[List[DeltaToolCall | ToolCall]] = None
     speculate_metrics: Optional[SpeculateMetrics] = None
 

fastdeploy/entrypoints/openai/serving_chat.py

@@ -614,10 +614,28 @@ async def chat_completion_full_generator(
                     request=request,
                 )
                 async for data in generator:
-                    if data.get("error_code", 200) != 200:
-                        raise ValueError("{}".format(data["error_msg"]))
                     idx = int(data["request_id"].split("_")[-1])
-                    # api_server_logger.debug(f"Client {request_id} received: {data}")
+                    if data.get("error_code", 200) != 200:
+                        # Error response - include already-generated tokens in the response
+                        if completion_token_ids[idx]:
+                            text = self.engine_client.data_processor.tokenizer.decode(
+                                completion_token_ids[idx], skip_special_tokens=True
+                            )
+                        else:
+                            text = ""
+                        data["outputs"] = {
+                            "text": text,
+                            "completion_tokens": text,
+                            "reasoning_content": "",
+                            "tool_calls": None,
+                            "reasoning_token_num": 0,
+                            "num_image_tokens": 0,
+                            "token_ids": [],
+                            "top_logprobs": None,
+                            "draft_top_logprobs": None,
+                        }
+                        data["metrics"] = data.get("metrics") or {}
+                        data["finished"] = True
                     previous_num_tokens[idx] += len(data["outputs"]["token_ids"])
                     completion_token_ids[idx].extend(data["outputs"]["token_ids"])
                     # The logprob for handling the response
@@ -804,6 +822,9 @@ async def _create_chat_completion_choice(
 
         if data.get("error_msg", None) is not None and "Aborted" in data["error_msg"]:
             finish_reason = "abort"
+
+        if data.get("error_msg", None) is not None and "PD Error" in data["error_msg"]:
+            finish_reason = "pd_reschedule"
         return ChatCompletionResponseChoice(
             index=idx,
             message=message,

fastdeploy/input/base_processor.py

@@ -236,6 +236,17 @@ def process_response_dict(self, response_dict, **kwargs):
 
         ``stream`` is read from ``kwargs`` (default: True).
         """
+        # Error responses (e.g., preemption) have outputs=None or error_code!=200.
+        # Skip token decoding and return as-is to let upstream error handling take over.
+        if isinstance(response_dict, dict):
+            outputs = response_dict.get("outputs")
+            error_code = response_dict.get("error_code", 200)
+        else:
+            outputs = getattr(response_dict, "outputs", None)
+            error_code = getattr(response_dict, "error_code", 200)
+        if outputs is None or error_code != 200:
+            return response_dict
+
         stream = kwargs.get("stream", True)
         if stream:
             return self.process_response_dict_streaming(response_dict, **kwargs)

fastdeploy/output/token_processor.py

@@ -528,8 +528,11 @@ def _recycle_resources(self, task_id, index, task, result=None, is_prefill=False
                         self.prefill_result_status[finished_task_id[0]] = finished_task_id[1]
                 if task_id in self.prefill_result_status:
                     if self.prefill_result_status[task_id] != "finished":
-                        result.error_code = 400
-                        result.error_message = f"{task_id} failed to {self.prefill_result_status[task_id]}"
+                        result.error_code = 501
+                        result.error_msg = (
+                            f"PD Error: prefill failed to send cache to decode, "
+                            f"{task_id}, {self.prefill_result_status[task_id]}"
+                        )
                     llm_logger.info(
                         f"wait for sending cache, request_id: {task_id}, cost seconds: {time.time()-start_time:.5f}"
                     )
@@ -737,7 +740,7 @@ def _process_batch_output(self):
         batch_result = list()
         # reschedule
         for i in range(batch):
-            if self.resource_manager.stop_flags[i]:
+            if self.resource_manager.stop_flags[i] or self.resource_manager.tasks_list[i] is None:
                 continue
 
             recovery_stop = False