perf(workflows): reuse S3 payload clients - #601
Open
lucasmrdt wants to merge 1 commit into
Open
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds an opt-in reusable S3 client for workflow payload offloading so a worker can reuse a single
aioboto3S3 client across many offload/reload calls instead of opening (and tearing down) a new client for every payload.BlobStorageConfig.reuse_client(defaultFalse) — fully backward compatible; the current per-context client behavior is unchanged unless a caller opts in.reuse_client=True,get_blob_storagepools oneS3BlobStorageper safe key(bucket_name, prefix, region_name, endpoint_url, access_key, secret_key), so repeated contexts for the same destination share one open client.S3BlobStoragenow tracks its client context explicitly (_client_context), returns early from__aenter__when a client is already open, and only tears down on__aexit__when it is not reusing. Real teardown lives inaclose(...).close_cached_blob_storages()closes every pooled client and is safe to call repeatedly (idempotent) — intended for worker shutdown.async with).Scope is intentionally minimal: only the reusable-client mechanism is added. No caching layer, no HEAD/existence-check changes, no other config fields.
Changes
src/mistralai/extra/workflows/encoding/config.py: addreuse_client: bool = False.src/mistralai/extra/workflows/encoding/storage/_s3.py: explicit client-context tracking, reuse guard,aclose, reuse-aware__aexit__.src/mistralai/extra/workflows/encoding/storage/blob_storage.py: keyed S3 client pool, reuse-aware construction,close_cached_blob_storages, safe enter/exit lifecycle.src/mistralai/extra/tests/test_workflow_payload_client_pool.py: new tests.All edited files are hand-owned (listed in
.genignore); no generated SDK surface is touched.Verification
test_workflow_payload_client_pool.py— defaultFalse, independent clients when disabled, single client reuse when enabled, distinct clients per distinct reuse config, repeatable cleanup, and enter/exit failure lifecycle (root error preserved, no cache leak).uv run python -m unittest discover -s src/mistralai/extra/tests -t src— 152 passed (CI path).uv run pytest tests/— 358 passed, 46 skipped.uv run ruff check src/mistralai/extra/— clean.uv run --all-extras mypy src/mistralai/extra/— clean.uv run --all-extras pyright src/mistralai/extra/— 0 errors.uv build— sdist + wheel build successfully.git diff --check— clean.aioboto3manual QA (no network) proving creation counts and cleanup:reuse_client=False→ 2 clients for 2 contexts;reuse_client=True→ 1 client kept open;close_cached_blob_storages()closes once and is idempotent; enter/exit failures surface the root error; repeated interruptions leave the pool stable.