Skip to content

[Fix] Add AgentLoop shutdown and pause/resume lifecycle hooks - #2135

Open
matrix72c wants to merge 4 commits into
InternLM:mainfrom
matrix72c:feat/sandbox-environment-design
Open

matrix72c wants to merge 4 commits into
InternLM:mainfrom
matrix72c:feat/sandbox-environment-design

Conversation

@matrix72c

@matrix72c matrix72c commented Oct 8, 2026 •

Copy link
Copy Markdown
Contributor

Problem

Resource-owning AgentLoops have no framework shutdown hook. When training finishes or raises, a loop can keep background work and external clients alive after the trainer has stopped. In the disaggregated trainer, a producer blocked on generation never observes the shutdown flag, so fit() can hang before any owner is closed.

Change

  • Add async AgentLoop.close(), a no-op by default. Both RL trainers close their training and evaluation managers when fit() exits, including on error. Managers forward close() to local loops, Ray actors and router workers; every loop is closed even if another one fails. Close failures are logged and do not replace the training result.
  • The disaggregated trainer cancels a producer that is still blocked on generation before closing loop owners. An idle producer still exits through the shutdown flag.
  • The built-in sandbox loop rejects new runs, cancels active runners and waits for their existing finally blocks to release their pools.
  • Add AgentLoop.prepare_pause() and AgentLoop.resume(), both no-ops by default. Managers call prepare_pause() before pause_generation() and resume() after continue_generation(), in both the colocated and disaggregated managers. Loops whose model calls come from an external harness use them to drain in-flight calls before the abort. If prepare_pause() raises, the manager does not abort inference.

The contract is documented in docs/en/rl/agent_loop_shutdown.md. Existing loops, Runner, Judger and AgentInSandboxLoopConfig usage are unchanged.

Test plan

python -m pytest -q \
  tests/rl/test_agent_loop_shutdown.py \
  tests/rl/test_rl_colocate_trainer.py \
  tests/rl/test_rl_disaggregated_trainer.py \
  tests/rl/test_multi_task_agent_loop_manager.py \
  tests/rl/test_producer.py \
  tests/rl/test_trajectory_logging.py

86 passed on top of current upstream/main. New tests cover normal/error/cancelled trainer exits, training and evaluation owners, local/actor/router dispatch, failure isolation, the sandbox Runner/Pool cleanup path, producer shutdown ordering, and pause/resume hook ordering. Ruff, docformatter and git diff --check pass; mypy reports no errors on changed lines. Tests use fake backends and do not exercise live sandbox deletion or GPU training.

@matrix72c
matrix72c force-pushed the feat/sandbox-environment-design branch from b08a5d5 to 2db9f71 Compare October 15, 2026 03:04
@matrix72c matrix72c changed the title [Docs] Propose sandbox agent loop environment interfaces [Refactor] Separate sandbox environment ownership from Runner execution Oct 15, 2026
@matrix72c
matrix72c marked this pull request as ready for review October 15, 2026 03:04
@matrix72c
matrix72c force-pushed the feat/sandbox-environment-design branch from 3225eaa to eee2016 Compare October 21, 2026 06:11
@matrix72c matrix72c changed the title [Refactor] Separate sandbox environment ownership from Runner execution [Fix] Close agent loops on RL trainer shutdown Oct 21, 2026
Agent loops that serve an external harness need to stop admitting model
calls before the manager aborts inference, and reopen admission after it
resumes; otherwise the abort cuts requests the harness is still waiting on.

Add optional prepare_pause/resume hooks to AgentLoop, RouterAgentLoop and
AgentLoopActor, and call them around pause_generation/continue_generation in
both the colocated and disaggregated managers. A failed drain now prevents
the global abort. Generalize the close helper into call_agent_loop_hook.
@matrix72c matrix72c changed the title [Fix] Close agent loops on RL trainer shutdown [Fix] Add AgentLoop shutdown and pause/resume lifecycle hooks Oct 26, 2026
Drop the cancellation-proof cleanup helpers and the error-chaining
machinery. The trainer closes managers through asyncio_run, which never
cancels the cleanup, so await_cleanup, gather_cleanup and CleanupError
guarded a path that does not occur. Close failures are now logged and no
longer replace the training result.

Also remove id()-based owner dedup and the optional-hook getattr guard,
which cannot trigger, keep a single closing check in the sandbox loop,
shrink the lifecycle doc to the extension contract, and revert
formatting changes to untouched test lines.

This branch has not been deployed

No deployments
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