Repository: vllm-project/aibrix
File: python/aibrix_kvcache/aibrix_kvcache/l2/connectors/infinistore.py (Python)

Code before:

    @contextmanager
    def _tcp_conn(self):
        assert self.config.connection_type == infinistore.TYPE_TCP
        conn = None
        try:
            conn = self.tcp_conns.get()
            yield conn
        finally:
            if conn is not None:
                self.tcp_conns.put(conn)

    @Status.capture_exception
    def open(self) -> Status:
        """Open a connection."""
        if self.config.connection_type == infinistore.TYPE_RDMA:
            if self.rdma_conn is None:
                self.rdma_conn = infinistore.InfinityConnection(self.config)
                self.rdma_conn.connect()
        else:
            for _ in range(self.tcp_conns.maxsize):
                tcp_conn = infinistore.InfinityConnection(self.config)
                tcp_conn.connect()
                self.tcp_conns.put(tcp_conn)
        return Status.ok()

    @Status.capture_exception
    def close(self) -> Status:
        """Close a connection."""
        if self.rdma_conn is not None:
            self.rdma_conn.close()
            self.rdma_conn = None

        if self.config.connection_type == infinistore.TYPE_TCP:
            while not self.tcp_conns.empty():
                tcp_conn = self.tcp_conns.get()
                if tcp_conn is not None:
                    tcp_conn.close()

        return Status.ok()

    @Status.capture_exception
    def register_slabs(self, slabs: List[torch.Tensor]) -> Status:
        assert self.rdma_conn is not None
        for slab in slabs:
            addr = slab.data_ptr()
            length = slab.numel()
            ret = self.rdma_conn.register_mr(addr, length)
            if ret != 0:
                return Status(StatusCodes.INVALID)
        return Status.ok()

Review comment (refactor):
seems the InfinityConnection abstraction is the connection so we have to manage the pool & connections outside? `self.tcp_conns.put(tcp_conn)`  

Code after:
        assert self.config.connection_type == infinistore.TYPE_TCP
        assert self.tcp_conns is not None
        conn = None
        try:
            conn = self.tcp_conns.get()
            yield conn
        finally:
            if conn is not None:
                self.tcp_conns.put(conn)

    @Status.capture_exception
    def open(self) -> Status:
        """Open a connection."""
        if self.config.connection_type == infinistore.TYPE_RDMA:
            if self.rdma_conn is None:
                self.rdma_conn = infinistore.InfinityConnection(self.config)
                self.rdma_conn.connect()
        else:
            assert hasattr(self._executor, "_max_workers")
            assert self.tcp_conns is None

            self.tcp_conns = Queue(
                maxsize=self._executor._max_workers  # type: ignore
            )
            for _ in range(self.tcp_conns.maxsize):
                tcp_conn = infinistore.InfinityConnection(self.config)
                tcp_conn.connect()
                self.tcp_conns.put(tcp_conn)
        return Status.ok()

    @Status.capture_exception
    def close(self) -> Status:
        """Close a connection."""
        if self.rdma_conn is not None:
            self.rdma_conn.close()
            self.rdma_conn = None

        if self.config.connection_type == infinistore.TYPE_TCP:
            if self.tcp_conns is not None:
                while not self.tcp_conns.empty():
                    tcp_conn = self.tcp_conns.get()
                    if tcp_conn is not None:
                        tcp_conn.close()
                del self.tcp_conns
                self.tcp_conns = None

        return Status.ok()

    @Status.capture_exception
    def register_slabs(self, slabs: List[torch.Tensor]) -> Status:
        assert self.rdma_conn is not None