fix

Author: niushengxiao

lightllm/server/core/objs/start_args_type.py

@@ -106,7 +106,7 @@ class StartArgs:
     pd_node_id: int = field(default=-1)
     enable_cpu_cache: bool = field(default=False)
     cpu_cache_storage_size: float = field(default=2)
-    cpu_cache_token_page_size: int = field(default=256)
+    cpu_cache_token_page_size: int = field(default=64)
     enable_disk_cache: bool = field(default=False)
     disk_cache_storage_size: float = field(default=10)
     # zmp ports

lightllm/server/multi_level_kv_cache/shm_objs.py

@@ -261,10 +261,11 @@ def init(self):
         self.value = -1
 
 
-def _create_shm(name: str, byte_size: int):
+def _create_shm(name: str, byte_size: int, auto_cleanup: bool = False):
     try:
         shm = shared_memory.SharedMemory(name=name, create=True, size=byte_size)
-        register_posix_shm_for_cleanup(name)
+        if auto_cleanup:
+            register_posix_shm_for_cleanup(name)
         logger.info(f"create lock shm {name}")
     except:
         shm = shared_memory.SharedMemory(name=name, create=False, size=byte_size)

lightllm/server/router/dynamic_prompt/shared_arr.py

@@ -5,7 +5,6 @@
 from multiprocessing import shared_memory
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.shm_utils import create_or_link_shm
-from lightllm.utils.auto_shm_cleanup import register_posix_shm_for_cleanup
 
 logger = init_logger(__name__)
 
@@ -15,27 +14,6 @@ def __init__(self, name, shape, dtype):
         dtype_byte_num = np.array([1], dtype=dtype).dtype.itemsize
         dest_size = np.prod(shape) * dtype_byte_num
         self.shm = create_or_link_shm(name, dest_size)
-        try:
-            shm = shared_memory.SharedMemory(name=name, create=True, size=dest_size)
-            logger.info(f"create shm {name}")
-            register_posix_shm_for_cleanup(name)
-        except:
-            shm = shared_memory.SharedMemory(name=name, create=False, size=dest_size)
-            logger.info(f"link shm {name}")
-
-        if shm.size != dest_size:
-            logger.info(f"size not same, unlink shm {name} and create again")
-            shm.unlink()
-            shm.close()
-            try:
-                shm = shared_memory.SharedMemory(name=name, create=True, size=dest_size)
-                logger.info(f"create shm {name}")
-                register_posix_shm_for_cleanup(name)
-            except Exception as e:
-                shm = shared_memory.SharedMemory(name=name, create=False, size=dest_size)
-                logger.info(f"error {str(e)} to link shm {name}")
-
-        self.shm = shm  # SharedMemory 对象一定要被持有，否则会被释放
         self.arr = np.ndarray(shape, dtype=dtype, buffer=self.shm.buf)
 
 

lightllm/server/router/model_infer/mode_backend/base_backend.py

@@ -367,9 +367,9 @@ def _read_reqs_buffer_and_init_reqs(self):
                 else:
                     assert False, f"error type {type(obj)}"
             if init_reqs:
-                self._init_reqs(reqs=init_reqs)
-                if self.args.enable_cpu_cache:
-                    self._fill_cpu_cache_to_reqs(req_ids=init_reqs)
+                req_ids = self._init_reqs(reqs=init_reqs)
+                if self.args.enable_cpu_cache and req_ids:
+                    self._fill_cpu_cache_to_reqs(req_ids=req_ids)
         return
 
     def _read_nixl_trans_io_buffer_and_update_req_status(self):
@@ -403,7 +403,6 @@ def _read_nixl_trans_io_buffer_and_update_req_status(self):
                                     is_master_in_dp=self.is_master_in_dp,
                                     nixl_prefill_chuncked_handle_func=None,
                                 )
-
         return
 
     # 一些可以复用的通用功能函数

lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py

@@ -35,14 +35,13 @@ def wait_for_init(self):
         if attach_shm_handle is not None:
             attach_shm_handle.wait()
 
-    def _compute_full_sequence_hash(self, req: InferReq):
+    def _compute_sequence_hash(self, req: InferReq):
+        # 综合考虑后只对prompt做缓存管理，不包含decode内容，这里与radix cache不一致
+        if not req.shm_req.token_hash_list.is_empty():
+            return req.shm_req.token_hash_list.get_all()
+
         input_tokens = req.shm_req.get_prompt_ids()
