fix(code_executors): stop ContainerCodeExecutor leaking via atexit#6392
Open
anxkhn wants to merge 1 commit into
Open
fix(code_executors): stop ContainerCodeExecutor leaking via atexit#6392anxkhn wants to merge 1 commit into
anxkhn wants to merge 1 commit into
Conversation
ContainerCodeExecutor.__init__ registered the bound method self.__cleanup_container with atexit. A bound method keeps a strong reference to the instance, so the process-global atexit registry retained every ContainerCodeExecutor ever created for the whole interpreter lifetime, and its long-lived Docker container was only stopped at exit rather than when the executor was discarded. A server that builds an executor per app/agent/session would grow both the atexit registry and the number of running containers without bound. Register the exit handler with a weakref.proxy instead, mirroring the existing idiom in bigquery_agent_analytics_plugin.py, so the handler no longer keeps the executor alive; a discarded executor (and its container reference) can be garbage collected. __cleanup_container becomes a staticmethod that guards ReferenceError for the case where the proxy is already dead at exit. Add regression tests asserting a dropped executor is not retained and that cleanup still stops and removes the container. 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.
Please ensure you have read the contribution guide before creating a pull request.
Link to Issue or Description of Change
1. Link to an existing issue (if applicable):
2. Or, if no issue exists, describe the change:
Problem:
ContainerCodeExecutor.__init__registers its cleanup withatexit.register(self.__cleanup_container). Becauseself.__cleanup_containeris a bound method, it holds a strong reference to the executor instance, and the
process-global
atexitregistry keeps that reference for the entire interpreterlifetime. As a result every
ContainerCodeExecutorthat is ever created isretained until the process exits, even after it is no longer used, and its
long-lived Docker container is only stopped at interpreter shutdown rather than
when the executor is discarded. In a long-running server that constructs an
executor per app/agent/session, both the
atexitregistry and the number ofrunning containers grow without bound.
Solution:
Register the exit handler through a
weakref.proxy(self)instead of a boundmethod, so the
atexitregistry no longer keeps the executor alive; a discardedexecutor (and its container reference) can be garbage collected normally, and
atexitstill stops the container at shutdown if the executor is still in use.__cleanup_containerbecomes a@staticmethodthat reads the executor throughthe proxy and guards
ReferenceError, so if the executor has already beencollected by the time the interpreter exits the handler is a safe no-op.
This mirrors the pattern already used in this repository for the same situation:
src/google/adk/plugins/bigquery_agent_analytics_plugin.pyregisters itsatexitcleanup withweakref.proxy(...)and guards the callback againstReferenceError.The change is limited to
container_code_executor.py(addimport weakref,switch the registration, convert the cleanup to a
ReferenceError-guardedstaticmethod) plus two regression tests. Behavior for a live executor is
unchanged: the container is still stopped and removed at exit.
Testing Plan
Unit Tests:
Added to
tests/unittests/code_executors/test_container_code_executor.py:test_executor_is_not_retained_by_atexit: constructs an executor (Dockermocked), drops it, forces a
gc.collect(), and asserts aweakref.refto itis
None, i.e. theatexitregistration no longer retains it. This testfails on the previous bound-method registration and passes with the fix.
test_cleanup_stops_and_removes_container: invokes the cleanup while theexecutor is still referenced and asserts the container is still stopped and
removed, confirming the exit behavior is preserved.
pytestresults (uv run --no-sync pytest, Python 3.11):Manual End-to-End (E2E) Tests:
This is an object-lifecycle fix with no user-facing behavior change, so it is
covered by unit tests rather than a manual
adk web/runnerflow. The leak canbe observed directly on the previous code with Docker mocked: after constructing
and deleting a
ContainerCodeExecutorand runninggc.collect(), aweakref.refto it is still alive (retained by theatexitregistry); with thefix the reference is cleared. Running
atexit._run_exitfuncs()after theexecutor has been collected completes without raising, confirming the
ReferenceErrorguard makes a dead-proxy callback a safe no-op at shutdown.Checklist
Additional context
The weakref approach is used deliberately rather than
atexit.unregister:atexit.unregistermatches on the exact registered(func, args, kwargs)andcannot precisely target a registration made with bound arguments, so a
weakref-based registration (observable via garbage collection) is the reliable
fix and is also the idiom already established in this repository.