rpyc and zmq use unix socket. (#653)

Author: hiworldwzj

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_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)
 

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():

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")
 

lightllm/server/router/model_infer/mode_backend/continues_batch/prefill_node_impl/prefill_impl.py

@@ -1,3 +1,4 @@
+import os
 import threading
 import torch
 import torch.multiprocessing as mp
@@ -29,8 +30,12 @@ def init_custom(self):
         self.lock_nccl_group = dist.new_group(backend="gloo")
         from .prefill_infer_rpyc import PDPrefillInferRpcServer
 
+        socket_path = f"/tmp/prefill_node_infer_rpyc_{self.pd_rpyc_port}"
+        if os.path.exists(socket_path):
+            os.remove(socket_path)
+
         t = ThreadedServer(
-            PDPrefillInferRpcServer(self), port=self.pd_rpyc_port, protocol_config={"allow_pickle": True}
+            PDPrefillInferRpcServer(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/prefill_node_impl/prefill_kv_move_manager.py

@@ -99,7 +99,10 @@ def __init__(self, args, info_queue: mp.Queue, mem_queues: List[mp.Queue]):
         self.infer_rpyc_objs: List[PDPrefillInferRpcServer] = []
         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/prefill_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 infer rpyc port: {port} ok")
         self.host_ip = get_hostname_ip()

lightllm/server/router/model_infer/model_rpc.py

@@ -1,6 +1,6 @@
 import asyncio
 import rpyc
-import torch
+import tempfile
 import torch.multiprocessing as mp
 from datetime import timedelta
 from typing import Dict, List, Tuple
@@ -243,7 +243,7 @@ async def get_max_total_token_num(self):
             return ans
 
 
-def _init_env(args, port, info_queue, mem_queue, router_lock):
+def _init_env(args, socket_path, info_queue, mem_queue, router_lock, success_event: mp.Event):
     import lightllm.utils.rpyc_fix_utils as _
 
     # 注册graceful 退出的处理
@@ -259,7 +259,10 @@ def _init_env(args, port, info_queue, mem_queue, router_lock):
 
     from rpyc.utils.server import ThreadedServer
 
-    t = ThreadedServer(ModelRpcServer(args, info_queue, mem_queue), port=port, protocol_config={"allow_pickle": True})
+    t = ThreadedServer(
+        ModelRpcServer(args, info_queue, mem_queue), socket_path=socket_path, protocol_config={"allow_pickle": True}
+    )
+    success_event.set()
     t.start()
     return
 
@@ -271,13 +274,18 @@ async def start_model_process(args, port, world_size, info_queue: mp.Queue, mem_
     if world_size == 1:
         return ModelRpcClient(ModelRpcServer(args, info_queue, mem_queue), world_size)
 
-    proc = mp.Process(target=_init_env, args=(args, port, info_queue, mem_queue, router_lock))
+    socket_path = tempfile.mktemp()
+    success_event = mp.Event()
+    proc = mp.Process(target=_init_env, args=(args, socket_path, info_queue, mem_queue, router_lock, success_event))
     proc.start()
-    await asyncio.sleep(2)
+    success_event.wait(timeout=40)
+
     repeat_count = 0
     while repeat_count < 20:
         try:
-            con = rpyc.connect("localhost", port, config={"allow_pickle": True})
+            from rpyc.utils.factory import unix_connect
+
+            con = unix_connect(socket_path, config={"allow_pickle": True})
             break
         except BaseException:
             await asyncio.sleep(1)