add vit mananger for vit-llm disaggr

Author: shihaobai

lightllm/server/api_cli.py

@@ -506,6 +506,17 @@ def make_argument_parser() -> argparse.ArgumentParser:
         default=0.03,
         help="""The interval of the schedule time, default is 30ms.""",
     )
+    parser.add_argument(
+        "--enable_remote_vit",
+        action="store_true",
+        help="Whether to enable remote vit for multimodal service.",
+    )
+    parser.add_argument(
+        "--remote_vit_port",
+        type=int,
+        default=12346,
+        help="The port number for the remote vit service.",
+    )
     # redis for vit llm disaggregation
     parser.add_argument(
         "--redis_port",

lightllm/server/api_server.py

@@ -13,7 +13,5 @@
         config_server_start(args)
     elif args.run_mode == "visual_only":
         visual_only_start(args)
-    elif args.run_mode == "llm_only":
-        llm_only_start(args)
     else:
         normal_or_p_d_start(args)

lightllm/server/api_start.py

@@ -5,7 +5,7 @@
 import subprocess
 import signal
 from lightllm.utils.net_utils import alloc_can_use_network_port, PortLocker
-from lightllm.utils.start_utils import process_manager, kill_recursive
+from lightllm.utils.start_utils import process_manager, kill_recursive, is_multimodal_mode
 from .metrics.manager import start_metric_manager
 from .embed_cache.manager import start_cache_manager
 from lightllm.utils.log_utils import init_logger
@@ -157,11 +157,13 @@ def check_and_set_args(args):
         assert args.mtp_draft_model_dir is None
         assert args.mtp_step == 0
 
+    args.enable_multimodal = is_multimodal_mode(args)
     # visual_only模式下才需要设置visual_embed_path
     if args.visual_embed_path is not None:
         assert (
             args.run_mode == "visual_only" or args.run_mode == "llm_only"
         ), "only visual_only or llm_only mode need visual_embed_path"
+
     # 检查GPU数量是否足够
     if args.visual_gpu_ids is None:
         args.visual_gpu_ids = list(range(args.visual_dp * args.visual_tp))
@@ -174,13 +176,11 @@ def check_and_set_args(args):
         args.visual_gpu_ids = args.visual_gpu_ids[:total_required_gpus]
 
     # 检查visual_nccl_port数量是否足够
-    if len(args.visual_nccl_ports) < args.visual_dp:
+    if args.visual_nccl_ports is not None and len(args.visual_nccl_ports) < args.visual_dp:
         raise ValueError(
             f"Not enough visual_nccl_ports specified. You need at least {args.visual_dp}, "
             f"but got ({len(args.visual_nccl_ports)})."
         )
-    else:
-        args.visual_nccl_ports = args.visual_nccl_ports[: args.visual_dp]
 
     if args.visual_dp <= 0:
         raise ValueError("visual_dp must be a positive integer.")
@@ -287,7 +287,6 @@ def normal_or_p_d_start(args):
     logger.info(f"all start args:{args}")
 
     ports_locker.release_port()
-
     if args.enable_multimodal:
         from .visualserver.manager import start_visual_process
 
@@ -381,105 +380,6 @@ def normal_or_p_d_start(args):
     return
 
 
-def llm_only_start(args):
-
-    check_and_set_args(args)
-    already_uesd_ports = [args.nccl_port, args.port]
-
-    # 提前锁定端口，防止在单个机器上启动多个实列的时候，要到模型启动的时候才能
-    # 捕获到端口设置冲突的问题
-    ports_locker = PortLocker(already_uesd_ports)
-    ports_locker.lock_port()
-
-    node_world_size = args.tp // args.nnodes
-    can_use_ports = alloc_can_use_network_port(num=4 + node_world_size, used_nccl_ports=already_uesd_ports)
-    logger.info(f"alloced ports: {can_use_ports}")
-    (
-        router_port,
-        detokenization_port,
-        detokenization_pub_port,
-        metric_port,
-    ) = can_use_ports[0:4]
-    can_use_ports = can_use_ports[4:]
-
-    # 将申请好的端口放入args参数中
-    args.router_port = router_port
-    args.detokenization_port = detokenization_port
-    args.detokenization_pub_port = detokenization_pub_port
-    args.metric_port = metric_port
-
-    # 申请在 p d 分离模式下，会用的端口
-    args.pd_node_infer_rpyc_ports = can_use_ports[0:node_world_size]
-    # p d 分离模式下用于标识节点的id
-    args.pd_node_id = uuid.uuid4().int
-    # p 节点用来建立torch kv 传输分布组的可用端口范围
-    args.pd_p_allowed_port_min = 20000
-    args.pd_p_allowed_port_max = 30000
-
-    # p d 分离模式下，decode节点的调度间隙是0
-    if args.run_mode == "decode":
-        args.router_max_wait_tokens = 0
-
-    send_and_receive_node_ip(args)  # 多机用于收发node ip
-    set_env_start_args(args)
-    logger.info(f"all start args:{args}")
-
-    ports_locker.release_port()
-
-    process_manager.start_submodule_processes(
-        start_funcs=[
-            start_metric_manager,
-        ],
-        start_args=[(metric_port, args)],
-    )
-
-    process_manager.start_submodule_processes(
-        start_funcs=[start_router_process, start_detokenization_process],
-        start_args=[
-            (args, router_port, detokenization_port, metric_port),
-            (args, detokenization_port, detokenization_pub_port),
-        ],
-    )
-
-    # 启动 gunicorn
-    command = [
-        "gunicorn",
-        "--workers",
-        f"{args.httpserver_workers}",
-        "--worker-class",
-        "uvicorn.workers.UvicornWorker",
-        "--bind",
-        f"{args.host}:{args.port}",
-        "--log-level",
-        "info",
-        "--access-logfile",
-        "-",
-        "--error-logfile",
-        "-",
-        "lightllm.server.api_http:app",
-        "--timeout",
-        f"{get_lightllm_gunicorn_time_out_seconds()}",
-        "--keep-alive",
-        f"{get_lightllm_gunicorn_keep_alive()}",
-    ]
-
-    # 启动子进程
-    http_server_process = subprocess.Popen(command)
-
-    if "s3://" in args.model_dir:
-        from lightllm.utils.petrel_helper import s3_model_clear
-
-        s3_model_clear(args.model_dir)
-
-    if args.health_monitor:
-        from lightllm.server.health_monitor.manager import start_health_check_process
-
-        process_manager.start_submodule_processes(start_funcs=[start_health_check_process], start_args=[(args,)])
-    setup_signal_handlers(http_server_process, process_manager)
-    http_server_process.wait()
-    return
-
-
 def pd_master_start(args):
     set_unique_server_name(args)
     if args.run_mode != "pd_master":

lightllm/server/core/objs/io_objs/group_req.py

@@ -4,12 +4,6 @@
 from ..req import Req
 
 
-@dataclass
-class VisualOnlyReqIndexes:
-    group_req_id: int
-    multimodal_params: MultimodalParams
-
-
 @dataclass
 class GroupReqIndexes:
     group_req_id: int

lightllm/server/httpserver/manager.py

@@ -81,10 +81,12 @@ def __init__(
                 )
 
         self.enable_multimodal = enable_multimodal
-        if self.enable_multimodal and self.args.run_mode != "llm_only":
+        if self.enable_multimodal:
             self.cache_client = rpyc.connect("localhost", cache_port, config={"allow_pickle": True})
-            self.send_to_visual = context.socket(zmq.PUSH)
-            self.send_to_visual.connect(f"{args.zmq_mode}127.0.0.1:{visual_port}")
+            # 初始化VIT连接管理器
+            from .vit_loop import VITConnectionManager
+
+            self.vit_manager = VITConnectionManager(args, context, visual_port)
 
         self.token_id_range_start = 100000000
         self.token_id_range_end = 2 ** 63 - 1
@@ -406,10 +408,7 @@ async def _encode(
                 ), "too many multimodal items!"
                 if multimodal_params.audios:
                     assert self.args.enable_multimodal_audio, "audio multimodal not enabled"
-                if self.args.run_mode == "llm_only":
-                    await self._get_image_embedding_from_afs(multimodal_params, sampling_params)
-                else:
-                    await self._alloc_multimodal_resources(multimodal_params, sampling_params)
+                await self._alloc_multimodal_resources(multimodal_params, sampling_params)
                 prompt_ids = self.tokenizer.encode(
                     prompt, multimodal_params, add_special_tokens=sampling_params.add_special_tokens
                 )
@@ -483,9 +482,9 @@ async def transfer_to_next_module(
         group_req_objs: Optional[GroupReqObjs] = None,
     ):
 
-        if self.pd_mode == NodeRole.P:
-            if self.enable_multimodal and self.args.run_mode != "llm_only":
-                self.send_to_visual.send_pyobj(
+        if self.pd_mode.is_P_or_NORMAL():
+            if self.enable_multimodal:
+                await self.vit_manager.send_to_vit(
                     group_req_objs.to_group_req_index(),
                     protocol=pickle.HIGHEST_PROTOCOL,
                 )
@@ -504,19 +503,6 @@ async def transfer_to_next_module(
             )
             return
 
-        if self.pd_mode == NodeRole.NORMAL or self.pd_mode == NodeRole.LLM_ONLY:
-            if self.enable_multimodal and self.args.run_mode != "llm_only":
-                self.send_to_visual.send_pyobj(
-                    group_req_objs.to_group_req_index(),
-                    protocol=pickle.HIGHEST_PROTOCOL,
-                )
-            else:
-                self.send_to_router.send_pyobj(
-                    group_req_objs.to_group_req_index(),
-                    protocol=pickle.HIGHEST_PROTOCOL,
-                )
-            return
-
         assert False, "dead code path"
         return
 

lightllm/server/httpserver/vit_loop.py

@@ -0,0 +1,136 @@
+import asyncio
+import zmq
+import zmq.asyncio
+import time
+import pickle
+from typing import Dict, List, Optional, Any
+from lightllm.utils.log_utils import init_logger
+import httpx
+import base64
+from dataclasses import dataclass
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class VIT_Obj:
+    node_id: int
+    host_ip_port: str
+
+    def to_log_str(self):
+        return f"VIT host_ip_port: {self.host_ip_port} node_id: {self.node_id}"
+
+
+class VITConnectionManager:
+    """VIT连接管理器"""
+
+    def __init__(self, args, context, local_visual_port: int):
+        self.args = args
+        self.context = context
+        self.local_visual_port = local_visual_port
+
+        self.send_to_visual = None
+        self.remote_vit_instances = []
+        self.current_vit_index = 0
+        self.remote_vit = args.enable_remote_vit
+        self.remote_vit_port = args.remote_vit_port
+
+        self._setup_vit_connections()
+
+    def _setup_vit_connections(self):
+        """
+        设置VIT连接，支持本地和远程VIT实例
+        支持多种连接模式：
+        1. 本地VIT实例 (默认)
+        2. 远程单个VIT实例
+        3. 远程多个VIT实例 (负载均衡)
+        """
+        if self.remote_vit:
+            # 远程VIT实例模式
+            self._setup_remote_vit_connections()
+        else:
+            self._setup_local_vit_connection()
+
+    def _setup_local_vit_connection(self):
+        self.send_to_visual = self.context.socket(zmq.PUSH)
+        self.send_to_visual.connect(f"{self.args.zmq_mode}127.0.0.1:{self.local_visual_port}")
+        logger.info(f"Connected to local VIT instance at {self.args.zmq_mode}127.0.0.1:{self.local_visual_port}")
+
+    def _setup_remote_vit_connections(self):
+        asyncio.create_task(self.vit_handle_loop())
+
+        # wait for remote vit instances
+        while True:
+            if len(self.remote_vit_instances) > 0:
+                break
+            time.sleep(1)
+
+    def _get_vit_instance(self):
+        """
+        获取下一个可用的VIT实例 (轮询负载均衡)
+        """
+        if not self.remote_vit:
+            return self.send_to_visual
+
+        # 简单的轮询负载均衡
+        index = (self.current_vit_index + 1) % len(self.remote_vit_instances)
+        self.current_vit_index = index
+        return self.remote_vit_instances[index]
+
+    async def send_to_vit(self, data, protocol=pickle.HIGHEST_PROTOCOL):
+        """
+        发送数据到VIT实例，支持本地和远程模式
+        """
+        instance = self._get_vit_instance()
+        try:
+            instance.send_pyobj(data, protocol=protocol)
+        except Exception as e:
+            logger.error(f"Failed to send to VIT instance {instance.host_ip_port}: {e}")
+            raise Exception(f"Failed to send to VIT instance {instance.host_ip_port}: {e}")
+
+    async def vit_handle_loop(self):
+        while True:
+            try:
+                id_to_vit_obj = await self._get_vit_objs()
+                logger.info(f"get vit_objs {id_to_vit_obj}")
+                for id, remote_instance in self.remote_vit_instances.items():
+                    if id not in id_to_vit_obj:
+                        try:
+                            remote_instance[id].close()
+                        except:
+                            pass
+                        self.remote_vit_instances.pop(id)
+                        logger.info(f"remote vit {id} closed")
+
+                for id, vit_obj in id_to_vit_obj.items():
+                    if id not in self.remote_vit_instances:
+                        self.remote_vit_instances[id] = self.context.socket(zmq.PUSH)
+                        self.remote_vit_instances[id].connect(
+                            f"tcp://{vit_obj.host_ip_port}:{self.args.remote_vit_port}"
+                        )
+                await asyncio.sleep(30)
+            except Exception as e:
+                logger.exception(str(e))
+                await asyncio.sleep(10)
+
+    async def _get_vit_objs(self) -> Optional[Dict[int, VIT_Obj]]:
+        """
+        get_vit_objs 主要负责从 config_server 获取所有的vit远程服务。
+        """
+        # 使用 config_server 服务来发现所有的 pd_master 节点。
+        uri = f"ws://{self.args.config_server_host}:{self.args.config_server_port}/registered_vit"
+
+        try:
+            async with httpx.AsyncClient() as client:
+                response = await client.get(uri)
+                if response.status_code == 200:
+                    base64data = response.json()["data"]
+                    id_to_vit_obj = pickle.loads(base64.b64decode(base64data))
+                    return id_to_vit_obj
+                else:
+                    logger.error(f"get pd_master_objs error {response.status_code}")
+                    return None
+        except Exception as e:
+            logger.exception(str(e))
+            await asyncio.sleep(10)
+            return None