Skip to content

perf: share Gemini clients instead of building one per prompt - #11

Open
JamesMcMinn wants to merge 1 commit into
mainfrom
reuse-gemini-client
Open

perf: share Gemini clients instead of building one per prompt#11
JamesMcMinn wants to merge 1 commit into
mainfrom
reuse-gemini-client

Conversation

@JamesMcMinn

Copy link
Copy Markdown
Contributor

Why

run_prompt and run_prompt_async each built a genai.Client, used it for one prompt, and closed it. Two costs:

  1. Construction isn't free — it resolves credentials and sets up a transport. PASTEL sends one prompt per sentence, so a 100-sentence checkworthiness request built 100 clients.
  2. close() doesn't give the memory back. Each live client holds native grpc/auth memory that isn't returned to the OS. Measured earlier in claims-analysis-api: /claimants at concurrency 8 ratcheted +127 MiB and held; an isolated serial loop of run_prompt ramped ~+20 MiB. This is the "memory leak forcing restarts" on the deployed API — the earlier fix added close(), 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):

  • sync: one client per project/location, shared across threads behind a lock (the sync client is thread-safe, and callers hit it from worker threads).
  • async: additionally keyed by the running event loop, because the async transport binds to the loop it's first used on. This is what the old comment said made sharing impossible — keying by loop is what makes it possible. A caller gathering many prompts under one loop now builds one client instead of one per prompt; a caller doing asyncio.run per 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 itself
  • close_clients() / aclose_clients() for shutdown hooks, tests, or when credentials change

Trade-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:

  • one client for repeated prompts; one per project/location; a different model on the same project/location reuses the client
  • async: one client per loop, and not shared between two asyncio.run calls
  • close_clients/aclose_clients both close and forget, so the next call rebuilds
  • get_async_client requires a running loop

Added tests/genai_utils/conftest.py with an autouse fixture that empties the caches between tests — needed now that clients outlive a single call, since tests patch genai.Client.

147 tests pass; black, isort, flake8 and mypy clean.

Follow-up

The largest remaining win for the checkworthiness path is PASTEL calling asyncio.run inside a worker thread per batch. If that became await-based on the caller's loop, the whole process would share one async client rather than one per batch.

`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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant