[fix]fix redis

Author: SangChengC

lightllm/models/qwen_vl/layer_infer/pre_layer_infer.py

@@ -1,3 +1,5 @@
+import rpyc
+import socket
 import torch
 import torch.distributed as dist
 
@@ -31,6 +33,8 @@ class LlamaMultimodalPreLayerInfer(LlamaPreLayerInfer):
     def __init__(self, network_config, mode):
         super().__init__(network_config, mode)
         self.args = get_env_start_args()
+        self.cache_client = rpyc.connect("localhost", self.args.cache_port, config={"allow_pickle": True})
+        self.cache_client._channel.stream.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
         return
 
     def context_forward(self, input_ids, infer_state: LlamaInferStateInfo, layer_weight: LlamaPreAndPostLayerWeight):
@@ -57,6 +61,7 @@ def context_forward(self, input_ids, infer_state: LlamaInferStateInfo, layer_wei
                     embed = read_afs(get_shm_name_embed(img["uuid"]), self.args.image_embed_dir)
                 else:
                     embed = read_shm(get_shm_name_embed(img["uuid"]))
+                self.cache_client.root.release([img["uuid"]])
                 img_weight.append(bytes2tensor(embed).cuda().reshape(img["token_num"], -1))
                 img_start_token_ids.append(img["token_id"])
                 img_token_lens.append(img["token_num"])

lightllm/server/embed_cache/impl/memory_cache_with_redis.py

@@ -32,10 +32,10 @@ def __init__(self, args) -> None:
     # llm 负责release
     def release(self, ids: list[int]) -> None:
         with self.lock:
-            for id_ in ids:
-                self._records[id_].ref -= 1
-                if self.redis_cache.query(str(id_)):
-                    self.redis_cache.decr(str(id_))
+            for id in ids:
+                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)
 
     # vit 负责set
@@ -44,27 +44,31 @@ def set_items_embed(self, ids: list[int]) -> None:
             for id in ids:
                 self.redis_cache.insert(str(id))
                 self._records[id].embed = True
-                self._records[id].ref -= 1  # vit端alloc之后ref+1 vit完成后ref-1
+                self._records[id].ref -= 1
+                self.redis_cache.decr(str(id))  # vit端alloc之后ref+1 vit完成后ref-1
 
-    def get_items_embed(self, ids: list[int]) -> list[Optional[bool]]:
+    def get_items_embed(self, ids: list[int], embeding_only: bool = False) -> list[Optional[bool]]:
         ret = []
         for id in ids:
-            exist = self.redis_cache.query(str(id))
+            if embeding_only:
+                exist = self.redis_cache.query(str(id))
+            else:
+                exist = self.redis_cache.query_and_incre(str(id))
             ret.append(exist)
             if exist:
                 self._records[id].embed = True
         return ret
 
-    def get_items_embed_and_incre(self, ids: list[int]) -> list[Optional[bool]]:
-        ret = []
-        for id in ids:
-            # if self.redis_cache.query(str(id)):
-            #     ret.append(True)
-            #     continue
-            # 避免重复的引用计数增加
-            if self._records[id].embed:
-                ret.append(True)
-                continue
-            self._records[id].embed = self.redis_cache.query_and_incre(str(id))
-            ret.append(self._records[id].embed)
-        return ret
+    # def get_items_embed_and_incre(self, ids: list[int]) -> list[Optional[bool]]:
+    #     ret = []
+    #     for id in ids:
+    #         # if self.redis_cache.query(str(id)):
+    #         #     ret.append(True)
+    #         #     continue
+    #         # 避免重复的引用计数增加
+    #         if self._records[id].embed:
+    #             ret.append(True)
+    #             continue
+    #         self._records[id].embed = self.redis_cache.query_and_incre(str(id))
+    #         ret.append(self._records[id].embed)
+    #     return ret

lightllm/server/embed_cache/impl/naive_memory_cache.py

