lmcache-cli 0.4.5.dev0__py3-none-any.whl

lmcache/v1/lookup_client/chunk_statistics_lookup_client.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# Standard
+from typing import Optional, Union
+import threading
+import time
+
+# Third Party
+import torch
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.observability import PrometheusLogger
+from lmcache.v1.config import LMCacheEngineConfig
+from lmcache.v1.lookup_client.abstract_client import LookupClientInterface
+from lmcache.v1.lookup_client.record_strategies import (
+    AsyncRecorder,
+    RecordStrategy,
+    create_record_strategy,
+)
+
+logger = init_logger(__name__)
+
+
+class ChunkStatisticsLookupClient(LookupClientInterface):
+    """Wrapper client that tracks chunk reuse statistics."""
+
+    def __init__(
+        self,
+        actual_lookup_client: LookupClientInterface,
+        config: LMCacheEngineConfig,
+    ):
+        self.actual_lookup_client = actual_lookup_client
+        self.lock = threading.RLock()
+        self.chunk_size = config.chunk_size
+        self.enabled = False
+        self.request_seen: set[str] = set()
+        self.lookup_time = 0.0
+        self.record_time = 0.0
+        self.check_exit_time = 0.0
+        self.statistics_start_time = 0.0
+        self.timeout_hours = config.chunk_statistics_auto_exit_timeout_hours
+        self.target_unique_chunks = (
+            config.chunk_statistics_auto_exit_target_unique_chunks
+        )
+        self.enable_auto_exit = (
+            self.timeout_hours > 0.0 or self.target_unique_chunks > 0
+        )
+        strategy: RecordStrategy = create_record_strategy(config)
+        self.recorder = AsyncRecorder(
+            strategy=strategy,
+            queue_capacity=config.get_extra_config_value(
+                "chunk_statistics_async_queue_capacity", 100000
+            ),
+            preprocess_in_caller=config.get_extra_config_value(
+                "chunk_statistics_async_preprocess_chunks", False
+            ),
+        )
+        self._setup_metrics()
+        if config.chunk_statistics_auto_start_statistics:
+            self.start_statistics()
+
+    def lookup_cache(self, lookup_id: str) -> Optional[int]:
+        return self.actual_lookup_client.lookup_cache(lookup_id)
+
+    def start_statistics(self) -> None:
+        with self.lock:
+            self.enabled = True
+            # Assign the start time while first recording
+            self.statistics_start_time = 0.0
+
+    def stop_statistics(self) -> None:
+        with self.lock:
+            self.enabled = False
+
+    def reset_statistics(self) -> None:
+        self.recorder.wait_for_completion(timeout=5.0)
+        with self.lock:
+            self.request_seen.clear()
+            self.recorder.reset()
+
+    def get_statistics(self) -> dict:
+        self.recorder.wait_for_completion(timeout=5.0)
+        with self.lock:
+            strategy_stats = self.recorder.get_statistics()
+            total_time = self.lookup_time + self.record_time + self.check_exit_time
+            overhead_time = self.record_time + self.check_exit_time
+            overhead_pct = (
+                (overhead_time / total_time * 100.0) if total_time > 0 else 0.0
+            )
+            result = {
+                "enabled": self.enabled,
+                "total_requests": len(self.request_seen),
+                "timing": {
+                    "lookup_time_seconds": self.lookup_time,
+                    "record_statistics_time_seconds": self.record_time,
+                    "check_exit_conditions_time_seconds": self.check_exit_time,
+                    "total_time_seconds": total_time,
+                    "overhead_time_seconds": overhead_time,
+                    "overhead_percentage": overhead_pct,
+                },
+                "total_chunks": strategy_stats.get("total_chunks", 0),
+                "unique_chunks": strategy_stats.get("unique_chunks", 0),
+                "duplicate_chunks": strategy_stats.get("duplicate_chunks", 0),
+                "reuse_rate": strategy_stats.get("reuse_rate", 0.0),
+                **{
+                    k: v
+                    for k, v in strategy_stats.items()
+                    if k in ("bloom_filter", "async_queue", "file_hash")
+                },
+            }
+            return result
+
+    def wait_for_async_processing(self, timeout: float = 5.0) -> bool:
+        return self.recorder.wait_for_completion(timeout)
+
+    def lookup(
+        self,
+        token_ids: Union[torch.Tensor, list[int]],
+        lookup_id: str,
+        request_configs: Optional[dict] = None,
+    ) -> Optional[int]:
+        start_time = time.time()
+        result = self.actual_lookup_client.lookup(
+            token_ids,
+            lookup_id,
+            request_configs,
+        )
+        lookup_elapsed = time.time() - start_time
+        with self.lock:
+            self.lookup_time += lookup_elapsed
+
+        if not self.enabled:
+            return result
+
+        with self.lock:
+            if lookup_id in self.request_seen:
+                return result
+            self.request_seen.add(lookup_id)
+
+        start_time = time.time()
+        self.recorder.record_async(token_ids, lookup_id)
+        record_elapsed = time.time() - start_time
+        with self.lock:
+            self.record_time += record_elapsed
+
+        start_time = time.time()
+        self._check_exit_conditions()
+        check_elapsed = time.time() - start_time
+        with self.lock:
+            self.check_exit_time += check_elapsed
+
+        return result
+
+    def clear_lookup_status(self, lookup_id: str) -> None:
+        self.actual_lookup_client.clear_lookup_status(lookup_id)
+
+    def supports_producer_reuse(self) -> bool:
+        return self.actual_lookup_client.supports_producer_reuse()
+
+    def close(self) -> None:
+        if self.enabled:
+            self.stop_statistics()
+        self.recorder.close()
+        self.actual_lookup_client.close()
+
+    def _check_exit_conditions(self) -> None:
+        if not self.enable_auto_exit:
+            return
+        if self.statistics_start_time == 0.0:
+            self.statistics_start_time = time.time()
+        stop_reason = None
+        if self.timeout_hours > 0.0:
+            elapsed_hours = (time.time() - self.statistics_start_time) / 3600.0
+            if elapsed_hours >= self.timeout_hours:
+                stop_reason = (
+                    f"Timeout: {elapsed_hours:.2f}h >= {self.timeout_hours:.2f}h"
+                )
+        if self.target_unique_chunks > 0:
+            unique = self.recorder.strategy.unique_chunks_count
+            if unique >= self.target_unique_chunks:
+                stop_reason = f"Target reached: {unique} >= {self.target_unique_chunks}"
+        if stop_reason:
+            self._trigger_stop(stop_reason)
+
+    def _trigger_stop(self, reason: str) -> None:
+        logger.warning("Auto-stop: %s", reason)
+        if self.enabled:
+            self.stop_statistics()
+
+    def _setup_metrics(self):
+        prometheus_logger = PrometheusLogger.GetInstanceOrNone()
+        if prometheus_logger is not None:
+            prometheus_logger.chunk_statistics_enabled.set_function(
+                lambda: 1.0 if self.enabled else 0.0
+            )
+            prometheus_logger.chunk_statistics_total_requests.set_function(
+                lambda: len(self.request_seen)
+            )
+            self.recorder.strategy.setup_metrics(prometheus_logger)

lmcache/v1/lookup_client/factory.py

@@ -0,0 +1,251 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import TYPE_CHECKING, Optional, Union
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.cache_engine import LMCacheEngine
+from lmcache.v1.config import LMCacheEngineConfig
+from lmcache.v1.lookup_client.abstract_client import LookupClientInterface
+from lmcache.v1.lookup_client.chunk_statistics_lookup_client import (
+    ChunkStatisticsLookupClient,
+)
+from lmcache.v1.lookup_client.hit_limit_lookup_client import HitLimitLookupClient
+from lmcache.v1.lookup_client.lmcache_lookup_client_bypass import (
+    LMCacheBypassLookupClient,
+)
+from lmcache.v1.lookup_client.mooncake_lookup_client import MooncakeLookupClient
+from lmcache.v1.metadata import LMCacheMetadata
+from lmcache.v1.rpc.zmq_transport import (
+    SocketParams,
+    ZmqReqRepClientTransport,
+    ZmqRouterServerTransport,
+)
+from lmcache.v1.rpc_utils import get_zmq_rpc_path_lmcache
+
+if TYPE_CHECKING:
+    # First Party
+    from lmcache.v1.lookup_client.lmcache_async_lookup_client import (
+        LMCacheAsyncLookupServer,
+    )
+    from lmcache.v1.lookup_client.lmcache_lookup_client import LMCacheLookupServer
+
+logger = init_logger(__name__)
+
+
+class LookupClientFactory:
+    """Factory for creating lookup clients and servers based on configuration."""
+
+    @staticmethod
+    def create_lookup_client(
+        config: LMCacheEngineConfig,
+        metadata: LMCacheMetadata,
+        lmcache_engine: Optional[LMCacheEngine] = None,
+    ) -> LookupClientInterface:
+        """
+        Create a lookup client based on the configuration.
+
+        Args:
+            config: The LMCache engine configuration
+            metadata: The LMCache engine metadata (includes engine_id,
+                world_size, kv_connector_extra_config)
+            lmcache_engine: Optional LMCacheEngine instance for
+                bypass lookup client
+
+        Returns:
+            A lookup client instance
+        """
+
+        client: LookupClientInterface
+        # Check if external_lookup_client is configured
+        if config.external_lookup_client is not None:
+            if config.enable_async_loading:
+                raise ValueError(
+                    "Asynchronous loading is not supported for external lookup clients."
+                )
+            client = LookupClientFactory._create_external_lookup_client(
+                config.external_lookup_client, config, metadata
+            )
+        else:
+            # First Party
+            from lmcache.v1.lookup_client.lmcache_async_lookup_client import (
+                LMCacheAsyncLookupClient,
+            )
+            from lmcache.v1.lookup_client.lmcache_lookup_client import (
+                LMCacheLookupClient,
+            )
+
+            # Check if bypass lookup is enabled and lmcache_engine is provided
+            if config.enable_scheduler_bypass_lookup and lmcache_engine is not None:
+                client = LMCacheBypassLookupClient(config, metadata, lmcache_engine)
+            elif config.enable_async_loading:
+                client = LMCacheAsyncLookupClient(config, metadata)
+            else:
+                transport = LookupClientFactory._create_zmq_client_transport(
+                    config, metadata
+                )
+                client = LMCacheLookupClient(config, metadata, transport)
+
+        if config.hit_miss_ratio is not None and 0 <= config.hit_miss_ratio <= 1:
+            client = HitLimitLookupClient(client, config)
+
+        # Wrap with ChunkStatisticsLookupClient if enabled
+        if config.enable_chunk_statistics:
+            client = ChunkStatisticsLookupClient(
+                client,
+                config,
+            )
+        return client
+
+    @staticmethod
+    def create_lookup_server(
+        lmcache_engine: LMCacheEngine,
+        metadata: LMCacheMetadata,
+    ) -> Optional[Union["LMCacheLookupServer", "LMCacheAsyncLookupServer"]]:
+        """
+        Create a lookup server based on the configuration.
+
+        Args:
+            lmcache_engine: The LMCache engine instance
+            metadata: The LMCache engine metadata (includes engine_id,
+                world_size, kv_connector_extra_config, worker_id)
+
+        Returns:
+            A lookup server instance, or None if no server should be created
+        """
+        config = lmcache_engine.config
+        assert isinstance(config, LMCacheEngineConfig), (
+            "LMCache v1 config is expected for lookup server and client"
+        )
+
+        lookup_server_worker_ids = config.get_lookup_server_worker_ids(
+            metadata.use_mla, metadata.world_size
+        )
+
+        if config.external_lookup_client is None and (
+            len(lookup_server_worker_ids) == 0
+            or metadata.worker_id in lookup_server_worker_ids
+        ):
+            # First Party
+            from lmcache.v1.lookup_client.lmcache_async_lookup_client import (
+                LMCacheAsyncLookupServer,
+            )
+            from lmcache.v1.lookup_client.lmcache_lookup_client import (
+                LMCacheLookupServer,
+            )
+
+            if config.enable_async_loading:
+                return LMCacheAsyncLookupServer(lmcache_engine, metadata)
+            else:
+                transport = LookupClientFactory._create_zmq_server_transport(metadata)
+                return LMCacheLookupServer(lmcache_engine, metadata, transport)
+
+        return None
+
+    @staticmethod
+    def _create_external_lookup_client(
+        external_lookup_uri: str,
+        config: LMCacheEngineConfig,
+        metadata: LMCacheMetadata,
+    ) -> LookupClientInterface:
+        """
+        Create an external lookup client based on the URI format.
+
+        Args:
+            external_lookup_uri: URI in format <scheme>://<address>
+            config: The LMCache engine configuration
+            metadata: The LMCache engine metadata
+
+        Returns:
+            A lookup client instance
+
+        Raises:
+            ValueError: If the URI format is unsupported
+        """
+        # Parse URI scheme and address
+        if "://" not in external_lookup_uri:
+            raise ValueError(
+                f"Invalid external lookup client URI format: {external_lookup_uri}. "
+                "Expected format: <scheme>://<address>"
+            )
+
+        scheme, address = external_lookup_uri.split("://", 1)
+
+        # Route to appropriate client based on scheme
+        if scheme == "mooncakestore":
+            return LookupClientFactory._create_mooncake_lookup_client(
+                address, config, metadata
+            )
+        else:
+            raise ValueError(
+                f"Unsupported external lookup client scheme: {scheme}. "
+                "Supported schemes: mooncakestore"
+            )
+
+    @staticmethod
+    def _create_mooncake_lookup_client(
+        master_address: str,
+        config: LMCacheEngineConfig,
+        metadata: LMCacheMetadata,
+    ) -> "MooncakeLookupClient":
+        """Create a MooncakeLookupClient instance."""
+        # First Party
+        from lmcache.v1.lookup_client.mooncake_lookup_client import (
+            MooncakeLookupClient,
+        )
+
+        return MooncakeLookupClient(config, metadata, master_address)
+
+    @staticmethod
+    def _create_zmq_client_transport(
+        config: LMCacheEngineConfig,
+        metadata: LMCacheMetadata,
+    ) -> ZmqReqRepClientTransport:
+        """Create a ZMQ REQ-REP client transport."""
+        kv_extra = metadata.kv_connector_extra_config or {}
+        rpc_port = kv_extra.get("lmcache_rpc_port", 0)
+        assert metadata.engine_id is not None, (
+            "engine_id is required for RPC communication"
+        )
+
+        lookup_ids = config.get_lookup_server_worker_ids(
+            metadata.use_mla, metadata.world_size
+        )
+        ranks = lookup_ids if len(lookup_ids) > 0 else list(range(metadata.world_size))
+
+        socket_params = [
+            SocketParams(
+                socket_path=get_zmq_rpc_path_lmcache(
+                    metadata.engine_id,
+                    "lookup",
+                    rpc_port,
+                    rank,
+                ),
+                rank=rank,
+            )
+            for rank in ranks
+        ]
+        return ZmqReqRepClientTransport(
+            socket_params=socket_params,
+            timeout_ms=config.lookup_timeout_ms,
+        )
+
+    @staticmethod
+    def _create_zmq_server_transport(
+        metadata: LMCacheMetadata,
+    ) -> ZmqRouterServerTransport:
+        """Create a ZMQ ROUTER server transport."""
+        kv_extra = metadata.kv_connector_extra_config or {}
+        rpc_port = kv_extra.get("lmcache_rpc_port", 0)
+        assert metadata.engine_id is not None, (
+            "engine_id is required for RPC communication"
+        )
+        socket_path = get_zmq_rpc_path_lmcache(
+            metadata.engine_id,
+            "lookup",
+            rpc_port,
+            metadata.worker_id,
+        )
+        return ZmqRouterServerTransport(
+            socket_path=socket_path,
+        )

