Connection pooling (#1071)

* Connection pooling in GRPC

* Undo auto formattings

* Make next_grpc_client() protected

* Prevent partial initialization of grpc clients on channel creation error

* also apply pool_size to rest client

* Fix using non clamped pool_size in httpx

* Review remarks

* fix: fix minor things (#1075)

---------

Co-authored-by: George <george.panchuk@qdrant.tech>

.gitignore

@@ -14,3 +14,4 @@ local_cache/*/*
 .python-version
 docs/source/examples/local_cache/*
 docs/source/examples/path/to/db/*
+.venv

qdrant_client/async_qdrant_client.py

@@ -72,8 +72,12 @@ class AsyncQdrantClient(AsyncQdrantFastembedMixin):
             Only use this if you can guarantee that you can resolve the thread safety outside QdrantClient.
         auth_token_provider: Callback function to get Bearer access token. If given, the function will be called before each request to get the token.
         check_compatibility: If `true` - check compatibility with the server version. Default: `true`
+        grpc_options: a mapping of gRPC channel options
+        cloud_inference: If `true` - do inference of `models.Document` and other models in Qdrant Cloud. Default: `False`.
+        local_inference_batch_size: inference batch size used by fastembed when using local inference with `models.Document` and other models.
+        pool_size: connection pool size, Default: None. Default value for gRPC connection pool is 3, rest default is
+            inherited from `httpx` (default: 100)
         **kwargs: Additional arguments passed directly into REST client initialization
-
     """
 
     def __init__(
@@ -97,6 +101,7 @@ class AsyncQdrantClient(AsyncQdrantFastembedMixin):
         cloud_inference: bool = False,
         local_inference_batch_size: Optional[int] = None,
         check_compatibility: bool = True,
+        pool_size: Optional[int] = None,
         **kwargs: Any,
     ):
         self._init_options = {
@@ -135,6 +140,7 @@ class AsyncQdrantClient(AsyncQdrantFastembedMixin):
                 grpc_options=grpc_options,
                 auth_token_provider=auth_token_provider,
                 check_compatibility=check_compatibility,
+                pool_size=pool_size,
                 **kwargs,
             )
             server_version = self._client.server_version

qdrant_client/async_qdrant_remote.py

@@ -54,6 +54,7 @@ from qdrant_client.uploader.uploader import BaseUploader
 
 class AsyncQdrantRemote(AsyncQdrantBase):
     DEFAULT_GRPC_TIMEOUT = 5
+    DEFAULT_GRPC_POOL_SIZE = 3
 
     def __init__(
         self,
@@ -71,6 +72,7 @@ class AsyncQdrantRemote(AsyncQdrantBase):
             Union[Callable[[], str], Callable[[], Awaitable[str]]]
         ] = None,
         check_compatibility: bool = True,
+        pool_size: Optional[int] = None,
         **kwargs: Any,
     ):
         super().__init__(**kwargs)
@@ -79,6 +81,10 @@ class AsyncQdrantRemote(AsyncQdrantBase):
         self._grpc_options = grpc_options or {}
         self._https = https if https is not None else api_key is not None
         self._scheme = "https" if self._https else "http"
+        self._pool_size: Optional[int] = None
+        if pool_size is not None:
+            pool_size = max(1, pool_size)
+            self._pool_size = pool_size
         self._prefix = prefix or ""
         if len(self._prefix) > 0 and self._prefix[0] != "/":
             self._prefix = f"/{self._prefix}"
@@ -115,6 +121,12 @@ class AsyncQdrantRemote(AsyncQdrantBase):
         if limits is None:
             if self._host in ["localhost", "127.0.0.1"]:
                 limits = httpx.Limits(max_connections=None, max_keepalive_connections=0)
+            elif self._pool_size is not None:
+                limits = httpx.Limits(max_connections=self._pool_size)
+        elif self._pool_size is not None:
+            raise ValueError(
+                f"`pool_size` and `limits` are mutually exclusive. `pool_size`: {pool_size}, `limit`: {limits}"
+            )
         http2 = kwargs.pop("http2", False)
         self._grpc_headers = []
         self._rest_headers = {k: v for (k, v) in kwargs.pop("metadata", {}).items()}
@@ -165,11 +177,12 @@ class AsyncQdrantRemote(AsyncQdrantBase):
         self.openapi_client: AsyncApis[AsyncApiClient] = AsyncApis(
             host=self.rest_uri, **self._rest_args
         )
-        self._grpc_channel = None
-        self._grpc_points_client: Optional[grpc.PointsStub] = None
-        self._grpc_collections_client: Optional[grpc.CollectionsStub] = None
-        self._grpc_snapshots_client: Optional[grpc.SnapshotsStub] = None
-        self._grpc_root_client: Optional[grpc.QdrantStub] = None
+        self._grpc_channel_pool: list[grpc.Channel] = []
+        self._grpc_points_client_pool: Optional[list[grpc.PointsStub]] = None
+        self._grpc_collections_client_pool: Optional[list[grpc.CollectionsStub]] = None
+        self._grpc_snapshots_client_pool: Optional[list[grpc.SnapshotsStub]] = None
+        self._grpc_root_client_pool: Optional[list[grpc.QdrantStub]] = None
+        self._grpc_client_next_index: int = 0
         self._closed: bool = False
         self.server_version = None
         if check_compatibility:
@@ -200,9 +213,10 @@ class AsyncQdrantRemote(AsyncQdrantBase):
         return self._closed
 
     async def close(self, grpc_grace: Optional[float] = None, **kwargs: Any) -> None:
-        if hasattr(self, "_grpc_channel") and self._grpc_channel is not None:
+        if hasattr(self, "_grpc_channel_pool") and len(self._grpc_channel_pool) > 0:
+            for channel in self._grpc_channel_pool:
                 try:
-                await self._grpc_channel.close(grace=grpc_grace)
+                    await channel.close(grace=grpc_grace)
                 except AttributeError:
                     show_warning(
                         message="Unable to close grpc_channel. Connection was interrupted on the server side",
@@ -232,11 +246,25 @@ class AsyncQdrantRemote(AsyncQdrantBase):
         )
         return (scheme, host, port, prefix)
 
+    def _get_grpc_pool_size(self) -> int:
+        """
+        Returns the pool size to use for GRPC connection pool.
+        This method should be preferred over accessing `self._pool_size` directly as it applies the
+        default value if no pool_size was provided.
+        """
+        if self._pool_size is not None:
+            return self._pool_size
+        else:
+            return self.DEFAULT_GRPC_POOL_SIZE
+
     def _init_grpc_channel(self) -> None:
         if self._closed:
             raise RuntimeError("Client was closed. Please create a new QdrantClient instance.")
-        if self._grpc_channel is None:
-            self._grpc_channel = get_channel(
+        try:
+            channel_pool = []
+            if len(self._grpc_channel_pool) == 0:
+                for _ in range(self._get_grpc_pool_size()):
+                    channel = get_channel(
                         host=self._host,
                         port=self._grpc_port,
                         ssl=self._https,
@@ -245,22 +273,41 @@ class AsyncQdrantRemote(AsyncQdrantBase):
                         compression=self._grpc_compression,
                         auth_token_provider=self._auth_token_provider,
                     )
+                    channel_pool.append(channel)
+                self._grpc_channel_pool = channel_pool
+        except Exception as e:
+            raise RuntimeError(f"Error initializing the grpc connection(s): {e}")
 
     def _init_grpc_points_client(self) -> None:
         self._init_grpc_channel()
-        self._grpc_points_client = grpc.PointsStub(self._grpc_channel)
+        self._grpc_points_client_pool = [
+            grpc.PointsStub(channel) for channel in self._grpc_channel_pool
+        ]
 
     def _init_grpc_collections_client(self) -> None:
         self._init_grpc_channel()
-        self._grpc_collections_client = grpc.CollectionsStub(self._grpc_channel)
+        self._grpc_collections_client_pool = [
+            grpc.CollectionsStub(channel) for channel in self._grpc_channel_pool
+        ]
 
     def _init_grpc_snapshots_client(self) -> None:
         self._init_grpc_channel()
-        self._grpc_snapshots_client = grpc.SnapshotsStub(self._grpc_channel)
+        self._grpc_snapshots_client_pool = [
+            grpc.SnapshotsStub(channel) for channel in self._grpc_channel_pool
+        ]
 
     def _init_grpc_root_client(self) -> None:
         self._init_grpc_channel()
-        self._grpc_root_client = grpc.QdrantStub(self._grpc_channel)
+        self._grpc_root_client_pool = [
+            grpc.QdrantStub(channel) for channel in self._grpc_channel_pool
+        ]
+
+    def _next_grpc_client(self) -> int:
+        current_index = self._grpc_client_next_index
+        self._grpc_client_next_index = (
+            self._grpc_client_next_index + 1
+        ) % self._get_grpc_pool_size()
+        return current_index
 
     @property
     def grpc_collections(self) -> grpc.CollectionsStub:
@@ -269,9 +316,10 @@ class AsyncQdrantRemote(AsyncQdrantBase):
         Returns:
             An instance of raw gRPC client, generated from Protobuf
         """
-        if self._grpc_collections_client is None:
+        if self._grpc_collections_client_pool is None:
             self._init_grpc_collections_client()
-        return self._grpc_collections_client
+        assert self._grpc_collections_client_pool is not None
+        return self._grpc_collections_client_pool[self._next_grpc_client()]
 
     @property
     def grpc_points(self) -> grpc.PointsStub:
@@ -280,9 +328,10 @@ class AsyncQdrantRemote(AsyncQdrantBase):
         Returns:
             An instance of raw gRPC client, generated from Protobuf
         """
-        if self._grpc_points_client is None:
+        if self._grpc_points_client_pool is None:
             self._init_grpc_points_client()
-        return self._grpc_points_client
+        assert self._grpc_points_client_pool is not None
+        return self._grpc_points_client_pool[self._next_grpc_client()]
 
     @property
     def grpc_snapshots(self) -> grpc.SnapshotsStub:
@@ -291,9 +340,10 @@ class AsyncQdrantRemote(AsyncQdrantBase):
         Returns:
             An instance of raw gRPC client, generated from Protobuf
         """
-        if self._grpc_snapshots_client is None:
+        if self._grpc_snapshots_client_pool is None:
             self._init_grpc_snapshots_client()
-        return self._grpc_snapshots_client
+        assert self._grpc_snapshots_client_pool is not None
+        return self._grpc_snapshots_client_pool[self._next_grpc_client()]
 
     @property
     def grpc_root(self) -> grpc.QdrantStub:
@@ -302,9 +352,10 @@ class AsyncQdrantRemote(AsyncQdrantBase):
         Returns:
             An instance of raw gRPC client, generated from Protobuf
         """
-        if self._grpc_root_client is None:
+        if self._grpc_root_client_pool is None:
             self._init_grpc_root_client()
-        return self._grpc_root_client
+        assert self._grpc_root_client_pool is not None
+        return self._grpc_root_client_pool[self._next_grpc_client()]
 
     @property
     def rest(self) -> AsyncApis[AsyncApiClient]:

qdrant_client/qdrant_client.py

@@ -73,8 +73,12 @@ class QdrantClient(QdrantFastembedMixin):
             Only use this if you can guarantee that you can resolve the thread safety outside QdrantClient.
         auth_token_provider: Callback function to get Bearer access token. If given, the function will be called before each request to get the token.
         check_compatibility: If `true` - check compatibility with the server version. Default: `true`
+        grpc_options: a mapping of gRPC channel options
+        cloud_inference: If `true` - do inference of `models.Document` and other models in Qdrant Cloud. Default: `False`.
+        local_inference_batch_size: inference batch size used by fastembed when using local inference with `models.Document` and other models.
+        pool_size: connection pool size, Default: None. Default value for gRPC connection pool is 3, rest default is
+            inherited from `httpx` (default: 100)
         **kwargs: Additional arguments passed directly into REST client initialization
-
     """
 
     def __init__(
@@ -98,6 +102,7 @@ class QdrantClient(QdrantFastembedMixin):
         cloud_inference: bool = False,
         local_inference_batch_size: Optional[int] = None,
         check_compatibility: bool = True,
+        pool_size: Optional[int] = None,
         **kwargs: Any,
     ):
         # Saving the init options to facilitate building AsyncQdrantClient from QdrantClient and vice versa.
@@ -142,6 +147,7 @@ class QdrantClient(QdrantFastembedMixin):
                 grpc_options=grpc_options,
                 auth_token_provider=auth_token_provider,
                 check_compatibility=check_compatibility,
+                pool_size=pool_size,
                 **kwargs,
             )
             server_version = self._client.server_version

qdrant_client/qdrant_remote.py

@@ -45,6 +45,7 @@ from qdrant_client.uploader.uploader import BaseUploader
 
 class QdrantRemote(QdrantBase):
     DEFAULT_GRPC_TIMEOUT = 5  # seconds
+    DEFAULT_GRPC_POOL_SIZE = 3
 
     def __init__(
         self,
@@ -62,6 +63,7 @@ class QdrantRemote(QdrantBase):
             Union[Callable[[], str], Callable[[], Awaitable[str]]]
         ] = None,
         check_compatibility: bool = True,
+        pool_size: Optional[int] = None,
         **kwargs: Any,
     ):
         super().__init__(**kwargs)
@@ -71,6 +73,13 @@ class QdrantRemote(QdrantBase):
         self._https = https if https is not None else api_key is not None
         self._scheme = "https" if self._https else "http"
 
+        # Pool size to use. This value should not be accessed directly; use _get_grpc_pool_size() instead.
+        self._pool_size: Optional[int] = None
+
+        if pool_size is not None:
+            pool_size = max(1, pool_size)  # Ensure pool_size is always > 0
+            self._pool_size = pool_size
+
         self._prefix = prefix or ""
         if len(self._prefix) > 0 and self._prefix[0] != "/":
             self._prefix = f"/{self._prefix}"
@@ -126,6 +135,14 @@ class QdrantRemote(QdrantBase):
                 # Disable keep-alive for local connections
                 # Cause in some cases, it may cause extra delays
                 limits = httpx.Limits(max_connections=None, max_keepalive_connections=0)
+            elif self._pool_size is not None:
+                # Set http connection pooling to `self._pool_size`, if no limits are specified.
+                limits = httpx.Limits(max_connections=self._pool_size)
+        elif self._pool_size is not None:
+            raise ValueError(
+                "`pool_size` and `limits` are mutually exclusive. "
+                f"`pool_size`: {pool_size}, `limit`: {limits}"
+            )
 
         http2 = kwargs.pop("http2", False)
         self._grpc_headers = []
@@ -193,11 +210,12 @@ class QdrantRemote(QdrantBase):
             **self._rest_args,
         )
 
-        self._grpc_channel = None
-        self._grpc_points_client: Optional[grpc.PointsStub] = None
-        self._grpc_collections_client: Optional[grpc.CollectionsStub] = None
-        self._grpc_snapshots_client: Optional[grpc.SnapshotsStub] = None
-        self._grpc_root_client: Optional[grpc.QdrantStub] = None
+        self._grpc_channel_pool: list[grpc.Channel] = []
+        self._grpc_points_client_pool: Optional[list[grpc.PointsStub]] = None
+        self._grpc_collections_client_pool: Optional[list[grpc.CollectionsStub]] = None
+        self._grpc_snapshots_client_pool: Optional[list[grpc.SnapshotsStub]] = None
+        self._grpc_root_client_pool: Optional[list[grpc.QdrantStub]] = None
+        self._grpc_client_next_index: int = 0  # The next index to use
 
         self._aio_grpc_points_client: Optional[grpc.PointsStub] = None
         self._aio_grpc_collections_client: Optional[grpc.CollectionsStub] = None
@@ -239,9 +257,10 @@ class QdrantRemote(QdrantBase):
         return self._closed
 
     def close(self, grpc_grace: Optional[float] = None, **kwargs: Any) -> None:
-        if hasattr(self, "_grpc_channel") and self._grpc_channel is not None:
+        if hasattr(self, "_grpc_channel_pool") and len(self._grpc_channel_pool) > 0:
+            for channel in self._grpc_channel_pool:
                 try:
-                self._grpc_channel.close()
+                    channel.close()
                 except AttributeError:
                     show_warning(
                         message="Unable to close grpc_channel. Connection was interrupted on the server side",
@@ -271,12 +290,28 @@ class QdrantRemote(QdrantBase):
         )
         return scheme, host, port, prefix
 
+    def _get_grpc_pool_size(self) -> int:
+        """
+        Returns the pool size to use for GRPC connection pool.
+        This method should be preferred over accessing `self._pool_size` directly as it applies the
+        default value if no pool_size was provided.
+        """
+
+        if self._pool_size is not None:
+            return self._pool_size
+        else:
+            return self.DEFAULT_GRPC_POOL_SIZE
+
     def _init_grpc_channel(self) -> None:
         if self._closed:
             raise RuntimeError("Client was closed. Please create a new QdrantClient instance.")
 
-        if self._grpc_channel is None:
-            self._grpc_channel = get_channel(
+        try:
+            channel_pool = []
+
+            if len(self._grpc_channel_pool) == 0:
+                for _ in range(self._get_grpc_pool_size()):
+                    channel = get_channel(
                         host=self._host,
                         port=self._grpc_port,
                         ssl=self._https,
@@ -287,22 +322,43 @@ class QdrantRemote(QdrantBase):
                         # but we can't check type here, since it'll get into async client as well
                         auth_token_provider=self._auth_token_provider,  # type: ignore
                     )
+                    channel_pool.append(channel)
+
+                # Apply the clients late to prevent half-initialized pools if a channel creation fails.
+                self._grpc_channel_pool = channel_pool
+        except Exception as e:
+            raise RuntimeError(f"Error initializing the grpc connection(s): {e}")
 
     def _init_grpc_points_client(self) -> None:
         self._init_grpc_channel()
-        self._grpc_points_client = grpc.PointsStub(self._grpc_channel)
+        self._grpc_points_client_pool = [
+            grpc.PointsStub(channel) for channel in self._grpc_channel_pool
+        ]
 
     def _init_grpc_collections_client(self) -> None:
         self._init_grpc_channel()
-        self._grpc_collections_client = grpc.CollectionsStub(self._grpc_channel)
+        self._grpc_collections_client_pool = [
+            grpc.CollectionsStub(channel) for channel in self._grpc_channel_pool
+        ]
 
     def _init_grpc_snapshots_client(self) -> None:
         self._init_grpc_channel()
-        self._grpc_snapshots_client = grpc.SnapshotsStub(self._grpc_channel)
+        self._grpc_snapshots_client_pool = [
+            grpc.SnapshotsStub(channel) for channel in self._grpc_channel_pool
+        ]
 
     def _init_grpc_root_client(self) -> None:
         self._init_grpc_channel()
-        self._grpc_root_client = grpc.QdrantStub(self._grpc_channel)
+        self._grpc_root_client_pool = [
+            grpc.QdrantStub(channel) for channel in self._grpc_channel_pool
+        ]
+
+    def _next_grpc_client(self) -> int:
+        current_index = self._grpc_client_next_index
+        self._grpc_client_next_index = (
+            self._grpc_client_next_index + 1
+        ) % self._get_grpc_pool_size()
+        return current_index
 
     @property
     def grpc_collections(self) -> grpc.CollectionsStub:
@@ -311,9 +367,10 @@ class QdrantRemote(QdrantBase):
         Returns:
             An instance of raw gRPC client, generated from Protobuf
         """
-        if self._grpc_collections_client is None:
+        if self._grpc_collections_client_pool is None:
             self._init_grpc_collections_client()
-        return self._grpc_collections_client
+        assert self._grpc_collections_client_pool is not None
+        return self._grpc_collections_client_pool[self._next_grpc_client()]
 
     @property
     def grpc_points(self) -> grpc.PointsStub:
@@ -322,9 +379,10 @@ class QdrantRemote(QdrantBase):
         Returns:
             An instance of raw gRPC client, generated from Protobuf
         """
-        if self._grpc_points_client is None:
+        if self._grpc_points_client_pool is None:
             self._init_grpc_points_client()
-        return self._grpc_points_client
+        assert self._grpc_points_client_pool is not None
+        return self._grpc_points_client_pool[self._next_grpc_client()]
 
     @property
     def grpc_snapshots(self) -> grpc.SnapshotsStub:
@@ -333,9 +391,10 @@ class QdrantRemote(QdrantBase):
         Returns:
             An instance of raw gRPC client, generated from Protobuf
         """
-        if self._grpc_snapshots_client is None:
+        if self._grpc_snapshots_client_pool is None:
             self._init_grpc_snapshots_client()
-        return self._grpc_snapshots_client
+        assert self._grpc_snapshots_client_pool is not None
+        return self._grpc_snapshots_client_pool[self._next_grpc_client()]
 
     @property
     def grpc_root(self) -> grpc.QdrantStub:
@@ -344,9 +403,10 @@ class QdrantRemote(QdrantBase):
         Returns:
             An instance of raw gRPC client, generated from Protobuf
         """
-        if self._grpc_root_client is None:
+        if self._grpc_root_client_pool is None:
             self._init_grpc_root_client()
-        return self._grpc_root_client
+        assert self._grpc_root_client_pool is not None
+        return self._grpc_root_client_pool[self._next_grpc_client()]
 
     @property
     def rest(self) -> SyncApis[ApiClient]:

tools/async_client_generator/transformers/remote/function_def_transformer.py

@@ -38,9 +38,10 @@ class RemoteFunctionDefTransformer(FunctionDefTransformer):
     def override_close() -> ast.stmt:
         code = """
 async def close(self, grpc_grace: Optional[float] = None, **kwargs: Any) -> None:
-    if hasattr(self, "_grpc_channel") and self._grpc_channel is not None:
+    if hasattr(self, "_grpc_channel_pool") and len(self._grpc_channel_pool) > 0:
+        for channel in self._grpc_channel_pool:
             try:
-            await self._grpc_channel.close(grace=grpc_grace)
+                await channel.close(grace=grpc_grace)
             except AttributeError:
                 show_warning(
                     message="Unable to close grpc_channel. Connection was interrupted on the server side",