perf(pd): reduce disaggregated serving latency

lightllm/server/detokenization/manager.py

@@ -1,3 +1,4 @@
+import os
 import uvloop
 import asyncio
 import setproctitle
@@ -20,6 +21,7 @@
 from lightllm.utils.shm_port_args import get_shm_port_args
 
 logger = init_logger(__name__)
+DETOKENIZATION_POLL_INTERVAL_S = float(os.getenv("LIGHTLLM_DETOKENIZATION_POLL_INTERVAL_S", "0.002"))
 
 
 class DeTokenizationManager:
@@ -93,7 +95,7 @@ def handle_loop(self):
                     logger.info(f"detokenize batch cost time {cost_time} ms")
 
                 if not exist_need_detoken:
-                    time.sleep(0.002)
+                    time.sleep(DETOKENIZATION_POLL_INTERVAL_S)
 
         except Exception as e:
             logger.exception(f"detoken process has exception {str(e)}")

lightllm/server/router/manager.py

@@ -1,3 +1,4 @@
+import os
 import time
 import uvloop
 import asyncio
@@ -52,6 +53,7 @@ def __init__(self, args: StartArgs):
         self.node_rank = args.node_rank
         self.dp_size = args.dp
         self.schedule_time_interval = args.schedule_time_interval  # 默认30ms 的调度周期
+        self.metric_gauge_step_interval = max(1, int(os.getenv("LIGHTLLM_ROUTER_GAUGE_STEP_INTERVAL", "1")))
         # 兼容多机纯tp的运行模式，这时候 1 // 2 == 0, 需要兼容
         self.dp_size_in_node = max(1, args.dp // self.nnodes)
         self.dp_world_size = self.world_size // self.dp_size
@@ -248,16 +250,17 @@ async def loop_for_fwd(
                     self.metric_client.gauge_set("lightllm_batch_pause_size", self._get_paused_req_num())
                 # pd decode mode need to update token_load more frequently
                 self.req_queue.update_token_load(self.running_batch, force_update=self.is_pd_decode_mode)
-                self.metric_client.gauge_set("lightllm_batch_current_size", len(self.running_batch.reqs))
-                self.metric_client.gauge_set("lightllm_num_running_reqs", len(self.running_batch.reqs))
-                self.metric_client.gauge_set("lightllm_queue_size", self.req_queue.get_wait_req_num())
-                self.metric_client.gauge_set(
-                    "lightllm_batch_current_max_tokens",
-                    int(
-                        sum(self.shared_token_load.get_dynamic_max_load(d_i) for d_i in range(self.dp_size_in_node))
-                        * self.max_total_token_num
-                    ),
-                )
+                if counter_count % self.metric_gauge_step_interval == 0:
+                    self.metric_client.gauge_set("lightllm_batch_current_size", len(self.running_batch.reqs))
+                    self.metric_client.gauge_set("lightllm_num_running_reqs", len(self.running_batch.reqs))
+                    self.metric_client.gauge_set("lightllm_queue_size", self.req_queue.get_wait_req_num())
+                    self.metric_client.gauge_set(
+                        "lightllm_batch_current_max_tokens",
+                        int(
+                            sum(self.shared_token_load.get_dynamic_max_load(d_i) for d_i in range(self.dp_size_in_node))
+                            * self.max_total_token_num
+                        ),
+                    )
             else:
                 self.req_queue.update_token_load(self.running_batch, force_update=True)
                 if counter_count % 300 == 0:
@@ -272,7 +275,14 @@ async def loop_for_fwd(
                             estimated_peak_token_count = self.shared_token_load.get_estimated_peak_token_count(dp_i)
                             logger.debug(f"dp_i {dp_i} estimated_peak_token_count: {estimated_peak_token_count} \n")
 
-            await asyncio.sleep(self._get_schedule_time_interval())
+            schedule_time_interval = self._get_schedule_time_interval()
+            if schedule_time_interval < 0.001:
+                # uvloop 会将亚毫秒级 sleep 向下取整为忙等待。
+                # 此处使用短暂的阻塞式 sleep 更准确，也能避免
+                # router/metrics 循环退化为占满 CPU 的轮询循环。
+                time.sleep(schedule_time_interval)
+            else:
+                await asyncio.sleep(schedule_time_interval)
 
     async def _step(self):
         """

lightllm/server/router/model_infer/mode_backend/pd/base_kv_move_manager.py

@@ -13,6 +13,7 @@
 from lightllm.utils.error_utils import log_exception
 
 logger = init_logger(__name__)
+PD_KV_RETURN_POLL_INTERVAL_S = float(os.getenv("LIGHTLLM_PD_KV_RETURN_POLL_INTERVAL_S", "0.01"))
 
 
 class BaseKVMoveManager:
@@ -77,7 +78,7 @@ def task_ret_upload_loop(self):
                     self.shm_pd_trans_io_buffer.set_ready()
                     break
                 else:
-                    time.sleep(0.01)
+                    time.sleep(PD_KV_RETURN_POLL_INTERVAL_S)
                     ret_objs.extend(self._collect_return_objects())
 
     def _collect_return_objects(self):

lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py

@@ -1,8 +1,9 @@
-import torch
+import os
 import time
 import inspect
 import threading
 import setproctitle
+import torch
 import torch.multiprocessing as mp
 import queue
 import pickle
@@ -24,6 +25,7 @@
 from lightllm.utils.process_check import start_parent_check_thread
 
 logger = init_logger(__name__)
+PD_KV_POLL_INTERVAL_S = float(os.getenv("LIGHTLLM_PD_KV_POLL_INTERVAL_S", "0.001"))
 
 
 def start_decode_trans_process(
@@ -337,7 +339,7 @@ def accept_peer_task_loop(
 
             self._check_tasks_time_out()
             if not notifies_dict:
-                time.sleep(0.001)
+                time.sleep(PD_KV_POLL_INTERVAL_S)
 
     def _check_tasks_time_out(self):
         with self.waiting_dict_lock:

lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py

@@ -1,8 +1,9 @@
-import torch
+import os
 import time
 import inspect
 import threading
 import setproctitle
+import torch
 import torch.multiprocessing as mp
 import queue
 import pickle
@@ -19,6 +20,7 @@
 
 
 logger = init_logger(__name__)
+PD_KV_POLL_INTERVAL_S = float(os.getenv("LIGHTLLM_PD_KV_POLL_INTERVAL_S", "0.001"))
 
 
 def start_prefill_trans_process(
@@ -283,7 +285,7 @@ def accept_decode_write_task_loop(self):
             self._check_tasks_time_out()
 
             if not notifies_dict:
-                time.sleep(0.001)
+                time.sleep(PD_KV_POLL_INTERVAL_S)
         return
 
     def _check_tasks_time_out(self):
@@ -328,7 +330,7 @@ def update_task_status_loop(
     ):
         while True:
             if len(self.waiting_dict) == 0:
-                time.sleep(0.001)
+                time.sleep(PD_KV_POLL_INTERVAL_S)
                 continue
 
             with self.waiting_dict_lock:
@@ -372,7 +374,7 @@ def update_task_status_loop(
                         trans_task.error_info = "time out in update_task_status_loop"
                         self.failed_queue.put(trans_task)
 
-            time.sleep(0.001)
+            time.sleep(PD_KV_POLL_INTERVAL_S)
 
     @log_exception
     def success_loop(self):