diff --git a/.gitignore b/.gitignore index a8960f2d..837ef9e6 100644 --- a/.gitignore +++ b/.gitignore @@ -14,3 +14,4 @@ local_cache/*/* .python-version docs/source/examples/local_cache/* docs/source/examples/path/to/db/* +.venv diff --git a/qdrant_client/async_qdrant_client.py b/qdrant_client/async_qdrant_client.py index 2501d15f..802e4468 100644 --- a/qdrant_client/async_qdrant_client.py +++ b/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 diff --git a/qdrant_client/async_qdrant_remote.py b/qdrant_client/async_qdrant_remote.py index 049ab4b5..87726927 100644 --- a/qdrant_client/async_qdrant_remote.py +++ b/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,17 +213,18 @@ 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: - try: - await self._grpc_channel.close(grace=grpc_grace) - except AttributeError: - show_warning( - message="Unable to close grpc_channel. Connection was interrupted on the server side", - category=UserWarning, - stacklevel=4, - ) - except RuntimeError: - pass + if hasattr(self, "_grpc_channel_pool") and len(self._grpc_channel_pool) > 0: + for channel in self._grpc_channel_pool: + try: + await channel.close(grace=grpc_grace) + except AttributeError: + show_warning( + message="Unable to close grpc_channel. Connection was interrupted on the server side", + category=UserWarning, + stacklevel=4, + ) + except RuntimeError: + pass try: await self.http.aclose() except Exception: @@ -232,35 +246,68 @@ 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( - host=self._host, - port=self._grpc_port, - ssl=self._https, - metadata=self._grpc_headers, - options=self._grpc_options, - compression=self._grpc_compression, - auth_token_provider=self._auth_token_provider, - ) + 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, + metadata=self._grpc_headers, + options=self._grpc_options, + 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]: diff --git a/qdrant_client/qdrant_client.py b/qdrant_client/qdrant_client.py index bb0a204f..ae73d696 100644 --- a/qdrant_client/qdrant_client.py +++ b/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 diff --git a/qdrant_client/qdrant_remote.py b/qdrant_client/qdrant_remote.py index 5b64912a..75f83d8a 100644 --- a/qdrant_client/qdrant_remote.py +++ b/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,15 +257,16 @@ 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: - try: - self._grpc_channel.close() - except AttributeError: - show_warning( - message="Unable to close grpc_channel. Connection was interrupted on the server side", - category=RuntimeWarning, - stacklevel=4, - ) + if hasattr(self, "_grpc_channel_pool") and len(self._grpc_channel_pool) > 0: + for channel in self._grpc_channel_pool: + try: + channel.close() + except AttributeError: + show_warning( + message="Unable to close grpc_channel. Connection was interrupted on the server side", + category=RuntimeWarning, + stacklevel=4, + ) try: self.openapi_client.close() @@ -271,38 +290,75 @@ 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( - host=self._host, - port=self._grpc_port, - ssl=self._https, - metadata=self._grpc_headers, - options=self._grpc_options, - compression=self._grpc_compression, - # sync get_channel does not accept coroutine functions, - # 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 - ) + 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, + metadata=self._grpc_headers, + options=self._grpc_options, + compression=self._grpc_compression, + # sync get_channel does not accept coroutine functions, + # 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]: diff --git a/tools/async_client_generator/transformers/remote/function_def_transformer.py b/tools/async_client_generator/transformers/remote/function_def_transformer.py index b8764f42..5d74028b 100644 --- a/tools/async_client_generator/transformers/remote/function_def_transformer.py +++ b/tools/async_client_generator/transformers/remote/function_def_transformer.py @@ -38,17 +38,18 @@ 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: - try: - await self._grpc_channel.close(grace=grpc_grace) - except AttributeError: - show_warning( - message="Unable to close grpc_channel. Connection was interrupted on the server side", - category=UserWarning, - stacklevel=4, - ) - except RuntimeError: - pass + if hasattr(self, "_grpc_channel_pool") and len(self._grpc_channel_pool) > 0: + for channel in self._grpc_channel_pool: + try: + await channel.close(grace=grpc_grace) + except AttributeError: + show_warning( + message="Unable to close grpc_channel. Connection was interrupted on the server side", + category=UserWarning, + stacklevel=4, + ) + except RuntimeError: + pass try: await self.http.aclose()