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

Code before:
                self.sockets[i].send_multipart(msg_buf, copy=False)

            # TODO(Jiayi): we can use zmq poll to optimize a bit
            for i in range(self.num_ranks):
                failed_rank = i
                resp = self.sockets[i].recv()
                result = int.from_bytes(resp, "big")
                results.append(result)
        except zmq.Again as e:
            logger.error(
                f"Timeout occurred for rank {failed_rank}, "
                f"recreating all sockets. Error: {str(e)}"
            )
            self._recreate_socket()
            return 0
        except zmq.ZMQError as e:
            logger.error(
                f"ZMQ error for rank {failed_rank}: {str(e)}, recreating all sockets"
            )
            self._recreate_socket()
            return 0

        assert len(results) == self.num_ranks
        if len(set(results)) > 1:
            logger.warning(
                f"Lookup results (number of hit tokens) differ "
                f"across (TP and PP) ranks: {results}."
            )
        # NOTE: it is possible that the number of hit tokens is different
        # across (TP and PP) ranks, so we can use the minimum value as the
        # number of hit tokens.
        num_hit_toks = min(results)
        self.reqs_status[lookup_id] = num_hit_toks

        return num_hit_toks

    def clear_lookup_status(self, lookup_id: str) -> None:
        self.reqs_status.pop(lookup_id, None)

    def supports_producer_reuse(self) -> bool:
        """Return True as LMCacheLookupClient supports producer kvcache reuse"""
        return True

    def close(self):
        for socket in self.sockets:
            try:
                socket.close(linger=0)
            except Exception as e:
                logger.warning(f"Error closing socket: {e}")

        try:

Review comment (bug):
nit: we can do similar updates here:
```suggestion
            logger.error(
                f"ZMQ error for rank %d: %s, recreating all sockets",
                failed_rank, e
            )
```

Code after:
        failed_rank = -1
        try:
            for i in range(self.num_ranks):
                failed_rank = i
                self.sockets[i].send_multipart(msg_buf, copy=False)

            # TODO(Jiayi): we can use zmq poll to optimize a bit
            for i in range(self.num_ranks):
                failed_rank = i
                resp = self.sockets[i].recv()
                result = int.from_bytes(resp, "big")
                results.append(result)
        except zmq.Again as e:
            logger.error(
                "Timeout occurred for rank %s, recreating all sockets. Error: %s",
                failed_rank,
                e,
            )
            self._recreate_socket()
            return 0
        except zmq.ZMQError as e:
            logger.error(
                "ZMQ error for rank %s: %s, recreating all sockets",
                failed_rank,
                e,
            )
            self._recreate_socket()
            return 0

        assert len(results) == self.num_ranks
        if len(set(results)) > 1:
            logger.warning(
                "Lookup results (number of hit tokens) differ "
                "across (TP and PP) ranks: %s.",
                results,
            )
        # NOTE: it is possible that the number of hit tokens is different
        # across (TP and PP) ranks, so we can use the minimum value as the
        # number of hit tokens.
        num_hit_toks = min(results)
        self.reqs_status[lookup_id] = num_hit_toks

        return num_hit_toks

    def clear_lookup_status(self, lookup_id: str) -> None:
        self.reqs_status.pop(lookup_id, None)

    def supports_producer_reuse(self) -> bool:
        """Return True as LMCacheLookupClient supports producer kvcache reuse"""
        return True