feat: add 'max_pool_size' parameter to Client and SessionManager

Signed-off-by: Abhishek Gaikwad <[email protected]>

python/aistore/sdk/client.py

@@ -24,22 +24,24 @@
 
 class Client:
     """
-    AIStore client for managing buckets, objects, ETL jobs
+    AIStore client for managing buckets, objects, and ETL jobs.
 
     Args:
-        endpoint (str): AIStore endpoint
+        endpoint (str): AIStore endpoint.
         skip_verify (bool, optional): If True, skip SSL certificate verification. Defaults to False.
-        ca_cert (str, optional): Path to a CA certificate file for SSL verification. If not provided, the
-            'AIS_CLIENT_CA' environment variable will be used. Defaults to None.
+        ca_cert (str, optional): Path to a CA certificate file for SSL verification. If not provided,
+            the 'AIS_CLIENT_CA' environment variable will be used. Defaults to None.
         client_cert (Union[str, Tuple[str, str], None], optional): Path to a client certificate PEM file
-            or a path pair (cert, key) for mTLS. If not provided, 'AIS_CRT' and 'AIS_CRT_KEY'
-            environment variables will be used. Defaults to None.
-        timeout (Union[float, Tuple[float, float], None], optional): Request timeout in seconds; a single float
-            for both connect/read timeouts (e.g., 5.0), a tuple for separate connect/read timeouts (e.g., (3.0, 10.0)),
+            or a tuple (cert, key) for mTLS. If not provided, 'AIS_CRT' and 'AIS_CRT_KEY' environment
+            variables will be used. Defaults to None.
+        timeout (Union[float, Tuple[float, float], None], optional): Request timeout in seconds.
+            Can be a single float (e.g., 5.0) for both connect/read timeouts, a tuple (e.g., (3.0, 10.0)),
             or None to disable timeout.
-        retry (urllib3.Retry, optional): Retry configuration object from the urllib3 library.
+        retry (urllib3.Retry, optional): Retry configuration object from the urllib3 library. Defaults to None.
         token (str, optional): Authorization token. If not provided, the 'AIS_AUTHN_TOKEN' environment variable
             will be used. Defaults to None.
+        max_pool_size (int, optional): Maximum number of connections per host in the connection pool.
+            Defaults to 10.
     """
 
     # pylint: disable=too-many-arguments
@@ -52,17 +54,20 @@ def __init__(
         timeout: Optional[Union[float, Tuple[float, float]]] = None,
         retry: Optional[Retry] = None,
         token: Optional[str] = None,
+        max_pool_size: int = 10,
     ):
         session_manager = SessionManager(
             retry=retry,
             ca_cert=ca_cert,
             client_cert=client_cert,
             skip_verify=skip_verify,
+            max_pool_size=max_pool_size,
         )
 
         # Check for token from arguments or environment variable
         if not token:
-            token = os.environ.get(AIS_AUTHN_TOKEN, None)
+            token = os.environ.get(AIS_AUTHN_TOKEN)
+
         self._request_client = RequestClient(
             endpoint=endpoint,
             session_manager=session_manager,

python/aistore/sdk/session_manager.py

@@ -27,14 +27,18 @@ class SessionManager:
         client_cert (Union[str, Tuple[str, str], None], optional): Path to a client certificate PEM file
             or a path pair (cert, key) for mTLS. If not provided, 'AIS_CRT' and 'AIS_CRT_KEY' environment
             variables will be used. Defaults to None.
+        max_pool_size (int, optional): Maximum number of connections per host in the connection pool.
+            Defaults to 10.
     """
 
+    # pylint: disable=too-many-arguments
     def __init__(
         self,
         retry: Retry = DEFAULT_RETRY,
         ca_cert: Optional[str] = None,
         skip_verify: bool = False,
         client_cert: Optional[Union[str, Tuple[str, str]]] = None,
+        max_pool_size: int = 10,
     ):
         self._retry = retry
         self._ca_cert = ca_cert
@@ -44,6 +48,7 @@ def __init__(
             key = os.getenv(AIS_CLIENT_KEY)
             client_cert = (cert, key) if cert and key else None
         self._client_cert = client_cert
+        self._max_pool_size = max_pool_size
         self._session_pool = {current_process().pid: self._create_session()}
 
     @property
@@ -104,6 +109,13 @@ def _create_session(self) -> Session:
         request_session = Session()
         request_session.cert = self._client_cert
         self._set_session_verification(request_session)
+
+        adapter = HTTPAdapter(
+            max_retries=self._retry,
+            pool_connections=self._max_pool_size,
+            pool_maxsize=self._max_pool_size,
+        )
+
         for protocol in (HTTP, HTTPS):
-            request_session.mount(protocol, HTTPAdapter(max_retries=self._retry))
+            request_session.mount(protocol, adapter)
         return request_session

python/tests/unit/sdk/test_client.py

@@ -34,16 +34,18 @@ def test_init_defaults(self, mock_request_client, mock_sm):
         )
 
     @cases(
-        (True, None, None, None, None, "dummy.token"),
-        (False, "ca_cert_location", None, None, None, None),
-        (False, None, "client_cert_location", None, None, None),
-        (False, None, None, 30.0, Retry(total=4), None),
-        (False, None, None, (10, 30.0), Retry(total=5, connect=2), "dummy.token"),
+        (True, None, None, None, None, "dummy.token", 50),
+        (False, "ca_cert_location", None, None, None, None, 10),
+        (False, None, "client_cert_location", None, None, None, None),
+        (False, None, None, 30.0, Retry(total=4), None, 100),
+        (False, None, None, (10, 30.0), Retry(total=5, connect=2), "dummy.token", 10),
     )
     @patch("aistore.sdk.client.SessionManager")
     @patch("aistore.sdk.client.RequestClient")
     def test_init(self, test_case, mock_request_client, mock_sm):
-        skip_verify, ca_cert, client_cert, timeout, retry, token = test_case
+        skip_verify, ca_cert, client_cert, timeout, retry, token, max_pool_size = (
+            test_case
+        )
         Client(
             self.endpoint,
             skip_verify=skip_verify,
@@ -52,12 +54,14 @@ def test_init(self, test_case, mock_request_client, mock_sm):
             timeout=timeout,
             retry=retry,
             token=token,
+            max_pool_size=max_pool_size,
         )
         mock_sm.assert_called_with(
             retry=retry,
             ca_cert=ca_cert,
             client_cert=client_cert,
             skip_verify=skip_verify,
+            max_pool_size=max_pool_size,
         )
         mock_request_client.assert_called_with(
             endpoint=self.endpoint,