perf(workflows): reuse S3 payload clients

src/mistralai/extra/tests/test_workflow_payload_client_pool.py

@@ -0,0 +1,193 @@
+"""Tests for the reusable S3 payload-offloading blob-storage client pool."""
+
+import unittest
+from typing import Any
+from unittest import mock
+
+from mistralai.extra.workflows.encoding.config import BlobStorageConfig, StorageProvider
+from mistralai.extra.workflows.encoding.storage import _s3 as s3_module
+from mistralai.extra.workflows.encoding.storage._s3 import S3BlobStorage
+from mistralai.extra.workflows.encoding.storage.blob_storage import (
+    _S3_STORAGE_CACHE,
+    close_cached_blob_storages,
+    get_blob_storage,
+)
+
+
+class _FakeS3Client:
+    def __init__(self, session: "_FakeSession") -> None:
+        self._session = session
+        self.entered = False
+        self.exited = False
+
+    async def __aenter__(self) -> "_FakeS3Client":
+        self.entered = True
+        self._session.client_enters += 1
+        return self
+
+    async def __aexit__(self, *_args: Any) -> None:
+        self.exited = True
+        self._session.client_exits += 1
+
+
+class _FakeSession:
+    def __init__(self) -> None:
+        self.clients_created = 0
+        self.client_enters = 0
+        self.client_exits = 0
+
+    def client(self, *_args: Any, **_kwargs: Any) -> _FakeS3Client:
+        self.clients_created += 1
+        return _FakeS3Client(self)
+
+
+def _s3_config(**overrides: Any) -> BlobStorageConfig:
+    values: dict[str, Any] = {
+        "storage_provider": StorageProvider.S3,
+        "bucket_name": "test-bucket",
+    }
+    values.update(overrides)
+    return BlobStorageConfig(**values)
+
+
+class ReuseClientConfigDefaultTest(unittest.TestCase):
+    def test_reuse_client_defaults_to_false(self) -> None:
+        assert "reuse_client" in BlobStorageConfig.model_fields
+        assert BlobStorageConfig().reuse_client is False
+        assert _s3_config().reuse_client is False
+
+
+class BlobStorageClientPoolTest(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self) -> None:
+        await close_cached_blob_storages()
+        self._session = _FakeSession()
+        patcher = mock.patch.object(
+            s3_module.aioboto3, "Session", return_value=self._session
+        )
+        self.addAsyncCleanup(close_cached_blob_storages)
+        self.addCleanup(patcher.stop)
+        patcher.start()
+
+    async def test_reuse_client_false_creates_independent_clients(self) -> None:
+        config = _s3_config(reuse_client=False)
+
+        async with get_blob_storage(config) as first:
+            assert isinstance(first, S3BlobStorage)
+            first_client = first._client
+        async with get_blob_storage(config) as second:
+            assert isinstance(second, S3BlobStorage)
+            second_client = second._client
+
+        assert self._session.clients_created == 2
+        assert self._session.client_enters == 2
+        # Each non-reusing context closes its own client on exit.
+        assert self._session.client_exits == 2
+        assert first_client is not second_client
+        assert first is not second
+        assert not _S3_STORAGE_CACHE
+
+    async def test_reuse_client_true_reuses_single_client(self) -> None:
+        config = _s3_config(reuse_client=True)
+
+        async with get_blob_storage(config) as first:
+            assert isinstance(first, S3BlobStorage)
+            first_client = first._client
+        async with get_blob_storage(config) as second:
+            assert isinstance(second, S3BlobStorage)
+            second_client = second._client
+
+        assert self._session.clients_created == 1
+        assert self._session.client_enters == 1
+        # The reused client is kept open across contexts.
+        assert self._session.client_exits == 0
+        assert first is second
+        assert first_client is second_client
+        assert len(_S3_STORAGE_CACHE) == 1
+
+    async def test_reuse_client_true_distinct_configs_get_distinct_clients(self) -> None:
+        async with get_blob_storage(_s3_config(reuse_client=True)) as a:
+            pass
+        async with get_blob_storage(
+            _s3_config(reuse_client=True, bucket_name="other-bucket")
+        ) as b:
+            pass
+
+        assert a is not b
+        assert self._session.clients_created == 2
+        assert len(_S3_STORAGE_CACHE) == 2
+
+    async def test_close_cached_blob_storages_closes_and_is_repeatable(self) -> None:
+        async with get_blob_storage(_s3_config(reuse_client=True)) as storage:
+            assert isinstance(storage, S3BlobStorage)
+
+        assert storage._client is not None
+        assert self._session.client_exits == 0
+
+        await close_cached_blob_storages()
+
+        assert self._session.client_exits == 1
+        assert storage._client is None
+        assert not _S3_STORAGE_CACHE
+
+        # Repeated cleanup is a no-op and must not raise or double-close.
+        await close_cached_blob_storages()
+        assert self._session.client_exits == 1
+        assert not _S3_STORAGE_CACHE
+
+
+class BlobStorageClientLifecycleFailureTest(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self) -> None:
+        await close_cached_blob_storages()
+        self.addAsyncCleanup(close_cached_blob_storages)
+
+    async def test_enter_failure_surfaces_root_error(self) -> None:
+        boom = RuntimeError("connect failed")
+
+        class _FailingEnterCtx:
+            async def __aenter__(self) -> Any:
+                raise boom
+
+            async def __aexit__(self, *_args: Any) -> None:
+                return None
+
+        class _FailingEnterSession:
+            def client(self, *_args: Any, **_kwargs: Any) -> _FailingEnterCtx:
+                return _FailingEnterCtx()
+
+        with mock.patch.object(
+            s3_module.aioboto3, "Session", return_value=_FailingEnterSession()
+        ):
+            with self.assertRaises(RuntimeError) as caught:
+                async with get_blob_storage(_s3_config(reuse_client=True)):
+                    pass
+
+        assert caught.exception is boom
+        # A failed enter must not leave a half-open client cached.
+        assert not _S3_STORAGE_CACHE
+
+    async def test_exit_failure_surfaces_root_error(self) -> None:
+        boom = RuntimeError("teardown failed")
+
+        class _FailingExitClient:
+            async def __aenter__(self) -> "_FailingExitClient":
+                return self
+
+            async def __aexit__(self, *_args: Any) -> None:
+                raise boom
+
+        class _FailingExitSession:
+            def client(self, *_args: Any, **_kwargs: Any) -> _FailingExitClient:
+                return _FailingExitClient()
+
+        with mock.patch.object(
+            s3_module.aioboto3, "Session", return_value=_FailingExitSession()
+        ):
+            with self.assertRaises(RuntimeError) as caught:
+                async with get_blob_storage(_s3_config(reuse_client=False)):
+                    pass
+
+        assert caught.exception is boom
+
+
+if __name__ == "__main__":
+    unittest.main()

