mirror of
https://github.com/qdrant/qdrant-client.git
synced 2026-07-23 11:11:01 -05:00
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:
1
.gitignore
vendored
1
.gitignore
vendored
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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]:
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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]:
|
||||||
|
|||||||
@@ -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()
|
||||||
|
|||||||
Reference in New Issue
Block a user