0911

Author: SangChengC

lightllm/models/vit/model.py

@@ -178,7 +178,8 @@ def encode(self, images: List[ImageItem]):
         for i, img in enumerate(images):
             if isinstance(img, ImageItem):
                 uuids.append(img.uuid)
-                image_data = read_shm(get_shm_name_data(img.uuid))
+                image_data = img._preload_data
+                # image_data = read_shm(get_shm_name_data(img.uuid))
                 image_data = Image.open(BytesIO(image_data))
                 t = self.load_image_func(image_data, max_num=img.extra_params["image_patch_max_num"])
                 img_tensors.append(t)

lightllm/server/embed_cache/impl/memory_cache_with_redis.py

@@ -36,21 +36,19 @@ def release(self, ids: list[int]) -> None:
                 self._records[id_].ref -= 1
                 if self.redis_cache.query(str(id_)):
                     self.redis_cache.decr(str(id_))
-                    print(self.redis_cache.stats(), flush=True)
+                    # print(self.redis_cache.stats(), flush=True)
 
     # vit 负责set
     def set_items_embed(self, ids: list[int]) -> None:
         with self.lock:
             for id in ids:
                 self.redis_cache.insert(str(id))
                 self._records[id].embed = True
-                self._records[id].ref -= 1
+                self._records[id].ref -= 1  # vit端alloc之后ref+1 vit完成后ref-1
 
     def get_items_embed(self, ids: list[int]) -> list[Optional[bool]]:
         ret = []
         for id in ids:
-            print(f"id is {id}")
-            print(f"self.redis_cache.query(str(id)) is {self.redis_cache.query(str(id))}")
             exist = self.redis_cache.query(str(id))
             ret.append(exist)
             if exist:

lightllm/server/embed_cache/impl/naive_memory_cache.py

@@ -78,9 +78,10 @@ def _clear(self, free_max_count: int):
                     free_shm(get_shm_name_data(id))
                 if record.embed:
                     # 仅vit释放掉afs里的, llm端不做释放
-                    if self.args.run_mode == "visual":
-                        free_afs(get_shm_name_embed(id), self.args.image_embed_dir)
-                    elif not self.args.enable_remote_vit:
+                    # if self.args.run_mode == "visual":
+                    #     free_afs(get_shm_name_embed(id), self.args.image_embed_dir)
+                    # elif not self.args.enable_remote_vit:
+                    if not self.args.run_mode == "visual":
                         free_shm(get_shm_name_embed(id))
                 del self._md5_to_record[record.md5sum]
                 del self._records[id]

lightllm/server/embed_cache/utils.py

@@ -8,6 +8,9 @@
 from pathlib import Path
 import multiprocessing.shared_memory as shm
 from lightllm.utils.envs_utils import get_env_start_args
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
 
 
 def tensor2bytes(t: torch.Tensor):
@@ -247,7 +250,7 @@ def _md5_to_afs_path(self, md5: str) -> str:
         """Convert md5 to AFS file path."""
         if not self.image_embed_dir:
             return None
-        filename = md5 + self.path_ext
+        filename = self.image_embed_dir + md5 + self.path_ext
         return filename
 
     def _delete_afs_files(self, victims: List[str]) -> None:
@@ -260,9 +263,9 @@ def _delete_afs_files(self, victims: List[str]) -> None:
                 file_path = self._md5_to_afs_path(md5)
                 if file_path and os.path.exists(file_path):
                     os.remove(file_path)
-                    print(f"Deleted AFS file: {file_path}")
+                    logger.debug(f"Deleted AFS file: {file_path}")
             except Exception as e:
-                print(f"Warning: Failed to delete AFS file for {md5}: {e}")
+                logger.debug(f"Warning: Failed to delete AFS file for {md5}: {e}")
 
     # ---------------- Lua scripts ----------------
     _INSERT_LUA = r"""
@@ -273,6 +276,7 @@ def _delete_afs_files(self, victims: List[str]) -> None:
 local md5 = ARGV[1]
 local capacity = tonumber(ARGV[2])
 
+local unpack = unpack or table.unpack
 local ref_key = ref_prefix .. md5
 if redis.call('GET', ref_key) then
   return {0}  -- Already exists
@@ -385,7 +389,7 @@ def _delete_afs_files(self, victims: List[str]) -> None:
     local now = redis.call('TIME')[1] * 1000
     redis.call('ZADD', zset, now, new_md5)
     
-    return {1, table.unpack(victims)}  -- success + victims
+    return {1, unpack(victims)}  -- success + victims
 else
     return {0}  -- 逐出失败，没有足够的候选
 end

lightllm/server/visualserver/manager.py

@@ -153,21 +153,28 @@ async def loop_for_fwd(self):
                     processing_group_reqs = []
                     images_need_infer = []
 
-    def _recv_reqs(self):
+    async def _recv_reqs(self):
         if self.remote_vit:
             recv_req: GroupReqIndexes = self.vit_receiver.recv_pyobj(zmq.NOBLOCK)
             # recv_req.multimodal_params.images[:]= [
             #     img for img in recv_req.multimodal_params.images
             #     if not self.cache_client.root.get_item_embed(img.uuid)  # embed已存在的被丢弃 , ref +1
             # ]
+            logger.info(f"Receive req {recv_req.group_req_id}, image_count:{len(recv_req.multimodal_params.images)}")
             uuids = [img.uuid for img in recv_req.multimodal_params.images]
             already_embed = self.cache_client.root.get_items_embed(uuids)
+            if all(already_embed):
+                return None
             token_nums = []
             for img, embed in zip(recv_req.multimodal_params.images, already_embed):
                 if not embed:
                     uuids.append(img.uuid)
                     token_nums.append(img.token_num)
-            self.cache_client.root.alloc(uuids, token_nums)
+            while True:
+                records = self.cache_client.root.alloc(uuids, token_nums)
+                if records is not None:
+                    break
+                await asyncio.sleep(0.1)
             return recv_req
         else:
             return self.vit_receiver.recv_pyobj(zmq.NOBLOCK)
@@ -179,11 +186,11 @@ 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 = await self._recv_reqs()
+                    if recv_req is None:
+                        continue
                     if isinstance(recv_req, GroupReqIndexes):
-                        # print(recv_req, flush=True)
                         self.waiting_reqs.append(recv_req)
-                        print(f"recv_req.multimodal_params is {recv_req.multimodal_params}")
                     else:
                         assert False, f"Error Req Inf {recv_req}"
                 self.visual_recv_max_count = min(self.visual_recv_max_count * 1.3, 256)
@@ -210,11 +217,21 @@ async def loop_for_fwd_visual_only(self):
                         images_need_infer.append(img)
 
                         if len(images_need_infer) == self.infer_batch_size:
+                            _t0 = time.perf_counter()
                             await self.infer_imgs(images_need_infer)
+                            logger.info(
+                                f"[visual] batch infer complete, image_count: {len(images_need_infer)}, "
+                                f"elapsed_time {(time.perf_counter()-_t0) * 1000}ms"
+                            )
                             images_need_infer = []
 
                     if len(images_need_infer) > 0:
+                        _t1 = time.perf_counter()
                         await self.infer_imgs(images_need_infer)
+                        logger.info(
+                            f"[visual] batch infer complete, image_count:{len(images_need_infer)}, "
+                            f"elapsed_time {(time.perf_counter()-_t1) * 1000}ms"
+                        )
                         images_need_infer = []
                     # 在这里release这个image，ref-1
                     logger.info(f"req-id {visual_req.group_req_id} has been release ok")

lightllm/server/visualserver/model_infer/model_rpc.py

@@ -50,7 +50,6 @@ def exposed_init_model(self, kvargs):
         self.dp_rank_id = kvargs["dp_rank_id"]
         self.tp_rank_id = kvargs["tp_rank_id"]
         kvargs["vit_rank_id"] = self.dp_rank_id * self.args.visual_tp + self.tp_rank_id
-        print(cache_port)
         self.cache_client = rpyc.connect("localhost", cache_port, config={"allow_pickle": True})
 
         init_vision_distributed_env(kvargs)
@@ -87,9 +86,7 @@ def exposed_init_model(self, kvargs):
             else:
                 raise Exception(f"can not support {self.model_type} now")
             self.model.load_model(weight_dir)
-            print("begin load model")
             self.model = self.model.cuda()
-            print("load model OK")
         except Exception as e:
             print("#" * 16)
             print("load model error:", str(e), e, type(e))
@@ -113,12 +110,11 @@ def exposed_encode(self, images: List[ImageItem]):
         all_img_embeds = all_img_embeds.to(torch.device("cpu"))
 
         if self.tp_rank_id == 0:
-            ready_flags = obtain(self.cache_client.root.get_items_embed(uuids))
-            print(f"ready_flags is {ready_flags}")
+            # ready_flags = obtain(self.cache_client.root.get_items_embed(uuids))
             ids_to_set = []
-            for i, ready in enumerate(ready_flags):
-                if ready:
-                    continue
+            for i, img in enumerate(images):
+                # if ready:
+                #     continue
                 uid = uuids[i]
                 start, end = valid_ids[i]
                 cur_embed_bytes = tensor2bytes(all_img_embeds[start:end])

lightllm/server/visualserver/vit_connect.py

@@ -137,7 +137,7 @@ 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)
-                    print(vit_obj.host_ip_port, self.args.remote_vit_port, flush=True)
+                    # 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
@@ -223,7 +223,6 @@ async def _wait_visual_embed_ready(self, req: GroupReqIndexes, timeout_seconds:
             return
 
         uuids = req.multimodal_params.get_all_uuids()
-        print(f"uuids is {uuids}")
 
         async def wait_for_embeds():
             while not all(self.cache_client.root.get_items_embed(uuids)):