fix vit transfer

Author: shihaobai

lightllm/server/httpserver/manager.py

@@ -134,8 +134,8 @@ async def _alloc_resource(self, items, md5sums, token_nums, datas):
                 uid_list.append(rec["id"])
 
             # # If enable the vit/audio-llm disaggregation, no need to cache the data in the memory of the server
-            # if self.enable_remote_vit:
-            #     return
+            if self.enable_remote_vit:
+                return
 
             ready_flags = obtain(self.cache_client.root.get_items_data(uid_list))
             update_data_ids = []
@@ -512,6 +512,7 @@ async def transfer_to_next_module(
                     group_req_objs.to_group_req_index(),
                     protocol=pickle.HIGHEST_PROTOCOL,
                 )
+
             if not self.enable_multimodal or self.args.enable_remote_vit:
                 self.send_to_router.send_pyobj(
                     group_req_objs.to_group_req_index(),

lightllm/server/multimodal_params.py

@@ -53,10 +53,7 @@ async def preload(self, request: Request):
 
     def read(self):
         assert self._preload_data is not None
-        ans = self._preload_data
-        # self._preload_data = None
-        # self._data = None
-        return ans
+        return self._preload_data
 
     def to_dict(self):
         ret = {}
@@ -112,10 +109,11 @@ async def preload(self, request: Request):
 
     def read(self):
         assert self._preload_data is not None
-        ans = self._preload_data
+        return self._preload_data
+
+    def free(self):
         self._preload_data = None
         self._data = None
-        return ans
 
     def to_dict(self):
         ret = {}
@@ -144,6 +142,12 @@ def __init__(
         self.audios = [AudioItem(**a) for a in audios]
         return
 
+    def free(self):
+        for image in self.images:
+            image.free()
+        for audio in self.audios:
+            audio.free()
+
     async def verify_and_preload(self, request: Request):
         for image in self.images:
             await image.preload(request)

lightllm/server/visualserver/manager.py

@@ -173,6 +173,7 @@ async def loop_for_netio_req(self):
                 for _ in range(self.visual_recv_max_count):
                     recv_req: GroupReqIndexes = self.vit_receiver.recv_pyobj(zmq.NOBLOCK)
                     if isinstance(recv_req, GroupReqIndexes):
+                        print(recv_req, flush=True)
                         self.waiting_reqs.append(recv_req)
                     else:
                         assert False, f"Error Req Inf {recv_req}"

lightllm/server/visualserver/vit_connect.py

@@ -131,7 +131,9 @@ def _update_vit_connections(self, id_to_vit_obj: Dict[int, VIT_Obj]):
             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}")
+                    print(vit_obj.host_ip_port, self.args.remote_vit_port, flush=True)
+                    ip, port = vit_obj.host_ip_port.split(":")
+                    socket.connect(f"tcp://{ip}:{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:
@@ -158,11 +160,14 @@ async def send_to_vit(self, data, protocol=pickle.HIGHEST_PROTOCOL):
         """
         instance = self._get_vit_instance()
         try:
+            print(instance, flush=True)
             instance.send_pyobj(data, protocol=protocol)
         except Exception as e:
             logger.error(f"Failed to send to VIT instance: {e}")
             raise Exception(f"Failed to send to VIT instance: {e}")
-
+        finally:
+            # 释放图片资源
+            data.multimodal_params.free()
         await self._wait_visual_embed_ready()
 
     async def vit_handle_loop(self):