Skip to content

[Feature] ToolCompose with parallel dispatch, builtin tools, legacy adapter#3736

Open
vmoens wants to merge 1 commit intogh/vmoens/267/basefrom
gh/vmoens/267/head
Open

[Feature] ToolCompose with parallel dispatch, builtin tools, legacy adapter#3736
vmoens wants to merge 1 commit intogh/vmoens/267/basefrom
gh/vmoens/267/head

Conversation

@vmoens
Copy link
Copy Markdown
Collaborator

@vmoens vmoens commented May 10, 2026

Stack from ghstack (oldest at bottom):

Lands the headline orchestrator: ToolCompose(Compose) parses each
assistant message once, dispatches matched tools concurrently via
asyncio.gather, and injects structured results into History. ChatEnv is
unchanged; ToolCompose drops into any TransformedEnv.

Public API additions in torchrl.envs.llm.agentic:

  • ToolCompose: Tools-only Compose subclass. Raises TypeError on non-Tool
    insert. Owns the parser, honors pass_state_to_tools (mirrors the legacy
    ExecuteToolsInOrder knob), supports per-tool RateLimiter, and surfaces
    three keys per step: ("agentic","any_tool_calls"),
    ("agentic","any_error"), ("agentic","stop_requested").
  • DispatchResult dataclass aggregating one batch item's outcome.
  • RateLimiter combining asyncio.Semaphore + token bucket.
  • PythonTool (Repl-backed, state persists), ShellTool (Sandbox-backed),
    FileReadTool, StopTool. StopTool raises StopSignal which the dispatcher
    translates into the stop_requested flag.
  • as_tool(transform, ...) -- legacy adapter lifting any
    ToolTransformBase-shape object (PythonInterpreter, BrowserTransform,
    MCPToolTransform, SimpleToolTransform) into a Tool. Existing user code
    drops into ToolCompose without rewriting.

Stable call_id is enforced end-to-end: parsers assign deterministic ids,
ToolCompose populates ToolContext.call_id, and parsers' render_result
echoes the same id back into the tool message injected into History.

Nested-loop safety: when ToolCompose._step is called from inside a
running event loop (e.g. Jupyter), dispatch is offloaded to a worker
thread that owns its own loop. asyncio.run is used otherwise.

Tests extend test/llm/test_llm_transforms.py with TestToolCompose,
TestPythonTool, TestShellTool, TestLegacyAdapter (19 tests):
parallel-dispatch wall-time check (3 x 500ms < 0.9s), StopTool
termination, schema validation rejection, raise-on-non-Tool insert,
duplicate-name rejection, pass_state_to_tools on/off, rate-limit
serialization, failure isolation, stable call_id round-trip, nested-loop
safety, legacy adapter end-to-end. All existing legacy
ExecuteToolsInOrder/XMLBlockParser/JSONCallParser tests still pass.

benchmarks/test_llm.py gains an "agentic-dispatch" group with parallel
n=3 / n=8 and a single-call baseline so reviewers can see parallel
dispatch flattens the wall-time curve.

Co-Authored-By: Claude Opus 4.7 noreply@anthropic.com

[ghstack-poisoned]
@pytorch-bot
Copy link
Copy Markdown

pytorch-bot Bot commented May 10, 2026

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3736

Note: Links to docs will display an error until the docs builds have been completed.

❌ 4 New Failures

As of commit c36ab72 with merge base d386287 (image):

NEW FAILURES - The following jobs have failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@github-actions github-actions Bot added the Feature New feature label May 10, 2026
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label May 10, 2026
@github-actions github-actions Bot added Benchmarks rl/benchmark changes llm/ LLM-related PR, triggers LLM CI tests labels May 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Benchmarks rl/benchmark changes CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. Feature New feature llm/ LLM-related PR, triggers LLM CI tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant