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>
This commit is contained in:
Jojii
2025-09-19 14:38:36 +02:00
committed by George Panchuk
parent 9a7a1d2d06
commit 8f8b1b6b66
6 changed files with 214 additions and 89 deletions

1
.gitignore vendored
View File

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

View File

@@ -72,8 +72,12 @@ class AsyncQdrantClient(AsyncQdrantFastembedMixin):
Only use this if you can guarantee that you can resolve the thread safety outside QdrantClient. 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. 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` 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 **kwargs: Additional arguments passed directly into REST client initialization
""" """
def __init__( def __init__(
@@ -97,6 +101,7 @@ class AsyncQdrantClient(AsyncQdrantFastembedMixin):
cloud_inference: bool = False, cloud_inference: bool = False,
local_inference_batch_size: Optional[int] = None, local_inference_batch_size: Optional[int] = None,
check_compatibility: bool = True, check_compatibility: bool = True,
pool_size: Optional[int] = None,
**kwargs: Any, **kwargs: Any,
): ):
self._init_options = { self._init_options = {
@@ -135,6 +140,7 @@ class AsyncQdrantClient(AsyncQdrantFastembedMixin):
grpc_options=grpc_options, grpc_options=grpc_options,
auth_token_provider=auth_token_provider, auth_token_provider=auth_token_provider,
check_compatibility=check_compatibility, check_compatibility=check_compatibility,
pool_size=pool_size,
**kwargs, **kwargs,
) )
server_version = self._client.server_version server_version = self._client.server_version

View File

@@ -54,6 +54,7 @@ from qdrant_client.uploader.uploader import BaseUploader
class AsyncQdrantRemote(AsyncQdrantBase): class AsyncQdrantRemote(AsyncQdrantBase):
DEFAULT_GRPC_TIMEOUT = 5 DEFAULT_GRPC_TIMEOUT = 5
DEFAULT_GRPC_POOL_SIZE = 3
def __init__( def __init__(
self, self,
@@ -71,6 +72,7 @@ class AsyncQdrantRemote(AsyncQdrantBase):
Union[Callable[[], str], Callable[[], Awaitable[str]]] Union[Callable[[], str], Callable[[], Awaitable[str]]]
] = None, ] = None,
check_compatibility: bool = True, check_compatibility: bool = True,
pool_size: Optional[int] = None,
**kwargs: Any, **kwargs: Any,
): ):
super().__init__(**kwargs) super().__init__(**kwargs)
@@ -79,6 +81,10 @@ class AsyncQdrantRemote(AsyncQdrantBase):
self._grpc_options = grpc_options or {} self._grpc_options = grpc_options or {}
self._https = https if https is not None else api_key is not None self._https = https if https is not None else api_key is not None
self._scheme = "https" if self._https else "http" 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 "" self._prefix = prefix or ""
if len(self._prefix) > 0 and self._prefix[0] != "/": if len(self._prefix) > 0 and self._prefix[0] != "/":
self._prefix = f"/{self._prefix}" self._prefix = f"/{self._prefix}"
@@ -115,6 +121,12 @@ class AsyncQdrantRemote(AsyncQdrantBase):
if limits is None: if limits is None:
if self._host in ["localhost", "127.0.0.1"]: if self._host in ["localhost", "127.0.0.1"]:
limits = httpx.Limits(max_connections=None, max_keepalive_connections=0) 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) http2 = kwargs.pop("http2", False)
self._grpc_headers = [] self._grpc_headers = []
self._rest_headers = {k: v for (k, v) in kwargs.pop("metadata", {}).items()} 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( self.openapi_client: AsyncApis[AsyncApiClient] = AsyncApis(
host=self.rest_uri, **self._rest_args host=self.rest_uri, **self._rest_args
) )
self._grpc_channel = None self._grpc_channel_pool: list[grpc.Channel] = []
self._grpc_points_client: Optional[grpc.PointsStub] = None self._grpc_points_client_pool: Optional[list[grpc.PointsStub]] = None
self._grpc_collections_client: Optional[grpc.CollectionsStub] = None self._grpc_collections_client_pool: Optional[list[grpc.CollectionsStub]] = None
self._grpc_snapshots_client: Optional[grpc.SnapshotsStub] = None self._grpc_snapshots_client_pool: Optional[list[grpc.SnapshotsStub]] = None
self._grpc_root_client: Optional[grpc.QdrantStub] = None self._grpc_root_client_pool: Optional[list[grpc.QdrantStub]] = None
self._grpc_client_next_index: int = 0
self._closed: bool = False self._closed: bool = False
self.server_version = None self.server_version = None
if check_compatibility: if check_compatibility:
@@ -200,17 +213,18 @@ class AsyncQdrantRemote(AsyncQdrantBase):
return self._closed return self._closed
async def close(self, grpc_grace: Optional[float] = None, **kwargs: Any) -> None: 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:
try: for channel in self._grpc_channel_pool:
await self._grpc_channel.close(grace=grpc_grace) try:
except AttributeError: await channel.close(grace=grpc_grace)
show_warning( except AttributeError:
message="Unable to close grpc_channel. Connection was interrupted on the server side", show_warning(
category=UserWarning, message="Unable to close grpc_channel. Connection was interrupted on the server side",
stacklevel=4, category=UserWarning,
) stacklevel=4,
except RuntimeError: )
pass except RuntimeError:
pass
try: try:
await self.http.aclose() await self.http.aclose()
except Exception: except Exception:
@@ -232,35 +246,68 @@ class AsyncQdrantRemote(AsyncQdrantBase):
) )
return (scheme, host, port, prefix) 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: def _init_grpc_channel(self) -> None:
if self._closed: if self._closed:
raise RuntimeError("Client was closed. Please create a new QdrantClient instance.") raise RuntimeError("Client was closed. Please create a new QdrantClient instance.")
if self._grpc_channel is None: try:
self._grpc_channel = get_channel( channel_pool = []
host=self._host, if len(self._grpc_channel_pool) == 0:
port=self._grpc_port, for _ in range(self._get_grpc_pool_size()):
ssl=self._https, channel = get_channel(
metadata=self._grpc_headers, host=self._host,
options=self._grpc_options, port=self._grpc_port,
compression=self._grpc_compression, ssl=self._https,
auth_token_provider=self._auth_token_provider, 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: def _init_grpc_points_client(self) -> None:
self._init_grpc_channel() 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: def _init_grpc_collections_client(self) -> None:
self._init_grpc_channel() 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: def _init_grpc_snapshots_client(self) -> None:
self._init_grpc_channel() 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: def _init_grpc_root_client(self) -> None:
self._init_grpc_channel() 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 @property
def grpc_collections(self) -> grpc.CollectionsStub: def grpc_collections(self) -> grpc.CollectionsStub:
@@ -269,9 +316,10 @@ class AsyncQdrantRemote(AsyncQdrantBase):
Returns: Returns:
An instance of raw gRPC client, generated from Protobuf 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() 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 @property
def grpc_points(self) -> grpc.PointsStub: def grpc_points(self) -> grpc.PointsStub:
@@ -280,9 +328,10 @@ class AsyncQdrantRemote(AsyncQdrantBase):
Returns: Returns:
An instance of raw gRPC client, generated from Protobuf 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() 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 @property
def grpc_snapshots(self) -> grpc.SnapshotsStub: def grpc_snapshots(self) -> grpc.SnapshotsStub:
@@ -291,9 +340,10 @@ class AsyncQdrantRemote(AsyncQdrantBase):
Returns: Returns:
An instance of raw gRPC client, generated from Protobuf 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() 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 @property
def grpc_root(self) -> grpc.QdrantStub: def grpc_root(self) -> grpc.QdrantStub:
@@ -302,9 +352,10 @@ class AsyncQdrantRemote(AsyncQdrantBase):
Returns: Returns:
An instance of raw gRPC client, generated from Protobuf 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() 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 @property
def rest(self) -> AsyncApis[AsyncApiClient]: def rest(self) -> AsyncApis[AsyncApiClient]:

View File

@@ -73,8 +73,12 @@ class QdrantClient(QdrantFastembedMixin):
Only use this if you can guarantee that you can resolve the thread safety outside QdrantClient. 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. 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` 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 **kwargs: Additional arguments passed directly into REST client initialization
""" """
def __init__( def __init__(
@@ -98,6 +102,7 @@ class QdrantClient(QdrantFastembedMixin):
cloud_inference: bool = False, cloud_inference: bool = False,
local_inference_batch_size: Optional[int] = None, local_inference_batch_size: Optional[int] = None,
check_compatibility: bool = True, check_compatibility: bool = True,
pool_size: Optional[int] = None,
**kwargs: Any, **kwargs: Any,
): ):
# Saving the init options to facilitate building AsyncQdrantClient from QdrantClient and vice versa. # Saving the init options to facilitate building AsyncQdrantClient from QdrantClient and vice versa.
@@ -142,6 +147,7 @@ class QdrantClient(QdrantFastembedMixin):
grpc_options=grpc_options, grpc_options=grpc_options,
auth_token_provider=auth_token_provider, auth_token_provider=auth_token_provider,
check_compatibility=check_compatibility, check_compatibility=check_compatibility,
pool_size=pool_size,
**kwargs, **kwargs,
) )
server_version = self._client.server_version server_version = self._client.server_version

