Repository: LMCache/LMCache
File: lmcache/v1/storage_backend/local_cpu_backend.py (Python)

Code before:
        self.cache_policy = get_cache_policy(config.cache_policy)
        self.hot_cache = self.cache_policy.init_mutable_mapping()

        self.use_hot = config.local_cpu
        self.lookup_server = lookup_server
        self.memory_allocator = memory_allocator
        self.lmcache_worker = lmcache_worker
        self.instance_id = config.lmcache_instance_id
        self.cpu_lock = threading.Lock()

        self.stream = torch.cuda.Stream()

        self.stats_monitor = LMCStatsMonitor.GetOrCreate()
        self.usage = 0

        self.layerwise = config.use_layerwise
        self.enable_blending = config.enable_blending

        # to help maintain suffix -> prefix order in the dict
        # assumption: only one request is looked up at a time
        # (only one worker per cache engine)
        self.keys_in_request: List[CacheEngineKey] = []
        self._setup_metrics()

    def _setup_metrics(self):
        Gauge(
            "lmcache:local_cpu_hot_cache_count",
            "The size of the hot_cache",
            ["uuid"],
            _labelvalues=[str(uuid.uuid4())],
        ).set_function(lambda: len(self.hot_cache))
        Gauge(
            "lmcache:local_cpu_keys_in_request_count",
            "The size of the keys_in_request",
            ["uuid"],
            _labelvalues=[str(uuid.uuid4())],
        ).set_function(lambda: len(self.keys_in_request))

    def __str__(self):
        return self.__class__.__name__

    def contains(self, key: CacheEngineKey, pin: bool = False) -> bool:
        with self.cpu_lock:
            if key not in self.hot_cache:
                return False
            if pin:
                self.hot_cache[key].pin()
                # vllm lookup sets pin to True
                self.keys_in_request.append(key)
            return True


Review comment (refactor):
Also, is it possible to add a helper function in `PrometheusLogger` for other modules to register metrics?

The goal of such a helper function is to make sure that the metric labels (right now include `model_id` and `worker_id`) can be correctly added. (See the function `PrometheusLogger._metadata_to_labels`)

Code after:

        self.use_hot = config.local_cpu
        self.lookup_server = lookup_server
        self.memory_allocator = memory_allocator
        self.lmcache_worker = lmcache_worker
        self.instance_id = config.lmcache_instance_id
        self.cpu_lock = threading.Lock()

        self.stream = torch.cuda.Stream()

        self.stats_monitor = LMCStatsMonitor.GetOrCreate()

        self.layerwise = config.use_layerwise
        self.enable_blending = config.enable_blending

        # to help maintain suffix -> prefix order in the dict
        # assumption: only one request is looked up at a time
        # (only one worker per cache engine)
        self.keys_in_request: List[CacheEngineKey] = []
        self._setup_metrics()

    def _setup_metrics(self):
        prometheus_logger = PrometheusLogger.GetInstanceOrNone()
        if prometheus_logger is not None:
            prometheus_logger.local_cpu_hot_cache_count.set_function(
                lambda: len(self.hot_cache)
            )
            prometheus_logger.local_cpu_keys_in_request_count.set_function(
                lambda: len(self.keys_in_request)
            )

    def __str__(self):
        return self.__class__.__name__

    def contains(self, key: CacheEngineKey, pin: bool = False) -> bool:
        with self.cpu_lock:
            if key not in self.hot_cache:
                return False
            if pin:
                self.hot_cache[key].pin()
                # vllm lookup sets pin to True
                self.keys_in_request.append(key)
            return True

    def touch_cache(self):
        # flip the order of the keys in the request
        with self.cpu_lock:
            for key in reversed(self.keys_in_request):
                self.cache_policy.update_on_hit(key, self.hot_cache)
            self.keys_in_request = []