[AsyncGRPO] Run sync tools on a thread pool and support async tools

[AmineDiro]

What does this PR do?

AsyncRolloutWorker ran tools inline on the asyncio event loop. A slow sync tool (sandbox,HTTP, code execution) froze every in-flight rollout and stopped the heartbeat. Async tools were rejected at init.

Now:

• Sync tools run on a thread pool sized to max_inflight_tasks, so the loop keeps dispatching
  and collecting vLLM requests while a tool runs.
• Async tools are awaited.
• Tool calls within one turn still run in order, since environment methods are stateful.

Related: #5446, #5444.

---

Note

Medium Risk
Changes concurrent tool execution in the live rollout worker; incorrect threading could cause races in stateful tools, though ordering per turn is preserved.

Overview
Async GRPO rollout no longer runs tools on the asyncio event loop or rejects coroutine tools at startup.

Synchronous tools are dispatched through a ThreadPoolExecutor sized to max_inflight_tasks, so slow blocking work (sandboxes, HTTP, etc.) does not stall concurrent vLLM rollouts or the worker heartbeat. Async tools are **await**ed directly. Tool calls in a single assistant turn still run sequentially in model order; the pool only allows up to max_inflight_tasks sync tools across rollouts at once, so shared-state tools must be thread-safe.

The worker shuts down the pool on exit. Docs add this behavior and tests cover sync/async/failure paths, call ordering, and non-blocking sync execution.

^{Reviewed by Cursor Bugbot for commit b56c754. Bugbot is set up for automated code reviews on this repo. Configure here.}

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

There are 2 total unresolved issues (including 1 from previous review).

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 2eba995. Configure here.}

[albertvillanova]

The approach is right and the failure mode is real: a blocking tool on the event loop stalls every in-flight rollout and the heartbeat, not just its own. A few things before this merges. Two of them are the open Bugbot threads, which I checked and think are both worth acting on.

1. The tests for this change were added and then removed. 117c4d3 drops the whole TestToolExecution class that 57385c2 introduced: the mixed sync/async ordering case (message order, failure counts, per-tool counters, latency rate) and the "slow sync tool does not block the loop" case. What is left is a stub that replaces _execute_tool_calls outright, so it only proves _generate_one awaits the call and asserts nothing about the new dispatch. Both tests are pure Python, no vLLM and no GPU, so they run in the normal suite. If the ticker test was dropped for flakiness, the threshold is the fragile part, not the idea: inline execution yields at most one tick, so asserting something like ticks > 2 keeps the signal with a lot more headroom. Could you restore them?

2. The pool is never shut down (https://github.com/huggingface/trl/pull/7175/changes#r3983308342). Bugbot marks this Low, but it is also a divergence from the sibling loop it points at, and that makes it worth fixing here rather than later: _HarnessRolloutLoop creates the same shaped pool and drains it in _run_loops, while _AsyncRolloutLoop now creates one and never does, with run() closing the loop right after _run_loops returns.

3. Async tools that are callable class instances are not detected (https://github.com/huggingface/trl/pull/7175/changes#r3980199670). This one is right and I would rate it above Medium, for a reason the thread does not give. inspect.iscoroutinefunction is False for an instance whose call is async def, so such a tool goes to the pool, which calls it, gets a coroutine back and never awaits it: the tool message content becomes "<coroutine object ...>", the call is counted as a success, and the only trace is a RuntimeWarning in the child process. That shape is not hypothetical here. Because everything is pickled into the rollout process, both the docs and the validation error in start() steer users toward a callable class instance precisely when a plain function will not do, so this PR documents async tool support and then silently breaks the tool shape the same page recommends. Bound async methods on an environment are detected correctly, so it is only this case.

4. The docs should mention thread safety. Until now every tool call in the process ran on the event loop, so tools were globally serialised and a tool with shared state was implicitly safe. They now run on up to max_inflight_tasks threads concurrently. That is the point of the change, but it is also a new requirement on user code, and the new paragraph only describes the benefit. One sentence saying sync tools must be thread-safe would save someone a very confusing debugging session.

One small suggestion: a word in the comment on why this is a dedicated pool rather than asyncio.to_thread, which is what the reward path uses. The default executor caps at min(32, cpu_count + 4), below max_inflight_tasks, and it is shared with reward scoring, so tools and rewards would contend for the same 32 slots and a burst of slow tools would back up the score queue. That reasoning is not obvious from the diff, and without it the next reader will try to collapse the two idioms.

[albertvillanova]

Thanks. The pool shutdown and the docs sentence are both right. Draining before the loop closes, with the queued calls cancelled, is a better stop than the one I described, and the thread-safety line says exactly what someone needs to know.

Two things left.

The first is my mistake, and it cost you a commit. Point 3 was wrong and you should not have been asked for it. The tool dict is built as {tool.__name__: tool}, and a callable class instance has no __name__, so it raises AttributeError when the dict is built and never reaches the dispatch at all. The new test passes only because it builds the tool dict by hand with string keys: a tool registered the ordinary way cannot take that shape. The __call__ clause therefore guards a case that cannot occur, and the test states support that does not exist. I would drop both, though it is your call. Apologies for the detour.

The second is the tool execution tests. test_every_tool_shape_runs_and_calls_keep_their_order is a good merge of what was there, but the other test did not come back, and it is the one that tests what this PR claims. Ordering and failure counting pass identically whether tools run on the pool or inline on the event loop, so nothing in the suite would catch a regression to inline execution. If it was dropped because the tick count was flaky, a coarser bound keeps the signal, since inline execution yields at most one tick.

[albertvillanova]

Both points are addressed and CI is green.

The restored test is the one that earns its place: awaiting a coroutine does not yield to the loop, so with inline execution the ticker never advances at all and the assertion fails. Anyone who moves tool execution back onto the event loop will hear about it from the suite rather than from a throughput graph.

Thanks for bearing with the detour on point 3, and for the clean shutdown.