8000 feat: add 'max_pool_size' parameter to Client and SessionManager · NVIDIA/aistore@8630b85 · GitHub
[go: up one dir, main page]
More Web Proxy on the site http://driver.im/
Skip to content

Commit 8630b85

Browse files
gaikwadabhishekgaikwadabhishek
committed
feat: add 'max_pool_size' parameter to Client and SessionManager
Signed-off-by: Abhishek Gaikwad <gaikwadabhishek1997@gmail.com>
1 parent 7ecba02 commit 8630b85

File tree

3 files changed

+38
-17
lines changed

3 files changed

+38
-17
lines changed

‎python/aistore/sdk/client.py

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -24,22 +24,24 @@
2424

2525
class Client:
2626
"""
27-
AIStore client for managing buckets, objects, ETL jobs
27+
AIStore client for managing buckets, objects, and ETL jobs.
2828
2929
Args:
30-
endpoint (str): AIStore endpoint
30+
endpoint (str): AIStore endpoint.
3131
skip_verify (bool, optional): If True, skip SSL certificate verification. Defaults to False.
32-
ca_cert (str, optional): Path to a CA certificate file for SSL verification. If not provided, the
33-
'AIS_CLIENT_CA' environment variable will be used. Defaults to None.
32+
ca_cert (str, optional): Path to a CA certificate file for SSL verification. If not provided,
33+
the 'AIS_CLIENT_CA' environment variable will be used. Defaults to None.
3434
client_cert (Union[str, Tuple[str, str], None], optional): Path to a client certificate PEM file
35-
or a path pair (cert, key) for mTLS. If not provided, 'AIS_CRT' and 'AIS_CRT_KEY'
36-
environment variables will be used. Defaults to None.
37-
timeout (Union[float, Tuple[float, float], None], optional): Request timeout in seconds; a single float
38-
for both connect/read timeouts (e.g., 5.0), a tuple for separate connect/read timeouts (e.g., (3.0, 10.0)),
35+
or a tuple (cert, key) for mTLS. If not provided, 'AIS_CRT' and 'AIS_CRT_KEY' environment
36+
variables will be used. Defaults to None.
37+
timeout (Union[float, Tuple[float, float], None], optional): Request timeout in seconds.
38+
Can be a single float (e.g., 5.0) for both connect/read timeouts, a tuple (e.g., (3.0, 10.0)),
3939
or None to disable timeout.
40-
retry (urllib3.Retry, optional): Retry configuration object from the urllib3 library.
40+
retry (urllib3.Retry, optional): Retry configuration object from the urllib3 library. Defaults to None.
4141
token (str, optional): Authorization token. If not provided, the 'AIS_AUTHN_TOKEN' environment variable
4242
will be used. Defaults to None.
43+
max_pool_size (int, optional): Maximum number of connections per host in the connection pool.
44+
Defaults to 10.
4345
"""
4446

4547
# pylint: disable=too-many-arguments
@@ -52,17 +54,20 @@ def __init__(
5254
timeout: Optional[Union[float, Tuple[float, float]]] = None,
5355
retry: Optional[Retry] = None,
5456
token: Optional[str] = None,
57+
max_pool_size: int = 10,
5558
):
5659
session_manager = SessionManager(
5760
retry=retry,
5861
ca_cert=ca_cert,
5962
client_cert=client_cert,
6063
skip_verify=skip_verify,
64+
max_pool_size=max_pool_size,
6165
)
6266

6367
# Check for token from arguments or environment variable
6468
if not token:
65-
token = os.environ.get(AIS_AUTHN_TOKEN, None)
69+
token = os.environ.get(AIS_AUTHN_TOKEN)
70+
6671
self._request_client = RequestClient(
6772
endpoint=endpoint,
6873
session_manager=session_manager,

‎python/aistore/sdk/session_manager.py

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,14 +27,18 @@ class SessionManager:
2727
client_cert (Union[str, Tuple[str, str], None], optional): Path to a client certificate PEM file
2828
or a path pair (cert, key) for mTLS. If not provided, 'AIS_CRT' and 'AIS_CRT_KEY' environment
2929
variables will be used. Defaults to None.
30+
max_pool_size (int, optional): Maximum number of connections per host in the connection pool.
31+
Defaults to 10.
3032
"""
3133

34+
# pylint: disable=too-many-arguments
3235
def __init__(
3336
self,
3437
retry: Retry = DEFAULT_RETRY,
3538
ca_cert: Optional[str] = None,
3639
skip_verify: bool = False,
3740
client_cert: Optional[Union[str, Tuple[str, str]]] = None,
41+
max_pool_size: int = 10,
3842
):
3943
self._retry = retry
4044
self._ca_cert = ca_cert
@@ -44,6 +48,7 @@ def __init__(
4448
key = os.getenv(AIS_CLIENT_KEY)
4549
client_cert = (cert, key) if cert and key else None
4650
self._client_cert = client_cert
51+
self._max_pool_size = max_pool_size
4752
self._session_pool = {current_process().pid: self._create_session()}
4853

4954
@property
@@ -104,6 +109,13 @@ def _create_session(self) -> Session:
104109
request_session = Session()
105110
request_session.cert = self._client_cert
106111
self._set_session_verification(request_session)
112+
113+
adapter = HTTPAdapter(
114+
max_retries=self._retry,
115+
pool_connections=self._max_pool_size,
116+
pool_maxsize=self._max_pool_size,
117+
)
118+
107119
for protocol in (HTTP, HTTPS):
108-
request_session.mount(protocol, HTTPAdapter(max_retries=self._retry))
120+
request_session.mount(protocol, adapter)
109121
return request_session

‎python/tests/unit/sdk/test_client.py

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -34,16 +34,18 @@ def test_init_defaults(self, mock_request_client, mock_sm):
3434
)
3535

3636
@cases(
37-
(True, None, None, None, None, "dummy.token"),
38-
(False, "ca_cert_location", None, None, None, None),
39-
(False, None, "client_cert_location", None, None, None),
40-
(False, None, None, 30.0, Retry(total=4), None),
41-
(False, None, None, (10, 30.0), Retry(total=5, connect=2), "dummy.token"),
37+
(True, None, None, None, None, "dummy.token", 50),
38+
(False, "ca_cert_location", None, None, None, None, 10),
39+
(False, None, "client_cert_location", None, None, None, None),
40+
(False, None, None, 30.0, Retry(total=4), None, 100),
41+
(False, None, None, (10, 30.0), Retry(total=5, connect=2), "dummy.token", 10),
4242
)
4343
@patch("aistore.sdk.client.SessionManager")
4444
@patch("aistore.sdk.client.RequestClient")
4545
def test_init(self, test_case, mock_request_client, mock_sm):
46-
skip_verify, ca_cert, client_cert, timeout, retry, token = test_case
46+
skip_verify, ca_cert, client_cert, timeout, retry, token, max_pool_size = (
47+
test_case
48+
)
4749
Client(
4850
self.endpoint,
4951
skip_verify=skip_verify,
@@ -52,12 +54,14 @@ def test_init(self, test_case, mock_request_client, mock_sm):
5254
timeout=timeout,
5355
retry=retry,
5456
token=token,
57+
max_pool_size=max_pool_size,
5558
)
5659
mock_sm.assert_called_with(
5760
retry=retry,
5861
ca_cert=ca_cert,
5962
client_cert=client_cert,
6063
skip_verify=skip_verify,
64+
max_pool_size=max_pool_size,
6165
)
6266
mock_request_client.assert_called_with(
6367
endpoint=self.endpoint,

0 commit comments

Comments
 (0)
0