Repository: LMCache/LMCache
File: lmcache/v1/lookup_client/lmcache_async_lookup_client.py (Python)

Code before:

        self.thread = threading.Thread(
            target=self.process_responses_from_workers, daemon=True
        )
        self.thread.start()

        # default backoff time
        self.lookup_backoff_time = 0.01
        if config.extra_config is not None:
            self.lookup_backoff_time = float(
                config.extra_config.get("lookup_backoff_time", self.lookup_backoff_time)
            )

    def lookup_cache(self, lookup_id: str) -> Optional[int]:
        with self.lock:
            return self.reqs_status.get(lookup_id, None)

    # TODO(Jiayi): Consider batching here
    def lookup(
        self,
        token_ids: Union[torch.Tensor, list[int]],
        lookup_id: str,
        request_configs: Optional[dict] = None,
        num_computed_tokens: int = 0,
    ) -> Optional[int]:
        # TODO num_computed_tokens is passed in
        # Due to complexity in preemption or eviction in async mode
        # We will handle skipping logic in future.

        # Check if any aborted lookups are finished, send cleanup messages
        self._cleanup_finished_aborted_lookups()

        with self.lock:
            # -1 indicates not found; None indicates ongoing.
            req_status = self.reqs_status.get(lookup_id, -1)
            if req_status is None:
                time.sleep(self.lookup_backoff_time)
                if (
                    time.time() - self.first_lookup_time[lookup_id]
                ) * 1000 > self.timeout_ms:
                    logger.warning(
                        (
                            "Request %s is still waiting for async lookup "
                            "after %d seconds, returning 0 lmcache cached tokens "
                            "so vllm can recompute"
                        ),
                        lookup_id,
                        self.timeout_ms // 1000,
                    )
                    self.first_lookup_time.pop(lookup_id, None)
                    return 0

Review comment (suggestion):
It would be nice if you could add your id to the todo comment.
Something like `TODO(zirui): num_computed_tokens is passed in`

Code after:
        self.running = True

        self.thread = threading.Thread(
            target=self.process_responses_from_workers, daemon=True
        )
        self.thread.start()

        # default backoff time
        self.lookup_backoff_time = 0.01
        if config.extra_config is not None:
            self.lookup_backoff_time = float(
                config.extra_config.get("lookup_backoff_time", self.lookup_backoff_time)
            )

    def lookup_cache(self, lookup_id: str) -> Optional[int]:
        """
        -1 means not found;
        None means ongoing;
        int >= 0 means number of hit tokens
        """
        # Check if any aborted lookups are finished, send cleanup messages
        self._cleanup_finished_aborted_lookups()

        with self.lock:
            if (req_status := self.reqs_status.get(lookup_id, -1)) == -1:
                self.reqs_status[lookup_id] = None
                self.first_lookup_time[lookup_id] = time.time()
            elif req_status is None:
                time.sleep(self.lookup_backoff_time)
                if (
                    time.time() - self.first_lookup_time[lookup_id]
                ) * 1000 > self.timeout_ms:
                    logger.warning(
                        (
                            "Request %s is still waiting for async lookup "
                            "after %d seconds, returning 0 lmcache cached tokens "
                            "so vllm can recompute"
                        ),
                        lookup_id,
                        self.timeout_ms // 1000,
                    )
                    self.first_lookup_time.pop(lookup_id, None)
                    return 0

            return req_status

    # TODO(Jiayi): Consider batching here
    def lookup(
        self,
        token_ids: Union[torch.Tensor, list[int]],
        lookup_id: str,