perf: share Gemini clients instead of building one per prompt - #11
Open
JamesMcMinn wants to merge 1 commit into
Open
perf: share Gemini clients instead of building one per prompt#11JamesMcMinn wants to merge 1 commit into
JamesMcMinn wants to merge 1 commit into
Conversation
`run_prompt` and `run_prompt_async` each built a `genai.Client`, used it once and closed it. Constructing one resolves credentials and sets up a transport, and every live client also holds native grpc/auth memory that `close()` does not return to the OS -- so in a long-running service a client per call reads as a memory leak. Measured previously in claims-analysis-api: `/claimants` at concurrency 8 ratcheted +127 MiB and held, and PASTEL builds one client per *sentence* (100 per 100-sentence request). Clients are now cached and reused, keyed by the project and location they were built for (the only construction arguments -- the model name is per request): - sync: one per project/location, shared across threads behind a lock. - async: additionally keyed by the running event loop, since the async transport binds to the loop it is first used on. A caller that gathers many prompts under one loop now builds one client instead of one per prompt; a caller using `asyncio.run` per batch gets one per batch. Loops are held weakly, so finished loops drop out of the cache on their own. `get_client`/`get_async_client` are public for callers that want the client itself, and `close_clients`/`aclose_clients` are there for shutdown hooks and tests. A conftest fixture empties the caches between tests, which tests that patch `genai.Client` now need. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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
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.
Why
run_promptandrun_prompt_asynceach built agenai.Client, used it for one prompt, and closed it. Two costs:close()doesn't give the memory back. Each live client holds native grpc/auth memory that isn't returned to the OS. Measured earlier inclaims-analysis-api:/claimantsat concurrency 8 ratcheted +127 MiB and held; an isolated serial loop ofrun_promptramped ~+20 MiB. This is the "memory leak forcing restarts" on the deployed API — the earlier fix addedclose(), which fixed the FD/connection leak but not this.The module comment already said the sync path was meant to share a client.
What
Clients are cached and reused, keyed by the project and location they were built for — the only arguments construction takes (the model name is per request):
asyncio.runper batch (which PASTEL does) gets one per batch instead of one per sentence. Loops are held weakly, so a finished loop's entry drops out on its own.Also public now:
get_client(model_config)/get_async_client(model_config)for callers that want the client itselfclose_clients()/aclose_clients()for shutdown hooks, tests, or when credentials changeTrade-off worth reviewing
Connections for an abandoned loop are now closed at garbage collection rather than deterministically at the end of each call.
aclose_clients()is there for callers that want to be explicit. I think this is the right trade: the previous behaviour was deterministic but paid 100 client constructions per request and still ratcheted native memory.Tests
Replaced the two tests that asserted per-call closing (that was the behaviour being changed) with:
asyncio.runcallsclose_clients/aclose_clientsboth close and forget, so the next call rebuildsget_async_clientrequires a running loopAdded
tests/genai_utils/conftest.pywith an autouse fixture that empties the caches between tests — needed now that clients outlive a single call, since tests patchgenai.Client.147 tests pass; black, isort, flake8 and mypy clean.
Follow-up
The largest remaining win for the checkworthiness path is PASTEL calling
asyncio.runinside a worker thread per batch. If that becameawait-based on the caller's loop, the whole process would share one async client rather than one per batch.