-        total_len = req.shm_req.input_len + req.shm_req.shm_cur_output_len
-        if total_len > req.shm_req.input_len:
-            full_sequence = req.shm_req.shm_prompt_ids.arr[:total_len].tolist()
-        else:
-            full_sequence = input_tokens
-        return compute_token_list_hash(full_sequence, self.args.cpu_cache_token_page_size)
+        return compute_token_list_hash(input_tokens, self.args.cpu_cache_token_page_size)
 
     def handle_finished_reqs(self, finished_reqs: List[InferReq]) -> List[InferReq]:
         """
@@ -98,9 +97,9 @@ def _start_kv_cache_offload_task(
     ) -> Optional["TransTask"]:
         with torch.cuda.stream(cpu_kv_cache_stream):
             if self.backend.is_master_in_dp:
-                all_token_hash_list = self._compute_full_sequence_hash(req)
+                token_hash_list = self._compute_sequence_hash(req)
                 block_size = req.cur_kv_len // self.args.cpu_cache_token_page_size
-                move_block_size = min(block_size, len(all_token_hash_list))
+                move_block_size = min(block_size, len(token_hash_list))
 
                 if move_block_size == 0:
                     dist.broadcast_object_list([0], group=self.gloo_group, group_src=0)
@@ -110,7 +109,7 @@ def _start_kv_cache_offload_task(
                 try:
                     self.cpu_cache_client.lock.acquire_sleep1ms()
                     page_list, ready_list = self.cpu_cache_client.allocate_pages(
-                        all_token_hash_list[:move_block_size],
+                        token_hash_list[:move_block_size],
                         disk_offload_enable=self.args.enable_disk_cache,
                     )
                 finally:

lightllm/utils/auto_shm_cleanup.py

@@ -2,6 +2,7 @@
 import ctypes
 import atexit
 import signal
+import threading
 import psutil
 from multiprocessing import shared_memory
 from typing import Set, Optional
@@ -11,6 +12,10 @@
 
 
 class AutoShmCleanup:
+    """
+    自动清理 System V 和 POSIX 共享内存
+    shared_memory.SharedMemory虽然有自动请理功能，但如果自动清理时仍有进程占用会清理失败，这里可做最后兜底清理
+    """
     def __init__(self):
         self.libc = None
         self._init_libc()
@@ -19,6 +24,7 @@ def __init__(self):
         self.registered_shm_ids = []
         # POSIX
         self.registered_posix_shm_names = []
+        self.signal_handlers_registered = False
         self._register_handlers_for_cleanup()
 
     def _init_libc(self):
@@ -34,8 +40,14 @@ def _init_libc(self):
 
     def _register_handlers_for_cleanup(self):
         atexit.register(self._cleanup)
+        self.register_signal_handlers()
+
+    def register_signal_handlers(self):
+        if self.signal_handlers_registered or not threading.current_thread() is threading.main_thread():
+            return
         for sig in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP):
             signal.signal(sig, self._signal_cleanup_handler)
+        self.signal_handlers_registered = True
 
     def _signal_cleanup_handler(self, signum, frame):
         self._cleanup()
@@ -110,6 +122,7 @@ def get_auto_cleanup() -> AutoShmCleanup:
     global _auto_cleanup
     if _auto_cleanup is None:
         _auto_cleanup = AutoShmCleanup()
+    _auto_cleanup.register_signal_handlers()
     return _auto_cleanup
 
 
@@ -118,4 +131,4 @@ def register_sysv_shm_for_cleanup(key: int, shmid: Optional[int] = None):
 
 
 def register_posix_shm_for_cleanup(name: str):
-    get_auto_cleanup().register_posix_shm(name)
+    get_auto_cleanup().register_posix_shm(name)

lightllm/utils/shm_utils.py

@@ -1,11 +1,12 @@
 from multiprocessing import shared_memory
 from filelock import FileLock
 from lightllm.utils.log_utils import init_logger
+from lightllm.utils.auto_shm_cleanup import register_posix_shm_for_cleanup
 
 logger = init_logger(__name__)
 
 
-def create_or_link_shm(name, expected_size, force_mode=None):
+def create_or_link_shm(name, expected_size, force_mode=None, auto_cleanup=False):
     """
     Args:
         name: name of the shared memory
@@ -26,15 +27,15 @@ def create_or_link_shm(name, expected_size, force_mode=None):
 
     if force_mode == "create":
         with FileLock(lock_name):
-            return _force_create_shm(name, expected_size)
+            return _force_create_shm(name, expected_size, auto_cleanup)
     elif force_mode == "link":
         return _force_link_shm(name, expected_size)
     else:
         with FileLock(lock_name):
-            return _smart_create_or_link_shm(name, expected_size)
+            return _smart_create_or_link_shm(name, expected_size, auto_cleanup)
 
 
-def _force_create_shm(name, expected_size):
+def _force_create_shm(name, expected_size, auto_cleanup):
     """强制创建新的共享内存"""
     try:
         existing_shm = shared_memory.SharedMemory(name=name)
@@ -45,6 +46,8 @@ def _force_create_shm(name, expected_size):
 
     # 创建新的共享内存
     shm = shared_memory.SharedMemory(name=name, create=True, size=expected_size)
+    if auto_cleanup:
+        register_posix_shm_for_cleanup(name)
     return shm
 
 
@@ -62,12 +65,12 @@ def _force_link_shm(name, expected_size):
         raise e
 
 
-def _smart_create_or_link_shm(name, expected_size):
+def _smart_create_or_link_shm(name, expected_size, auto_cleanup):
     """优先连接，不存在则创建"""
     try:
         shm = _force_link_shm(name=name, expected_size=expected_size)
         return shm
     except:
         pass
 
-    return _force_create_shm(name=name, expected_size=expected_size)
+    return _force_create_shm(name=name, expected_size=expected_size, auto_cleanup=auto_cleanup)