src/mistralai/extra/workflows/encoding/config.py

@@ -13,6 +13,7 @@ class StorageProvider(str, Enum):
 class BlobStorageConfig(BaseModel):
     storage_provider: StorageProvider = StorageProvider.S3
     prefix: Optional[str] = None
+    reuse_client: bool = False
 
     # Azure settings
     container_name: Optional[str] = None

src/mistralai/extra/workflows/encoding/storage/_s3.py

@@ -17,14 +17,17 @@ def __init__(
         endpoint_url: str | None = None,
         aws_access_key_id: str | None = None,
         aws_secret_access_key: str | None = None,
+        reuse_client: bool = False,
     ):
         self.bucket_name = bucket_name
         self.prefix = prefix or ""
         self.region_name = region_name
         self.endpoint_url = endpoint_url
         self.aws_access_key_id = aws_access_key_id
         self.aws_secret_access_key = aws_secret_access_key
+        self.reuse_client = reuse_client
         self._session: aioboto3.Session | None = None
+        self._client_context: Any = None
         self._client: Any = None
 
     def _get_full_key(self, key: str) -> str:
@@ -35,6 +38,8 @@ def _get_full_key(self, key: str) -> str:
         return f"{self.prefix}/{key}"
 
     async def __aenter__(self) -> "S3BlobStorage":
+        if self._client is not None:
+            return self
         self._session = aioboto3.Session()
         assert self._session is not None
         kwargs: dict[str, Any] = {}
@@ -47,14 +52,21 @@ async def __aenter__(self) -> "S3BlobStorage":
         if self.aws_secret_access_key:
             kwargs["aws_secret_access_key"] = self.aws_secret_access_key
 
-        self._client = self._session.client("s3", **kwargs)
-        self._client = await self._client.__aenter__()
+        self._client_context = self._session.client("s3", **kwargs)
+        self._client = await self._client_context.__aenter__()
         return self
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        if self._client:
-            await self._client.__aexit__(exc_type, exc_val, exc_tb)
+        if not self.reuse_client:
+            await self.aclose(exc_type, exc_val, exc_tb)
+
+    async def aclose(
+        self, exc_type: Any = None, exc_val: Any = None, exc_tb: Any = None
+    ) -> None:
+        if self._client_context is not None:
+            await self._client_context.__aexit__(exc_type, exc_val, exc_tb)
         self._session = None
