From ba978347f45cc762d1afb67971f9611c4b49f7f0 Mon Sep 17 00:00:00 2001 From: Lucas Marandat Date: Thu, 30 Jul 2026 17:35:44 +0200 Subject: [PATCH] perf(workflows): reuse S3 payload clients Add an opt-in reusable S3 client for workflow payload offloading. BlobStorageConfig gains reuse_client (default False), preserving the current per-context client behavior. When reuse_client is True, get_blob_storage pools one S3BlobStorage per safe credential/location key so a worker reuses a single aioboto3 client across offload calls instead of opening one per payload. S3BlobStorage tracks its client context explicitly, skips re-entering an already-open client, and only tears down on __aexit__ when not reusing; aclose performs the real teardown. close_cached_blob_storages closes every pooled client and is safe to call repeatedly. A failed client enter never leaves a half-open entry in the pool, and context exit still forwards the original exception. Tests cover the default, independent clients when disabled, single client reuse when enabled, repeatable cleanup, and enter/exit failure lifecycle. --- .../test_workflow_payload_client_pool.py | 193 ++++++++++++++++++ .../extra/workflows/encoding/config.py | 1 + .../extra/workflows/encoding/storage/_s3.py | 20 +- .../encoding/storage/blob_storage.py | 93 +++++++-- 4 files changed, 282 insertions(+), 25 deletions(-) create mode 100644 src/mistralai/extra/tests/test_workflow_payload_client_pool.py diff --git a/src/mistralai/extra/tests/test_workflow_payload_client_pool.py b/src/mistralai/extra/tests/test_workflow_payload_client_pool.py new file mode 100644 index 00000000..e4854fe4 --- /dev/null +++ b/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() diff --git a/src/mistralai/extra/workflows/encoding/config.py b/src/mistralai/extra/workflows/encoding/config.py index d54ad9fe..d6adba36 100644 --- a/src/mistralai/extra/workflows/encoding/config.py +++ b/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 diff --git a/src/mistralai/extra/workflows/encoding/storage/_s3.py b/src/mistralai/extra/workflows/encoding/storage/_s3.py index 4a0ce063..83422694 100644 --- a/src/mistralai/extra/workflows/encoding/storage/_s3.py +++ b/src/mistralai/extra/workflows/encoding/storage/_s3.py @@ -17,6 +17,7 @@ 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 "" @@ -24,7 +25,9 @@ def __init__( 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: diff --git a/src/mistralai/extra/workflows/encoding/storage/blob_storage.py b/src/mistralai/extra/workflows/encoding/storage/blob_storage.py index e267b979..7d78f58b 100644 --- a/src/mistralai/extra/workflows/encoding/storage/blob_storage.py +++ b/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", +]