View File

@@ -45,6 +45,7 @@ from qdrant_client.uploader.uploader import BaseUploader
class QdrantRemote(QdrantBase): class QdrantRemote(QdrantBase):
DEFAULT_GRPC_TIMEOUT = 5 # seconds DEFAULT_GRPC_TIMEOUT = 5 # seconds
DEFAULT_GRPC_POOL_SIZE = 3
def __init__( def __init__(
self, self,
@@ -62,6 +63,7 @@ class QdrantRemote(QdrantBase):
Union[Callable[[], str], Callable[[], Awaitable[str]]] Union[Callable[[], str], Callable[[], Awaitable[str]]]
] = None, ] = None,
check_compatibility: bool = True, check_compatibility: bool = True,
pool_size: Optional[int] = None,
**kwargs: Any, **kwargs: Any,
): ):
super().__init__(**kwargs) 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._https = https if https is not None else api_key is not None
self._scheme = "https" if self._https else "http" 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 "" self._prefix = prefix or ""
if len(self._prefix) > 0 and self._prefix[0] != "/": if len(self._prefix) > 0 and self._prefix[0] != "/":
self._prefix = f"/{self._prefix}" self._prefix = f"/{self._prefix}"
@@ -126,6 +135,14 @@ class QdrantRemote(QdrantBase):
# Disable keep-alive for local connections # Disable keep-alive for local connections
# Cause in some cases, it may cause extra delays # Cause in some cases, it may cause extra delays
limits = httpx.Limits(max_connections=None, max_keepalive_connections=0) 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) http2 = kwargs.pop("http2", False)
self._grpc_headers = [] self._grpc_headers = []
@@ -193,11 +210,12 @@ class QdrantRemote(QdrantBase):
**self._rest_args, **self._rest_args,
) )
self._grpc_channel = None self._grpc_channel_pool: list[grpc.Channel] = []
self._grpc_points_client: Optional[grpc.PointsStub] = None self._grpc_points_client_pool: Optional[list[grpc.PointsStub]] = None
self._grpc_collections_client: Optional[grpc.CollectionsStub] = None self._grpc_collections_client_pool: Optional[list[grpc.CollectionsStub]] = None
self._grpc_snapshots_client: Optional[grpc.SnapshotsStub] = None self._grpc_snapshots_client_pool: Optional[list[grpc.SnapshotsStub]] = None
self._grpc_root_client: Optional[grpc.QdrantStub] = 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_points_client: Optional[grpc.PointsStub] = None
self._aio_grpc_collections_client: Optional[grpc.CollectionsStub] = None self._aio_grpc_collections_client: Optional[grpc.CollectionsStub] = None
@@ -239,15 +257,16 @@ class QdrantRemote(QdrantBase):
return self._closed return self._closed
def close(self, grpc_grace: Optional[float] = None, **kwargs: Any) -> None: 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:
try: for channel in self._grpc_channel_pool:
self._grpc_channel.close() try:
except AttributeError: channel.close()
show_warning( except AttributeError:
message="Unable to close grpc_channel. Connection was interrupted on the server side", show_warning(
category=RuntimeWarning, message="Unable to close grpc_channel. Connection was interrupted on the server side",
stacklevel=4, category=RuntimeWarning,
) stacklevel=4,
)
try: try:
self.openapi_client.close() self.openapi_client.close()
@@ -271,38 +290,75 @@ class QdrantRemote(QdrantBase):
) )
return scheme, host, port, prefix 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: def _init_grpc_channel(self) -> None:
if self._closed: if self._closed:
raise RuntimeError("Client was closed. Please create a new QdrantClient instance.") raise RuntimeError("Client was closed. Please create a new QdrantClient instance.")
if self._grpc_channel is None: try:
self._grpc_channel = get_channel( channel_pool = []
host=self._host,
port=self._grpc_port, if len(self._grpc_channel_pool) == 0:
ssl=self._https, for _ in range(self._get_grpc_pool_size()):
metadata=self._grpc_headers, channel = get_channel(
options=self._grpc_options, host=self._host,
compression=self._grpc_compression, port=self._grpc_port,
# sync get_channel does not accept coroutine functions, ssl=self._https,
# but we can't check type here, since it'll get into async client as well metadata=self._grpc_headers,
auth_token_provider=self._auth_token_provider, # type: ignore 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: def _init_grpc_points_client(self) -> None:
self._init_grpc_channel() 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: def _init_grpc_collections_client(self) -> None:
self._init_grpc_channel() 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: def _init_grpc_snapshots_client(self) -> None:
self._init_grpc_channel() 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: def _init_grpc_root_client(self) -> None:
self._init_grpc_channel() 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 @property
def grpc_collections(self) -> grpc.CollectionsStub: def grpc_collections(self) -> grpc.CollectionsStub:
@@ -311,9 +367,10 @@ class QdrantRemote(QdrantBase):
Returns: Returns:
An instance of raw gRPC client, generated from Protobuf 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() 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 @property
def grpc_points(self) -> grpc.PointsStub: def grpc_points(self) -> grpc.PointsStub:
@@ -322,9 +379,10 @@ class QdrantRemote(QdrantBase):
Returns: Returns:
An instance of raw gRPC client, generated from Protobuf 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() 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 @property
def grpc_snapshots(self) -> grpc.SnapshotsStub: def grpc_snapshots(self) -> grpc.SnapshotsStub:
@@ -333,9 +391,10 @@ class QdrantRemote(QdrantBase):
Returns: Returns:
An instance of raw gRPC client, generated from Protobuf 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() 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 @property
def grpc_root(self) -> grpc.QdrantStub: def grpc_root(self) -> grpc.QdrantStub:
@@ -344,9 +403,10 @@ class QdrantRemote(QdrantBase):
Returns: Returns:
An instance of raw gRPC client, generated from Protobuf 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() 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 @property
def rest(self) -> SyncApis[ApiClient]: def rest(self) -> SyncApis[ApiClient]:

View File

@@ -38,17 +38,18 @@ class RemoteFunctionDefTransformer(FunctionDefTransformer):
def override_close() -> ast.stmt: def override_close() -> ast.stmt:
code = """ code = """
async def close(self, grpc_grace: Optional[float] = None, **kwargs: Any) -> None: 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:
try: for channel in self._grpc_channel_pool:
await self._grpc_channel.close(grace=grpc_grace) try:
except AttributeError: await channel.close(grace=grpc_grace)
show_warning( except AttributeError:
message="Unable to close grpc_channel. Connection was interrupted on the server side", show_warning(
category=UserWarning, message="Unable to close grpc_channel. Connection was interrupted on the server side",
stacklevel=4, category=UserWarning,
) stacklevel=4,
except RuntimeError: )
pass except RuntimeError:
pass
try: try:
await self.http.aclose() await self.http.aclose()