[fix]batch rpyc call in multimodal

Author: niushengxiao

lightllm/models/whisper/whisper_audio.py

@@ -190,7 +190,7 @@ def encode(self, audio_items: List[AudioItem]):
         audio_lens_after_cnn = np.array(audio_lens_after_cnn, dtype=np.int32)
         audio_token_num = (audio_lens_after_cnn - 2) // 2 + 1
 
-        ready_audio = self.cache_client.root.get_items_data(uuids)
+        ready_audio = self.cache_client.root.get_items_embed(uuids)
         ids_to_set = []
         for i, ready in enumerate(ready_audio):
             if not ready:

lightllm/server/embed_cache/impl/naive_memory_cache.py

@@ -2,7 +2,7 @@
 import threading
 import dataclasses
 import requests
-from typing import Union
+from typing import Union, Optional
 import torch
 import time
 from collections import deque
@@ -87,41 +87,42 @@ def _clear(self):
                 if deleted >= max_delete:
                     break
 
-    def alloc(self, md5sum_list: list[str], token_num_list: list[int]) -> list[dict]:
-        results = []
+    def alloc(self, md5sum_list: list[str], token_num_list: list[int]) -> Optional[list[dict]]:
+        now = time.time()
         with self.lock:
+            new_md5s = [m for m in md5sum_list if m not in self._md5_to_record]
+            new_needed = len(new_md5s)
+
+            if self.occupied + new_needed > self.capacity:
+                self._clear()
+            if self.occupied + new_needed > self.capacity:
+                return None
+
+            results: list[dict] = []
             for md5sum, token_num in zip(md5sum_list, token_num_list):
-                t = time.time()
-                if md5sum not in self._md5_to_record:
-                    if self.occupied >= self.capacity:
-                        self._clear()
-                        if self.occupied >= self.capacity:
-                            results.append(None)
-                            continue
-                    id = uuid.uuid1()
-                    id = id.int
+                if md5sum in self._md5_to_record:
+                    rec = self._md5_to_record[md5sum]
+                    rec.visittime = now
+                    rec.ref += 1
+                else:
+                    uid_int = uuid.uuid1().int
                     self._check_and_set_new_id_range(token_num)
-                    record = Record(
-                        id=id,
+                    rec = Record(
+                        id=uid_int,
                         md5sum=md5sum,
                         ref=1,
                         data=False,
                         embed=False,
-                        createtime=t,
-                        visittime=t,
+                        createtime=now,
+                        visittime=now,
                         token_id=self.token_id_range_start,
                         token_num=token_num,
                     )
                     self.token_id_range_start += token_num
-                    self._records[id] = record
-                    self._md5_to_record[md5sum] = record
+                    self._records[uid_int] = rec
+                    self._md5_to_record[md5sum] = rec
                     self.occupied += 1
-                # cache hit
-                else:
-                    record = self._md5_to_record[md5sum]
-                    record.visittime = t
-                    record.ref += 1
-                results.append({"id": record.id, "token_id": record.token_id, "token_num": record.token_num})
+                results.append({"id": rec.id, "token_id": rec.token_id, "token_num": rec.token_num})
         return results
 
     def release(self, ids: list[int]) -> None:

lightllm/server/embed_cache/manager.py

@@ -1,7 +1,7 @@
 import rpyc
 import uuid
 import inspect
-from typing import Union
+from typing import Union, Optional
 from lightllm.utils.graceful_utils import graceful_registry
 from lightllm.server.embed_cache.impl.naive_memory_cache import InMemoryCache
 from rpyc.utils.classic import obtain
@@ -22,7 +22,7 @@ def on_disconnect(self, conn):
         # (to finalize the service, if needed)
         pass
 
-    def exposed_alloc(self, md5sum_list: list[str], token_num_list: list[int]) -> dict:
+    def exposed_alloc(self, md5sum_list: list[str], token_num_list: list[int]) -> Optional[list[dict]]:
         md5sum_list = obtain(md5sum_list)
         token_num_list = obtain(token_num_list)
         record = self._impl.alloc(md5sum_list, token_num_list)

lightllm/server/httpserver/manager.py

@@ -114,53 +114,34 @@ def __init__(
         self.latest_success_infer_time_mark.set_value(int(time.time()))
         return
 
-    async def _alloc_resource(self, items, md5sums, tokens_nums, datas):
+    async def _alloc_resource(self, items, md5sums, token_nums, datas):
         wait_time = 1
-        pending_idx = list(range(len(items)))
-        while pending_idx:
-            sub_md5sum = [md5sums[i] for i in pending_idx]
-            sub_tokens_num = [tokens_nums[i] for i in pending_idx]
-
-            records = self.cache_client.root.alloc(sub_md5sum, sub_tokens_num)
+        while True:
+            records = self.cache_client.root.alloc(md5sums, token_nums)
 
-            if all(record is None for record in records):
+            if records is None:
                 await asyncio.sleep(wait_time)
                 wait_time = min(wait_time + 2, 9)
                 continue
 
-            next_pending = []  # record为None，安排在下一轮
-            uids_to_check = []  # record存在，本轮处理
-            uid_to_idx = {}  # uid → 原items下标
+            uid_list = []
+            for item, rec in zip(items, records):
+                item.uuid = rec["id"]
+                item.token_id = rec["token_id"]
+                item.token_num = rec["token_num"]
+                uid_list.append(rec["id"])
 
-            for local_pos, record in enumerate(records):
-                global_pos = pending_idx[local_pos]
+            ready_flags = self.cache_client.root.get_items_data(uid_list)
+            need_write = []
 
-                if record is None:
-                    next_pending.append(global_pos)
-                    continue
+            for uid, ready, data in zip(uid_list, ready_flags, datas):
+                if not ready:
+                    create_shm(get_shm_name_data(uid), data)
+                    need_write.append(uid)
 
-                uid = record["id"]
-                uid_to_idx[uid] = global_pos
-                uids_to_check.append(uid)
-
-                item = items[global_pos]
-                item.uuid = uid
-                item.token_id = record["token_id"]
-                item.token_num = record["token_num"]
-
-            if uids_to_check:
-                ready_flags = self.cache_client.root.get_items_data(uids_to_check)
-                need_write = []
-
-                for uid, ready in zip(uids_to_check, ready_flags):
-                    if not ready:
-                        idx = uid_to_idx[uid]
-                        create_shm(get_shm_name_data(uid), datas[idx])
-                        need_write.append(uid)
-                if need_write:
-                    self.cache_client.root.set_items_data(need_write)
-            pending_idx = next_pending
-        return
+            if need_write:
+                self.cache_client.root.set_items_data(need_write)
+            return
 
     async def _alloc_multimodal_resources(self, multimodal_params: MultimodalParams, sampling_params: SamplingParams):
         # 只有 P 和 NORMAL 节点需要真的管理多模态资源