fix(models): stop CompletionsHTTPClient leaking its httpx client via atexit#6379
Open
anxkhn wants to merge 2 commits into
Open
fix(models): stop CompletionsHTTPClient leaking its httpx client via atexit#6379anxkhn wants to merge 2 commits into
anxkhn wants to merge 2 commits into
Conversation
…atexit CompletionsHTTPClient._client registered its per-instance httpx.AsyncClient with atexit.register(self._cleanup_client, client), leaving a permanent strong reference to the client and its connection pool in the process-global atexit registry. Nothing ever called atexit.unregister, so neither close(), aclose(), nor __del__ could release the client before interpreter exit. A long-running deployment that creates many CompletionsHTTPClient instances therefore grew the atexit registry and retained httpx clients without bound. Register the cleanup callback against weakref.proxy(client) so the registry no longer pins the client, mirroring the existing pattern in bigquery_agent_analytics_plugin, and unregister the per-instance callback in close()/aclose() once the client is explicitly closed. Tolerate a dead proxy in _cleanup_client, and retain the fire-and-forget aclose() task in a module-level set so it is not garbage collected while pending. Add regression tests covering the register/unregister pairing and asserting the client is releasable after close(). Signed-off-by: Anas Khan <83116240+anxkhn@users.noreply.github.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.
Link to Issue or Description of Change
1. Link to an existing issue (if applicable):
CONTRIBUTING ("you may instead describe the bug or feature directly within the
PR description, following the structure of our issue templates").
2. Or, if no issue exists, describe the change:
Problem:
CompletionsHTTPClient._client(the OpenAI-compatible HTTP client used byApigeeLlm, insrc/google/adk/models/apigee_llm.py) registers its per-instancehttpx.AsyncClientfor cleanup withatexit.register(self._cleanup_client, client)on first access.atexitholds its callbacks (and their arguments) forthe entire life of the process, so this keeps a strong reference to the
AsyncClientand its connection pool in the process-globalatexitregistry.Nothing ever calls
atexit.unregister, soclose(),aclose(), and__del__cannot actually release the client before interpreter exit. A long-running
deployment that constructs many
CompletionsHTTPClient/ApigeeLlminstancestherefore grows the
atexitregistry and leaks one open connection pool perinstance, without bound.
There is also a smaller latent issue in
_cleanup_client: it schedulesloop.create_task(client.aclose())without keeping a reference to the resultingtask, so the task can be garbage collected before it finishes ("Task was
destroyed but it is pending").
Steps to Reproduce:
pip install google-adkcredentials or network calls are needed):
Observed Behavior (before this change):
Every client (and its connection pool) is pinned by the
atexitregistry evenafter an explicit
close(), an owner drop, and a full GC.Expected Behavior:
After
close()(or once the owner is dropped and finalized), the client shouldbe releasable and its connection pool freed; the process should not accumulate
one retained
AsyncClientper constructed instance.Solution:
weakref.proxy(client)so theatexitregistry no longer keeps the client alive. This mirrors the pattern already
used in
bigquery_agent_analytics_plugin.py(
atexit.register(self._atexit_cleanup, weakref.proxy(batch_processor))).(
self._atexit_callback = partial(self._cleanup_client, weakref.proxy(client)))and unregister exactly that callback in
close()/aclose(). A per-instancepartialis required because_cleanup_clientis a shared@staticmethod, andatexit.unregister(func)removes every registration offunc; unregisteringthe unique partial removes only this instance's handler and leaves other
instances untouched.
_cleanup_client(catchReferenceError), matchingthe same in-repo precedent.
aclose()task in a module-level set until itcompletes, so it is not garbage collected while pending.
Observed Behavior (after this change):
The same snippet reports
0/5alive. The del-only path (drop the owner withoutcalling
close(), relying on__del__+ the non-pinning proxy) also reports0/5.Testing Plan
Unit Tests:
Added
tests/unittests/models/test_completions_http_client.pyregression testscovering: the register/unregister pairing on
close()andaclose(), the no-opwhen no client was ever created, that the client is releasable (weakref goes to
None) afterclose()+ GC, that cleanup tolerates a deadweakref.proxy, andthat the pending
aclose()task is retained until it finishes.Manual End-to-End (E2E) Tests:
The reproduction snippet above is the manual check:
5/5retained clients beforethe change,
0/5after. It needs no Apigee setup because it only drives the HTTPclient's construction/cleanup lifecycle.
Checklist
Additional context
The correct non-pinning pattern already exists in this codebase
(
bigquery_agent_analytics_plugin.py); this change bringsCompletionsHTTPClientin line with it.