[misc] feat: trace rollout generation and tool calls using weave #2345

chenhaiq · 2025-07-03T09:29:20Z

What does this PR do?

Provide rollout generation and tool calls details in wandb weave to help debugging agentic RL.

2 new interfaces:

rollout_trace_attr contextmanager: used to mark sample_index、step、rollout_n and experience name for a trajectory.
rollout_trace_op decorator：mark the method to trace. It must be a method of an instance.

related issue #2188

Checklist Before Starting

Search for similar PRs. Paste at least one query link here: ...
Format the PR title as [{modules}] {type}: {description} (This will be checked by the CI)
- {modules} include fsdp, megatron, sglang, vllm, rollout, trainer, ci, training_utils, recipe, hardware, deployment, ray, worker, single_controller, misc, perf, model, algo, env, tool, ckpt, doc, data
- If this PR involves multiple modules, separate them with , like [megatron, fsdp, doc]
- {type} is in feat, fix, refactor, chore, test
- If this PR breaks any API (CLI arguments, config, function signature, etc.), add [BREAKING] to the beginning of the title.
- Example: [BREAKING][fsdp, megatron] feat: dynamic batching

Test

API and Usage Example

options:
+trainer.rollout_trace.backend=weave: only wandb weave is support in this PR. Leave the reset of trace tool to the community.
+trainer.rollout_trace.token2text=False: whether append decoded text in result of run method.

High-Level Design

n/a

Specific Changes

Only works for async rollout from agent loop. No effect for sync rollout.

Checklist Before Submitting

Important

Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.

Read the Contribute Guide.
Apply pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
Add / Update the documentation.
Add unit or end-to-end test(s) to the CI workflow to cover all the code. If not feasible, explain why: ...
Once your PR is ready for CI, send a message in the ci-request channel in the verl Slack workspace.

zhaochenyang20 · 2025-07-03T15:56:14Z

@chenhaiq Haiquan, could you help to check wether this is also supported in SGLang? I am checking with our team also. Thanks!

chenhaiq · 2025-07-04T01:20:41Z

@chenhaiq Haiquan, could you help to check wether this is also supported in SGLang? I am checking with our team also. Thanks!

sure，I will support sglang

chenhaiq · 2025-07-07T03:45:50Z

@chenhaiq Haiquan, could you help to check wether this is also supported in SGLang? I am checking with our team also. Thanks!

It supports sglang async mode, please have a code review.

wuxibin89 · 2025-07-08T06:40:00Z

verl/experimental/agent_loop/agent_loop.py

+            rollout_n += 1
+        else:
+            rollout_n = 0
+        trajectory_info.append({"step": step, "sample_index": index[i], "rollout_n": rollout_n})


what is the meaning of rollout_n?

rollout_n is used to associate trace entry with actor_rollout_ref.rollout.n.
For example, when actor_rollout_ref.rollout.n=2, the _run_agent_loop method will be executed 2 times.
Both of the 2 _run_agent_loop calls have their own trace entry. Setting a rollout_n is used to tell which time it is traced from.

wuxibin89 · 2025-07-08T06:55:11Z

verl/utils/rollout_trace.py

+        yield
+
+
+def rollout_trace_op(func):


Seems like this function can trace every function, include tooling, agent loop run, ...?

It can trace any method in an class instance.

eric-haibin-lin · 2025-07-08T16:45:42Z

docs/faq/faq.rst

+.. code:: bash
+
+  +trainer.rollout_trace.token2text=True # default to False for better performance
+


@chenhaiq could u add the visualization result from this PR to the documnetation page as well? please upload the image to https://github.com/eric-haibin-lin/verl-community/tree/main/docs

chenhaiq · 2025-07-09T10:40:40Z

corresponding feature in mlflow:

…cengine#2345) ### What does this PR do? Provide rollout generation and tool calls details in wandb weave to help debugging agentic RL. 2 new interfaces: 1. rollout_trace_attr contextmanager: used to mark sample_index、step、rollout_n and experience name for a trajectory. 2. rollout_trace_op decorator：mark the method to trace. It must be a method of an instance. related issue volcengine#2188 ### Checklist Before Starting - [X] Search for similar PRs. Paste at least one query link here: ... - [X] Format the PR title as `[{modules}] {type}: {description}` (This will be checked by the CI) - `{modules}` include `fsdp`, `megatron`, `sglang`, `vllm`, `rollout`, `trainer`, `ci`, `training_utils`, `recipe`, `hardware`, `deployment`, `ray`, `worker`, `single_controller`, `misc`, `perf`, `model`, `algo`, `env`, `tool`, `ckpt`, `doc`, `data` - If this PR involves multiple modules, separate them with `,` like `[megatron, fsdp, doc]` - `{type}` is in `feat`, `fix`, `refactor`, `chore`, `test` - If this PR breaks any API (CLI arguments, config, function signature, etc.), add `[BREAKING]` to the beginning of the title. - Example: `[BREAKING][fsdp, megatron] feat: dynamic batching` ### Test <img width="1910" alt="截屏2025-07-03 下午4 09 58" src="https://github.com/user-attachments/assets/ff30bbca-f9c8-434f-a3c2-0e333d16fa68" /> <img width="1895" alt="截屏2025-07-03 下午4 11 27" src="https://github.com/user-attachments/assets/0b9ed8db-58a7-4769-88fb-bda204dc9fc8" /> ### API and Usage Example options: +trainer.rollout_trace.backend=weave: only wandb weave is support in this PR. Leave the reset of trace tool to the community. +trainer.rollout_trace.token2text=False: whether append decoded text in result of run method. ### High-Level Design n/a ### Specific Changes Only works for async rollout from agent loop. No effect for sync rollout. ### Checklist Before Submitting > [!IMPORTANT] > Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review. - [x] Read the [Contribute Guide](https://github.com/volcengine/verl?tab=readme-ov-file#contribution-guide). - [x] Apply [pre-commit checks](https://github.com/volcengine/verl?tab=readme-ov-file#code-linting-and-formatting): `pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always` - [x] Add / Update [the documentation](https://github.com/volcengine/verl/tree/main/docs). - [x] Add unit or end-to-end test(s) to [the CI workflow](https://github.com/volcengine/verl/tree/main/.github/workflows) to cover all the code. If not feasible, explain why: ... - [x] Once your PR is ready for CI, send a message in [the `ci-request` channel](https://verl-project.slack.com/archives/C091TCESWB1) in [the `verl` Slack workspace](https://join.slack.com/t/verl-project/shared_invite/zt-3855yhg8g-CTkqXu~hKojPCmo7k_yXTQ).

