Skip to content

fix(code_executors): stop ContainerCodeExecutor leaking via atexit#6392

Open
anxkhn wants to merge 1 commit into
google:mainfrom
anxkhn:fix/container-executor-atexit-leak
Open

fix(code_executors): stop ContainerCodeExecutor leaking via atexit#6392
anxkhn wants to merge 1 commit into
google:mainfrom
anxkhn:fix/container-executor-atexit-leak

Conversation

@anxkhn

@anxkhn anxkhn commented Jul 15, 2026

Copy link
Copy Markdown
Contributor

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

  • N/A (no existing issue; described below per the issue template structure).

2. Or, if no issue exists, describe the change:

Problem:

ContainerCodeExecutor.__init__ registers its cleanup with
atexit.register(self.__cleanup_container). Because self.__cleanup_container
is a bound method, it holds a strong reference to the executor instance, and the
process-global atexit registry keeps that reference for the entire interpreter
lifetime. As a result every ContainerCodeExecutor that is ever created is
retained 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 atexit registry and the number of
running containers grow without bound.

Solution:

Register the exit handler through a weakref.proxy(self) instead of a bound
method, so the atexit registry no longer keeps the executor alive; a discarded
executor (and its container reference) can be garbage collected normally, and
atexit still stops the container at shutdown if the executor is still in use.
__cleanup_container becomes a @staticmethod that reads the executor through
the proxy and guards ReferenceError, so if the executor has already been
collected 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.py registers its
atexit cleanup with weakref.proxy(...) and guards the callback against
ReferenceError.

The change is limited to container_code_executor.py (add import weakref,
switch the registration, convert the cleanup to a ReferenceError-guarded
staticmethod) plus two regression tests. Behavior for a live executor is
unchanged: the container is still stopped and removed at exit.

Testing Plan

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

Added to tests/unittests/code_executors/test_container_code_executor.py:

  • test_executor_is_not_retained_by_atexit: constructs an executor (Docker
    mocked), drops it, forces a gc.collect(), and asserts a weakref.ref to it
    is None, i.e. the atexit registration no longer retains it. This test
    fails on the previous bound-method registration and passes with the fix.
  • test_cleanup_stops_and_removes_container: invokes the cleanup while the
    executor is still referenced and asserts the container is still stopped and
    removed, confirming the exit behavior is preserved.

pytest results (uv run --no-sync pytest, Python 3.11):

tests/unittests/code_executors/test_container_code_executor.py .... [100%]
4 passed

tests/unittests/code_executors/  ->  80 passed

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/runner flow. The leak can
be observed directly on the previous code with Docker mocked: after constructing
and deleting a ContainerCodeExecutor and running gc.collect(), a
weakref.ref to it is still alive (retained by the atexit registry); with the
fix the reference is cleared. Running atexit._run_exitfuncs() after the
executor has been collected completes without raising, confirming the
ReferenceError guard makes a dead-proxy callback a safe no-op at shutdown.

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end. (Covered by unit tests; see above.)
  • Any dependent changes have been merged and published in downstream modules. (None.)

Additional context

The weakref approach is used deliberately rather than atexit.unregister:
atexit.unregister matches on the exact registered (func, args, kwargs) and
cannot 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.

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>
@adk-bot adk-bot added the tools [Component] This issue is related to tools label Jul 15, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

tools [Component] This issue is related to tools

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants