[0828]fix vit embed

Author: SangChengC

lightllm/server/embed_cache/impl/naive_memory_cache.py

@@ -103,7 +103,7 @@ def alloc(self, md5sum_list: list[str], token_num_list: list[int]) -> Optional[l
                     rec.visittime = now
                     rec.ref += 1
                 else:
-                    uid_int = int(md5sum, 16)
+                    uid_int = md5sum
                     self._check_and_set_new_id_range(token_num)
                     rec = Record(
                         id=uid_int,

lightllm/server/embed_cache/utils.py

@@ -183,6 +183,11 @@ def insert(self, md5: str) -> Tuple[bool, List[str]]:
             self._release_lock()
             raise e
 
+    def query(self, md5: str) -> bool:
+        """Quert if md5 exists."""
+        self._wait_if_eviction()
+        return bool(self.r.exists(self.ref_prefix + md5))
+
     def query_and_incre(self, md5: str) -> bool:
         """Query if md5 exists and increment ref_count if found."""
         self._wait_if_eviction()

lightllm/server/httpserver/manager.py

@@ -117,10 +117,16 @@ def __init__(
         self.latest_success_infer_time_mark.set_value(int(time.time()))
         return
 
-    async def _alloc_resource(self, items, md5sums, token_nums, datas):
+    async def _alloc_resource(self, items, uuids, token_nums, datas):
 
         while True:
-            records = obtain(self.cache_client.root.alloc(md5sums, token_nums))
+            # 检查这个图片在redis总是否已经存在
+            # embed_exists = obtain(self.cache_client.root.get_items_embed(uuids))
+            # for exist in embed_exists:
+            #     if exist:
+            #         continue
+            # else:
+            records = obtain(self.cache_client.root.alloc(uuids, token_nums))
 
             if records is None:
                 await asyncio.sleep(0.1)
@@ -156,14 +162,15 @@ async def _alloc_multimodal_resources(self, multimodal_params: MultimodalParams,
             # 如果不加任何锁，假如请求1和请求2都有6张图片，而cache_capacity为10，
             # 那么如果某一时刻shm中存在请求1的5张图和请求2的5张图，将会资源竞争产生死锁。
             async with self._resource_lock:
-                items, md5sums, tokens_nums, datas = [], [], [], []
+                items, uuids, tokens_nums, datas = [], [], [], []
                 for img in multimodal_params.images:
                     self.tokenizer.init_imageitem_extral_params(img, multimodal_params, sampling_params)
                     data = img.read()
                     # must after init_imageitem_extral_params
                     token_num = self.tokenizer.get_image_token_length(img)
                     md5sum = "{}_{}".format(hashlib.md5(data).hexdigest(), img.patch_num)
-                    md5sums.append(md5sum)
+                    uuid = int(md5sum, 16)
+                    uuids.append(uuid)
                     tokens_nums.append(token_num)
                     datas.append(data)
                     items.append(img)
@@ -175,12 +182,13 @@ async def _alloc_multimodal_resources(self, multimodal_params: MultimodalParams,
                         hashlib.md5(data).hexdigest(),
                         hashlib.md5(pickle.dumps(audio.extra_params, protocol=4)).hexdigest(),
                     )
-                    md5sums.append(md5sum)
+                    uuid = int(md5sum, 16)
+                    uuids.append(uuid)
                     tokens_nums.append(token_num)
                     datas.append(data)
                     items.append(audio)
 
-                await self._alloc_resource(items, md5sums, tokens_nums, datas)
+                await self._alloc_resource(items, uuids, tokens_nums, datas)
         return
 
     async def _release_multimodal_resources(self, multimodal_params: MultimodalParams):

lightllm/server/visualserver/manager.py

@@ -153,16 +153,22 @@ async def loop_for_fwd(self):
                     processing_group_reqs = []
                     images_need_infer = []
 
-    # def _recv_reqs(self):
-    #     if self.remote_vit:
-    #         recv_req: GroupReqIndexes = self.recv_from_httpserver.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
-    #         ]
-    #         return recv_req
-    #     else:
-    #         return self.recv_from_httpserver.recv_pyobj(zmq.NOBLOCK)
+    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
+            # ]
+            uuids = []
+            token_nums = []
+            for img in recv_req.multimodal_params.images:
+                uuids.append(img.uuid)
+                token_nums.append(img.token_num)
+            self.cache_client.root.alloc(uuids, token_nums)
+            return recv_req
+        else:
+            return self.vit_receiver.recv_pyobj(zmq.NOBLOCK)
 
     async def loop_for_netio_req(self):
         if not hasattr(self, "visual_recv_max_count"):
@@ -171,7 +177,7 @@ async def loop_for_netio_req(self):
         while True:
             try:
                 for _ in range(self.visual_recv_max_count):
-                    recv_req: GroupReqIndexes = self.vit_receiver.recv_pyobj(zmq.NOBLOCK)
+                    recv_req: GroupReqIndexes = self._recv_reqs()
                     if isinstance(recv_req, GroupReqIndexes):
                         print(recv_req, flush=True)
                         self.waiting_reqs.append(recv_req)