chenhaiq added 5 commits July 2, 2025 19:43

[misc] feat: trace llm generate and toolcall in rollout

1c84b12

Merge branch 'main' into rollout_trace

21ccdef

[misc] feat: trace rollout generation and toolcalls using weave

478f609

Merge branch 'main' into rollout_trace

bf83688

add doc and rename test

d8da588

chenhaiq requested review from PeterSH6, eric-haibin-lin, tongyx361, vermouth1992 and zhaochenyang20 as code owners July 3, 2025 09:29

chenhaiq requested review from SwordFaith and wuxibin89 and removed request for PeterSH6, eric-haibin-lin, tongyx361, vermouth1992 and zhaochenyang20 July 3, 2025 09:29

chenhaiq added 4 commits July 3, 2025 17:58

add test dependance

7e74ceb

fix cpu ut

56691e5

add license

dd09dee

fix error in vllm sync mode

c952790

chenhaiq added 3 commits July 4, 2025 09:37

fix error in test_basic_agent_loop.py

47b244a

Merge branch 'main' into rollout_trace

6297f0d

fix import

39e496f

chenhaiq requested a review from eric-haibin-lin July 7, 2025 06:20

Merge branch 'main' into rollout_trace

0d2d0e6

wuxibin89 reviewed Jul 8, 2025

View reviewed changes

wuxibin89 approved these changes Jul 8, 2025

View reviewed changes

wuxibin89 merged commit ec4433c into volcengine:main Jul 8, 2025
54 of 57 checks passed

eric-haibin-lin reviewed Jul 8, 2025

View reviewed changes

shuyhere mentioned this pull request Jul 22, 2025

[fix] rollout: Fix rollout_trace_op not tracing decoded prompt and response #2690

Open

7 tasks

Jakhongir0103 mentioned this pull request Jul 25, 2025

[rollout] feat: extend agent loop to multimodal modals #2754

Closed

7 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[misc] feat: trace rollout generation and tool calls using weave #2345

[misc] feat: trace rollout generation and tool calls using weave #2345

Uh oh!

chenhaiq commented Jul 3, 2025

Uh oh!

zhaochenyang20 commented Jul 3, 2025

Uh oh!

chenhaiq commented Jul 4, 2025

Uh oh!

chenhaiq commented Jul 7, 2025

Uh oh!

wuxibin89 Jul 8, 2025

Uh oh!

chenhaiq Jul 8, 2025

Uh oh!

wuxibin89 Jul 8, 2025

Uh oh!

chenhaiq Jul 8, 2025

Uh oh!

Uh oh!

eric-haibin-lin Jul 8, 2025

Uh oh!

chenhaiq commented Jul 9, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		.. code:: bash

		+trainer.rollout_trace.token2text=True # default to False for better performance

[misc] feat: trace rollout generation and tool calls using weave #2345

[misc] feat: trace rollout generation and tool calls using weave #2345

Uh oh!

Conversation

chenhaiq commented Jul 3, 2025

What does this PR do?

Checklist Before Starting

Test

API and Usage Example

High-Level Design

Specific Changes

Checklist Before Submitting

Uh oh!

zhaochenyang20 commented Jul 3, 2025

Uh oh!

chenhaiq commented Jul 4, 2025

Uh oh!

chenhaiq commented Jul 7, 2025

Uh oh!

wuxibin89 Jul 8, 2025

Choose a reason for hiding this comment

Uh oh!

chenhaiq Jul 8, 2025

Choose a reason for hiding this comment

Uh oh!

wuxibin89 Jul 8, 2025

Choose a reason for hiding this comment

Uh oh!

chenhaiq Jul 8, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

eric-haibin-lin Jul 8, 2025

Choose a reason for hiding this comment

Uh oh!

chenhaiq commented Jul 9, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants