Merge branch 'main' into reduce_fix

Author: shihaobai

lightllm/common/basemodel/basemodel.py

@@ -166,7 +166,14 @@ def _check_mem_size(self):
         return
 
     def _init_req_manager(self):
-        self.req_manager = ReqManager(self.max_req_num, self.max_seq_length, self.mem_manager)
+        create_max_seq_len = 0
+
+        if self.batch_max_tokens is not None:
+            create_max_seq_len = max(create_max_seq_len, self.batch_max_tokens)
+        if self.max_seq_length is not None:
+            create_max_seq_len = max(create_max_seq_len, self.max_seq_length)
+
+        self.req_manager = ReqManager(self.max_req_num, create_max_seq_len, self.mem_manager)
         return
 
     def _init_infer_layer(self):

lightllm/server/api_cli.py

@@ -13,6 +13,12 @@ def make_argument_parser() -> argparse.ArgumentParser:
     )
     parser.add_argument("--host", type=str, default="127.0.0.1")
     parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument(
+        "--zmq_mode",
+        type=str,
+        default="ipc:///tmp/",
+        help="use socket mode or ipc mode, only can be set in ['tcp://', 'ipc:///tmp/']",
+    )
 
     parser.add_argument(
         "--pd_master_ip",

lightllm/server/api_models.py

@@ -1,6 +1,6 @@
 import time
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, field_validator
 from typing import Dict, List, Optional, Union, Literal
 import uuid
 
@@ -9,7 +9,7 @@ class ChatCompletionRequest(BaseModel):
     # The openai api native parameters
     model: str
     messages: List[Dict[str, str]]
-    function_call: Optional[str] = 'none'
+    function_call: Optional[str] = "none"
     temperature: Optional[float] = 1
     top_p: Optional[float] = 1.0
     n: Optional[int] = 1
@@ -52,6 +52,10 @@ class ChatCompletionResponse(BaseModel):
     choices: List[ChatCompletionResponseChoice]
     usage: UsageInfo
 
+    @field_validator("id", mode="before")
+    def ensure_id_is_str(cls, v):
+        return str(v)
+
 
 class DeltaMessage(BaseModel):
     role: Optional[str] = None
@@ -70,3 +74,7 @@ class ChatCompletionStreamResponse(BaseModel):
     created: int = Field(default_factory=lambda: int(time.time()))
     model: str
     choices: List[ChatCompletionStreamResponseChoice]
+
+    @field_validator("id", mode="before")
+    def ensure_id_is_str(cls, v):
+        return str(v)

lightllm/server/api_start.py

@@ -4,7 +4,7 @@
 from lightllm.server import TokenLoad
 from .api_lightllm import lightllm_generate, lightllm_generate_stream
 from .api_tgi import tgi_generate_impl, tgi_generate_stream_impl
-from lightllm.utils.net_utils import alloc_can_use_network_port
+from lightllm.utils.net_utils import alloc_can_use_network_port, PortLocker
 from lightllm.utils.start_utils import start_submodule_processes
 from .metrics.manager import start_metric_manager
 from .embed_cache.manager import start_cache_manager
@@ -27,6 +27,15 @@ def normal_or_p_d_start(g_objs):
     if args.run_mode not in ["normal", "prefill", "decode"]:
         return
 
+    assert args.zmq_mode in ["tcp://", "ipc:///tmp/"]
+
+    # 确保单机上多实列不冲突
+    if args.zmq_mode == "ipc:///tmp/":
+        zmq_mode = f"{args.zmq_mode}_{str(args.nccl_port)}_"
+        args.zmq_mode = None  # args 的参数不能直接设置，只能先设置None，再设置才能成功
+        args.zmq_mode = zmq_mode
+        logger.info(f"zmq mode head: {args.zmq_mode}")
+
     if args.use_tgi_api:
         g_objs.g_generate_func = tgi_generate_impl
         g_objs.g_generate_stream_func = tgi_generate_stream_impl
@@ -117,9 +126,18 @@ def normal_or_p_d_start(g_objs):
         assert args.data_type in ["fp16", "float16", "bf16", "bfloat16", "fp32", "float32"]
 
     already_uesd_ports = args.visual_nccl_ports + [args.nccl_port, args.port]
+    if args.run_mode == "decode":
+        already_uesd_ports = args.visual_nccl_ports + [args.nccl_port, args.port, args.pd_decode_rpyc_port]
+
+    # 提前锁定端口，防止在单个机器上启动多个实列的时候，要到模型启动的时候才能
+    # 捕获到端口设置冲突的问题
+    ports_locker = PortLocker(already_uesd_ports)
+    ports_locker.lock_port()
+
     can_use_ports = alloc_can_use_network_port(
         num=6 + args.tp + args.tp + args.visual_dp * args.visual_tp, used_nccl_ports=already_uesd_ports
     )
+    logger.info(f"alloced ports: {can_use_ports}")
     router_port, detokenization_port, httpserver_port, visual_port, cache_port, metric_port = can_use_ports[0:6]
     model_rpc_ports = can_use_ports[6 : 6 + args.tp]
     can_use_ports = can_use_ports[6 + args.tp :]
@@ -144,6 +162,8 @@ def normal_or_p_d_start(g_objs):
 
     logger.info(f"all start args:{args}")
 
+    ports_locker.release_port()
+
     if args.enable_multimodal:
         start_submodule_processes(
             start_funcs=[

lightllm/server/detokenization/manager.py

@@ -33,10 +33,10 @@ def __init__(
         self.args = args
         context = zmq.asyncio.Context(2)
         self.recv_from_router = context.socket(zmq.PULL)
-        self.recv_from_router.bind(f"tcp://127.0.0.1:{detokenization_port}")
+        self.recv_from_router.bind(f"{args.zmq_mode}127.0.0.1:{detokenization_port}")
 
         self.send_to_httpserver = context.socket(zmq.PUSH)
-        self.send_to_httpserver.connect(f"tcp://127.0.0.1:{httpserver_port}")
+        self.send_to_httpserver.connect(f"{args.zmq_mode}127.0.0.1:{httpserver_port}")
 
         self.tokenizer = get_tokenizer(model_weightdir, tokenizor_mode, trust_remote_code=trust_remote_code)
         self.all_special_ids = set(self.tokenizer.all_special_ids)

lightllm/server/httpserver/manager.py

@@ -45,16 +45,16 @@ def __init__(
         self.args = args
         context = zmq.asyncio.Context(2)
         self.send_to_router = context.socket(zmq.PUSH)
-        self.send_to_router.connect(f"tcp://127.0.0.1:{router_port}")
+        self.send_to_router.connect(f"{args.zmq_mode}127.0.0.1:{router_port}")
 
         self.enable_multimodal = enable_multimodal
         if self.enable_multimodal:
             self.cache_client = rpyc.connect("localhost", cache_port)
             self.send_to_visual = context.socket(zmq.PUSH)
-            self.send_to_visual.connect(f"tcp://127.0.0.1:{visual_port}")
+            self.send_to_visual.connect(f"{args.zmq_mode}127.0.0.1:{visual_port}")
 
         self.recv_from_detokenization = context.socket(zmq.PULL)
-        self.recv_from_detokenization.bind(f"tcp://127.0.0.1:{httpserver_port}")
+        self.recv_from_detokenization.bind(f"{args.zmq_mode}127.0.0.1:{httpserver_port}")
 
         self.tokenizer = get_tokenizer(args.model_dir, args.tokenizer_mode, trust_remote_code=args.trust_remote_code)
 
@@ -67,6 +67,7 @@ def __init__(
         assert self.pd_mode in [NodeRole.P, NodeRole.D, NodeRole.NORMAL]
         self.id_gen = ReqIDGenerator()
         self.first_time_costs = MovingAverage()
+        self.per_token_costs = MovingAverage()
         # 有的模型的vocab size 读取tokenizer和config.json中不一致
         self.vocab_size = max(get_vocab_size(args.model_dir), self.tokenizer.vocab_size)
 
@@ -340,6 +341,7 @@ async def _wait_to_token_package(
                             pass
                         total_cost_time_ms = (time.time() - start_time) * 1000
                         mean_per_token_cost_time_ms = (total_cost_time_ms - first_token_cost_ms) / out_token_counter
+                        self.per_token_costs.add(mean_per_token_cost_time_ms)
                         x_request_id = request.headers.get("X-Request-Id", "")
                         x_session_id = request.headers.get("X-Session-Id", "")
                         prompt_cache_len = metadata.pop("prompt_cache_len", 0)
@@ -441,6 +443,7 @@ async def timer_to_pd_master(self):
                         await asyncio.sleep(3)
                         if log_count % 5 == 0:
                             logger.info(f"mean first cost: {self.first_time_costs.average()} ms")
+                            logger.info(f"mean per token cost: {self.per_token_costs.average()} ms")
 
             except Exception as e:
                 logger.error("connetion to pd_master has error")

lightllm/server/httpserver_for_pd_master/manager.py

@@ -41,7 +41,9 @@ def __init__(
         self.id_to_event: Dict[int, asyncio.Event] = {}
         self.session = None
         self.first_time_costs = MovingAverage()
-        self.create_session_costs = MovingAverage()
+        self.prefill_create_session_costs = MovingAverage()
+        self.decode_create_session_costs = MovingAverage()
+        self.per_token_costs = MovingAverage()
         return
 
     async def register_pd(self, pd_info_json):
@@ -181,7 +183,7 @@ async def fetch_stream(
             req = await self._to_req_info(prompt, sampling_params, multimodal_params)
             create_start_time = time.time()
             async with self.session.post(p_node.to_llm_url(), json=req) as response:
-                self.create_session_costs.add((time.time() - create_start_time) * 1000)
+                self.prefill_create_session_costs.add((time.time() - create_start_time) * 1000)
                 if response.status == 200:
                     async for line in response.content:
                         line = line.decode("utf-8").strip()
@@ -217,7 +219,9 @@ async def fetch_stream(
             sampling_params.suggested_dp_index = event.upkv_status.dp_index
 
             req = await self._to_req_info(prompt_ids, sampling_params, multimodal_params)
+            create_start_time = time.time()
             async with self.session.post(d_node.to_llm_url(), json=req) as response:
+                self.decode_create_session_costs.add((time.time() - create_start_time) * 1000)
                 if response.status == 200:
                     async for line in response.content:
                         line = line.decode("utf-8").strip()
@@ -269,6 +273,7 @@ async def _wait_to_token_package(
 
         total_cost_time_ms = (time.time() - start_time) * 1000
         mean_per_token_cost_time_ms = (total_cost_time_ms - first_token_cost_ms) / out_token_counter
+        self.per_token_costs.add(mean_per_token_cost_time_ms)
         x_request_id = request.headers.get("X-Request-Id", "")
         x_session_id = request.headers.get("X-Session-Id", "")
         prompt_cache_len = metadata.pop("prompt_cache_len", 0)
@@ -312,5 +317,7 @@ async def handle_loop(self):
             # 可以做一个定时任务
             await asyncio.sleep(20)
             logger.info(f"mean first cost: {self.first_time_costs.average()} ms")
-            logger.info(f"create_session_costs: {self.create_session_costs.average()} ms")
+            logger.info(f"prefill mean create_session_costs: {self.prefill_create_session_costs.average()} ms")
+            logger.info(f"decode mean create_session_costs: {self.decode_create_session_costs.average()} ms")
+            logger.info(f"mean per token cost: {self.per_token_costs.average()} ms")
         return

lightllm/server/router/manager.py

@@ -66,10 +66,10 @@ def __init__(self, args, router_port, detokenization_port, model_rpc_ports, metr
 
         context = zmq.asyncio.Context(2)
         self.recv_from_httpserver = context.socket(zmq.PULL)
-        self.recv_from_httpserver.bind(f"tcp://127.0.0.1:{router_port}")
+        self.recv_from_httpserver.bind(f"{args.zmq_mode}127.0.0.1:{router_port}")
 
         self.send_to_detokenization = context.socket(zmq.PUSH)
-        self.send_to_detokenization.connect(f"tcp://127.0.0.1:{detokenization_port}")
+        self.send_to_detokenization.connect(f"{args.zmq_mode}127.0.0.1:{detokenization_port}")
         self.model_rpc_ports = model_rpc_ports
 
         self.is_splitfuse_mode = args.splitfuse_mode
@@ -283,14 +283,15 @@ async def _step(self):
                 self.running_batch = new_batch
                 await self._prefill_batch(self.running_batch)
                 self._filter_runing_batch()
-                self.has_wait_tokens = 0
+                self.has_wait_tokens = self.max_wait_tokens
             return
 
         # 有运行请求，但是已经到了可以调度新的请求合并推理的时机
         if self.has_wait_tokens >= self.max_wait_tokens:
             new_mini_batch = self.req_queue.generate_new_batch(self.running_batch)
             self.has_wait_tokens = 0
             if new_mini_batch is not None:
+                self.has_wait_tokens = self.max_wait_tokens
                 self.stats_tool.count_prompt_tokens(new_mini_batch)
                 await self._prefill_batch(new_mini_batch)
                 if not new_mini_batch.is_clear():
@@ -426,6 +427,9 @@ def _update_init_status_to_batch(self, batch: Batch, req_to_req_status):
 
     def _update_out_status_to_batch(self, batch: Batch, req_to_out_status):
         new_batch_decode_need_tokens = [0 for _ in range(self.dp_size)]  # 只有在 splitfuse 模式下有意义
+
+        start_time = 0
+        # extral_info 字段如果推理后端输入时间标记, 则用来评估序列化所占用的时间, 主要用于调试时使用
         for req_id, (
             req_status,
             cur_kv_len,
@@ -434,6 +438,8 @@ def _update_out_status_to_batch(self, batch: Batch, req_to_out_status):
             finish_status_value,
             extral_info,
         ) in req_to_out_status.items():
+            if extral_info is not None:
+                start_time = max(start_time, extral_info)
             req: Req = batch.id_to_reqs[req_id]
             req.req_status = req_status
             req.cur_kv_len = cur_kv_len
@@ -446,6 +452,9 @@ def _update_out_status_to_batch(self, batch: Batch, req_to_out_status):
             new_batch_decode_need_tokens[req_dp_index] += req.get_decode_need_tokens()
 
         batch.batch_decode_need_tokens = new_batch_decode_need_tokens
+        rpyc_cost_time = (time.time() - start_time) * 1000
+        if 8 <= rpyc_cost_time <= 1000:
+            logger.warning(f"rpyc use too much time {rpyc_cost_time} ms, batch_size {len(req_to_out_status)}")
         return
 
     def _can_decode(self, batch: Batch):

lightllm/server/router/model_infer/mode_backend/continues_batch/decode_node_impl/decode_impl.py

@@ -1,3 +1,4 @@
+import os
 import torch
 import torch.multiprocessing as mp
 import torch.distributed as dist
@@ -30,7 +31,13 @@ def init_custom(self):
         self.lock_nccl_group = dist.new_group(backend="gloo")
         from .decode_infer_rpyc import PDDecodeInferRpcServer
 
-        t = ThreadedServer(PDDecodeInferRpcServer(self), port=self.pd_rpyc_port, protocol_config={"allow_pickle": True})
+        socket_path = f"/tmp/decode_node_infer_rpyc_{self.pd_rpyc_port}"
+        if os.path.exists(socket_path):
+            os.remove(socket_path)
+
+        t = ThreadedServer(
+            PDDecodeInferRpcServer(self), socket_path=socket_path, protocol_config={"allow_pickle": True}
+        )
         threading.Thread(target=lambda: t.start(), daemon=True).start()
         return
 

lightllm/server/router/model_infer/mode_backend/continues_batch/decode_node_impl/decode_kv_move_manager.py

@@ -87,7 +87,10 @@ def __init__(self, args, info_queue: mp.Queue, mem_queues: List[mp.Queue]):
         self.infer_rpyc_objs: List[PDDecodeInferRpcServer] = []
         self.node_id_to_trans_obj: Dict[str, TransProcessObj] = {}
         for port in self.args.pd_tp_infer_rpyc_ports:
-            con = retry(max_attempts=20, wait_time=2)(rpyc.connect)("localhost", port, config={"allow_pickle": True})
+            socket_path = f"/tmp/decode_node_infer_rpyc_{port}"
+            from rpyc.utils.factory import unix_connect
+
+            con = retry(max_attempts=20, wait_time=2)(unix_connect)(socket_path, config={"allow_pickle": True})
             self.infer_rpyc_objs.append(con.root)
             logger.info(f"rpyc connect to port: {port} ok")