[Bug fix] Test td cache messager (#3242)

* support disable cache task in decode node

* fix busg

* Update engine.py

* Update expert_service.py

* Update splitwise_connector.py

* Optimize log for debug

* Optimize log for debug

* fix bug

---------

Co-authored-by: ltd0924 <[email protected]>
Co-authored-by: ltd0924 <[email protected]>

Author: rainyfly

fastdeploy/cache_manager/cache_messager.py

@@ -17,8 +17,9 @@
 import argparse
 import json
 import math
-import time
 import threading
+import time
+
 import numpy as np
 import paddle
 
@@ -196,7 +197,9 @@ def __init__(
 
         self.gpu_id = gpu_id
         self.cache_info = dict()
-        self.rank_id = self.rank + local_data_parallel_id * self.nranks  # align with engine worker rank (paddle.distributed.launch)
+        self.rank_id = (
+            self.rank + local_data_parallel_id * self.nranks
+        )  # align with engine worker rank (paddle.distributed.launch)
 
         connect_rdma_thread = threading.Thread(target=self._handle_connect_task)
         connect_rdma_thread.daemon = True
@@ -284,7 +287,7 @@ def prefill_layerwise_send_cache_thread(self):
                 if not self.cache_info:
                     time.sleep(0.001)
                     continue
-                logger.info(f"prefilled_layer_idx: {prefilled_layer_idx}, prefilled_step_idx: {prefilled_step_idx}")
+                logger.debug(f"prefilled_layer_idx: {prefilled_layer_idx}, prefilled_step_idx: {prefilled_step_idx}")
                 for req_id, item in list(self.cache_info.items()):
                     if "status" not in item:
                         continue
@@ -364,7 +367,7 @@ def prefill_layerwise_send_cache_thread(self):
 
         except Exception as e:
             logger.info(f"prefill layerwise send cache thread has exception: {e}")
-            
+
     def _handle_connect_task(self):
         while True:
             try:
@@ -465,7 +468,8 @@ def main():
 if __name__ == "__main__":
 
     args = parse_args()
-    logger = get_logger("cache_messager", "cache_messager.log")
+    rank_id = args.rank + args.local_data_parallel_id * args.mp_num
+    logger = get_logger("cache_messager", f"cache_messager_rank{rank_id}.log")
 
     logger.info("create cache messager...")
     logger.info(f"{args}")

fastdeploy/engine/engine.py

@@ -113,6 +113,8 @@ def __init__(self, cfg):
 
         self.start_queue_service()
 
+        self.enable_decode_cache_task = envs.FD_ENABLE_CACHE_TASK == "1"
+
         if envs.ENABLE_V1_KVCACHE_SCHEDULER:
             self.resource_manager = ResourceManagerV1(
                 cfg.max_num_seqs, cfg, cfg.tensor_parallel_size, cfg.splitwise_role
@@ -630,11 +632,15 @@ def receiver_loop():
                                         if self.resource_manager.is_resource_sufficient(task.prompt_token_ids_len):
                                             self.insert_tasks([task])
                                         else:
+                                            if not self.enable_decode_cache_task:
+                                                task.error_msg = "Not enough resources"
                                             new_waiting.append(task)
-
                                     if new_waiting:
-                                        self.waiting_requests.extend(new_waiting)
-                                        llm_logger.info(f"Added {len(new_waiting)} tasks to waiting queue")
+                                        if not self.enable_decode_cache_task:
+                                            self.split_connector.send_cache_infos(new_waiting, -1)
+                                        else:
+                                            self.waiting_requests.extend(new_waiting)
+                                            llm_logger.info(f"Added {len(new_waiting)} tasks to waiting queue")
 
                     else:
                         time.sleep(0.001)
@@ -805,6 +811,22 @@ def insert_tasks(self, tasks, current_id=-1, allocated=False):
 
         for task in tasks:
             start_span_request("DEQUEUE", task, trace.SpanKind.CONSUMER)
+            if self.cfg.splitwise_role != "mixed":
+                status, msg = self.split_connector.check_decode_allocated(task)
+                if not status:
+                    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,
+                            )
+                        ]
+                    )
+                    tasks.remove(task)
+                    continue
             if task.sampling_params.bad_words is not None:
                 task.sampling_params.update_from_tokenizer(self.data_processor.tokenizer)
 
@@ -1020,7 +1042,6 @@ def _exit_sub_services(self):
             except Exception as e:
                 print(f"Error extracting sub services: {e}")
 
-
         for worker_queue in self.engine_worker_queue_server:
             worker_queue.cleanup()
         if hasattr(self, "send_response_server") and self.send_response_server is not None:

fastdeploy/engine/expert_service.py

@@ -26,6 +26,7 @@
 
 import numpy as np
 
+from fastdeploy.engine.request import RequestOutput
 from fastdeploy.engine.resource_manager import ResourceManager
 from fastdeploy.inter_communicator import EngineWorkerQueue
 from fastdeploy.metrics.metrics import main_process_metrics
@@ -34,6 +35,7 @@
 from fastdeploy.splitwise.splitwise_connector import SplitwiseConnector
 from fastdeploy.utils import EngineError, console_logger, envs, get_logger, llm_logger
 
+
 class ExpertService:
     """
     Engine class responsible for managing the Large Language Model (LLM) operations.
@@ -146,7 +148,7 @@ def start(
 
         # Start TokenProcessor thread
         os.environ["INFERENCE_MSG_QUEUE_ID"] = str(local_data_parallel_id + int(self.cfg.engine_worker_queue_port))
-
+        self.enable_decode_cache_task = envs.FD_ENABLE_CACHE_TASK
         self.token_processor.run()
 
         self.cfg.init_cache_info()
@@ -262,11 +264,15 @@ def receiver_loop():
                                         if self.resource_manager.is_resource_sufficient(task.prompt_token_ids_len):
                                             self.insert_tasks([task])
                                         else:
+                                            if not self.enable_decode_cache_task:
+                                                task.error_msg = "Not enough resources"
                                             new_waiting.append(task)
-
                                     if new_waiting:
-                                        self.waiting_requests.extend(new_waiting)
-                                        self.llm_logger.info(f"Added {len(new_waiting)} tasks to waiting queue")
+                                        if not self.enable_decode_cache_task:
+                                            self.split_connector.send_cache_infos(new_waiting, -1)
+                                        else:
+                                            self.waiting_requests.extend(new_waiting)
+                                            self.llm_logger.info(f"Added {len(new_waiting)} tasks to waiting queue")
 
                     else:
                         time.sleep(0.001)
@@ -310,8 +316,24 @@ def insert_tasks(self, tasks, current_id=-1, allocated=False):
         if not isinstance(tasks, list):
             tasks = [tasks]
 
-        for item in tasks:
-            item.schedule_start_time = time.time()
+        for task in tasks:
+            if self.cfg.splitwise_role != "mixed":
+                status, msg = self.split_connector.check_decode_allocated(task)
+                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,
+                            )
+                        ]
+                    )
+                    tasks.remove(task)
+                    continue
+            task.schedule_start_time = time.time()
 
         available_batch = np.sum(self.resource_manager.stop_flags)
         if len(tasks) > available_batch:

fastdeploy/envs.py

@@ -90,6 +90,8 @@
     "FD_ZMQ_CONTROL_CMD_SERVER_PORTS": lambda: os.getenv("FD_ZMQ_CONTROL_CMD_SERVER_PORTS", "8202"),
     # Whether to use PLUGINS.
     "FD_PLUGINS": lambda: None if "FD_PLUGINS" not in os.environ else os.environ["FD_PLUGINS"].split(","),
+    # Whether to enable cache task in decode node
+    "FD_ENABLE_CACHE_TASK": lambda: os.getenv("FD_ENABLE_CACHE_TASK", "1"),
 }