feat: refine

Author: niushengxiao

lightllm/server/api_cli.py

@@ -537,4 +537,10 @@ def make_argument_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--disk_cache_storage_size", type=float, default=10, help="""The capacity of disk cache. GB used."""
     )
+    parser.add_argument(
+        "--disk_cache_dir",
+        type=str,
+        default=None,
+        help="""Directory used to persist disk cache data. Defaults to a temp directory when not set.""",
+    )
     return parser

lightllm/server/core/objs/atomic_lock.py

@@ -29,9 +29,13 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
     # acquire_sleep1ms 和 release 是某些特定场景下主动使用进行锁获取的操作函数
     def acquire_sleep1ms(self):
+        last_log_time = time.monotonic()
         with atomics.atomicview(buffer=self.shm.buf, atype=atomics.INT) as a:
             while not a.cmpxchg_weak(0, 1):
-                logger.warning("acquire_sleep1ms wait for 1ms")
+                now = time.monotonic()
+                if now - last_log_time >= 0.1:
+                    logger.warning("acquire_sleep1ms wait for 100ms")
+                    last_log_time = now
                 time.sleep(0.001)
                 pass
 

lightllm/server/core/objs/start_args_type.py

@@ -111,6 +111,7 @@ class StartArgs:
     cpu_cache_token_page_size: int = field(default=64)
     enable_disk_cache: bool = field(default=False)
     disk_cache_storage_size: float = field(default=10)
+    disk_cache_dir: Optional[str] = field(default=None)
     # zmp ports
     router_port: int = field(default=None)
     detokenization_port: int = field(default=None)

lightllm/server/httpserver/manager.py

@@ -614,15 +614,21 @@ async def _wait_to_token_package(
                             f"disk_prompt_cache_ratio:{disk_prompt_cache_ratio} "
                             f"mtp_avg_token_per_step:{mtp_avg_token_per_step} "
                         )
+                        if prompt_cache_len > 0:
+                            logger.info(
+                                f"[gpu cache hit] "
+                                f"prompt_cache_len:{prompt_cache_len} "
+                                f"prompt_cache_ratio:{prompt_cache_ratio} "
+                            )
                         if cpu_prompt_cache_len > 0:
                             logger.info(
-                                f"blueswhen "
+                                f"[cpu cache hit] "
                                 f"cpu_prompt_cache_len:{cpu_prompt_cache_len} "
                                 f"cpu_prompt_cache_ratio:{cpu_prompt_cache_ratio} "
                             )
                         if disk_prompt_cache_len > 0:
                             logger.info(
-                                f"blueswhen "
+                                f"[disk cache hit] "
                                 f"disk_prompt_cache_len:{disk_prompt_cache_len} "
                                 f"disk_prompt_cache_ratio:{disk_prompt_cache_ratio} "
                             )

lightllm/server/multi_level_kv_cache/cpu_cache_client.py

@@ -1,7 +1,7 @@
 import ctypes
 import torch
 import numpy as np
-from lightllm.utils.envs_utils import get_env_start_args, get_unique_server_name
+from lightllm.utils.envs_utils import get_env_start_args, get_unique_server_name, get_disk_cache_prompt_limit_length
 from typing import List, Optional, Tuple
 from lightllm.utils.log_utils import init_logger
 from .shm_objs import ShmDict, ShmLinkedList, _LinkedListItem, IntList
@@ -38,10 +38,12 @@ def __init__(self, only_create_meta_data: bool, init_shm_data: bool):
         return
 
     @staticmethod
+    # 负数编码，用于标记一个page index是一个offload group的第一个page
     def _encode_offload_head(page_index: int) -> int:
         return -(page_index + 1)
 
     @staticmethod
+    # 解码恢复page index，并返回该page index是否是一个offload group的第一个page
     def _decode_offload_value(value: int) -> Tuple[int, bool]:
         if value < 0:
             return -(value + 1), True
@@ -126,6 +128,19 @@ def update_pages_status_to_ready(
                     assert cur_page.ref_count > 0
                     cur_page.ref_count -= 1
 
+        # 控制prompt长度，较短的prompt不进行disk offload
+        limit_length = get_disk_cache_prompt_limit_length()
+        if (
+            disk_offload_enable
+            and offload_candidates
+            and len(page_list) * self.args.cpu_cache_token_page_size < limit_length
+        ):
+            logger.info(
+                f"skip disk offload for small page, " f"length = {len(page_list) * self.args.cpu_cache_token_page_size}"
+            )
+            self.mark_pages_recyclable(page_list=offload_candidates)
+            return
+
         if disk_offload_enable and offload_candidates:
             for idx, page_index in enumerate(offload_candidates):
                 if idx == 0:
@@ -225,11 +240,19 @@ def recycle_pages(self, page_list: List[int]):
             if page_index == -1:
                 continue
             cur_page: _CpuPageStatus = self.page_items.get_item_by_index(page_index)
-            cur_page.del_self_from_list()
-            if not cur_page.is_empty() and cur_page.hash_key != 0:
+
+            if cur_page.ref_count != 0:
+                if cur_page.status == cur_page.LOADING and cur_page.ref_count == 1:
+                    cur_page.ref_count = 0
+                else:
+                    continue
+
+            if cur_page.hash_key != 0:
                 existing_index = self.page_hash_dict.get(cur_page.hash_key)
-                if existing_index is not None:
+                if existing_index is not None and existing_index == cur_page.self_index:
                     self.page_hash_dict.remove(cur_page.hash_key)
+
+            cur_page.del_self_from_list()
             cur_page.hash_key = 0
             cur_page.status = cur_page.EMPTY
             cur_page.ref_count = 0

lightllm/server/multi_level_kv_cache/disk_cache_worker.py

@@ -6,13 +6,20 @@
 from typing import List, Optional
 
 import torch
-
-from cache import PyLocalCacheService, PyState
 from lightllm.utils.envs_utils import get_unique_server_name
 from lightllm.utils.log_utils import init_logger
 
 logger = init_logger(__name__)
 
+try:
+    from cache import PyLocalCacheService, PyState
+except ImportError as e:
+    logger.error(
+        "Failed to import LightMem library. Please install it first.\n"
+        "You can install it by building from source: https://github.com/ModelTC/LightMem"
+    )
+    raise ImportError("LightMem library is required for disk cache functionality") from e
+
 
 @dataclass
 class _PagePayload:
@@ -23,21 +30,28 @@ class _PagePayload:
 class DiskCacheWorker:
     """Background worker that offloads CPU KV pages to disk using kvcache."""
 
-    def __init__(self, disk_cache_storage_size: float, cpu_cache_client):
+    def __init__(
+        self,
+        disk_cache_storage_size: float,
+        cpu_cache_client,
+        disk_cache_dir: Optional[str] = None,
+    ):
         self.cpu_cache_client = cpu_cache_client
         self._pages_all_idle = False
 
         assert disk_cache_storage_size > 0
         storage_size = int(disk_cache_storage_size * (1024 ** 3))
-        num_shard = 32
-        num_worker = 32
+        num_shard = 64
+        num_worker = 48
+        max_concurrent_write_tasks = 16
 
-        cache_dir = os.getenv("LIGHTLLM_DISK_CACHE_DIR")
+        cache_dir = disk_cache_dir
         if not cache_dir:
             cache_dir = os.path.join(tempfile.gettempdir(), f"lightllm_disk_cache_{get_unique_server_name()}")
         os.makedirs(cache_dir, exist_ok=True)
         cache_file = os.path.join(cache_dir, "cache_file")
 
+        self.max_concurrent_write_tasks = max_concurrent_write_tasks
         self._page_major_tensor = self._prepare_tensor(cpu_cache_client.cpu_kv_cache_tensor)
 
         self.service = PyLocalCacheService(
@@ -49,7 +63,7 @@ def __init__(self, disk_cache_storage_size: float, cpu_cache_client):
         )
 
         logger.info(
-            "blueswhen disk cache worker initialized: dir=%s size_bytes=%d shards=%d workers=%d pages_per_block=%d",
+            "disk cache worker initialized: dir=%s size_bytes=%d shards=%d workers=%d pages_per_block=%d",
             cache_dir,
             storage_size,
             num_shard,
@@ -63,35 +77,15 @@ def _prepare_tensor(self, tensor: torch.Tensor) -> torch.Tensor:
 
     def run(self) -> None:
         while True:
-            time.sleep(0.01)
+            time.sleep(0.1)
             payload_groups = self._gather_offload_payloads()
-            # self._log_idle_once()
             if not payload_groups:
                 continue
             for payloads in payload_groups:
                 if not payloads:
                     continue
                 self._persist_pages_to_disk(payloads)
 
-    def _log_idle_once(self) -> int:
-        locked_pages = 0
-        self.cpu_cache_client.lock.acquire_sleep1ms()
-        try:
-            for page_idx in range(self.cpu_cache_client.page_num):
-                page_item = self.cpu_cache_client.page_items.get_item_by_index(page_idx)
-                if not page_item.is_ready_recycle() or page_item.ref_count != 0:
-                    locked_pages += 1
-        finally:
-            self.cpu_cache_client.lock.release()
-
-        if locked_pages == 0:
-            if not self._pages_all_idle:
-                logger.info("blueswhen all cpu cache pages are idle and ready to reuse")
-            self._pages_all_idle = True
-        else:
-            self._pages_all_idle = False
-        return locked_pages
-
     def _gather_offload_payloads(self) -> List[List[_PagePayload]]:
         self.cpu_cache_client.lock.acquire_sleep1ms()
         try:
@@ -109,6 +103,7 @@ def _gather_offload_payloads(self) -> List[List[_PagePayload]]:
         finally:
             self.cpu_cache_client.lock.release()
 
+    # 数据写入磁盘
     def _persist_pages_to_disk(self, payloads: List[_PagePayload]) -> None:
         if not payloads:
             return
@@ -120,16 +115,21 @@ def _persist_pages_to_disk(self, payloads: List[_PagePayload]) -> None:
         kv_indexer = torch.tensor(page_indexes, dtype=torch.int32, device="cpu")
         query_result = self.service.query(tokens)
         if not all(query_result):
+            # 限制写入并发量，给读取操作留资源
+            while (
+                self.service.active_threads("r") and self.service.active_threads("w") >= self.max_concurrent_write_tasks
+            ):
+                time.sleep(0.001)
+
             task = self.service.create(tokens=tokens, kv_page_indexer=kv_indexer, mode="w")
-            while not task.ready():
+            # 数据安全即可结束等待，无需写入完成
+            while not task.data_safe():
                 time.sleep(0.001)
 
         self.cpu_cache_client.lock.acquire_sleep1ms()
         self.cpu_cache_client.update_pages_status_to_ready_recycle(page_list=page_indexes, deref=True)
         self.cpu_cache_client.lock.release()
 
-        # self._log_idle_once()
-
     def blocks_exist(self, tokens: List[int], start_pos: int = 0) -> bool:
         if not tokens or start_pos < 0 or start_pos >= len(tokens):
             return False
@@ -141,12 +141,18 @@ def blocks_exist(self, tokens: List[int], start_pos: int = 0) -> bool:
             return False
         return all(query_result[block_start:block_end])
 
+    # 从磁盘读取数据到内存
     def load_pages(self, tokens: List[int], page_indexes: List[int], start_pos: int = 0) -> bool:
         if not tokens or not page_indexes or len(tokens) != len(page_indexes):
             return False
         if start_pos < 0 or start_pos >= len(tokens):
             return False
 
+        # 检测当前是否有写操作在进行，若有则跳过本次load请求，暂时不用
+        # if self.service.active_threads("w") > 0:
+        #     logger.warning("disk cache worker is busy writing, skip load_pages")
+        #     return False
+
         kv_indexer = torch.tensor(page_indexes, dtype=torch.int32, device="cpu")
         task = self.service.create(tokens=tokens, kv_page_indexer=kv_indexer, mode="r", start_pos=start_pos)
         while not task.ready():

lightllm/server/multi_level_kv_cache/manager.py

@@ -14,7 +14,6 @@
 from lightllm.server.core.objs.io_objs import GroupReqIndexes
 from lightllm.utils.graceful_utils import graceful_registry
 from .cpu_cache_client import CpuKvCacheClient
-from .disk_cache_worker import DiskCacheWorker
 from lightllm.utils.log_utils import init_logger
 
 logger = init_logger(__name__)
@@ -45,9 +44,12 @@ def __init__(
         self.disk_cache_worker = None
         self.disk_cache_thread = None
         if self.args.enable_disk_cache:
+            from .disk_cache_worker import DiskCacheWorker
+
             self.disk_cache_worker = DiskCacheWorker(
                 disk_cache_storage_size=self.args.disk_cache_storage_size,
                 cpu_cache_client=self.cpu_cache_client,
+                disk_cache_dir=self.args.disk_cache_dir,
             )
             self.disk_cache_thread = threading.Thread(target=self.disk_cache_worker.run, daemon=True)
             self.disk_cache_thread.start()
@@ -71,8 +73,8 @@ def _handle_group_req_cpu_cache_match(self, group_req_indexes: GroupReqIndexes,
         if current_time - start_time >= self.cpu_cache_time_out:
             self.send_to_router.send_pyobj(group_req_indexes, protocol=pickle.HIGHEST_PROTOCOL)
             logger.warning(
-                f"blueswhen cpu cache match time out {current_time - start_time}s, "
-                "group_req_id: {group_req_indexes.group_req_id}"
+                f"cpu cache match time out {current_time - start_time}s, "
+                f"group_req_id: {group_req_indexes.group_req_id}"
             )
             return
 

lightllm/utils/envs_utils.py

@@ -194,3 +194,8 @@ def enable_radix_tree_timer_merge() -> bool:
 @lru_cache(maxsize=None)
 def get_radix_tree_merge_update_delta() -> int:
     return int(os.getenv("LIGHTLMM_RADIX_TREE_MERGE_DELTA", 6000))
+
+
+@lru_cache(maxsize=None)
+def get_disk_cache_prompt_limit_length():
+    return int(os.getenv("LIGHTLLM_DISK_CACHE_PROMPT_LIMIT_LENGTH", 10000))

requirements.txt

@@ -87,4 +87,8 @@ librosa==0.11.0
 cuda_bindings==12.9.0
 orjson==3.11.2
 setproctitle==1.3.6
-xxhash==3.6.0
+xxhash==3.6.0
+torchvision==0.23.0
+interegular
+partial_json_parser
+websockets