[Feature] support clear data (#3601)

* [Feature] support clear data

* update

* fix

* fix

* fix

* fix

* fix

* fix

* fix

Author: ltd0924

fastdeploy/engine/common_engine.py

@@ -934,6 +934,18 @@ def start_cache_service(self, device_ids, ipc_signal_suffix):
     def check_and_free_block_tables(self):
         self.resource_manager.check_and_free_block_tables()
 
+    def clear_data(self):
+        try:
+            llm_logger.info("Clear Data: Start")
+            self.token_processor.clear_data()
+            self.engine_worker_queue.clear_data()
+            self.zmq_server.req_dict.clear()
+            llm_logger.info("Clear Data: Successfully")
+            return True
+        except Exception as e:
+            llm_logger.error(f"Clear data error: {e}")
+            return False
+
     def _exit_sub_services(self):
         """
         exit sub services

fastdeploy/engine/sched/resource_manager_v1.py

@@ -752,3 +752,7 @@ def finish_requests(self, request_ids: Union[str, Iterable[str]]):
                         del self.req_dict[req_id]
         except Exception as e:
             llm_logger.error(f"finish_request err: {e}, {str(traceback.format_exc())}")
+
+    def clear_data(self):
+        self.waiting: deque[Request] = deque()
+        self.to_be_rescheduled_request_id_set = set()

fastdeploy/entrypoints/engine_client.py

@@ -346,3 +346,6 @@ def clear_load_weight(self, timeout=300):
             return False, "clear model weight timeout"
         time.sleep(1)
         return True, ""
+
+    def check_model_weight_status(self):
+        return self.model_weights_status_signal.value[0] < 0

fastdeploy/entrypoints/openai/api_server.py

@@ -495,6 +495,8 @@ def reset_scheduler():
 
     if llm_engine is None:
         return Response("Engine not loaded", status_code=500)
+
+    llm_engine.engine.clear_data()
     llm_engine.engine.scheduler.reset()
     return Response("Scheduler Reset Successfully", status_code=200)
 

fastdeploy/entrypoints/openai/serving_chat.py

@@ -220,6 +220,8 @@ async def chat_completion_stream_generator(
                 decoder_base_url=self.tokenizer_base_url,
             )
             while num_choices > 0:
+                if self.engine_client.check_model_weight_status():
+                    raise ValueError("Engine is clearing model weight")
                 try:
                     response = await asyncio.wait_for(response_queue.get(), timeout=10)
                     current_waiting_time = 0
@@ -435,6 +437,14 @@ async def chat_completion_full_generator(
                 decoder_base_url=self.tokenizer_base_url,
             )
             while True:
+                if self.engine_client.check_model_weight_status():
+                    return ErrorResponse(
+                        error=ErrorInfo(
+                            message="Model weight cleared",
+                            code=ErrorCode.INVALID_VALUE,
+                            type=ErrorType.INVALID_REQUEST_ERROR,
+                        )
+                    )
                 try:
                     response = await asyncio.wait_for(response_queue.get(), timeout=10)
                     current_waiting_time = 0
@@ -523,6 +533,7 @@ async def chat_completion_full_generator(
 
         if final_res.get("error_msg") is not None and "Recover" in final_res["error_msg"]:
             choice.finish_reason = "recover_stop"
+
         choices.append(choice)
 
         num_prompt_tokens = len(prompt_token_ids)

fastdeploy/entrypoints/openai/serving_completion.py

@@ -227,6 +227,14 @@ async def completion_full_generator(
             completion_batched_token_ids = [[] for _ in range(num_choices)]
             current_waiting_time = 0
             while num_choices > 0:
+                if self.engine_client.check_model_weight_status():
+                    return ErrorResponse(
+                        error=ErrorInfo(
+                            message="Model weight cleared",
+                            code=ErrorCode.INVALID_VALUE,
+                            type=ErrorType.INVALID_REQUEST_ERROR,
+                        )
+                    )
                 try:
                     response = await asyncio.wait_for(response_queue.get(), timeout=10)
                     current_waiting_time = 0
@@ -281,7 +289,6 @@ async def completion_full_generator(
             return res
         except Exception as e:
             api_server_logger.error(f"Error in completion_full_generator: {e}", exc_info=True)
-            raise
         finally:
             self.engine_client.semaphore.release()
             if dealer is not None:
@@ -360,6 +367,8 @@ async def completion_stream_generator(
             )
             current_waiting_time = 0
             while num_choices > 0:
+                if self.engine_client.check_model_weight_status():
+                    raise ValueError("Engine is clearing model weight")
                 try:
                     response = await asyncio.wait_for(response_queue.get(), timeout=10)
                     current_waiting_time = 0
@@ -447,6 +456,7 @@ async def completion_stream_generator(
                         choices[-1].finish_reason = self.calc_finish_reason(
                             request.max_tokens, output_tokens[idx], output, tool_called[idx]
                         )
+
                     send_idx = output.get("send_idx")
                     # 只有当 send_idx 明确为 0 时才记录日志
                     if send_idx == 0 and not request.return_token_ids:

fastdeploy/inter_communicator/engine_worker_queue.py

@@ -503,6 +503,13 @@ def get_disaggregated_tasks(self):
         llm_logger.debug("get tasks from queue success")
         return item
 
+    def clear_data(self):
+        self.lock.acquire()
+        self.tasks[:] = list()
+        self.client_read_flag[:] = [1] * self.num_client
+        self.lock.release()
+        llm_logger.info("clear data for engine worker queue")
+
     def cleanup(self):
         """
         Exit the worker queue gracefully.

fastdeploy/output/token_processor.py

@@ -677,6 +677,31 @@ def _record_speculative_decoding_mertics(self, accept_num):
                     single_head_acceptance_rate
                 )
 
+    def clear_data(self):
+        if envs.ENABLE_V1_KVCACHE_SCHEDULER:
+            self.resource_manager.clear_data()
+        for i in range(self.cfg.max_num_seqs):
+            if self.resource_manager.stop_flags[i]:
+                continue
+            task = self.resource_manager.tasks_list[i]
+            result = RequestOutput(
+                request_id=task.request_id,
+                outputs=CompletionOutput(
+                    index=i,
+                    send_idx=self.tokens_counter[task.request_id],
+                    token_ids=task.eos_token_ids,
+                    draft_token_ids=[],
+                ),
+                finished=True,
+                metrics=RequestMetrics(
+                    arrival_time=time.time(),
+                    request_start_time=task.arrival_time,
+                ),
+            )
+            is_prefill = task.disaggregate_info is not None and task.disaggregate_info["role"] == "prefill"
+            self._recycle_resources(task.request_id, i, task, result, is_prefill)
+            llm_logger.warning(f"clear data for task {task.request_id}")
+
 
 class WarmUpTokenProcessor(TokenProcessor):
     """

fastdeploy/rl/dynamic_weight_manager.py

@@ -259,6 +259,7 @@ def check_model_weights_status(model_weights_status, model_runner, pid):
                 model_runner.update_parameters(pid)
             elif model_weights_status.value[0] == -1:
                 logger.info("infer engine stopped! start to clear checkpoint...")
+                model_runner.clear_requests()
                 model_runner.clear_parameters(pid)
 
             while True:

fastdeploy/worker/gcu_model_runner.py

@@ -1232,6 +1232,10 @@ def clear_parameters(self, pid):
         paddle.device.cuda.empty_cache()
         self.dynamic_weight_manager._log_memory("dynamic weight manager clear all memory")
 
+    def clear_requests(self):
+        """Dynamic model loader use to clear requests use for RL"""
+        self.share_inputs["stop_flags"][:] = True
+
     def update_parameters(self, pid):
         """ " Dynamic model loader use to update parameters use for RL"""
         self.dynamic_weight_manager.update_parameters(pid)