fix vit manager

Author: shihaobai

lightllm/server/api_cli.py

@@ -353,7 +353,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
         "--visual_nccl_ports",
         nargs="+",
         type=int,
-        default=[29500],
+        default=None,
         help="List of NCCL ports to build a distributed environment for Vit, e.g., 29500 29501 29502",
     )
     parser.add_argument(

lightllm/server/api_http.py

@@ -95,21 +95,6 @@ def set_args(self, args):
             )
         elif args.run_mode == "visual":
             self.metric_client = MetricClient(args.metric_port)
-        elif args.run_mode == "llm_only":
-            init_tokenizer(args)  # for openai api
-            SamplingParams.load_generation_cfg(args.model_dir)
-            self.metric_client = MetricClient(args.metric_port)
-            self.httpserver_manager = HttpServerManager(
-                args,
-                router_port=args.router_port,
-                cache_port=None,
-                detokenization_pub_port=args.detokenization_pub_port,
-                visual_port=None,
-                enable_multimodal=args.enable_multimodal,
-                metric_port=args.metric_port,
-            )
-            dp_size_in_node = max(1, args.dp // args.nnodes)  # 兼容多机纯tp的运行模式，这时候 1 // 2 == 0, 需要兼容
-            self.shared_token_load = TokenLoad(f"{get_unique_server_name()}_shared_token_load", dp_size_in_node)
         else:
             init_tokenizer(args)  # for openai api
             SamplingParams.load_generation_cfg(args.model_dir)
@@ -365,9 +350,10 @@ async def shutdown():
 @app.on_event("startup")
 async def startup_event():
     logger.info("server start up")
+    if g_objs.httpserver_manager is None:
+        return
     loop = asyncio.get_event_loop()
     g_objs.set_args(get_env_start_args())
-    if g_objs.args.run_mode != "visual":
-        loop.create_task(g_objs.httpserver_manager.handle_loop())
+    loop.create_task(g_objs.httpserver_manager.handle_loop())
     logger.info(f"server start up ok, loop use is {asyncio.get_event_loop()}")
     return

lightllm/server/api_start.py

@@ -208,9 +208,9 @@ def check_and_set_args(args):
 def normal_or_p_d_start(args):
 
     check_and_set_args(args)
-    already_uesd_ports = args.visual_nccl_ports + [args.nccl_port, args.port]
+    already_uesd_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]
+        already_uesd_ports = [args.nccl_port, args.port, args.pd_decode_rpyc_port]
 
     # 提前锁定端口，防止在单个机器上启动多个实列的时候，要到模型启动的时候才能
     # 捕获到端口设置冲突的问题
@@ -219,7 +219,7 @@ def normal_or_p_d_start(args):
 
     node_world_size = args.tp // args.nnodes
     can_use_ports = alloc_can_use_network_port(
-        num=7 + node_world_size + args.visual_dp * args.visual_tp, used_nccl_ports=already_uesd_ports
+        num=7 + node_world_size + args.visual_dp * args.visual_tp + args.visual_dp, used_nccl_ports=already_uesd_ports
     )
     logger.info(f"alloced ports: {can_use_ports}")
     (
@@ -239,6 +239,9 @@ def normal_or_p_d_start(args):
         can_use_ports = can_use_ports[args.visual_tp :]
         visual_model_tp_ports.append(tp_ports_for_dp)
 
+    args.visual_nccl_ports = can_use_ports[0 : args.visual_dp]
+    can_use_ports = can_use_ports[args.visual_dp :]
+
     # 将申请好的端口放入args参数中
     args.router_port = router_port
     args.detokenization_port = detokenization_port
@@ -436,7 +439,6 @@ def visual_start(args):
         metric_port,
     ) = can_use_ports[0:5]
     can_use_ports = can_use_ports[5:]
-    print(cache_port)
 
     visual_model_tp_ports = []
     for _ in range(args.visual_dp):

lightllm/server/config_server/api_http.py

@@ -8,7 +8,8 @@
 from typing import Dict, List
 from fastapi.responses import JSONResponse
 from lightllm.utils.log_utils import init_logger
-from ..pd_io_struct import PD_Master_Obj, Visual_Server_Obj
+from lightllm.server.visualserver.vit_connect import VIT_Obj
+from ..pd_io_struct import PD_Master_Obj
 from .nccl_tcp_store import start_tcp_store_server
 from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.utils.process_check import start_parent_check_thread
@@ -18,7 +19,7 @@
 app = FastAPI()
 
 registered_pd_master_objs: Dict[str, PD_Master_Obj] = {}
-registered_visual_server_obj: Dict[str, Visual_Server_Obj] = {}
+registered_visual_server_objs: Dict[str, VIT_Obj] = {}
 registered_pd_master_obj_lock = Lock()
 registered_visual_server_obj_lock = Lock()
 
@@ -73,15 +74,15 @@ async def websocket_endpoint(websocket: WebSocket):
     return
 
 
-@app.websocket("/visual_server_register")
+@app.websocket("/visual_register")
 async def visual_websocket_endpoint(websocket: WebSocket):
     await websocket.accept()
     client_ip, client_port = websocket.client
     logger.info(f"ws connected from IP: {client_ip}, Port: {client_port}")
-    registered_visual_server_obj: Visual_Server_Obj = pickle.loads(await websocket.receive_bytes())
+    registered_visual_server_obj: VIT_Obj = pickle.loads(await websocket.receive_bytes())
     logger.info(f"recieved registered_visual_server_obj {registered_visual_server_obj}")
     with registered_visual_server_obj_lock:
-        registered_visual_server_obj_lock[registered_visual_server_obj.node_id] = registered_visual_server_obj
+        registered_visual_server_objs[registered_visual_server_obj.node_id] = registered_visual_server_obj
 
     try:
         while True:
@@ -93,7 +94,7 @@ async def visual_websocket_endpoint(websocket: WebSocket):
     finally:
         logger.error(f"registered_visual_server_obj {registered_visual_server_obj} removed")
         with registered_visual_server_obj_lock:
-            registered_visual_server_obj.pop(registered_visual_server_obj.node_id, None)
+            registered_visual_server_objs.pop(registered_visual_server_obj.node_id, None)
     return
 
 
@@ -105,10 +106,10 @@ async def get_registered_objects():
         return {"data": base64_encoded}
 
 
-@app.get("/registered_visual_server_objects")
+@app.get("/registered_visual_objects")
 async def get_vit_registered_objects():
     with registered_visual_server_obj_lock:
-        serialized_data = pickle.dumps(registered_visual_server_obj)
+        serialized_data = pickle.dumps(registered_visual_server_objs)
         base64_encoded = base64.b64encode(serialized_data).decode("utf-8")
         return {"data": base64_encoded}
 

lightllm/server/httpserver/manager.py

@@ -697,6 +697,9 @@ async def handle_loop(self):
 
             asyncio.create_task(pd_handle_loop(self))
 
+        if self.enable_multimodal:
+            asyncio.create_task(self.vit_manager.vit_handle_loop())
+
         while True:
             try:
                 await asyncio.wait_for(self.recv_from_detokenization.recv_pyobj(), timeout=0.05)

lightllm/server/pd_io_struct.py

@@ -73,15 +73,6 @@ def to_log_str(self):
         return f"PD_MASTER host_ip_port: {self.host_ip_port} node_id: {self.node_id}"
 
 
-@dataclass
-class Visual_Server_Obj:
-    node_id: int
-    host_ip_port: str
-
-    def to_log_str(self):
-        return f"Visual_Server host_ip_port: {self.host_ip_port} node_id: {self.node_id}"
-
-
 @dataclass
 class UpKVStatus:
     type: str = "kv_move_status"

lightllm/server/visualserver/manager.py

@@ -53,11 +53,11 @@ def __init__(
     def _setup_connections(self):
         context = zmq.Context(2)
         if self.remote_vit:
-            self.recv_from_remote_llm = context.socket(zmq.PULL)
-            self.recv_from_remote_llm.bind(f"tcp://*:{self.args.remote_vit_port}")
+            self.vit_receiver = context.socket(zmq.PULL)
+            self.vit_receiver.bind(f"tcp://*:{self.args.remote_vit_port}")
         else:
-            self.recv_from_httpserver = context.socket(zmq.PULL)
-            self.recv_from_httpserver.bind(f"{self.args.zmq_mode}127.0.0.1:{self.visual_port}")
+            self.vit_receiver = context.socket(zmq.PULL)
+            self.vit_receiver.bind(f"{self.args.zmq_mode}127.0.0.1:{self.visual_port}")
             self.send_to_next_module = context.socket(zmq.PUSH)  # router or audio server (if --enable_multimodal_audio)
             self.send_to_next_module.connect(f"{self.args.zmq_mode}127.0.0.1:{self.next_module_port}")
         self.cache_client = rpyc.connect("localhost", self.cache_port, config={"allow_pickle": True})
@@ -153,7 +153,7 @@ async def loop_for_fwd(self):
 
     def _recv_reqs(self):
         if self.remote_vit:
-            recv_req: GroupReqIndexes = self.recv_from_httpserver.recv_pyobj(zmq.NOBLOCK)
+            recv_req: GroupReqIndexes = self.vit_receiver.recv_pyobj(zmq.NOBLOCK)
             for img in recv_req.multimodal_params.images:
                 image_patch = self.tokenizer.get_image_patch_func(img)
                 data = img._preload_data
@@ -164,7 +164,7 @@ def _recv_reqs(self):
                 self.cache_client.root.set_items_data([md5])
             return recv_req
         else:
-            return self.recv_from_httpserver.recv_pyobj(zmq.NOBLOCK)
+            return self.vit_receiver.recv_pyobj(zmq.NOBLOCK)
 
     async def loop_for_netio_req(self):
         if not hasattr(self, "visual_recv_max_count"):
@@ -173,7 +173,7 @@ async def loop_for_netio_req(self):
         while True:
             try:
                 for _ in range(self.visual_recv_max_count):
-                    recv_req: GroupReqIndexes = self._recv_reqs()
+                    recv_req: GroupReqIndexes = self.vit_receiver.recv_pyobj(zmq.NOBLOCK)
                     if isinstance(recv_req, GroupReqIndexes):
                         self.waiting_reqs.append(recv_req)
                     else:
@@ -182,6 +182,9 @@ async def loop_for_netio_req(self):
             except zmq.ZMQError:
                 # 当队列已经开始清空的时候，将一次接受数量下调
                 self.visual_recv_max_count = 64
+            except Exception as e:
+                logger.exception(f"Error in loop_for_netio_req: {e}")
+                raise e
             await asyncio.sleep(0.01)
 
     # code for visual only mode
@@ -249,9 +252,6 @@ def handle_exception(loop, context):
     loop = asyncio.new_event_loop()
     loop.set_exception_handler(handle_exception)
     asyncio.set_event_loop(loop)
-    if args.run_mode == "visual":
-        loop.create_task(visualserver.loop_for_fwd_visual_only())
-    else:
-        loop.create_task(visualserver.loop_for_fwd())
+    create_forward_loop(args, visualserver, loop)
     loop.run_until_complete(visualserver.loop_for_netio_req())
     return

lightllm/server/visualserver/register_loop.py

@@ -20,7 +20,7 @@ async def register_loop(args):
     while True:
 
         try:
-            uri = f"ws://{args.config_server_host}:{args.config_server_port}/visual_server_register"
+            uri = f"ws://{args.config_server_host}:{args.config_server_port}/visual_register"
             async with websockets.connect(uri, max_queue=(2048 * 1024, 2048 * 1023)) as websocket:
 
                 sock = websocket.transport.get_extra_info("socket")
@@ -33,7 +33,7 @@ async def register_loop(args):
 
                 while True:
                     await websocket.send("heartbeat")
-                    await asyncio.sleep(60)
+                    await asyncio.sleep(40)
 
         except Exception as e:
             logger.error("connetion to config_server has error")

lightllm/server/visualserver/vit_connect.py

@@ -57,6 +57,7 @@ def _setup_local_vit_connection(self):
         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):
+        print("_setup_remote_vit_connections", "fdakpgdakgjadpgkjadk")
         asyncio.create_task(self.vit_handle_loop())
 
         # wait for remote vit instances
@@ -89,6 +90,7 @@ async def send_to_vit(self, data, protocol=pickle.HIGHEST_PROTOCOL):
             raise Exception(f"Failed to send to VIT instance {instance.host_ip_port}: {e}")
 
     async def vit_handle_loop(self):
+        print("vit_handle_loop", "fdakpgdakgjadpgkjadk")
         while True:
             try:
                 id_to_vit_obj = await self._get_vit_objs()
@@ -118,8 +120,8 @@ 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"
-
+        uri = f"ws://{self.args.config_server_host}:{self.args.config_server_port}/registered_visual_objects"
+        print("uri", uri)
         try:
             async with httpx.AsyncClient() as client:
                 response = await client.get(uri)