+        self._client_context = None
         self._client = None
 
     async def upload_blob(self, key: str, content: bytes) -> str:

src/mistralai/extra/workflows/encoding/storage/blob_storage.py

@@ -1,12 +1,18 @@
 from __future__ import annotations
 
+import sys
 from abc import ABC, abstractmethod
 from contextlib import asynccontextmanager
-from typing import Any, AsyncGenerator
+from typing import TYPE_CHECKING, Any, AsyncGenerator
 
 from mistralai.extra.workflows.encoding.config import BlobStorageConfig, StorageProvider
 from mistralai.extra.exceptions import WorkflowPayloadOffloadingException
 
+if TYPE_CHECKING:
+    from ._s3 import S3BlobStorage
+
+_S3_STORAGE_CACHE: dict[tuple[str | None, ...], "S3BlobStorage"] = {}
+
 
 class BlobNotFoundError(Exception):
     """Raised when a blob is not found in storage."""
@@ -59,6 +65,7 @@ async def get_blob_storage(
     """
     storage: BlobStorage
     prefix = blob_storage_config.prefix
+    cache_key: tuple[str | None, ...] | None = None
 
     if blob_storage_config.storage_provider == StorageProvider.AZURE:
         try:
@@ -116,30 +123,74 @@ async def get_blob_storage(
             raise WorkflowPayloadOffloadingException(
                 "bucket_name is required for S3 blob storage"
             )
-        storage = S3BlobStorage(
-            bucket_name=blob_storage_config.bucket_name,
-            prefix=prefix,
-            region_name=blob_storage_config.region_name,
-            endpoint_url=blob_storage_config.endpoint_url,
-            aws_access_key_id=(
-                blob_storage_config.aws_access_key_id.get_secret_value()
-                if blob_storage_config.aws_access_key_id
-                else None
-            ),
-            aws_secret_access_key=(
-                blob_storage_config.aws_secret_access_key.get_secret_value()
-                if blob_storage_config.aws_secret_access_key
-                else None
-            ),
+        access_key = (
+            blob_storage_config.aws_access_key_id.get_secret_value()
+            if blob_storage_config.aws_access_key_id
+            else None
+        )
+        secret_key = (
+            blob_storage_config.aws_secret_access_key.get_secret_value()
+            if blob_storage_config.aws_secret_access_key
+            else None
         )
+        if blob_storage_config.reuse_client:
+            cache_key = (
+                blob_storage_config.bucket_name,
+                prefix,
+                blob_storage_config.region_name,
+                blob_storage_config.endpoint_url,
+                access_key,
+                secret_key,
+            )
+            storage = _S3_STORAGE_CACHE.get(cache_key) or S3BlobStorage(
+                bucket_name=blob_storage_config.bucket_name,
+                prefix=prefix,
+                region_name=blob_storage_config.region_name,
+                endpoint_url=blob_storage_config.endpoint_url,
+                aws_access_key_id=access_key,
+                aws_secret_access_key=secret_key,
+                reuse_client=True,
+            )
+            _S3_STORAGE_CACHE[cache_key] = storage
+        else:
+            storage = S3BlobStorage(
+                bucket_name=blob_storage_config.bucket_name,
+                prefix=prefix,
+                region_name=blob_storage_config.region_name,
+                endpoint_url=blob_storage_config.endpoint_url,
+                aws_access_key_id=access_key,
+                aws_secret_access_key=secret_key,
+                reuse_client=False,
+            )
 
     else:
         raise ValueError(
             f"Unsupported storage provider: {blob_storage_config.storage_provider}"
         )
 
-    async with storage as blob_storage_instance:
-        yield blob_storage_instance
-
-
-__all__ = ["BlobStorage", "BlobNotFoundError", "get_blob_storage"]
+    try:
+        entered = await storage.__aenter__()
+    except BaseException:
+        # A pooled client that never finished entering must not stay cached.
+        if cache_key is not None:
+            _S3_STORAGE_CACHE.pop(cache_key, None)
+        raise
+    try:
+        yield entered
+    finally:
+        await storage.__aexit__(*sys.exc_info())
+
+
+async def close_cached_blob_storages() -> None:
+    storages = list(_S3_STORAGE_CACHE.values())
+    _S3_STORAGE_CACHE.clear()
+    for storage in storages:
+        await storage.aclose()
+
+
+__all__ = [
+    "BlobStorage",
+    "BlobNotFoundError",
+    "close_cached_blob_storages",
+    "get_blob_storage",
+]