fix: reuse Gemini client instead of building one per call#7
Conversation
Each run_prompt/run_prompt_async call constructed a fresh genai.Client and never closed it. The client owns httpx connection pools, so under sustained load the abandoned pools accumulated and process memory grew steadily. - Cache the sync client per (project, location) and reuse it; the sync client is thread-safe and the set of pairs is tiny, so this is a bounded process-lifetime singleton. - run_prompt now calls the sync API directly, dropping the per-call asyncio.run event loop as well. - run_prompt_async still builds a client per call (its async transport binds to the running loop, which is often short-lived), but now closes it in a finally so the pool is released. - Extract _build_request / _parse_response so the sync and async paths share request building and response handling. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
dearden
left a comment
dearden
left a comment
There was a problem hiding this comment.
Seems fine. I am a little confused why it fixed it differently for sync and async (would finally not work for both?) but I don't object - all makes sense. Had one little question about TTL which may be unfounded.
| ) | ||
| try: | ||
| response = await client.aio.models.generate_content(**request) | ||
| finally: |
There was a problem hiding this comment.
I don't object to how it's doing either, but how come it didn't do both async and sync the same way (with finally?)
|
|
||
|
|
||
| @lru_cache(maxsize=None) | ||
| def _get_client(project: str, location: str) -> genai.Client: |
There was a problem hiding this comment.
This makes sense, but I am wondering - if the claims analysis API is always running, will the clients not become stale without some kind of timeout. Shouldn't we set some kind of TTL?
There was a problem hiding this comment.
It uses a httpx connection pool which should be fine as it should handle removing idle connections behind the scenes.
|
I have just realised that this is literally the worst possible description for what this PR does now, but alas, the button has been pressed. |
2 participants
Problem
run_prompt/run_prompt_asyncconstructed a freshgenai.Clienton every call and never closed it. The client ownshttpxconnection pools, so under sustained load the abandoned pools accumulate and process memory climbs steadily over time. This was surfaced investigating a slow memory leak in the claims-analysis-api, where every Gemini-backed task (claim matching, claimant detection, checkworthy keyword filter, pastel) routes through these functions.Fix
(project, location)via_get_clientand reuse it. The sync client is thread-safe and the set of pairs is effectively one, so this is a bounded process-lifetime singleton.run_promptnow calls the syncmodels.generate_contentAPI directly — also dropping the per-callasyncio.runevent loop.run_prompt_asyncstill builds a client per call (its async transport binds to the running event loop, which is often short-lived — e.g. driven underasyncio.run— so it can't be safely shared), but now closes it in afinallyso the pool is released rather than leaked._build_request/_parse_responseso the sync and async paths share request building and response handling with no behavioural drift.Verification
uv run pytest→ 102 passed (addedtest_run_prompt_reuses_cached_clientandtest_run_prompt_async_closes_client)uv run ruff check/format→ cleanuv run mypy src/genai_utils/gemini.py→ cleanFollow-up
run_prompt_asyncstill creates a client per call (correctness over the cross-event-loop hazard). If pastel's fan-out makes that churn show up in latency, a follow-up could reuse one client across the gathered calls within a single request.🤖 Generated with Claude Code