@@ -144,5 +144,5 @@ def set_items_embed(self, ids: list[int]) -> None:
         for id_ in ids:
             self._records[id_].embed = True
 
-    def get_items_embed(self, ids: list[int]) -> list[Optional[bool]]:
+    def get_items_embed(self, ids: list[int], embeding_only: bool = False) -> list[Optional[bool]]:
         return [self._records.get(id_).embed if id_ in self._records else False for id_ in ids]

lightllm/server/embed_cache/manager.py

@@ -49,9 +49,9 @@ def exposed_set_items_embed(self, ids: list[int]) -> None:
         ids = obtain(ids)
         return self._impl.set_items_embed(ids)
 
-    def exposed_get_items_embed(self, ids: list[int]) -> list[bool]:
+    def exposed_get_items_embed(self, ids: list[int], embeding_only: bool = False) -> list[bool]:
         ids = obtain(ids)
-        return self._impl.get_items_embed(ids)
+        return self._impl.get_items_embed(ids, embeding_only)
 
 
 def get_cache_manager(args):

lightllm/server/embed_cache/utils.py

@@ -118,7 +118,7 @@ def __init__(
         self,
         redis_url: str = "redis://localhost:6379/0",
         capacity: int = 50000,
-        evict_fraction: float = 0.2,
+        evict_fraction: float = 0.1,
         key_prefix: str = "md5:",
         image_embed_dir: str = None,
         path_ext: str = "-embed",
@@ -128,7 +128,7 @@ def __init__(
         - capacity: max count of md5 entries allowed in Redis
         - evict_fraction: fraction to evict when inserting a NEW md5 and at capacity
         - image_embed_dir: base directory for image embed files (e.g., "/afs/embeds")
-        - path_ext: file extension for embed files (default: ".embed")
+        - path_ext: file extension for embed files (default: "-embed")
         """
         if not (0.0 <= evict_fraction <= 1.0):
             raise ValueError("evict_fraction must be 0..1")
@@ -152,7 +152,7 @@ def __init__(
         self._evict_and_insert_script = self.r.register_script(self._EVICT_AND_INSERT_LUA)
 
     def insert(self, md5: str) -> Tuple[bool, List[str]]:
-        """Insert a new md5 with default ref_count=0. May trigger LRU eviction."""
+        """Insert a new md5 with default ref_count=1. May trigger LRU eviction."""
         # 等待任何正在进行的逐出操作
         self._wait_if_eviction()
 
@@ -176,16 +176,20 @@ def insert(self, md5: str) -> Tuple[bool, List[str]]:
                     success = bool(evict_res[0])
                     victims = evict_res[1:] if len(evict_res) > 1 else []
 
-                    # 删除被逐出md5对应的AFS文件
-                    if victims and self.image_embed_dir:
-                        self._delete_afs_files(victims)
-
-                    return success, victims
+                    if success:
+                        # 删除被逐出md5对应的AFS文件
+                        if victims and self.image_embed_dir:
+                            self._delete_afs_files(victims)
+                        return True, victims
+                    else:
+                        # 逐出失败，短暂退避后重试
+                        time.sleep(0.01)
+                        return self.insert(md5)
                 finally:
                     self._release_lock()
             else:
                 # 等待锁释放后重试
-                time.sleep(0.1)
+                time.sleep(0.01)
                 return self.insert(md5)
         except Exception as e:
             self._release_lock()
@@ -199,7 +203,6 @@ def query(self, md5: str) -> bool:
     def query_and_incre(self, md5: str) -> bool:
         """Query if md5 exists and increment ref_count if found."""
         self._wait_if_eviction()
-
         res = self._query_incre_script(
             keys=[self.zset_key, self.ref_prefix],
             args=[md5],
@@ -228,6 +231,11 @@ def stats(self) -> dict:
             "evict_fraction": self.evict_fraction,
         }
 
+    def get_ref(self, md5: str) -> int | None:
+        self._wait_if_eviction()
+        val = self.r.get(self.ref_prefix + md5)
+        return int(val) if val is not None else None
+
     def _wait_if_eviction(self) -> None:
         max_wait = 30
         start_time = time.time()
@@ -284,8 +292,8 @@ def _delete_afs_files(self, victims: List[str]) -> None:
 
 local size = redis.call('ZCARD', zset)
 if size < capacity then
-  -- Insert with ref_count=0
-  redis.call('SET', ref_key, 0)
+  -- Insert with ref_count=1
+  redis.call('SET', ref_key, 1)
   local now = redis.call('TIME')[1] * 1000
   redis.call('ZADD', zset, now, md5)
   return {0}  -- Success, no eviction
@@ -332,17 +340,16 @@ def _delete_afs_files(self, victims: List[str]) -> None:
 
 --ref 递减到 0 时保留键，只更新计数与 LRU
 local rc = tonumber(val) - 1
-if rc < 0 then
-  rc = 0
-end
-
+if rc < 0 then rc = 0 end
 redis.call('SET', ref_key, rc)
 
--- 更新 LRU 时间戳（最近释放的条目更不容易被立即逐出）
-local now = redis.call('TIME')[1] * 1000
-redis.call('ZADD', zset, now, md5)
+if rc > 0 then
+  -- 只有仍被引用时才更新 LRU
+  local now = redis.call('TIME')[1] * 1000
+  redis.call('ZADD', zset, now, md5)
+end
 
-return {rc, 0}  -- 未删除
+return {rc, 0}
 """
 
     _EVICT_AND_INSERT_LUA = r"""
@@ -354,43 +361,64 @@ def _delete_afs_files(self, victims: List[str]) -> None:
 local capacity = tonumber(ARGV[2])
 local evict_fraction = tonumber(ARGV[3])
 
--- 计算需要逐出的数量
-local need = math.max(1, math.floor(capacity * evict_fraction + 0.5))
+local unpack = unpack or table.unpack
+
+-- helper: now millis
+local function now_ms()
+  local t = redis.call('TIME')
+  return t[1] * 1000 + math.floor(t[2] / 1000)
+end
+
+local new_ref_key = ref_prefix .. new_md5
+
+-- If already exists, treat as a hit: bump ref_count and refresh LRU
+local cur = redis.call('GET', new_ref_key)
+if cur then
+  local rc = tonumber(cur) + 1
+  redis.call('SET', new_ref_key, rc)
+  redis.call('ZADD', zset, now_ms(), new_md5)
+  return {1}  -- success, no victims
+end
+
+-- If not at capacity, just insert
+local size = redis.call('ZCARD', zset)
+if size < capacity then
+  redis.call('SET', new_ref_key, 1)
+  redis.call('ZADD', zset, now_ms(), new_md5)
+  return {1}  -- success, no victims
+end
+
+-- At capacity: try to evict up to max_try items with rc==0, but success if at least 1 is freed
+local max_try = math.max(1, math.floor(size * evict_fraction + 0.5))
 local victims = {}
+local freed = 0
 
--- 获取所有键并按LRU排序
+-- Scan from LRU (smallest score) to MRU
 local all_keys = redis.call('ZRANGE', zset, 0, -1, 'WITHSCORES')
 local i = 1
-
--- 查找引用计数为0的键作为逐出候选
-while #victims < need and i <= #all_keys do
-    local md5 = all_keys[i]
-    local ref_key = ref_prefix .. md5
-    local rc = redis.call('GET', ref_key)
-    
-    if rc and tonumber(rc) <= 0 then
-        table.insert(victims, md5)
-    end
-    i = i + 2  -- 跳过分数
+while freed < 1 and i <= #all_keys and #victims < max_try do
+  local md5 = all_keys[i]
+  local ref_key = ref_prefix .. md5
+  local v = redis.call('GET', ref_key)
+  if v and tonumber(v) <= 0 then
+    table.insert(victims, md5)
+    freed = freed + 1
+  end
+  i = i + 2  -- skip score
 end
 
--- 如果找到足够的候选，执行逐出
-if #victims >= need then
-    -- 删除受害者
-    for _, v in ipairs(victims) do
-        local ref_key = ref_prefix .. v
-        redis.call('DEL', ref_key)
-        redis.call('ZREM', zset, v)
-    end
-    
-    -- 插入新的md5
-    local ref_key = ref_prefix .. new_md5
-    redis.call('SET', ref_key, 0)
-    local now = redis.call('TIME')[1] * 1000
-    redis.call('ZADD', zset, now, new_md5)
-    
-    return {1, unpack(victims)}  -- success + victims
+if freed >= 1 then
+  -- delete victims
+  for _, v in ipairs(victims) do
+    redis.call('DEL', ref_prefix .. v)
+    redis.call('ZREM', zset, v)
+  end
+  -- insert new
+  redis.call('SET', new_ref_key, 1)
+  redis.call('ZADD', zset, now_ms(), new_md5)
+  return {1, unpack(victims)}
 else
-    return {0}  -- 逐出失败，没有足够的候选
+  -- no zero-ref items found
+  return {0}
 end
 """

lightllm/server/httpserver/manager.py

@@ -122,12 +122,6 @@ def __init__(
     async def _alloc_resource(self, items, uuids, token_nums, datas):
 
         while True:
-            # 检查这个图片在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:
@@ -212,8 +206,8 @@ async def _release_multimodal_resources(self, multimodal_params: MultimodalParam
                         audio.uuid = None
                         audio.token_id = None
                         audio.token_num = None
-                if ids_to_release:
-                    self.cache_client.root.release(ids_to_release)
+                # if ids_to_release:
+                #     self.cache_client.root.release(ids_to_release)
         return
 
     def tokens(self, prompt, multimodal_params, samping_params: SamplingParams, kwargs=None):
@@ -370,7 +364,7 @@ async def generate(
             # 对于还没有形成正式请求对象管理的多模态资源，需要单独自己释放
             # 已经放入到 req_id_to_out_inf 中的请求对象，由统一的回收循环
             # 进行回收。
-            if group_request_id not in self.req_id_to_out_inf and self.args.run_mode != "llm_only":
+            if group_request_id not in self.req_id_to_out_inf:
                 await self._release_multimodal_resources(multimodal_params)
             await self.abort(group_request_id)
             raise e
@@ -410,7 +404,7 @@ async def get_image_embeding(
             visual_req_status = GroupReqObjs(group_request_id, multimodal_params, None, start_time)
 
             await self.transfer_to_next_module_or_node(
-                None, sampling_params, original_multimodal_params, visual_req_status
+                None, sampling_params, original_multimodal_params, visual_req_status, embeding_only=True
             )
 
         except Exception as e:
@@ -513,6 +507,7 @@ async def transfer_to_next_module_or_node(
         sampling_params: SamplingParams,
         original_multimodal_params: MultimodalParams,
         group_req_objs: Optional[GroupReqObjs] = None,
+        embeding_only: Optional[bool] = False,
     ):
         # 多节点纯tp 运行模式下，master 节点需要将请求转发给slave节点.
         if self.is_multinode_tp_master:
@@ -522,19 +517,21 @@ async def transfer_to_next_module_or_node(
                     protocol=pickle.HIGHEST_PROTOCOL,
                 )
 
-        await self.transfer_to_next_module(group_req_objs)
+        await self.transfer_to_next_module(group_req_objs, embeding_only)
         return
 
     async def transfer_to_next_module(
         self,
         group_req_objs: Optional[GroupReqObjs] = None,
+        embeding_only: Optional[bool] = False,
     ):
 
         if self.pd_mode.is_P_or_NORMAL():
             if self.enable_multimodal:
                 await self.vit_manager.send_to_vit(
                     group_req_objs.to_group_req_index(),
                     protocol=pickle.HIGHEST_PROTOCOL,
+                    embeding_only=embeding_only,
                 )
 
             if not self.enable_multimodal or self.args.enable_remote_vit: