update healh check

Author: shihaobai

lightllm/server/router/manager.py

@@ -29,6 +29,7 @@
 from lightllm.utils.log_utils import init_logger, log_time_ready
 from lightllm.server.router.token_load import TokenLoad
 from lightllm.server.metrics.manager import MetricClient
+from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt
 from lightllm.common.basemodel.infer_lock import g_router_lock
 from lightllm.common.mem_manager import ReadOnlyStaticsMemoryManager
 from lightllm.utils.graceful_utils import graceful_registry
@@ -71,6 +72,11 @@ def __init__(self, args, router_port, detokenization_port, metric_port):
             self.shared_token_load.set_logical_max_load(0.0, dp_index)
             self.shared_token_load.set_dynamic_max_load(0.0, dp_index)
 
+        # The timemark of the latest inference(prefill/decode) which is used to check the health status of the system.
+        # If the timemark is not updated for a pre-set time, a prob request will be sent to the backend.
+        self.latest_req_inference_timemark = SharedInt(f"{get_unique_server_name()}_latest_req_inference_timemark")
+        self.latest_req_inference_timemark.set_value(int(time.time()))
+
         self.pause_strategy = Fcfs()
         self.running_batch: Batch = None
         self.eos_id = args.eos_id
@@ -389,6 +395,8 @@ async def _prefill_batch(self, batch: Batch):
         self.metric_client.histogram_observe(
             "lightllm_batch_inference_duration_bucket", time.time() - start_time, "prefill"
         )
+        # update inference timemark
+        self.latest_req_inference_timemark.set_value(int(time.time()))
         return
 
     async def _decode_batch(self, batch: Batch):
@@ -404,6 +412,8 @@ async def _decode_batch(self, batch: Batch):
         self.metric_client.histogram_observe(
             "lightllm_batch_inference_duration_bucket", time.time() - start_time, "decode"
         )
+        # update inference timemark
+        self.latest_req_inference_timemark.set_value(int(time.time()))
         return
 
     async def _pause_reqs(self, pasue_reqs):

lightllm/utils/health_check.py

@@ -1,13 +1,16 @@
 import os
+import time
 import asyncio
 import numpy as np
 from dataclasses import dataclass
 from lightllm.server.core.objs import SamplingParams
 from lightllm.server.multimodal_params import MultimodalParams
 from lightllm.server.httpserver.manager import HttpServerManager
+from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt
 from fastapi import Request
 from lightllm.server.req_id_generator import ReqIDGenerator
 from lightllm.utils.log_utils import init_logger
+from lightllm.utils.envs_utils import get_unique_server_name
 
 logger = init_logger(__name__)
 
@@ -24,6 +27,7 @@ class HealthObj:
     _failure_threshold: int = int(os.getenv("HEALTH_FAILURE_THRESHOLD", 3))
     timeout: int = int(os.getenv("HEALTH_TIMEOUT", 100))
     dynamic_timeout: int = int(os.getenv("HEALTH_TIMEOUT", 100))
+    latest_req_inferece_timemark = SharedInt(f"{get_unique_server_name()}_latest_req_inference_timemark")
 
     def begin_check(self):
         self._is_health_checking = True
@@ -48,12 +52,17 @@ def is_health(self):
     def is_checking(self):
         return self._is_health_checking
 
+    def has_latest_inference(self):
+        last_timemark = self.latest_req_inferece_timemark.get_value()
+        time_diff = time.time() - last_timemark
+        return time_diff < self.timeout
+
 
 health_obj = HealthObj()
 
 
 async def health_check(args, httpserver_manager: HttpServerManager, request: Request):
-    if health_obj.is_checking():
+    if health_obj.is_checking() or health_obj.has_latest_inference():
         return health_obj.is_health()
     health_obj.begin_check()
     try: