fix llm remote vit init

Author: shihaobai

lightllm/server/api_http.py

@@ -350,10 +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.httpserver_manager is None:
+        return
     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/core/objs/req.py

@@ -214,9 +214,6 @@ def can_release(self):
         # 只有管理节点有一个引用
         ref_count_ok = self.ref_count == 1
         can_released_mark = self.can_released_mark
-        print(f"self.is_aborted is {self.is_aborted}")
-        print(f"self.finish_status.is_finished() is {self.finish_status.is_finished()}")
-        print(f"self.ref_count is {self.ref_count}")
         if self.is_aborted and can_released_mark and ref_count_ok:
             return True
 

lightllm/server/httpserver/manager.py

@@ -512,7 +512,7 @@ async def transfer_to_next_module(
                     group_req_objs.to_group_req_index(),
                     protocol=pickle.HIGHEST_PROTOCOL,
                 )
-            else:
+            if not self.enable_multimodal or self.args.enable_remote_vit:
                 self.send_to_router.send_pyobj(
                     group_req_objs.to_group_req_index(),
                     protocol=pickle.HIGHEST_PROTOCOL,

lightllm/server/visualserver/register_loop.py

@@ -26,7 +26,7 @@ async def register_loop(args):
                 sock = websocket.transport.get_extra_info("socket")
                 sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
 
-                vit_obj = VIT_Obj(node_id=args.visual_node_id, host_ip_port=f"{host_ip}:{args.port}")
+                vit_obj = VIT_Obj(node_id=args.visual_node_id, host_ip_port=f"{host_ip}:{args.remote_vit_port}")
 
                 await websocket.send(pickle.dumps(vit_obj))
                 logger.info(f"Sent registration vit_obj: {vit_obj}")

lightllm/server/visualserver/vit_connect.py

@@ -30,7 +30,7 @@ def __init__(self, args, context, local_visual_port: int):
         self.local_visual_port = local_visual_port
 
         self.send_to_visual = None
-        self.remote_vit_instances = []
+        self.remote_vit_instances = {}
         self.current_vit_index = 0
         self.remote_vit = args.enable_remote_vit
         self.remote_vit_port = args.remote_vit_port
@@ -42,8 +42,7 @@ def _setup_vit_connections(self):
         设置VIT连接，支持本地和远程VIT实例
         支持多种连接模式：
         1. 本地VIT实例 (默认)
-        2. 远程单个VIT实例
-        3. 远程多个VIT实例 (负载均衡)
+        2. 远程多个VIT实例 (负载均衡)
         """
         if self.remote_vit:
             # 远程VIT实例模式
@@ -57,14 +56,86 @@ 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())
+        """
+        初始化远程VIT连接，同步获取初始实例
+        """
+        logger.info("Setting up remote VIT connections...")
 
-        # wait for remote vit instances
-        while True:
-            if len(self.remote_vit_instances) > 0:
-                break
+        self._sync_init_vit_instances()
+
+        retry_count = 0
+        max_retries = 30  # 最多等待30秒
+        while len(self.remote_vit_instances) == 0 and retry_count < max_retries:
+            logger.info(f"Waiting for VIT instances... (attempt {retry_count + 1}/{max_retries})")
             time.sleep(1)
+            retry_count += 1
+            self._sync_init_vit_instances()
+
+        if len(self.remote_vit_instances) == 0:
+            logger.warning("No VIT instances available after initialization")
+        else:
+            logger.info(f"Successfully connected to {len(self.remote_vit_instances)} VIT instances")
+
+    def _sync_init_vit_instances(self):
+        """
+        同步初始化VIT实例连接
+        """
+        try:
+            # 使用同步方式获取VIT实例
+            vit_objs = self._sync_get_vit_objs()
+            if vit_objs:
+                self._update_vit_connections(vit_objs)
+        except Exception as e:
+            logger.error(f"Failed to initialize VIT instances: {e}")
+
+    def _sync_get_vit_objs(self) -> Optional[Dict[int, VIT_Obj]]:
+        """
+        同步获取VIT实例信息
+        """
+        import requests
+
+        uri = f"http://{self.args.config_server_host}:{self.args.config_server_port}/registered_visual_objects"
+        try:
+            response = requests.get(uri, timeout=10)
+            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"Failed to get VIT instances: {response.status_code}")
+                return None
+        except Exception as e:
+            logger.error(f"Error getting VIT instances: {e}")
+            return None
+
+    def _update_vit_connections(self, id_to_vit_obj: Dict[int, VIT_Obj]):
+        """
+        更新VIT连接，添加新的连接，关闭失效的连接
+        """
+        # 关闭不再存在的连接
+        closed_ids = []
+        for id, remote_instance in self.remote_vit_instances.items():
+            if id not in id_to_vit_obj:
+                try:
+                    remote_instance.close()
+                except:
+                    pass
+                closed_ids.append(id)
+                logger.info(f"Closed VIT connection {id}")
+
+        for id in closed_ids:
+            self.remote_vit_instances.pop(id)
+
+        # 建立新的连接
+        for id, vit_obj in id_to_vit_obj.items():
+            if id not in self.remote_vit_instances:
+                try:
+                    socket = self.context.socket(zmq.PUSH)
+                    socket.connect(f"tcp://{vit_obj.host_ip_port}:{self.args.remote_vit_port}")
+                    self.remote_vit_instances[id] = socket
+                    logger.info(f"Connected to VIT instance {id} at {vit_obj.host_ip_port}")
+                except Exception as e:
+                    logger.error(f"Failed to connect to VIT instance {id}: {e}")
 
     def _get_vit_instance(self):
         """
@@ -73,10 +144,13 @@ def _get_vit_instance(self):
         if not self.remote_vit:
             return self.send_to_visual
 
+        if len(self.remote_vit_instances) == 0:
+            raise Exception("No available VIT instances")
+
         # 简单的轮询负载均衡
         index = (self.current_vit_index + 1) % len(self.remote_vit_instances)
         self.current_vit_index = index
-        return self.remote_vit_instances[index]
+        return list(self.remote_vit_instances.values())[index]
 
     async def send_to_vit(self, data, protocol=pickle.HIGHEST_PROTOCOL):
         """
@@ -86,42 +160,32 @@ async def send_to_vit(self, data, protocol=pickle.HIGHEST_PROTOCOL):
         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}")
+            logger.error(f"Failed to send to VIT instance: {e}")
+            raise Exception(f"Failed to send to VIT instance: {e}")
+
+        await self._wait_visual_embed_ready()
 
     async def vit_handle_loop(self):
-        print("vit_handle_loop", "fdakpgdakgjadpgkjadk")
+        """
+        异步VIT连接管理循环，由外部启动
+        """
+        logger.info("Starting VIT connection management loop")
         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}"
-                        )
+                id_to_vit_obj = await self._async_get_vit_objs()
+                if id_to_vit_obj:
+                    logger.debug(f"Retrieved {len(id_to_vit_obj)} VIT instances")
+                    self._update_vit_connections(id_to_vit_obj)
                 await asyncio.sleep(30)
             except Exception as e:
-                logger.exception(str(e))
+                logger.exception(f"Error in VIT handle loop: {e}")
                 await asyncio.sleep(10)
 
-    async def _get_vit_objs(self) -> Optional[Dict[int, VIT_Obj]]:
+    async def _async_get_vit_objs(self) -> Optional[Dict[int, VIT_Obj]]:
         """
-        get_vit_objs 主要负责从 config_server 获取所有的vit远程服务。
+        异步获取VIT实例信息
         """
-        # 使用 config_server 服务来发现所有的 pd_master 节点。
         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)
@@ -130,9 +194,14 @@ async def _get_vit_objs(self) -> Optional[Dict[int, VIT_Obj]]:
                     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}")
+                    logger.error(f"Failed to get VIT instances: {response.status_code}")
                     return None
         except Exception as e:
-            logger.exception(str(e))
-            await asyncio.sleep(10)
+            logger.exception(f"Error getting VIT instances: {e}")
             return None
+
+    async def _wait_visual_embed_ready(self):
+        """
+        等待VIT实例的embed准备好
+        """
+        await asyncio.sleep(10)