lmcache/v1/lookup_client/hit_limit_lookup_client.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Optional, Union
+
+# Third Party
+import torch
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.config import LMCacheEngineConfig
+from lmcache.v1.lookup_client.abstract_client import LookupClientInterface
+
+logger = init_logger(__name__)
+
+
+"""
+HitLimitLookupClient now is used for test, when lookup is called, cal the cache hit,
+- if the cache hit <= (1 - hit_miss_ratio), direct return the result
+- if the cache hit > (1 - hit_miss_ratio), re-compute the result by hit_miss_ratio
+"""
+
+
+class HitLimitLookupClient(LookupClientInterface):
+    def __init__(
+        self, actual_lookup_client: LookupClientInterface, config: LMCacheEngineConfig
+    ):
+        assert config.hit_miss_ratio is not None and 0 <= config.hit_miss_ratio <= 1
+        self.actual_lookup_client = actual_lookup_client
+        self.config = config
+        logger.info(
+            "create HitLimitLookupClient succeed, "
+            "the hit ratio upper is %s, chunk size is %s",
+            1 - config.hit_miss_ratio,
+            config.chunk_size,
+        )
+
+    def lookup_cache(self, lookup_id: str) -> Optional[int]:
+        return self.actual_lookup_client.lookup_cache(lookup_id)
+
+    def lookup(
+        self,
+        token_ids: Union[torch.Tensor, list[int]],
+        lookup_id: str,
+        request_configs: Optional[dict] = None,
+    ) -> Optional[int]:
+        # get real hit tokens
+        result = self.actual_lookup_client.lookup(
+            token_ids,
+            lookup_id,
+            request_configs,
+        )
+        if result is not None:
+            total_tokens_length = len(token_ids)
+            assert result <= total_tokens_length
+            current_hit_ratio = 0.0
+            if total_tokens_length > 0:
+                current_hit_ratio = result / total_tokens_length
+            # limit the hit tokens
+            hit_ratio_upper = 1 - self.config.hit_miss_ratio
+            if current_hit_ratio > hit_ratio_upper:
+                origin_result = result
+                # align to chunk size
+                chunk_size = self.config.chunk_size
+                new_result = (
+                    int(total_tokens_length * hit_ratio_upper)
+                    // chunk_size
+                    * chunk_size
+                )
+                # check again
+                result = min(result, new_result)
+                logger.debug(
+                    f"hit ratio upper: {hit_ratio_upper} is smaller than "
+                    f"the real hit ratio {current_hit_ratio}, "
+                    f"the origin result is {origin_result}, "
+                    f"the new result is {new_result}, the final result is {result}"
+                )
+        return result
+
+    def clear_lookup_status(self, lookup_id: str) -> None:
+        self.actual_lookup_client.clear_lookup_status(lookup_id)
+
+    def supports_producer_reuse(self) -> bool:
+        return self.actual_lookup_client.supports_producer_reuse()
+
+    def close(self) -> None:
+        self.actual_lookup_client.close()