add global req_id alloc.

Author: hiworldwzj

lightllm/server/config_server/api_http.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, Query
 from threading import Lock
 from typing import Dict
 from fastapi.responses import JSONResponse
@@ -7,13 +7,17 @@
 import base64
 import pickle
 import os
+import requests
 
 logger = init_logger(__name__)
 app = FastAPI()
 
 registered_pd_master_objs:Dict[str, PD_Master_Obj] = {}
 registered_pd_master_obj_lock = Lock()
 
+global_req_id = 0
+global_req_id_lock = Lock()
+
 
 @app.get("/liveness")
 @app.post("/liveness")
@@ -41,7 +45,6 @@ async def websocket_endpoint(websocket: WebSocket):
     logger.info(f"ws connected from IP: {client_ip}, Port: {client_port}")
     registered_pd_master_obj: PD_Master_Obj = pickle.loads(await websocket.receive_bytes())
     logger.info(f"recieved registered_pd_master_obj {registered_pd_master_obj}")
-    
     with registered_pd_master_obj_lock:
         registered_pd_master_objs[registered_pd_master_obj.node_id] = registered_pd_master_obj
 
@@ -64,3 +67,28 @@ async def get_registered_objects():
         serialized_data = pickle.dumps(registered_pd_master_objs)
         base64_encoded = base64.b64encode(serialized_data).decode('utf-8')
         return {"data": base64_encoded}
+
+@app.get("/allocate_global_unique_id_range")
+async def allocate_global_id_range():
+    """
+    Allocate a global ID range for the requesting client without requiring parameters.
+
+    Returns:
+        dict: A dictionary containing the start and end of the allocated ID range.
+
+    Example HTTP client usage:
+    ```python
+    response = requests.get("http://<server_address>/allocate_global_unique_id_range")
+    print(response.json())
+    ```
+    """
+    global global_req_id
+    range_size = 800000
+    with global_req_id_lock:
+        if req_id + range_size > 2**63 - 1:
+            req_id = 0    
+        start_id = req_id
+        req_id += range_size
+        end_id = req_id
+
+    return {"start_id": start_id, "end_id": end_id}

lightllm/server/core/objs/start_args_type.py

@@ -15,6 +15,8 @@ class StartArgs:
     )
     pd_master_ip: str = field(default="127.0.0.1")
     pd_master_port: int = field(default=1212)
+    config_server_host: str = field(default=None)
+    config_server_port: int = field(default=None)
     pd_decode_rpyc_port: int = field(default=42000)
     model_name: str = field(default="default_model_name")
     model_dir: Optional[str] = field(default=None)

lightllm/server/req_id_generator.py

@@ -1,5 +1,9 @@
-import threading
+import time
+import requests
 import numpy as np
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
 
 # 可以支持的最大 beam 参数上限，为了让生成的请求的group_req_id 和 sub_req_id 可以有直接的计算映射关系
 # id 生成器，只会以 MAX_BEST_OF 的间隔生成id 作为 group_req_id, (sub_req_id // MAX_BEST_OF * MAX_BEST_OF) 即可
@@ -12,15 +16,44 @@ class ReqIDGenerator:
     def __init__(self):
         from lightllm.server.core.objs.atomic_lock import AtomicShmLock
         from lightllm.server.core.objs.shm_array import ShmArray
-        from lightllm.utils.envs_utils import get_unique_server_name
-
-        self.current_id = ShmArray(f"{get_unique_server_name()}_req_id_gen", (1,), dtype=np.int64)
+        from lightllm.utils.envs_utils import get_unique_server_name, get_env_start_args
+        
+        self.args = get_env_start_args()
+        self.use_config_server = self.args.config_server_host and self.args.config_server_port
+        self.current_id = ShmArray(f"{get_unique_server_name()}_req_id_gen", (2,), dtype=np.int64)
         self.current_id.create_shm()
         self.current_id.arr[0] = 0
+        self.current_id.arr[1] = 0
         self.lock = AtomicShmLock(f"{get_unique_server_name()}_req_id_gen_lock")
 
+    def _check_and_set_new_id_range(self):
+        need_update_range = self.current_id.arr[0] + MAX_BEST_OF >= self.current_id.arr[1]
+        if need_update_range:
+            if not self.use_config_server:
+                self.current_id.arr[0] = MAX_BEST_OF
+                self.current_id.arr[1] = np.iinfo(np.int64).max
+            else:
+                while True:
+                    try:
+                        url = f"http://{self.args.config_server_host}:{self.args.config_server_port}/allocate_global_unique_id_range"
+                        response = requests.get(url)
+                        if response.status_code == 200:
+                            id_range = response.json()
+                            logger.info(f"get new id range {id_range}")
+                            # 保证id满足倍乘关系
+                            self.current_id.arr[0] = (id_range["start_id"] // MAX_BEST_OF + 1) * MAX_BEST_OF
+                            self.current_id.arr[1] = id_range["end_id"]
+                            assert self.current_id.arr[0] + MAX_BEST_OF < self.current_id.arr[1], f"get id range error {self.current_id.arr[0]} {self.current_id.arr[1]}"
+                            return
+                        else:
+                            raise RuntimeError(f"Failed to fetch ID range from config server: {response.status_code}")
+                    except BaseException as e:
+                        logger.exception(str(e))
+                        time.sleep(3)
+
     def generate_id(self):
         with self.lock:
+            self._check_and_set_new_id_range()
             id = self.current_id.arr[0]
             self.current_id.arr[0] += MAX_BEST_OF
         return id

lightllm/utils/envs_utils.py

@@ -10,7 +10,10 @@
 
 
 def set_unique_server_name(args):
-    os.environ["LIGHTLLM_UNIQUE_SERVICE_NAME_ID"] = str(args.nccl_port) + "_" + str(args.node_rank)
+    if args.run_mode == "pd_master":
+        os.environ["LIGHTLLM_UNIQUE_SERVICE_NAME_ID"] = str(args.port) + "_pd_master"
+    else:
+        os.environ["LIGHTLLM_UNIQUE_SERVICE_NAME_ID"] = str(args.nccl_port) + "_" + str(args.node_rank)
     return
 
 

lightllm/utils/health_check.py

@@ -10,14 +10,12 @@
 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
+from lightllm.utils.envs_utils import get_unique_server_name, get_env_start_args
 
 logger = init_logger(__name__)
 
 
 _g_health_req_id_gen = ReqIDGenerator()
-_g_health_req_id_gen.generate_id()
-
 
 @dataclass
 class HealthObj:
@@ -78,7 +76,13 @@ async def health_check(args, httpserver_manager: HttpServerManager, request: Req
         sampling_params = SamplingParams()
         sampling_params.init(tokenizer=httpserver_manager.tokenizer, **sample_params_dict)
         sampling_params.verify()
-        sampling_params.group_request_id = -_g_health_req_id_gen.generate_id()  # health monitor 的 id 是负的
+        
+        if get_env_start_args().run_mode == "pd_master":
+            # Since the id assigned by pd master needs to be passed to prefill and decode nodes for inference, 
+            # a normal request id is required instead of a negative id.
+            sampling_params.group_request_id = _g_health_req_id_gen.generate_id()
+        else:
+            sampling_params.group_request_id = -_g_health_req_id_gen.generate_id()  # health monitor 的 id 是负的
         multimodal_params_dict = request_dict.get("multimodal_params", {})
         multimodal_params = MultimodalParams(**multimodal_params_dict)
         results_generator = httpserver_manager.generate(