[PD Disaggregation] Add timestamp for analyzing splitwise deployment (#5317)

* Add timestamp for analyzing splitwise deployment

* up

* up

* up

* up

* up

* up

* fix format

* fix

Author: juncaipeng

docs/online_serving/README.md

@@ -231,6 +231,8 @@ mm_hashes: Optional[list] = None
 # Hash values for multimodal (e.g., image/audio) inputs, used for verification or tracking.
 # Default None indicates no multimodal input or hash validation required.
 
+collect_metrics: Optional[bool] = False
+# Whether to return metrics information, for performance analysis or debugging (default is False, meaning no metrics are returned).
 ```
 
 ### Differences in Return Fields

docs/zh/online_serving/README.md

@@ -224,6 +224,8 @@ logits_processors_args: Optional[Dict] = None
 mm_hashes: Optional[list] = None
 # 多模态（multimodal）输入的哈希值列表，用于验证或跟踪输入内容（如图像、音频等）。默认 None 表示无多模态输入或无需哈希验证。
 
+collect_metrics: Optional[bool] = False
+# 是否返回生成过程中的指标信息，用于性能分析或调试（默认 False 表示不返回）。
 ```
 
 ### 返回字段差异

fastdeploy/demo/offline_demo.py

@@ -17,11 +17,9 @@
 from fastdeploy.engine.sampling_params import SamplingParams
 from fastdeploy.entrypoints.llm import LLM
 
-model_name_or_path = "/workspace/ERNIE-4.5-0.3B-Paddle"
-
-# 超参设置
-sampling_params = SamplingParams(temperature=0.1, max_tokens=30, prompt_logprobs=100)
-llm = LLM(model=model_name_or_path, tensor_parallel_size=1, enable_prefix_caching=False)
-output = llm.generate(prompts="who are you？", use_tqdm=True, sampling_params=sampling_params)
+model_name_or_path = "PaddlePaddle/ERNIE-4.5-0.3B-Paddle"
+sampling_params = SamplingParams(temperature=0.1, max_tokens=30)
+llm = LLM(model=model_name_or_path)
+output = llm.generate(prompts="who are you?", use_tqdm=True, sampling_params=sampling_params)
 
 print(output)

fastdeploy/engine/async_llm.py

@@ -402,7 +402,7 @@ async def add_request(
 
         try:
             request = Request.from_dict(prompt)
-            request.llm_engine_recv_req_timestamp = time.time()
+            request.metrics.scheduler_recv_req_time = time.time()
 
             # Check if already preprocessed by AsyncEngineClient
             is_preprocessed = prompt.get("_preprocessed", False)
@@ -419,7 +419,7 @@ async def add_request(
             request.need_prefill_tokens = prompt_token_ids_len
 
             if not is_preprocessed:
-                request.preprocess_start_time = arrival_time
+                request.metrics.preprocess_start_time = arrival_time
                 input_ids_len = request.prompt_token_ids_len
 
                 request.set(
@@ -448,7 +448,7 @@ async def add_request(
                     llm_logger.error(error_msg)
                     raise EngineError(error_msg, error_code=400)
 
-                request.preprocess_end_time = time.time()
+                request.metrics.preprocess_end_time = time.time()
 
             # Register output queue first, then add request
             await self.output_processor.register_request(request_id, output_queue)

fastdeploy/engine/common_engine.py

@@ -362,7 +362,6 @@ def insert_tasks(self, tasks: List[Request], current_id=-1):
             tasks.remove(tmp_task)
 
         for item in tasks:
-            item.schedule_start_time = time.time()
             trace_print(LoggingEventName.RESOURCE_ALLOCATE_START, item.request_id, getattr(item, "user", ""))
         available_batch = np.sum(self.resource_manager.stop_flags)
         if len(tasks) > available_batch:
@@ -400,7 +399,7 @@ def insert_tasks(self, tasks: List[Request], current_id=-1):
         if not is_decode:
             self.llm_logger.info(f"Tasks are sent to engine, req_ids={req_ids}")
             for task in tasks:
-                task.inference_start_time = time.time()
+                task.metrics.inference_start_time = time.time()
                 trace_print(LoggingEventName.RESOURCE_ALLOCATE_END, task.request_id, getattr(task, "user", ""))
                 trace_print(LoggingEventName.REQUEST_SCHEDULE_END, task.request_id, getattr(task, "user", ""))
                 trace_print(LoggingEventName.INFERENCE_START, task.request_id, getattr(task, "user", ""))
@@ -415,7 +414,7 @@ def insert_tasks(self, tasks: List[Request], current_id=-1):
     def _insert_prefilled_requests(self, request_outputs: List[RequestOutput]):
         """
         Decode insert prefilled requests into engine worker queue.
-        Used in v1_kvcache_scheduler.
+        Used in v0_kvcache_scheduler.
         Args:
             request_outputs: a list of RequestOutput sent by prefill instance
         """
@@ -437,6 +436,10 @@ def _insert_prefilled_requests(self, request_outputs: List[RequestOutput]):
 
             cur_req.prompt_token_ids[0] = req_out.outputs.token_ids[0]
             cur_req.num_cached_tokens = req_out.num_cached_tokens
+            req_out.metrics.decode_recv_req_time = cur_req.metrics.decode_recv_req_time
+            req_out.metrics.decode_preallocate_req_time = cur_req.metrics.decode_preallocate_req_time
+            cur_req.metrics = req_out.metrics
+            cur_req.metrics.decode_inference_start_time = time.time()
             if self.cfg.speculative_config.method in ["mtp"] and self.cfg.scheduler_config.splitwise_role == "decode":
                 cur_req.draft_token_ids = copy.deepcopy(req_out.outputs.draft_token_ids)
 
@@ -644,6 +647,7 @@ def _schedule_request_to_worker(self):
                     batch=num_prefill_batch,
                 )
                 for task in tasks:
+                    task.metrics.engine_get_req_time = time.time()
                     trace_print(LoggingEventName.REQUEST_QUEUE_END, task.request_id, getattr(task, "user", ""))
                 if len(tasks) == 0:
                     time.sleep(0.001)
@@ -706,7 +710,7 @@ def _fetch_request():
                     batch=num_prefill_batch,
                 )
                 for task in tasks:
-                    task.schedule_start_time = time.time()
+                    task.metrics.engine_get_req_time = time.time()
                     trace_print(LoggingEventName.REQUEST_QUEUE_END, task.request_id, getattr(task, "user", ""))
 
                 if self.cfg.scheduler_config.splitwise_role == "decode":
@@ -732,7 +736,10 @@ def _fetch_request():
                             # assure can allocate block ids in P
                             while not self.resource_manager.preallocate_resource_in_p(task):
                                 time.sleep(0.005)
-                            self.llm_logger.debug(f"P has allocated resources for request: {task.request_id}")
+                            self.llm_logger.debug(
+                                f"P has allocated resources and then ask D resource for request: {task.request_id}"
+                            )
+                            task.metrics.ask_decode_resource_start_time = time.time()
                             while True:
                                 self.split_connector.send_splitwise_tasks([task], task.idx)
                                 status, msg = self.split_connector.check_decode_allocated(task)
@@ -742,39 +749,39 @@ def _fetch_request():
                                     )
                                     time.sleep(0.05)
                                 else:
+                                    task.metrics.ask_decode_resource_finish_time = time.time()
                                     break
                             self.llm_logger.debug(f"D has allocated resource for request: {task.request_id}")
                     else:
                         for task in tasks:
                             # assure can allocate block ids in P
                             while not self.resource_manager.preallocate_resource_in_p(task):
                                 time.sleep(0.005)
-                            self.llm_logger.debug(f"P has allocated resources for request: {task.request_id}")
+
+                            self.llm_logger.debug(
+                                f"P has allocated resources and then ask D resource for req_id: {task.request_id}"
+                            )
+                            task.metrics.ask_decode_resource_start_time = time.time()
                             self.split_connector.send_splitwise_tasks([task], task.idx)
 
                         for task in tasks:
-                            if self.cfg.scheduler_config.splitwise_role != "mixed":
-                                # assure fetch block ids from D
-                                status, msg = self.split_connector.check_decode_allocated(task)
-                                if not status:
-                                    self.llm_logger.error(
-                                        f"D failed to allocate resource for request {task.request_id}, message: {msg}."
-                                    )
-                                    self.scheduler.put_results(
-                                        [
-                                            RequestOutput(
-                                                request_id=task.request_id,
-                                                finished=True,
-                                                error_code=500,
-                                                error_msg=msg,
-                                            )
-                                        ]
-                                    )
-                                    need_delete_tasks.append(task)
-                                    continue
-                                else:
-                                    self.llm_logger.debug(f"D has allocated resource for request: {task.request_id}")
-
+                            # assure fetch block ids from D
+                            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}.")
+                                self.scheduler.put_results(
+                                    [
+                                        RequestOutput(
+                                            request_id=task.request_id,
+                                            finished=True,
+                                            error_code=500,
+                                            error_msg=msg,
+                                        )
+                                    ]
+                                )
+                                need_delete_tasks.append(task)
+                                continue
                     for tmp_task in need_delete_tasks:
                         tasks.remove(tmp_task)
                         # release resource in P
@@ -822,6 +829,7 @@ def _fetch_request():
                 # Fetch requests and add them to the scheduling queue
                 if tasks:
                     for task in tasks:
+                        task.metrics.add_req_to_resource_manager_time = time.time()
                         trace_print(
                             LoggingEventName.RESOURCE_ALLOCATE_START, task.request_id, getattr(task, "user", "")
                         )
@@ -895,6 +903,11 @@ def _fetch_request():
                                 LoggingEventName.REQUEST_SCHEDULE_END, task.request_id, getattr(task, "user", "")
                             )
                             trace_print(LoggingEventName.INFERENCE_START, task.request_id, getattr(task, "user", ""))
+                        if isinstance(task, Request):
+                            if self.cfg.scheduler_config.splitwise_role == "decode":
+                                task.metrics.decode_inference_start_time = time.time()
+                            else:
+                                task.metrics.inference_start_time = time.time()
                     self.engine_worker_queue.put_tasks((tasks, self.resource_manager.real_bsz))
 
                 # 4. Response error tasks
@@ -962,7 +975,7 @@ def _insert_zmq_task_to_scheduler(self):
                     err_msg = None
                     try:
                         request = Request.from_dict(data)
-                        request.llm_engine_recv_req_timestamp = time.time()
+                        request.metrics.scheduler_recv_req_time = time.time()
                         start_span("ENQUEUE_ZMQ", data, trace.SpanKind.PRODUCER)
                         main_process_metrics.requests_number.inc()
                         trace_print(LoggingEventName.PREPROCESSING_END, data["request_id"], data.get("user", ""))
@@ -1132,6 +1145,8 @@ def _fetch_requests():
                     self.llm_logger.debug(
                         f"D has received tasks to preallocate resource for tasks: {[task.request_id for task in tasks]}"
                     )
+                    for task in tasks:
+                        task.metrics.decode_recv_req_time = time.time()
                     allocate_resource_requests.extend(tasks)
                 elif isinstance(tasks[0], RequestOutput):
                     self.llm_logger.debug(
@@ -1141,6 +1156,7 @@ def _fetch_requests():
                         tasks = [tasks]
                     for task in tasks:
                         task.finished = False
+                        task.metrics.decode_recv_first_token_time = time.time()
                     prefilled_request_ouputs.extend(tasks)
 
         def _process_allocate_resource_requests():
@@ -1150,6 +1166,8 @@ def _process_allocate_resource_requests():
 
                 if envs.ENABLE_V1_KVCACHE_SCHEDULER:
                     if self.resource_manager.preallocate_resource_in_d(task):
+                        task.metrics.decode_preallocate_req_time = time.time()
+                        self.llm_logger.info(f"Resource available, processing task {task.request_id}")
                         self.split_connector.send_cache_info_to_prefill([task])
                         self.llm_logger.debug(f"D has successfully sent cache infos for task {task.request_id}")
                         processed_indices.append(idx)
@@ -1158,6 +1176,7 @@ def _process_allocate_resource_requests():
                     if self.resource_manager.is_resource_sufficient(task.prompt_token_ids_len):
                         self.llm_logger.debug(f"D Resource available, processing task {task.request_id}")
                         self.insert_tasks([task])
+                        task.metrics.decode_preallocate_req_time = time.time()
                         processed_indices.append(idx)
                         is_success = True
 

fastdeploy/engine/engine.py

@@ -256,13 +256,13 @@ def add_requests(self, task, sampling_params=None, **kwargs):
         if sampling_params is not None:
             task.update(asdict(sampling_params))
         request = Request.from_dict(task)
-        request.llm_engine_recv_req_timestamp = time.time()
+        request.metrics.scheduler_recv_req_time = time.time()
         llm_logger.info(f"Receive request {request}")
         if sampling_params is not None:
             if sampling_params.temperature is not None and abs(sampling_params.temperature) < 1e-06:
                 sampling_params.temperature = 1e-06
             request.sampling_params = sampling_params
-        request.preprocess_start_time = time.time()
+        request.metrics.preprocess_start_time = time.time()
         chat_template_kwargs = kwargs.get("chat_template_kwargs") or {}
         chat_template_kwargs["chat_template"] = kwargs.get("chat_template")
         kwargs["chat_template_kwargs"] = chat_template_kwargs
@@ -324,7 +324,8 @@ def add_requests(self, task, sampling_params=None, **kwargs):
                 llm_logger.error(err_msg)
                 raise EngineError(err_msg, error_code=400)
 
-        request.preprocess_end_time = time.time()
+        request.metrics.preprocess_end_time = time.time()
+        request.metrics.scheduler_recv_req_time = time.time()
         self.engine.scheduler.put_requests([request])
         llm_logger.info(f"Cache task with request_id ({request.get('request_id')})")
         llm_logger.debug(f"cache task: {request}")