From 9f7ed8f9e5ce7f2d031363bafdbea19e7e04d859 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=80=E8=BF=9F?= Date: Thu, 27 Aug 2026 18:11:48 +0800 Subject: [PATCH] [tracing] squashed: sample-flow trace, progress panel and trajectory instrumentation for v0+v1 trainer paths Squashed commits (eeca5b74..HEAD) on verl-tracing-v1: eeca5b74 [tracing] feat: add sample-flow trace, training progress and trajectory instrumentation 60d8a58b [tracing] feat: extend sample-flow trace and trajectory dump to V1 trainer path 36f1321b [tracing] feat: align tracing with latest verl-horizon and extend to V1 trainer path 968855f0 [tracing] chore: rename NGPUS to NPROC_PER_NODE in smoke scripts; multinode defaults USE_V1=True 0725d470 [tracing] fix: suppress sample-flow trace and trajectory dump for validation samples c7962df4 [tracing] fix: multinode script ray wait loop and cross-node NCCL/RoCEv2 config 857f6309 [tracing] feat: expose progress dir via /tmp/verl_progress_path.txt e9ca8012 [tracing] feat: add progress panel lines and restructure trajectory JSONL 2a478db0 [tracing] chore: upgrade smoke test scripts to Qwen3.8-27B and enrich multiturn task b72a913e [tracing] feat: rename AgentLoopMetrics.compute_score to reward_score and store trajectory reward 549594d0 [tracing] fix: address review findings in trace pipeline and smoke scripts --- .qoder/skills/verl-trace-instrument/SKILL.md | 281 +++++++++++ .../TRACE_SMOKE_NEW_IMAGE_GUIDE.md | 230 +++++++++ .../verl-trace-instrument/insertion_points.md | 447 ++++++++++++++++++ .../insertion_points_v1.md | 310 ++++++++++++ .../verl-trace-instrument/progress_points.md | 199 ++++++++ .../trace_data_formats.md | 239 ++++++++++ .../verl-trace-instrument/trace_module.py | 313 ++++++++++++ .../trajectory_points.md | 293 ++++++++++++ .../trajectory_store_module.py | 136 ++++++ .../trajectory_tracker_module.py | 226 +++++++++ .../trajectory_types_module.py | 321 +++++++++++++ scripts/coding_tool.py | 50 ++ scripts/prep_gsm8k_coding.py | 59 +++ scripts/run_trace_smoke_test.sh | 241 ++++++++++ scripts/run_trace_smoke_test_7b.sh | 230 +++++++++ scripts/run_trace_smoke_test_7b_multinode.sh | 313 ++++++++++++ scripts/run_trace_smoke_test_multiturn.sh | 265 +++++++++++ verl/experimental/agent_loop/agent_loop.py | 258 +++++++--- .../agent_loop/single_turn_agent_loop.py | 32 ++ .../agent_loop/tool_agent_loop.py | 55 +++ verl/experimental/separation/ray_trainer.py | 53 ++- verl/single_controller/ray/base.py | 2 + verl/trainer/constants_ppo.py | 5 + verl/trainer/ppo/ray_trainer.py | 121 +++++ verl/trainer/ppo/v1/agent_loop_tq.py | 41 +- verl/trainer/ppo/v1/trainer_base.py | 137 ++++++ verl/trainer/ppo/v1/trainer_colocate_async.py | 7 + verl/trainer/ppo/v1/trainer_separate_async.py | 7 + verl/trainer/ppo/v1/trainer_sync.py | 7 + verl/utils/tracing/__init__.py | 43 ++ verl/utils/tracing/trace.py | 427 +++++++++++++++++ verl/utils/tracing/trajectory_store.py | 136 ++++++ verl/utils/tracing/trajectory_tracker.py | 226 +++++++++ verl/utils/tracing/trajectory_types.py | 321 +++++++++++++ .../sglang_rollout/async_sglang_server.py | 13 + 35 files changed, 5976 insertions(+), 68 deletions(-) create mode 100644 .qoder/skills/verl-trace-instrument/SKILL.md create mode 100644 .qoder/skills/verl-trace-instrument/TRACE_SMOKE_NEW_IMAGE_GUIDE.md create mode 100644 .qoder/skills/verl-trace-instrument/insertion_points.md create mode 100644 .qoder/skills/verl-trace-instrument/insertion_points_v1.md create mode 100644 .qoder/skills/verl-trace-instrument/progress_points.md create mode 100644 .qoder/skills/verl-trace-instrument/trace_data_formats.md create mode 100644 .qoder/skills/verl-trace-instrument/trace_module.py create mode 100644 .qoder/skills/verl-trace-instrument/trajectory_points.md create mode 100644 .qoder/skills/verl-trace-instrument/trajectory_store_module.py create mode 100644 .qoder/skills/verl-trace-instrument/trajectory_tracker_module.py create mode 100644 .qoder/skills/verl-trace-instrument/trajectory_types_module.py create mode 100644 scripts/coding_tool.py create mode 100644 scripts/prep_gsm8k_coding.py create mode 100644 scripts/run_trace_smoke_test.sh create mode 100755 scripts/run_trace_smoke_test_7b.sh create mode 100755 scripts/run_trace_smoke_test_7b_multinode.sh create mode 100755 scripts/run_trace_smoke_test_multiturn.sh create mode 100644 verl/utils/tracing/__init__.py create mode 100644 verl/utils/tracing/trace.py create mode 100644 verl/utils/tracing/trajectory_store.py create mode 100644 verl/utils/tracing/trajectory_tracker.py create mode 100644 verl/utils/tracing/trajectory_types.py diff --git a/.qoder/skills/verl-trace-instrument/SKILL.md b/.qoder/skills/verl-trace-instrument/SKILL.md new file mode 100644 index 00000000..c9c1b525 --- /dev/null +++ b/.qoder/skills/verl-trace-instrument/SKILL.md @@ -0,0 +1,281 @@ +--- +name: verl-trace-instrument +description: Automatically instruments the verl codebase (v0.9.x; v0 RayPPOTrainer + AgentLoop and v1 TaskRunnerV1 + TransferQueue paths) with three tracing facilities aligned with the latest verl-horizon implementation - sample-flow tracing (batch trace_flow tracking each (uid, traj_uid) through the full RL pipeline, CSV under VERL_TRACE_DIR/sample/), training progress tracking for colocated, separation and v1 trainer-mode paths ([horizon raw lines] stdout + hourly progress/ log files via log_progress), and per-trajectory JSONL dumps (Turn/Trajectory/TurnRecord/TrajectoryTracker/TrajectoryStore). Use when the user asks to add trace/tracing/插桩/trace_flow/训练进度展示/trajectory监控 to verl, enable sample flow tracking, or mentions VERL_TRACE_DIR. +--- + +# verl Trace 自动插桩(对齐最新 verl-horizon,覆盖 v0 + v1 trainer) + +为 verl(v0.9.x)自动插入三类追踪代码,同时覆盖两条训练主路径: + +- **v0 路径**:RayPPOTrainer + AgentLoop(`trainer.use_v1=False`); +- **v1 路径(默认)**:TaskRunnerV1 + TransferQueue + (`verl/trainer/ppo/v1/`,trainer_mode = sync / colocate_async / separate_async), + agent loop 侧复用共享的 AgentLoopWorker 代码,traj_uid ≡ session_id。 + +1. **样本流转 trace**(trace_flow,**batch API**):`uid`/`traj_uid` 既接受单个字符串 + 也接受序列(list/ndarray),函数内部自动 fan-out、标量广播;支持 `global_step` + 参数。追踪每个 `(uid, traj_uid)` 样本从数据加载到参数更新的完整链路。 + `[FLOW_TRACE]` stdout 与 CSV(`{trace_dir}/sample/`,cat=sample)**均由 + `VERL_TRACE_DIR` 统一门控**(未设置时两者都静默)。 +2. **训练进度跟踪(训推一体 + 训推分离)**:与 verl-horizon 一致的裸格式行, + 两条路径各用其原生格式:训推一体无前缀(如 `start compute reward`、 + `start generation`),训推分离带 `disagg: ` 前缀(如 `disagg: start compute + reward`)且有 `Trainer: successfully load global batch@N ...` fetch 行; + stdout 常开,设置 `VERL_TRACE_DIR` 后同步逐行落盘 + `{trace_dir}/progress/{YYYYMMDD-HH}.log`(按小时轮转)。 + 覆盖前端 11 个训练阶段中的 10 个(`update prompt2accs` 为 horizon + 异步过采样特有,无实现不插桩)。 +3. **Trajectory 监控**(TrajectoryStore,**horizon schema**): + `Turn`/`Trajectory`/`TurnRecord`/`TrajectoryTracker` 数据模型;每次 + `server_manager.generate` 采集一个 TurnRecord(completion text/ids、 + finish_reason、timing,tool 轮次经 `patch_observation` 记录 observation); + agent loop 完成后在 `AgentLoopWorker._run_agent_loop` 中央点位 build + dump + 到 `{trace_dir}/trajectory/trajectory_{instance}_{ts}.jsonl`(100MB 轮转、 + 自动裁剪 token_ids 大字段)。 + +## 铁律(必须遵守) + +1. **零侵入**:只插入 trace 调用与必要的 import;**严禁**修改任何业务逻辑、 + 控制流、数据结构或函数签名。允许的数据新增仅限:`traj_uid` 字段(插入点 2 + 注入与插入点 2b 丢失兜底重注入,必须成对)、`trace_traj_key`(经 kwargs 传给 + run(),供 TurnRecord 采集挂载)。 +2. **插桩代码绝不中断训练**:样本字段读取一律用 `.get(..., 兜底值)` 防御式写法; + dump/trace 内部可选字段用 getattr 兜底;**序列化前必须把 pydantic 对象转为 + dict**(`AgentLoopOutput.metrics` 是 `AgentLoopMetrics`,直接传 `json.dumps` + 会抛 not JSON serializable,用 `model_dump()` 强转);引擎侧取出的 + `stop_reason` 用 `str()` 保护后写入 `Turn.finish_reason`。 +3. **锚点定位,不依赖行号**:每个插入点给出锚点代码片段。先 grep 定位锚点,锚点未命中时 + 语义搜索等价代码并**重新确认上下文**,仍找不到则报告该点跳过,不得猜测插入。 +4. **幂等**:插入前检查目标位置附近是否已有 `trace_flow` 调用,已有则跳过该点。 +5. **import 统一在模块顶部**:`from verl.utils.tracing import ...`,不在函数体内局部 import。 + +## 工作流 + +复制此清单并跟踪进度: + +``` +插桩进度: +- [ ] Step 1: 环境检查 +- [ ] Step 2: 创建 trace 工具模块(4 个文件) +- [ ] Step 3: 注入 Ray worker 环境变量(两处) +- [ ] Step 4: 应用 17+1 个 v0 样本流转插入点(batch 风格) +- [ ] Step 4a: 应用 14 个 v1 样本流转插入点 + v1 进度插桩(32 个 log_progress) +- [ ] Step 4b: 应用 v0 训练进度插桩点(P1-P11 + 心跳 + worker 侧 CSV) +- [ ] Step 4c: 应用 trajectory 采集与中央 dump +- [ ] Step 5: 验证 +- [ ] Step 6: 输出报告 +``` + +### Step 1: 环境检查 + +1. 确认仓库根(存在 `verl/trainer/ppo/ray_trainer.py`)。 +2. 版本适配检查:确认存在 `verl/experimental/agent_loop/agent_loop.py` 与 + `verl/workers/engine/base.py`(本 skill 适配 v0.9.x engine 架构)。 + 若存在 `verl/trainer/disaggregated/core.py`,说明是旧版 disaggregated + 架构,**停止**并告知用户本 skill 不适用。 +3. 幂等检查:`grep -rn "trace_flow(" verl/ | wc -l`,若已有插桩,逐点比对后只补缺失点。 + +### Step 2: 创建 trace 工具模块 + +1. 将 [trace_module.py](trace_module.py) 完整内容写入 `verl/utils/tracing/trace.py` + (含 batch 版 `trace_flow`、`log_progress` 及其按小时落盘逻辑)。 +2. 将 [trajectory_types_module.py](trajectory_types_module.py)、 + [trajectory_store_module.py](trajectory_store_module.py)、 + [trajectory_tracker_module.py](trajectory_tracker_module.py) 原样写入 + `verl/utils/tracing/trajectory_types.py`、`trajectory_store.py`、`trajectory_tracker.py`。 +3. 创建 `verl/utils/tracing/__init__.py`,导出: + `trace_flow, log_progress, get_trace_dir, + set_trace_dir, suppress_sample_trace, Turn, Trajectory, TurnRecord, + TrajectoryStore, TrajectoryTracker`。 + +关键语义(勿改):`VERL_TRACE_DIR` 统一门控 stdout 与文件输出(优先级 +`set_trace_dir()` > 环境变量);CSV 写到 `{trace_dir}/sample/{hostname}-pid{pid}.csv`, +列顺序 `job_id,cat,sample_id,sample_status,timestamp,content`(表头行不写入文件, +由消费方如 SLS 采集器 csv_headers 提供);后台线程异步批量落盘 +(每批 ≤64 行 flush),队列满丢弃;`sample_id = f"{uid}/{traj_uid}" if str(traj_uid).isdigit() +else f"{uid}/sample"`;`sample_status` 仅 `TRAIN_UPDATE`/`FILTER_TRAJECTORIES` 为 +`COMPLETE`; +进度 stdout 日志同步落盘 `{trace_dir}/progress/{YYYYMMDD-HH}.log`(按小时轮转); +发现文件 `/tmp/verl_trace_sample_path.txt`(trace.py 初始化时)、 +`/tmp/verl_trace_trajectory_path.txt`(TrajectoryStore 初始化时)与 +`/tmp/verl_progress_path.txt`(首次 progress 落盘时)。 +**验证样本静默**:验证(eval)样本永不进训练段,若被 trace 会呈现大量永远 +RUNNING 的"未完成"样本;`trace.py` 提供 `suppress_sample_trace()`(asyncio +任务级 contextvar 开关),在每个样本任务入口(v0/共享 `_run_agent_loop`、v1 +`_run_prompt`)对 `validate=True` 置位后全链路 trace_flow 自动静默,任务外的 +manager/worker 批量分发点用局部 `validate` 显式 gate,trajectory +register/dump 对验证样本一并跳过。 + +### Step 3: 注入 Ray worker 环境变量 + +两处都要做,保证 WorkerDict worker 与裸 `@ray.remote` actor(AgentLoopWorker)都可见: + +1. `verl/single_controller/ray/base.py`:锚点 `env_vars.update(worker_env)` 之后插入 + `if os.environ.get("VERL_TRACE_DIR"): env_vars["VERL_TRACE_DIR"] = os.environ["VERL_TRACE_DIR"]`。 +2. `verl/trainer/constants_ppo.py` 的 `get_ppo_ray_runtime_env()`:在 call-time 转发段 + 追加(**仅在已设置时转发**,不得像 PYTHONHASHSEED 那样默认 "0"): + +```python + for key in ("VERL_TRACE_DIR", "VERL_PROGRESS_HEARTBEAT_SEC"): + if os.environ.get(key): + runtime_env["env_vars"][key] = os.environ[key] +``` + +### Step 4: 应用 17+1 个样本流转插入点(batch 风格) + +逐个处理,每个点的锚点、插入位置与完整代码见 [insertion_points.md](insertion_points.md)。 +顺序执行:先 ray_trainer.py(含点 2 与点 2b 重注入兜底,二者必须成对),再 +agent_loop.py、single_turn_agent_loop.py、tool_agent_loop.py。 +**批量点位(1/2/3/12-16)直接把整个 uid/traj_uid 数组传给 trace_flow, +禁止在调用侧写 for 循环**;driver 侧点位传 `global_step=self.global_steps`, +agent_loop 侧传 `batch.meta_info.get("global_steps", -1)`。 + +### Step 4a: 应用 v1 trainer 插桩点(样本流转 + 进度) + +锚点与完整代码见 [insertion_points_v1.md](insertion_points_v1.md)。v1 主路径 +`TaskRunnerV1 + TransferQueue`(`verl/trainer/ppo/v1/`): + +- 样本流转 14 点:trainer_base.py(V1-1/V1-2 + 训练段 V1-3~V1-8,统一经模块级 + trace-only 辅助函数 `_trace_ids_from_batch` 从 TQ key `{uid}_{session_id}_{index}` + 批量解析,padding 样本跳过)、agent_loop_tq.py(V1-9~V1-11,V1-11 注入 + `traj_uid=session_id` 后共享 agent-loop 点位 5-11 与 trajectory dump 自动生效)、 + 3 个 trainer mode 的 SYNC_PARAMS(V1-12); +- 进度 32 点:sync/colocate_async 用一体版无前缀行(含 `start/end generation`, + 对应 ReplayBuffer 采样等待),separate_async 用 `disagg: ` 前缀 + + `Trainer: try to/successfully load global batch@N` fetch 行;banner/init 行/ + save checkpoint/eval/record metrics 与 v0 同格式,前缀由 + `self.trainer_mode == 'separate_async'` 决定。 + +### Step 4b: 应用训练进度插桩点(训推一体 + 训推分离) + +锚点与代码见 [progress_points.md](progress_points.md): +- fetch batch:`_logged_train_dataloader()` 包装生成器(**仅分离版**,horizon 的 + `Trainer: try to/successfully load global batch@N` 行); +- 10 个阶段的 start/end 对(reward/compute log prob/ + compute ref log prob/compute adv/update actor/record metrics/ + sync params/eval/save checkpoint):一体版无前缀、分离版 `disagg: ` 前缀; + 一体版另含 `start/end generation`; +- 无 heartbeat;step banner 按 horizon 格式输出(一体版含 Epoch/consumed + batches,分离版仅 Step)。 + +### Step 4c: 应用 trajectory 采集与中央 dump + +锚点与代码见 [trajectory_points.md](trajectory_points.md): +- `AgentLoopBase.__init__` 挂 `self._traj_tracker = TrajectoryTracker()`; +- `AgentLoopWorker.__init__` 在 `get_trace_dir()` 为真时创建 + `TrajectoryStore({trace_dir}/trajectory, instance_id=os.getpid())` 并写发现文件; +- `_run_agent_loop`:run() 前 `register_trajectory`(经 SimpleNamespace 适配 item), + `trace_traj_key` 经 kwargs 传入 run(),run() 返回后立即 `kwargs.pop`(防止泄漏进 `_compute_score` 的 reward DataProto);postprocess 后 `build_trajectory` + + `write`(agent_metrics 先 `model_dump()` 强转);`unregister_trajectory` 在成功路径的 + dump 块 `finally` 与失败路径(run/postprocess 抛异常)的 `except ... raise` 中都会执行; +- single_turn/tool agent loop:每次 `server_manager.generate` 包一个 `TurnRecord` + (fill completion 字段 → `stamp_end()` → `append_turn_record`);tool 轮次在 + messages.extend 后用 `patch_observation` 记录 observation 时序。 +未设置 `VERL_TRACE_DIR` 时全链路静默跳过(store 为 None,不 register)。 + +涉及文件汇总: + +| 文件 | 样本插入点 | 进度插桩点 | trajectory | +|---|---|---|---| +| `verl/trainer/ppo/ray_trainer.py` | 1, 2, 2b(重注入兜底), 12-17 | P1-P11 | - | +| `verl/experimental/agent_loop/agent_loop.py` | 3, 4, 5, 11 | rollout 心跳 | tracker+store+register/build/dump | +| `verl/experimental/agent_loop/single_turn_agent_loop.py` | 6, 7 | - | TurnRecord 采集 | +| `verl/experimental/agent_loop/tool_agent_loop.py` | 8, 9, 10(可选) | - | TurnRecord + patch_observation | +| `verl/experimental/separation/ray_trainer.py` | - | P2-P10 复用 | - | +| `verl/trainer/ppo/v1/trainer_base.py` | V1-1, V1-2, V1-3~V1-8 | init/banner/fetch/阶段对 26 处 | - | +| `verl/trainer/ppo/v1/agent_loop_tq.py` | V1-9, V1-10, V1-11(traj_uid 注入) | - | 复用共享 dump | +| `verl/trainer/ppo/v1/trainer_sync.py` | V1-12 | sync params 对 | - | +| `verl/trainer/ppo/v1/trainer_colocate_async.py` | V1-12 | sync params 对 | - | +| `verl/trainer/ppo/v1/trainer_separate_async.py` | V1-12 | disagg sync params 对 | - | +| `verl/single_controller/ray/base.py` | 环境变量注入 | - | - | +| `verl/trainer/constants_ppo.py` | runtime_env 转发(v1 裸 actor 依赖此处) | - | - | + +### Step 5: 验证 + +```bash +# 1. 语法检查(所有改动文件) +python -m py_compile verl/utils/tracing/*.py verl/trainer/ppo/ray_trainer.py \ + verl/trainer/ppo/v1/*.py \ + verl/experimental/agent_loop/agent_loop.py \ + verl/experimental/agent_loop/single_turn_agent_loop.py \ + verl/experimental/agent_loop/tool_agent_loop.py \ + verl/trainer/constants_ppo.py verl/single_controller/ray/base.py + +# 2. 插桩点计数(期望 trace_flow = 31:v0 17 + v1 14) +grep -rn "trace_flow(" verl/ --include="*.py" | grep -v "utils/tracing" | wc -l + +# 2b. 进度点计数(期望 log_progress = 82:v0 50(一体 27 + separation 23)+ v1 32) +grep -rn "log_progress(" verl/ --include="*.py" | grep -v "utils/tracing" | wc -l + +# 3. import 检查 +grep -rln "from verl.utils.tracing import" verl/ + +# 4. 工具模块自检(含批量 trace_flow / TrajectoryStore schema) +VERL_TRACE_DIR=/tmp/_verl_trace_selfcheck python - <<'EOF' +import glob, json, time +from verl.utils.tracing import (trace_flow, log_progress, + Trajectory, TrajectoryStore, TurnRecord) +trace_flow(uid=['u1', 'u2'], traj_uid=['0', '1'], from_module='T', to_module='T', + stage='DATA_PREPROCESS', detail='self_check', global_step=0) +log_progress('step', 'start', step=0, detail='self_check') +rec = TurnRecord([1, 2, 3]); rec.stamp_end() +store = TrajectoryStore('/tmp/_verl_trace_selfcheck/trajectory', 0) +store.write(Trajectory(request_id='0', prompt_uid='u1', turns=[rec.turn])); store.close() +time.sleep(0.5) +d = json.loads(open(glob.glob('/tmp/_verl_trace_selfcheck/trajectory/*.jsonl')[0]).readline()) +assert d['request_id'] == '0' and 'turns' in d +assert glob.glob('/tmp/_verl_trace_selfcheck/progress/*.log') +print('SELF_CHECK_OK') +EOF +``` + +任何一步失败:修复后重新验证,直到全部通过。 +端到端冒烟:`bash scripts/run_trace_smoke_test.sh`(8xA100 GRPO,3 step, +默认 v1 sync 路径,USE_V1=False 回退 v0;覆盖全部 stage 与 trajectory schema 自检)。 + +### Step 6: 输出报告 + +向用户报告: + +``` +## verl trace 插桩报告(对齐最新 verl-horizon,v0 + v1) +- 已创建: verl/utils/tracing/{trace.py,trajectory_types.py,trajectory_store.py,trajectory_tracker.py,__init__.py} +- 样本流转: N/31 个插桩点(v0 17 + v1 14;batch trace_flow + global_step) +- 训练进度: 10/11 前端阶段(v0 两路径 + v1 三 trainer_mode 各一套; + 一体版/v1 sync/colocate_async 无前缀 + generation 对, + 分离版/v1 separate_async disagg: 前缀 + fetch 行;prompt2accs 无实现不插桩), + horizon 裸格式行 + 小时轮转 progress/ log +- trajectory 监控: TurnRecord 采集(single_turn/tool)+ 中央 dump(_run_agent_loop) +- 验证: py_compile 通过 / 计数 / 自检输出 +- 使用: export VERL_TRACE_DIR=/path/to/trace_dir 后启动训练 + CSV 输出: $VERL_TRACE_DIR/sample/{hostname}-pid{pid}.csv(cat=sample) + 进度 log: $VERL_TRACE_DIR/progress/{YYYYMMDD-HH}.log(按小时轮转,horizon 裸格式) + 轨迹 dump: $VERL_TRACE_DIR/trajectory/trajectory_{instance}_{ts}.jsonl + stdout: [FLOW_TRACE] 样本流转(受 VERL_TRACE_DIR 门控)/ + horizon 裸格式进度行(常开;一体版无前缀、分离版 disagg: 前缀) + 心跳间隔: 已移除 heartbeat(与 horizon 一致) + 发现文件: /tmp/verl_trace_sample_path.txt、/tmp/verl_trace_trajectory_path.txt +``` + +## 已知坑(实施与自检时注意) + +1. `AgentLoopOutput.metrics` 是 pydantic `AgentLoopMetrics`:dump 前必须 + `model_dump()` 强转,否则每条轨迹都抛 not JSON serializable 且被 try/except + 吞掉,表现为 trajectory JSONL 全为 0 字节(到 Ray worker 日志 + `/tmp/ray/session_latest/logs/worker-*.err` grep "TrajectoryStore dump failed" 定位)。 +2. 自检脚本 `head -1` 传多个文件会输出 `==> 文件名 <==` 头行导致 JSON 解析失败, + 只取单个文件抽查。 +3. 训练结束退出期 Ray 清理会 SIGKILL DataLoader 多进程 worker,torch signal + handler 报 "DataLoader worker is killed by signal: Killed"(训练结果已完整, + 属退出期无害错误);冒烟脚本默认 `data.dataloader_num_workers=0` 规避。 + +## 参考文档 + +- 样本流转插桩点锚点与代码:[insertion_points.md](insertion_points.md)(v0)、 + [insertion_points_v1.md](insertion_points_v1.md)(v1,含 v1 进度插桩) +- 训练进度插桩点锚点与代码:[progress_points.md](progress_points.md) +- trajectory 采集与 dump 锚点与代码:[trajectory_points.md](trajectory_points.md) +- 工具模块源码:[trace_module.py](trace_module.py)、 + [trajectory_types_module.py](trajectory_types_module.py)、 + [trajectory_store_module.py](trajectory_store_module.py)、 + [trajectory_tracker_module.py](trajectory_tracker_module.py) diff --git a/.qoder/skills/verl-trace-instrument/TRACE_SMOKE_NEW_IMAGE_GUIDE.md b/.qoder/skills/verl-trace-instrument/TRACE_SMOKE_NEW_IMAGE_GUIDE.md new file mode 100644 index 00000000..2782650d --- /dev/null +++ b/.qoder/skills/verl-trace-instrument/TRACE_SMOKE_NEW_IMAGE_GUIDE.md @@ -0,0 +1,230 @@ +# 在新镜像(如 nvgc / NGC PyTorch 系镜像)中运行 Trace 冒烟测试指南 + +适用脚本(仓库分支 `verl-tracing-v1`): + +- `scripts/run_trace_smoke_test.sh` —— 单机 8×A100,GRPO Qwen2.5-0.5B(最快冒烟) +- `scripts/run_trace_smoke_test_7b.sh` —— 单机 8×A100,GRPO Qwen3.8-27B +- `scripts/run_trace_smoke_test_multiturn.sh` —— 单机 8×A100,Qwen3.8-27B 多轮 tool_agent(coding 任务) +- `scripts/run_trace_smoke_test_7b_multinode.sh` —— N 机 × 8×A100,Qwen3.8-27B 多机版 + +四个脚本验证同一套追踪输出:样本流转 CSV(`$VERL_TRACE_DIR/sample/`)、 +训练进度日志(`$VERL_TRACE_DIR/progress/{YYYYMMDD-HH}.log`,horizon 裸格式行)、 +轨迹 JSONL(`$VERL_TRACE_DIR/trajectory/`)。 + +下文所有依赖版本均为**已实际跑通冒烟的验证组合**,照装即可。 + +--- + +## 1. 硬件 / 系统要求 + +| 项 | 单机版 | 多机版(每节点) | +|---|---|---| +| GPU | 8×A100-80GB(40GB 版显存不足) | 同单机 | +| 驱动 | ≥550.x(支持 CUDA 12.8/12.9 runtime) | 同单机 | +| 网络 | - | 节点间 `RAY_PORT`(默认 6379)+ Ray/NCCL 高端口段双向互通 | +| 存储 | 模型 ~54GB + 数据 ~10MB + trace 输出 | **建议** NFS 等共享存储,模型/数据/`TRACE_DIR` 各节点同路径;否则需每节点预置模型与数据,trace 输出分散需自行汇总 | + +## 2. 基础镜像选择 + +- 任何 **Python ≥3.10 + CUDA ≥12.8 + 预装 PyTorch** 的镜像均可(nvgc、NGC + `nvcr.io/nvidia/pytorch:*` 等)。验证组合基于 **Python 3.12 + torch 2.9.1+cu128**。 +- 镜像自带的 torch **不要重装**,先确认: + ```bash + python3 -c "import torch; print(torch.__version__, torch.cuda.is_available())" + ``` +- 若镜像无 torch,先装 `pip install torch==2.9.1 --index-url https://download.pytorch.org/whl/cu128`。 + +## 3. 依赖安装(验证版本) + +Ubuntu 24.04 系统 Python 受 PEP668 保护,需加 `--break-system-packages`(或自建 venv)。 + +```bash +# 第一步:除 transformers 外的全部依赖 +pip install --break-system-packages \ + "ray[default]==2.56.1" \ + "datasets==4.4.2" \ + "sglang==0.5.10.post1" \ + "flashinfer_python==0.6.7.post3" \ + "flashinfer_cubin==0.6.7.post3" \ + "torch_memory_saver==0.0.9.post1" \ + "TransferQueue>=0.1.9" \ + "modelscope==1.39.1" \ + "omegaconf==2.3.1" \ + "tensordict==0.10.0" \ + "hydra-core==1.3.5" \ + "codetiming==1.4.0" \ + "accelerate==1.14.0" \ + "peft==0.20.0" \ + "pyarrow>=19.0.0" \ + "pandas==2.2.3" \ + "numpy==1.26.4" \ + "pylatexenc==2.10" \ + "uvicorn==0.52.1" \ + "fastapi==0.141.1" \ + "cachetools" + +# 第二步:强制升级 transformers(必须在第一步之后单独执行) +# sglang 0.5.10.post1 元数据钉死 transformers==5.3.0,与 Qwen3.8-27B 要求的 +# >=5.8.0 冲突;两步合装会 ResolutionImpossible,只能先装 sglang 再强制升级 +# (pip 会提示 sglang 元数据不兼容的警告,运行时已验证无碍)。 +pip install --break-system-packages "transformers==5.10.4" +``` + +说明: + +- `sglang 0.5.10.post1` 会自动带入 `sgl-kernel 0.3.21`、`sglang-kernel 0.4.1`(已验证); + `flashinfer_python` / `flashinfer_cubin` 需显式装 `0.6.7.post3`,两者版本必须一致。 +- `torch_memory_saver` 必须 `>=0.0.9.post1`:0.0.9 的 `pause()` 在初始化 + `sleep_replicas` 时报 `cudaError 1 (invalid argument)` 直接拖死 sglang + scheduler(报 `ActorUnavailableError`,见 §9)。 +- `transformers 5.x` 需配套 `huggingface_hub>=1.x`、`tokenizers>=0.22`(pip 自动解析, + 勿手动降级 transformers 到 4.x,模型 Qwen3.8 系列需要 5.x 的建模代码)。 +- `TransferQueue` 为 v1 trainer(脚本默认 `USE_V1=True`)必需,直接 pip 安装即可 + (公共源最高 0.1.9,满足 verl 的 >=0.1.9 要求;内部源有 0.1.10.dev0 亦可)。 + 或设 `USE_V1=False` 回退 v0 路径(无需该包)。 +- `ray[default]` 含 dashboard,多机脚本的 `ray nodes/status` 依赖 ray CLI。 +- **flash_attn 可选**:脚本默认 `SGLANG_ATTENTION_BACKEND=triton`(A100 不支持 fa3), + 不装 flash_attn 也能跑;若要装(源码编译约 10-20 分钟): + `pip install --break-system-packages flash-attn==2.8.3 --no-build-isolation`。 +- **不要** `pip install -e .` 或 `pip install -r requirements.txt`: + requirements 钉了 `TransferQueue==0.1.8`,与验证组合(0.1.10.dev0)冲突; + 脚本通过 `PYTHONPATH=<仓库根>` 直接 import verl,无需安装 verl 包。 +- **vLLM 后端未验证**(本环境未装 vllm,默认 sglang)。如需: + `INFER_BACKEND=vllm` 并自行安装匹配 torch2.9/cu128 的 vllm 版本。 +- 离线/内网环境:pip 加镜像源参数;`modelscope` 也可换成预置快照(见 §5)。 + +## 4. 获取代码 + +```bash +git clone -b verl-tracing-v1 <仓库地址> verl && cd verl +``` + +多机版要求**每个节点**都有该仓库(建议同路径);仓库内含脚本依赖的 +`examples/data_preprocess/gsm8k.py`、`scripts/prep_gsm8k_coding.py`、 +`scripts/coding_tool.py` 与 tracing 插桩代码。 + +## 5. 模型与数据(可预置,也可交给脚本自动兜底) + +脚本默认兜底链:本地 ModelScope 快照 → modelscope 在线下载 → HF 直连。 + +- 0.5B 脚本模型(Qwen2.5-0.5B-Instruct,~1GB)预置路径: + `${MODELSCOPE_CACHE:-$HOME/modelscope_cache}/models/Qwen--Qwen2.5-0.5B-Instruct/snapshots/master` +- 7b / multiturn / multinode 脚本模型(Qwen3.8-27B,~54GB)预置路径: + `${MODELSCOPE_CACHE:-$HOME/modelscope_cache}/models/Qwen--Qwen3.8-27B/snapshots/master` +- 任意脚本均可 `MODEL_PATH=/path/to/model` 显式指定,跳过兜底链。 +- 数据(GSM8K parquet)默认 `$HOME/data/gsm8k`(multiturn 另派生 + `$HOME/data/gsm8k_coding`),缺失时脚本自动生成 + (需 `datasets`;HF 不可达时经 modelscope 下载原始数据)。 + +## 6. 单机运行 + +```bash +cd verl +bash scripts/run_trace_smoke_test_7b.sh # 全默认 +# 或按需覆盖: +TRACE_DIR=/data/trace_7b TOTAL_STEPS=3 \ +SGLANG_ATTENTION_BACKEND=triton \ +bash scripts/run_trace_smoke_test_7b.sh +``` + +常用环境变量: + +| 变量 | 默认 | 说明 | +|---|---|---| +| `MODEL_PATH` | 自动兜底 | 模型路径 | +| `TRACE_DIR` | `<仓库根>/test/trace_7b` | trace 输出目录 | +| `INFER_BACKEND` | `sglang` | `vllm` 未验证 | +| `SGLANG_ATTENTION_BACKEND` | `triton` | A100 必用 triton(fa3 不支持) | +| `USE_V1` | `True` | v1 TaskRunnerV1 + TransferQueue 是打点主路径;需先装 TransferQueue,设 `False` 回退 v0 RayPPOTrainer | +| `DATALOADER_NUM_WORKERS` | `0` | 避免训练结束 Ray 清理 SIGKILL 多进程 worker | +| `TRAIN_BATCH_SIZE` / `ROLLOUT_N` / `TOTAL_STEPS` | 64 / 4 / 3 | 规模参数 | + +multiturn 脚本另有 `MAX_ASSISTANT_TURNS`(默认 8)、`MAX_TOOL_RESPONSE_LENGTH`(默认 1024)。 + +默认跑 3 个 step,单 step 数分钟级;结束自动打印自检结果(见 §8)。 + +## 7. 多机运行 + +每台机器各执行一次同一脚本;容器调度器通常已注入 `RANK`/`WORLD_SIZE`/`MASTER_ADDR`。 + +```bash +# 主节点(rank 0) +RANK=0 WORLD_SIZE=2 bash scripts/run_trace_smoke_test_7b_multinode.sh +# 其余节点(rank ≥ 1,MASTER_ADDR 必填为主节点 IP) +RANK=1 WORLD_SIZE=2 MASTER_ADDR=<主节点IP> bash scripts/run_trace_smoke_test_7b_multinode.sh +``` + +流程与分工: + +1. 数据准备:rank0 生成/下载,worker 仅校验可见(非共享存储会 fail fast,需先每节点预置)。 +2. 组网:rank0 `ray start --head` 并等待 `WORLD_SIZE` 个节点 ALIVE(超时 + `RAY_JOIN_TIMEOUT_SEC`=600s);worker `ray start --address=...` 后**阻塞等待**, + 不跑训练,head 结束 `ray stop` 后自动退出。 +3. 训练:仅 head 提交 `trainer.nnodes=$WORLD_SIZE`;`TRAIN_BATCH_SIZE` 默认 + `64×节点数` 保持每卡负载与单机一致。 +4. 自检:仅 head 执行;共享存储下覆盖全集群输出,否则需先汇总各节点 + `TRACE_DIR`(非共享时输出分散在本机磁盘)。 +5. 清理:`STOP_RAY_AFTER_DONE=1`(默认)时 head `ray stop` 关整集群。 + +注意:脚本会在 `ray start` 前 `unset RANK WORLD_SIZE`,防止调度器变量泄漏进 +Ray daemon 被裸 actor 错误继承——手动包装时也不要把这两个变量带进 ray 进程。 + +多机额外变量:`RAY_PORT`(6379)、`RAY_JOIN_TIMEOUT_SEC`、`STOP_RAY_AFTER_DONE`。 + +## 8. 结果验证 + +脚本结尾自动自检并打印;也可手动检查 `$TRACE_DIR` 目录结构: + +``` +$TRACE_DIR/ +├── sample/ # {hostname}-pid{pid}.csv,样本流转(cat=sample) +├── progress/ # {YYYYMMDD-HH}.log,按小时轮转的训练进度 +└── trajectory/ # trajectory_{instance}_{ts}.jsonl,轨迹 dump +``` + +进度日志样例(v1 训推一体无前缀;v1 另有 TransferQueue 的 `Task ... is completed` +行与 `Step N rollout started` 行): + +``` +Total training steps: 3 +'train_batch_size': 64, +'rollout_n': 4, +'ppo_mini_batch_size': 32, +'num_minibatches': 8, +'total_training_steps': 3 +==================== Epoch 0, consumed 0 batches, Step 1 ==================== +Step 1 rollout started +start generation +end generation +start compute reward +end compute reward +... +``` + +自检期望: + +- `sample/` 各 stage 覆盖:DATA_PREPROCESS / DATA_DISPATCH / AGENT_LOOP / + LLM_GENERATE / TRAIN_REWARD / TRAIN_LOGPROB / TRAIN_REF_LOGPROB / + TRAIN_ADVANTAGE / TRAIN_UPDATE / SYNC_PARAMS(v1 另有 ROLLOUT_REWARD / TRAIN_SAMPLE); +- `progress/` 行数 > 0,阶段行成对(start/end),step 2 触发 eval; +- `trajectory/` 条数 ≈ `TRAIN_BATCH_SIZE × ROLLOUT_N × TOTAL_STEPS`(单机 768)。 + v1 格式:顶层含 prompt_uid/sample_index/raw_message,turns 在 raw_message JSON 内; + v0 格式:顶层直接含 request_id/prompt_uid/sample_index/turns/initial_prompt_text。 + 两者均不含 `initial_prompt_token_ids`(已裁剪); +- multiturn 脚本额外关注 turns > 1 的多轮轨迹占比(自检自动统计)。 + +## 9. 常见问题 + +| 现象 | 处理 | +|---|---| +| `error: externally-managed-environment` | pip 加 `--break-system-packages` 或用 venv | +| huggingface.co 不可达 | 安装 `modelscope` 走自动兜底;或按 §5 预置快照 | +| sglang 报 fa3 / flashinfer 错 | `SGLANG_ATTENTION_BACKEND=triton`(脚本默认);确认 `flashinfer_python` 与 `flashinfer_cubin` 同为 `0.6.7.post3` | +| 脚本预检报缺 `transfer_queue` | `pip install --break-system-packages "TransferQueue>=0.1.9"`(见 §3);或 `USE_V1=False` 回退 v0 | +| 初始化报 `ActorUnavailableError`,sglang 日志有 `[torch_memory_saver.cpp] cudaError error: 1 ... func=pause` | `torch_memory_saver` 版本低于 0.0.9.post1,升级:`pip install --break-system-packages "torch_memory_saver==0.0.9.post1"` | +| 报 `model type qwen3_5 ... Transformers does not recognize` | transformers 低于 5.8,升级到 `5.10.4`(见 §3 第二步) | +| v1 训练中 step 组批卡住,worker 日志报 `ZMQ recv timeout ... storage unit ... MessagePack data is malformed` | TransferQueue storage unit 反序列化崩溃(已知偶发):重跑一般可恢复;频繁复现时升级 TransferQueue 或换 `transfer_queue.backend.storage_backend=MooncakeStore` | +| 训练结束报 `DataLoader worker ... killed by signal` | 脚本已默认 `dataloader_num_workers=0`;若仍出现,确认未手动调大该值 | +| 多机等待节点超时 | 检查 `MASTER_ADDR`/防火墙 6379 与高端口段;确认各节点 `RANK`/`WORLD_SIZE` 正确 | +| 显存不足 | 降 `actor_rollout_ref.rollout.gpu_memory_utilization`(7b/multiturn 默认 0.5、0.5B 版 0.4)、降 `ppo_max_token_len_per_gpu`;ref 已默认 `param_offload=True` | +| `import verl` 失败 | 必须在仓库根目录运行脚本(脚本自动设置 PYTHONPATH) | diff --git a/.qoder/skills/verl-trace-instrument/insertion_points.md b/.qoder/skills/verl-trace-instrument/insertion_points.md new file mode 100644 index 00000000..1063dee5 --- /dev/null +++ b/.qoder/skills/verl-trace-instrument/insertion_points.md @@ -0,0 +1,447 @@ +# 插入点锚点清单(verl v0.8.x/0.9.x) + +每个插入点包含:**锚点**(用于 grep 定位的原始代码,勿改动)、**插入位置**、**插入代码**。 +插入代码的缩进必须与锚点所在行一致。所有文件需在模块顶部追加 +`from verl.utils.tracing import trace_flow`(若已有则跳过)。 + +**全局消歧规则**:部分锚点在 `verl/experimental/separation/ray_trainer.py`、 +`verl/trainer/ppo/v1/` 等变体中存在同文副本。本 skill 只改主路径: +ray_trainer.py 的锚点必须落在 `verl/trainer/ppo/ray_trainer.py` 的 `RayPPOTrainer.fit()` 内, +agent_loop.py 的锚点必须落在 `verl/experimental/agent_loop/agent_loop.py` +(而非 `verl/trainer/ppo/v1/agent_loop_tq.py`)。变体文件的插桩见 tracer_plan.md 第六节,默认不做。 + +批量样本 trace 的通用写法(对齐最新 verl-horizon 的 batch API): +**直接把整个 uid/traj_uid 数组传给 trace_flow,禁止在调用侧写 for 循环**, +函数内部自动 fan-out、一侧为标量时自动广播;driver 侧点位传 +`global_step=self.global_steps`,agent_loop 侧传 `batch.meta_info.get("global_steps", -1)`。 +从 `batch.non_tensor_batch` 取 `uid` / `traj_uid`; +对可能不存在 `traj_uid` 的位置,用 `batch.non_tensor_batch.get("traj_uid", np.array(["N/A"] * len(batch)))`。 +**训练段(插入点 12-16)一律用该 `.get` 兜底写法,禁止直接索引 `['traj_uid']`**: +插桩代码绝不允许抛异常中断训练,且流式 reward 模式下 rollout 后处理会丢弃透传的 +non_tensor 字段(见插入点 2b)。 + +--- + +## 插入点 1: uid 生成(DATA_PREPROCESS / uid_generated) + +**文件**: `verl/trainer/ppo/ray_trainer.py` — `RayPPOTrainer.fit()` + +**锚点**(uid 数组生成语句,3 行): +```python + batch.non_tensor_batch["uid"] = np.array( + [str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object + ) +``` + +**插入位置**: 锚点之后(batch 风格,标量 'N/A' 自动广播): + +```python + trace_flow(uid=batch.non_tensor_batch["uid"], traj_uid='N/A', + from_module='DataLoader', to_module='RayPPOTrainer', + stage='DATA_PREPROCESS', detail='uid_generated', + global_step=self.global_steps) +``` + +--- + +## 插入点 2: traj_uid 注入(DATA_PREPROCESS / traj_uid_generated_by_repeat) + +**文件**: `verl/trainer/ppo/ray_trainer.py` — `RayPPOTrainer.fit()` + +**锚点**: +```python + gen_batch_output = gen_batch.repeat(repeat_times=rollout_n, interleave=True) +``` + +**插入位置**: 锚点之后。这是唯一允许新增数据字段的点(interleave 模式下组内序号 `i % rollout_n`)。 + +```python + gen_batch_output.non_tensor_batch["traj_uid"] = np.array( + [str(i % rollout_n) for i in range(len(gen_batch_output))], dtype=object + ) + trace_flow(uid=gen_batch_output.non_tensor_batch["uid"], + traj_uid=gen_batch_output.non_tensor_batch["traj_uid"], + from_module='RayPPOTrainer', to_module='RayPPOTrainer', + stage='DATA_PREPROCESS', detail='traj_uid_generated_by_repeat', + global_step=self.global_steps) +``` + +> 说明:若 rollout 后处理透传了输入 non_tensor_batch,后续 `batch = batch.repeat(...)` 与 +> `batch.union(gen_batch_output)` 会自动携带 traj_uid 进入训练段。但**流式 reward 模式** +> (存在 reward_loop_worker_handles)下 `AgentLoopWorker._postprocess` 会重建 +> non_tensor_batch 并丢弃 traj_uid(uid 因原 batch 保留而幸存),因此必须配合插入点 2b。 + +--- + +## 插入点 2b: traj_uid 重注入兜底(防流式 reward 模式丢失) + +**文件**: `verl/trainer/ppo/ray_trainer.py` — `RayPPOTrainer.fit()` + +**锚点**(rollout 输出切片与 `__do_sample__` 清理,在 `marked_timer("gen")` 块外): +```python + gen_batch_output = combined_gen_output.slice(0, num_sampled_prompts) + if "__do_sample__" in gen_batch_output.non_tensor_batch: + gen_batch_output.pop(non_tensor_batch_keys=["__do_sample__"]) +``` +(注意:不同版本中该区域可能还有 REMAX 基线分支/`del combined_gen_batch` 等语句, +均以 `slice(0, num_sampled_prompts)` + `__do_sample__` pop 两行为锚点定位。) + +**插入位置**: 锚点之后。缺失时按与插入点 2 完全相同的 interleave 规则重建 +(rollout 输入输出顺序一一对应,数组内容一致,`union` 无冲突): + +```python + # traj_uid may be dropped by rollout postprocess (streaming reward mode); + # re-inject with the same interleave pattern as point 2 so downstream + # trace points can resolve it after batch.union(gen_batch_output). + if "traj_uid" not in gen_batch_output.non_tensor_batch: + gen_batch_output.non_tensor_batch["traj_uid"] = np.array( + [str(i % rollout_n) for i in range(len(gen_batch_output))], dtype=object + ) +``` + +> 说明:这是插入点 2 同一 `traj_uid` 字段的兜底重注入,属铁律 1 允许的数据新增范围; +> `rollout_n` 变量在 fit() 循环作用域内已由插入点 2 上文定义。 + +--- + +## 插入点 3: rollout 分发(DATA_DISPATCH / dispatched_to_agent_loop_worker) + +**文件**: `verl/experimental/agent_loop/agent_loop.py` — `AgentLoopManager.generate_sequences()` + +**锚点**(注意:同文锚点在 `verl/trainer/ppo/v1/agent_loop_tq.py` 也有一处, +只取 `verl/experimental/agent_loop/agent_loop.py` 内的一处): +```python + chunkes = prompts.chunk(len(self.agent_loop_workers)) +``` + +**插入位置**: 锚点之后、`outputs = await asyncio.gather(...)` 之前。 +验证批次(`meta_info["validate"]=True`)不打样本 trace(永不进训练段, +否则统计上会呈现大量"未完成"样本)。 + +```python + if not prompts.meta_info.get("validate", False): + for _worker_idx, _chunk in enumerate(chunkes): + _traj_uids = _chunk.non_tensor_batch.get("traj_uid", np.array(["N/A"] * len(_chunk))) + trace_flow(uid=_chunk.non_tensor_batch["uid"], traj_uid=_traj_uids, + from_module='AgentLoopManager', to_module=f'AgentLoopWorker_{_worker_idx}', + stage='DATA_DISPATCH', detail='dispatched_to_agent_loop_worker', + global_step=prompts.meta_info.get("global_steps", -1)) +``` + +确认文件已 import numpy(上游已有 `import numpy as np`)。 + +--- + +## 插入点 4: 单样本任务开始(AGENT_LOOP / agent_task_started) + +**文件**: `verl/experimental/agent_loop/agent_loop.py` — `AgentLoopWorker.generate_sequences()` + +**锚点**(for 循环内第一行): +```python + kwargs = {k: v[i] for k, v in batch.non_tensor_batch.items() if k != "__do_sample__"} +``` + +**插入位置**: 锚点之后、`tasks.append(...)` 之前(用方法内已有的局部 +`validate` gate,验证样本跳过)。 + +```python + if not validate: + trace_flow(uid=kwargs.get('uid', 'unknown'), traj_uid=kwargs.get('traj_uid', 'N/A'), + from_module='AgentLoopWorker', + to_module=f"AgentLoop({kwargs.get('agent_name', 'single_turn')})", + stage='AGENT_LOOP', detail='agent_task_started') +``` + +--- + +## 插入点 5: 轨迹完成(AGENT_LOOP / agent_run_completed)+ 验证样本任务级静默 + +**文件**: `verl/experimental/agent_loop/agent_loop.py` — `AgentLoopWorker._run_agent_loop()` + +**前置**(方法体第一行,`with rollout_trace_attr(` 之前):每个样本的 +`_run_agent_loop` 是独立 asyncio.Task(contextvars 隔离),在任务入口对验证样本 +调用 `suppress_sample_trace()`,该样本调用链上的全部 trace_flow(本点 5、 +点 6-11 的 LLM_GENERATE / ROLLOUT_REWARD 等)自动静默,无需逐点加判断: + +```python + # validation samples never enter the training stages; silence their + # sample-flow trace for this sample task's whole call chain + # (LLM_GENERATE / ROLLOUT_REWARD / ... — see suppress_sample_trace) + if trajectory.get("validate", False): + suppress_sample_trace() +``` + +(import 一并从 `verl.utils.tracing` 导入 `suppress_sample_trace`。) + +**锚点**: +```python + output: AgentLoopOutput = await agent_loop.run(sampling_params, **kwargs) +``` + +**插入位置**: 锚点之后、`return await self._agent_loop_postprocess(...)` 之前。 + +```python + trace_flow(uid=kwargs.get('uid', 'unknown'), traj_uid=kwargs.get('traj_uid', 'N/A'), + from_module=f'AgentLoop({agent_name})', to_module='AgentLoopWorker', + stage='AGENT_LOOP', detail='agent_run_completed') +``` + +--- + +## 插入点 6/7: 单轮 LLM 请求前后(LLM_GENERATE / sending|received) + +**文件**: `verl/experimental/agent_loop/single_turn_agent_loop.py` — `SingleTurnAgentLoop.run()` +(run() 的 kwargs 直接含 uid/traj_uid) + +**锚点**: +```python + with simple_timer("generate_sequences", metrics): +``` + +**插入位置与代码**: `with` 块之前插入 sending,`with` 块(含 `output = await self.server_manager.generate(...)`)整体结束后插入 received: + +```python + trace_flow(uid=kwargs.get('uid', 'unknown'), traj_uid=kwargs.get('traj_uid', 'N/A'), + from_module='SingleTurnAgentLoop', to_module='RolloutReplica', + stage='LLM_GENERATE', detail='sending_generate_request') + with simple_timer("generate_sequences", metrics): + output: TokenOutput = await self.server_manager.generate(...) + trace_flow(uid=kwargs.get('uid', 'unknown'), traj_uid=kwargs.get('traj_uid', 'N/A'), + from_module='RolloutReplica', to_module='SingleTurnAgentLoop', + stage='LLM_GENERATE', detail='received_generate_response') +``` + +--- + +## 插入点 8/9: 多轮 LLM 请求前后(LLM_GENERATE / sending|received) + +**文件**: `verl/experimental/agent_loop/tool_agent_loop.py` + +分两步: + +**步骤 A** — 挂载 trace 字段。锚点(`run()` 内 AgentData 构建完成后的状态机起点): +```python + # State machine loop + state = AgentState.PENDING +``` +在其前插入(纯数据挂载,trace_traj_key 供 trajectory TurnRecord 采集): +```python + agent_data.trace_uid = kwargs.get('uid', 'unknown') + agent_data.trace_traj_uid = kwargs.get('traj_uid', 'N/A') + agent_data.trace_traj_key = kwargs.get('trace_traj_key') +``` + +**步骤 B** — generate 前后。锚点(`_handle_generating_state` 内): +```python + with simple_timer("generate_sequences", agent_data.metrics): +``` +与插入点 6/7 相同模式,uid/traj_uid 从 `agent_data` 取: +```python + trace_flow(uid=getattr(agent_data, 'trace_uid', 'unknown'), + traj_uid=getattr(agent_data, 'trace_traj_uid', 'N/A'), + from_module='ToolAgentLoop', to_module='RolloutReplica', + stage='LLM_GENERATE', detail='sending_generate_request') + with simple_timer("generate_sequences", agent_data.metrics): + output: TokenOutput = await self.server_manager.generate(...) + trace_flow(uid=getattr(agent_data, 'trace_uid', 'unknown'), + traj_uid=getattr(agent_data, 'trace_traj_uid', 'N/A'), + from_module='RolloutReplica', to_module='ToolAgentLoop', + stage='LLM_GENERATE', detail='received_generate_response') +``` + +--- + +## 插入点 10(可选): 工具执行(AGENT_TOOL / tool_executed) + +**文件**: `verl/experimental/agent_loop/tool_agent_loop.py` — `_call_tool()` +(方法签名含 `agent_data: AgentData`,可直接取步骤 A 挂载的 trace 字段; +选在方法末尾插桩可同时覆盖 FunctionTool 与 BaseTool 两条执行路径) + +**锚点**(`_call_tool` 的最后一行 return): +```python + return ToolResponse(**tool_response_kwargs), tool_reward, res +``` + +**插入位置**: 锚点之前。 + +```python + trace_flow(uid=getattr(agent_data, 'trace_uid', 'unknown'), + traj_uid=getattr(agent_data, 'trace_traj_uid', 'N/A'), + from_module='Tool', to_module='ToolAgentLoop', + stage='AGENT_TOOL', detail=f'tool_executed_name={tool_call.name}') + return ToolResponse(**tool_response_kwargs), tool_reward, res +``` + +--- + +## 插入点 11: 流式 reward 完成(ROLLOUT_REWARD / reward_score_computed) + +**文件**: `verl/experimental/agent_loop/agent_loop.py` — `AgentLoopWorker._compute_score()` + +**锚点**: +```python + final_output.extra_fields["reward_extra_info"] = result["reward_extra_info"] +``` + +**插入位置**: 锚点之后(仍在 `with simple_timer("compute_score", timing):` 块内)。 + +```python + trace_flow(uid=kwargs.get('uid', 'unknown'), traj_uid=kwargs.get('traj_uid', 'N/A'), + from_module='RewardLoopWorker', to_module='AgentLoopWorker', + stage='ROLLOUT_REWARD', + detail=f"reward_score_computed_score={result['reward_score']}") +``` + +--- + +## 插入点 12: 训练侧 reward 汇总(TRAIN_REWARD / reward_extracted) + +**文件**: `verl/trainer/ppo/ray_trainer.py` — `fit()` 的 `marked_timer("reward")` 块内 + +**锚点**(同文副本存在于 separation/v1 变体,只取 `verl/trainer/ppo/ray_trainer.py` +`fit()` 内、位于 `marked_timer("reward"...)` 块内的一处): +```python + reward_tensor, reward_extra_infos_dict = extract_reward(batch) +``` + +**插入位置**: 锚点之后。 + +```python + _traj_uids = batch.non_tensor_batch.get("traj_uid", np.array(["N/A"] * len(batch))) + trace_flow(uid=batch.non_tensor_batch['uid'], traj_uid=_traj_uids, + from_module='RayPPOTrainer', to_module='RewardManager', + stage='TRAIN_REWARD', detail='reward_extracted', + global_step=self.global_steps) +``` + +--- + +## 插入点 13: old_log_prob(TRAIN_LOGPROB / old_log_prob_computed) + +**文件**: `verl/trainer/ppo/ray_trainer.py` — `marked_timer("old_log_prob")` 块内 + +**锚点**(同文副本存在于 separation 变体,只取 `verl/trainer/ppo/ray_trainer.py` +`marked_timer("old_log_prob")` 块内的一处;注意区分方法定义 `def _compute_old_log_prob`): +```python + old_log_prob, old_log_prob_mfu = self._compute_old_log_prob(batch) +``` + +**插入位置**: 锚点之后。 + +```python + _traj_uids = batch.non_tensor_batch.get("traj_uid", np.array(["N/A"] * len(batch))) + trace_flow(uid=batch.non_tensor_batch['uid'], traj_uid=_traj_uids, + from_module='ActorRolloutRefWorker', to_module='RayPPOTrainer', + stage='TRAIN_LOGPROB', detail='old_log_prob_computed', + global_step=self.global_steps) +``` + +--- + +## 插入点 14: ref_log_prob(TRAIN_REF_LOGPROB / ref_log_prob_computed) + +**文件**: `verl/trainer/ppo/ray_trainer.py` — `marked_timer(str(Role.RefPolicy))` 块内 + +**锚点**(同文副本存在于 separation 变体及 `_compute_ref_log_prob` 方法内部, +只取 `marked_timer(str(Role.RefPolicy))` 块内的调用处): +```python + ref_log_prob = self._compute_ref_log_prob(batch) +``` + +**插入位置**: 锚点之后、`batch = batch.union(ref_log_prob)` 之前。 + +```python + _traj_uids = batch.non_tensor_batch.get("traj_uid", np.array(["N/A"] * len(batch))) + trace_flow(uid=batch.non_tensor_batch['uid'], traj_uid=_traj_uids, + from_module='RefPolicyWorker', to_module='RayPPOTrainer', + stage='TRAIN_REF_LOGPROB', detail='ref_log_prob_computed', + global_step=self.global_steps) +``` + +--- + +## 插入点 15: advantage(TRAIN_ADVANTAGE / advantage_computed) + +**文件**: `verl/trainer/ppo/ray_trainer.py` — `marked_timer("adv")` 块内 + +**锚点**(`compute_advantage` 多行调用的最后一个实参,定位其后闭合括号): +```python + config=self.config.algorithm, + ) +``` + +**插入位置**: `compute_advantage(...)` 调用结束后的下一行(仍在 adv timer 块内)。 +若该锚点在文件中不唯一,用 `batch = compute_advantage(` 向上定位确认。 + +```python + _traj_uids = batch.non_tensor_batch.get("traj_uid", np.array(["N/A"] * len(batch))) + trace_flow(uid=batch.non_tensor_batch['uid'], traj_uid=_traj_uids, + from_module='RayPPOTrainer', to_module='RayPPOTrainer', + stage='TRAIN_ADVANTAGE', detail='advantage_computed', + global_step=self.global_steps) +``` + +--- + +## 插入点 16: actor 更新(TRAIN_UPDATE / actor_parameters_updated) + +**文件**: `verl/trainer/ppo/ray_trainer.py` — `marked_timer("update_actor")` 块内 + +**锚点**(同文副本存在于 separation 变体,只取 `marked_timer("update_actor")` 块内的一处): +```python + actor_output = self._update_actor(batch) +``` + +**插入位置**: 锚点之后。 + +```python + _traj_uids = batch.non_tensor_batch.get("traj_uid", np.array(["N/A"] * len(batch))) + trace_flow(uid=batch.non_tensor_batch['uid'], traj_uid=_traj_uids, + from_module='ActorRolloutRefWorker', to_module='RayPPOTrainer', + stage='TRAIN_UPDATE', detail='actor_parameters_updated', + global_step=self.global_steps) +``` + +--- + +## 插入点 17: 权重同步(SYNC_PARAMS / weights_synced) + +**文件**: `verl/trainer/ppo/ray_trainer.py` + +**锚点**(注意:`update_weights` 在文件中出现多次,必须取 `marked_timer("update_weights"` 块内的一处): +```python + with marked_timer("update_weights", timing_raw, color="red"): + self.checkpoint_manager.update_weights(self.global_steps) +``` + +**插入位置**: `self.checkpoint_manager.update_weights(self.global_steps)` 之后。 + +```python + trace_flow(uid='*', traj_uid='*', from_module='CheckpointEngineManager', + to_module='RolloutReplica', stage='SYNC_PARAMS', + detail=f'weights_synced_global_step={self.global_steps}', + global_step=self.global_steps) +``` + +--- + +## 完成标准 + +- 15 个必选点(1-9, 11-17)全部插入,插入点 2b 重注入兜底必须与点 2 一同插入, + 可选点 10 视作用域情况决定。 +- 批量点位(1/2/3/12-16)均为单次 trace_flow 数组调用,无调用侧 for 循环。 +- 所有改动文件 `python -m py_compile` 通过。 +- `grep -rn "trace_flow(" verl/ --include="*.py" | grep -v "utils/tracing"` 计数 = 17。 + +--- + +rl-tracing Step 视图说明(已调整) + +step 级 rollout/train 窗口的起止**不再走 sample-flow CSV**(曾短暂以 +`ROLLOUT/rollout_started|rollout_finished`、`TRAIN_UPDATE/update_started` +uid/sample 行实现,已移除),改为 progress 行输出,见 progress_points.md +的「step 级窗口行」章节。 + +保留的相关改动:`AGENT_LOOP` 的 `agent_task_started` / `agent_run_completed` +(agent_loop.py)携带 `global_step`(分别取 `batch.meta_info['global_steps']` +与 `trajectory['step']`),供流转分析按 step 归因逐轨迹事件。 diff --git a/.qoder/skills/verl-trace-instrument/insertion_points_v1.md b/.qoder/skills/verl-trace-instrument/insertion_points_v1.md new file mode 100644 index 00000000..abe507dd --- /dev/null +++ b/.qoder/skills/verl-trace-instrument/insertion_points_v1.md @@ -0,0 +1,310 @@ +# V1 插入点锚点清单(verl v1 trainer:TaskRunnerV1 + TransferQueue) + +本清单是 [insertion_points.md](insertion_points.md)(v0)的 V1 对照版,已对齐最新 +verl-horizon 实现(batch `trace_flow` + `global_step` + `log_progress` 裸格式进度行)。 +V1 主路径为 `verl/trainer/main_ppo.py::TaskRunnerV1` → +`verl/trainer/ppo/v1/trainer_base.py::PPOTrainer` +(`trainer_mode` = sync / colocate_async / separate_async)+ +`verl/trainer/ppo/v1/agent_loop_tq.py`(AgentLoopManagerTQ / AgentLoopWorkerTQ)。 + +**V1 与 V0 的关键差异(决定插桩位置)**: + +1. 样本经 **TransferQueue(TQ)kv store** 流转,而非 DataProto 直接返回。prompt 以 `uid` 为 + key 注册(`_submit_batch_to_rollout`),轨迹输出以 **`{uid}_{session_id}_{index}`** 为 key + 写入(`AgentLoopWorkerTQ._agent_loop_postprocess`),训练段由 `ReplayBuffer.sample()` + 按 key 采样出 `KVBatchMeta(partition_id, keys, tags)`。 +2. **traj_uid ≡ session_id**:V1 没有 v0 的 `batch.repeat(rollout_n)` 注入;`rollout.n` 展开 + 发生在 `AgentLoopWorkerTQ._run_prompt()` 的 `for i in range(n)` 循环,`session_id=i` + 即组内序号。插入点 V1-11 负责把 `traj_uid=str(session_id)` 注入 agent loop kwargs。 +3. 训练段 trace 从 **key 解析** `(uid, traj_uid)`:模块级 trace-only 辅助函数 + `_trace_ids_from_batch(batch)` 对 `zip(batch.keys, batch.tags)` 做 + `key.rsplit("_", 2)` + `isdigit()` 校验,失败回退 `(key, "N/A")`,padding 样本 + (tag 含 `is_padding=True`)一律跳过;返回 `(uids, traj_uids)` 两个列表,直接整体 + 传给 batch 版 `trace_flow`(禁止调用侧写 for 循环)。 +4. 所有 driver 侧点位传 `global_step=self.global_steps`;agent loop 侧从 + prompt/chunk 里防御式取 `global_steps`(取不到传 -1)。 + +**V0 插桩点在 V1 下自动生效的部分(不要重复插)**: + +- `AgentLoopWorker._run_agent_loop()`(v0 点 5)与 `_compute_score()`(v0 点 11, + ROLLOUT_REWARD)被 `AgentLoopWorkerTQ` 继承,`kwargs.get('traj_uid')` 由 V1-11 + 注入后可正确解析; +- `single_turn_agent_loop.py` / `tool_agent_loop.py` 的点 6-10(LLM_GENERATE 等)与 + trajectory 采集(TurnRecord)/中央 dump(`_run_agent_loop` 内 register/build/write) + 是 v0/v1 共用代码,无需改动;trajectory dump 在 V1 下同样生效。 + +**环境变量传递**:V1 的 agent loop worker 是裸 `@ray.remote` actor(不经 +WorkerGroup),v0 Step 3(`single_controller/ray/base.py`)覆盖不到;V1 靠 +`verl/trainer/constants_ppo.py::get_ppo_ray_runtime_env()` 末尾的 call-time 转发段 +(`VERL_TRACE_DIR` / `VERL_PROGRESS_HEARTBEAT_SEC`,仅在已设置时转发,见 V1-E)。 + +涉及文件需在模块顶部追加 +`from verl.utils.tracing import log_progress, trace_flow`(仅用到其一时只导入所需)。 + +--- + +## 辅助函数: `_trace_ids_from_batch`(trainer_base.py 模块级) + +**文件**: `verl/trainer/ppo/v1/trainer_base.py` — `logger.setLevel(...)` 之后、 +`_tq_supports_checkpoint()` 之前。 + +```python +def _trace_ids_from_batch(batch) -> tuple[list[str], list[str]]: + """Trace-only helper: parse (uid, traj_uid) pairs from TransferQueue + trajectory keys ``{uid}_{session_id}_{index}`` via rsplit with isdigit + validation; padding samples (``is_padding`` tag) are skipped.""" + _uids, _traj_uids = [], [] + for _k, _tag in zip(batch.keys, batch.tags): + if not isinstance(_tag, dict) or _tag.get("is_padding", False): + continue + _parts = str(_k).rsplit("_", 2) + if len(_parts) == 3 and _parts[1].isdigit() and _parts[2].isdigit(): + _uids.append(_parts[0]) + _traj_uids.append(_parts[1]) + else: + _uids.append(str(_k)) + _traj_uids.append("N/A") + return _uids, _traj_uids +``` + +--- + +## 插入点 V1-1: uid 生成(DATA_PREPROCESS / uid_generated) + +**文件**: `verl/trainer/ppo/v1/trainer_base.py` — `PPOTrainer._fetch_one_gen_batch()` + +**锚点**: +```python + batch_dict["uid"] = np.array([str(uuid.uuid4()) for _ in range(len(batch_dict["raw_prompt"]))], dtype=object) +``` + +**插入位置**: 锚点之后、`return tu.get_tensordict(batch_dict)` 之前(batch 风格)。 + +```python + trace_flow(uid=batch_dict["uid"], traj_uid='N/A', from_module='DataLoader', + to_module='PPOTrainerV1', stage='DATA_PREPROCESS', detail='uid_generated', + global_step=self.global_steps) +``` + +--- + +## 插入点 V1-2: prompt 注册与分发(DATA_DISPATCH / dispatched_to_agent_loop_manager) + +**文件**: `verl/trainer/ppo/v1/trainer_base.py` — `PPOTrainer._submit_batch_to_rollout()` + +**锚点**: `tq.kv_batch_put(...)` if/else 之后、 +`self.agent_loop_manager.generate_sequences(batch)` 之前。 + +```python + trace_flow(uid=[str(getattr(_uid, "data", _uid)) for _uid in batch["uid"]], traj_uid='N/A', + from_module='PPOTrainerV1', to_module='AgentLoopManagerTQ', + stage='DATA_DISPATCH', detail='dispatched_to_agent_loop_manager', + global_step=self.global_steps) +``` + +注:TensorDict 取出的 uid 可能是 `NonTensorData`,用 `getattr(_uid, "data", _uid)` 解包。 + +--- + +## 插入点 V1-3 ~ V1-8: 训练段 trace(`PPOTrainer._step_once()`) + +**文件**: `verl/trainer/ppo/v1/trainer_base.py`。每个点均为「锚点语句之后紧跟」, +统一模板(``/``/``/`` 见下表): + +```python + _t_uids, _t_trajs = _trace_ids_from_batch(batch) + trace_flow(uid=_t_uids, traj_uid=_t_trajs, + from_module='', to_module='', + stage='', detail='', + global_step=self.global_steps) +``` + +| 点 | 锚点(之后插入) | STAGE | from → to | detail | +|---|---|---|---|---| +| V1-3 | `batch, off_policy_metrics = self.replay_buffer.sample(...)` 后的 `metrics.update(off_policy_metrics)` | TRAIN_SAMPLE | ReplayBuffer → PPOTrainerV1 | sampled_from_replay_buffer | +| V1-4 | `batch = self._compute_reward_colocate(batch, metrics=metrics)`(仅 colocated reward 分支) | TRAIN_REWARD | PPOTrainerV1 → RewardLoopManager | reward_extracted | +| V1-5 | `batch = self._compute_old_log_prob(batch, metrics=metrics)` | TRAIN_LOGPROB | ActorRolloutWorker → PPOTrainerV1 | old_log_prob_computed | +| V1-6 | `batch = self._compute_ref_log_prob(batch, metrics=metrics)`(`use_reference_policy` 分支) | TRAIN_REF_LOGPROB | RefPolicyWorker → PPOTrainerV1 | ref_log_prob_computed | +| V1-7 | `batch = self._compute_advantage(batch, metrics=metrics)` | TRAIN_ADVANTAGE | PPOTrainerV1 → PPOTrainerV1 | advantage_computed | +| V1-8 | `batch = self._update_actor(batch, metrics=metrics)`(critic_warmup 分支) | TRAIN_UPDATE | ActorRolloutWorker → PPOTrainerV1 | actor_parameters_updated | + +--- + +## 插入点 V1-9: manager→worker 分发(DATA_DISPATCH / dispatched_to_agent_loop_worker) + +**文件**: `verl/trainer/ppo/v1/agent_loop_tq.py` — `AgentLoopManagerTQ.generate_sequences()` + +**锚点**: `chunkes = prompts.chunk(len(self.agent_loop_workers))` 之后、`ray.get([...])` 之前。 +验证批次(`prompts["validate"]=True`)不打样本 trace。 + +```python + _val = prompts.get("validate", False) + _val = bool(getattr(_val, "data", _val)) + for _worker_idx, _chunk in enumerate(chunkes): + _uids = _chunk.get("uid", None) + if _uids is not None and not _val: + _gs = _chunk.get("global_steps", -1) + _gs = getattr(_gs, "data", _gs) + trace_flow(uid=[str(getattr(_uids[i], "data", _uids[i])) for i in range(len(_chunk))], + traj_uid='N/A', + from_module='AgentLoopManagerTQ', to_module=f'AgentLoopWorkerTQ_{_worker_idx}', + stage='DATA_DISPATCH', detail='dispatched_to_agent_loop_worker', + global_step=int(_gs) if isinstance(_gs, (int, np.integer)) else -1) +``` + +注:需要模块顶部 `import numpy as np`。 + +--- + +## 插入点 V1-10: 单样本任务开始(AGENT_LOOP / agent_task_started) + +**文件**: `verl/trainer/ppo/v1/agent_loop_tq.py` — `AgentLoopWorkerTQ.generate_sequences()` + +**锚点**: per-sample 循环内 `task = asyncio.create_task(self._run_prompt(...))` 之前 +(“fire-and-forget” 注释处;用方法内已有的局部 `validate` gate,验证样本跳过)。 + +```python + if not validate: + trace_flow(uid=str(prompt.get('uid', 'unknown')), traj_uid='N/A', + from_module='AgentLoopWorkerTQ', + to_module=f"AgentLoop({prompt.get('agent_name', 'single_turn')})", + stage='AGENT_LOOP', detail='agent_task_started', + global_step=int(prompt.get('global_steps', -1))) +``` + +--- + +## 插入点 V1-11: session 展开 + traj_uid 注入(AGENT_LOOP / session_started) + +**文件**: `verl/trainer/ppo/v1/agent_loop_tq.py` — `AgentLoopWorkerTQ._run_prompt()` + +**前置**(方法体第一行):每条 prompt 的 `_run_prompt` 是独立 asyncio.Task, +在任务入口对验证样本调用 `suppress_sample_trace()`(contextvars 会复制进下方 +创建的 session 子任务),该 prompt 全链路 trace_flow(含本点 V1-11 与共享 +agent-loop 点位)自动静默: + +```python + # validation samples never enter the training stages; silence their + # sample-flow trace for this prompt task and all session sub-tasks + # (contextvars are copied into tasks created below) + if trajectory.get("validate", False): + suppress_sample_trace() +``` + +(import 一并从 `verl.utils.tracing` 导入 `suppress_sample_trace`。) + +**锚点**: `for i in range(n):` 循环体内、`task = asyncio.create_task(self._run_agent_loop(...))` 之前。 + +```python + trace_flow(uid=str(prompt.get('uid', 'unknown')), traj_uid=str(i), + from_module='AgentLoopWorkerTQ', to_module='AgentLoopWorkerTQ', + stage='AGENT_LOOP', + detail=f'session_started_session_id={i}_rollout_n={n}', + global_step=int(prompt.get('global_steps', -1))) + # traj_uid injection: v1 uses session_id as the per-trajectory id + # (TQ key format {uid}_{session_id}_{index}); merge into a copy so a + # pre-existing "traj_uid" in the prompt batch never causes a + # duplicate-keyword crash and wins over the derived value. + _session_prompt = {"traj_uid": str(i), **prompt} +``` + +并把 `self._run_agent_loop(...)` 的 `**prompt` 改为 `**_session_prompt` +(这是 V1 唯一允许的数据新增:`traj_uid` 字段,供共用 agent loop 点位与 trajectory 使用)。 + +--- + +## 插入点 V1-12: 权重同步(SYNC_PARAMS / weights_synced) + +三个 trainer mode 的 `on_step_end()` 各插一处,锚点均为 +`checkpoint_manager.update_weights(self.global_steps)`(separate_async 为 +`standalone_checkpoint_manager`)之后: + +| 文件 | 进度行前缀 | +|---|---| +| `verl/trainer/ppo/v1/trainer_sync.py` | 无(colocated) | +| `verl/trainer/ppo/v1/trainer_colocate_async.py` | 无(colocated) | +| `verl/trainer/ppo/v1/trainer_separate_async.py` | `disagg: ` | + +```python + log_progress('start sync params') # separate_async: 'disagg: start sync params' + with marked_timer("update_weights", self.timing_raw, color="red"): + self.checkpoint_manager.update_weights(self.global_steps) + trace_flow(uid='*', traj_uid='*', from_module='CheckpointEngineManager', + to_module='RolloutReplica', stage='SYNC_PARAMS', + detail=f'weights_synced_global_step={self.global_steps}', + global_step=self.global_steps) + log_progress('end sync params') # separate_async: 'disagg: end sync params' +``` + +--- + +## V1 训练进度插桩(log_progress,`trainer_base.py`) + +前缀规则与 v0 一致:**sync / colocate_async 为训推一体(无前缀)**, +**separate_async 为训推分离(`disagg: ` 前缀 + fetch 行)**。在 `fit()` 与 +`_step_once()` 开头各计算一次 +`_pp = 'disagg: ' if self.trainer_mode == 'separate_async' else ''`。 + +`fit()`: + +- init 行(`on_train_begin()` 之后):`Total training steps: N` + 5 行 dict 风格 + config dump(`'train_batch_size': ...,` 等,与 v0 完全一致); +- 每 step banner(while 循环体开头):一体版 + `==== Epoch {e}, consumed {(global_steps-1) % steps_per_epoch} batches, Step {n} ====`, + 分离版 `==== Step {n} ====`; +- `start/end save checkpoint`(save_freq 分支)、`start/end eval`(test_freq 分支)、 + `start record metrics`(`self._compute_metrics(...)` 之前)/ + `end record metrics`(`self.logger.log(...)` + dapo 落表之后)。 + +`_step_once()`: + +- replay_buffer.sample 前后:一体版 `start generation` / `end generation`;分离版 + `Trainer: try to load global batch@{N}, chunk[0/0] from replay_queue` / + `Trainer: successfully load global batch@{N}, chunk[0/0] from replay_queue` + (V1 中「取全局 batch」即 ReplayBuffer 采样); +- `start/end compute reward`(仅 colocated reward 分支)、`start/end compute log prob`、 + `start/end compute ref log prob`、`start/end compute values`、`start/end compute adv`、 + `start/end update actor`,均按 `_pp` 加前缀。 + +--- + +## 插入点 V1-E: trace 环境变量注入 Ray runtime_env + +**文件**: `verl/trainer/constants_ppo.py` — `get_ppo_ray_runtime_env()` 尾部 +(与 v0 Step 3 第 2 处相同,若已存在则跳过): + +```python + for key in ("VERL_TRACE_DIR", "VERL_PROGRESS_HEARTBEAT_SEC"): + if os.environ.get(key): + runtime_env["env_vars"][key] = os.environ[key] +``` + +--- + +## 完成标准 + +```bash +# V1 侧 trace_flow 计数(期望 14:trainer_base 8 + agent_loop_tq 3 + 3 个 mode 各 1) +grep -c "trace_flow(" verl/trainer/ppo/v1/*.py | awk -F: '{s+=$2} END {print s}' + +# V1 侧 log_progress 计数(期望 32:trainer_base 26 + 3 个 mode 各 2) +grep -c "log_progress(" verl/trainer/ppo/v1/*.py | awk -F: '{s+=$2} END {print s}' + +# 全仓库合计(v0 17 + v1 14 = 31;log_progress v0 50 + v1 32 = 82) +grep -rn "trace_flow(" verl/ --include="*.py" | grep -v "utils/tracing" | wc -l +grep -rn "log_progress(" verl/ --include="*.py" | grep -v "utils/tracing" | wc -l + +python3 -m py_compile verl/trainer/ppo/v1/*.py +``` + +--- + +rl-tracing Step 视图说明(已调整) + +step 级 rollout/train 窗口的起止**不再走 sample-flow CSV**(曾短暂以 +`ROLLOUT/rollout_finished`、`TRAIN_UPDATE/update_started` 打点实现,已移除), +改为 progress 行输出,见 progress_points.md 的「step 级窗口行」章节。 + +保留的相关改动:`AGENT_LOOP` 两个点的 `global_step` 由共享的 agent_loop.py +提供(`agent_task_started` 取 `batch.meta_info['global_steps']`, +`agent_run_completed` 取 `trajectory['step']`),v1 无需单独处理。 diff --git a/.qoder/skills/verl-trace-instrument/progress_points.md b/.qoder/skills/verl-trace-instrument/progress_points.md new file mode 100644 index 00000000..32106020 --- /dev/null +++ b/.qoder/skills/verl-trace-instrument/progress_points.md @@ -0,0 +1,199 @@ +# 训练进度展示插桩点清单(TRAIN_PROGRESS,对齐 verl-horizon 裸日志格式) + +进度日志**内容与 verl-horizon 原始 print 一致**,两条路径各用其原生格式: + +- **训推一体(v0 `RayPPOTrainer.fit`)**:horizon 一体版格式,**无前缀**, + 如 `start compute reward` / `start generation`(horizon 一体版为单侧 start, + 此处补成 start/end 对以便前端计算耗时); +- **训推分离(`SeparateRayPPOTrainer`)**:horizon 分离版格式,**`disagg: ` 前缀**, + 如 `disagg: start compute reward`,另有 fetch 行 + `Trainer: try to/successfully load global batch@N, chunk[0/0] from replay_queue`。 + +示例: + +``` +Trainer: try to load global batch@3, chunk[0/0] from replay_queue # 仅分离版 +Trainer: successfully load global batch@3, chunk[0/0] from replay_queue # 仅分离版 +start compute reward # 一体版 +end compute reward +disagg: start compute reward # 分离版 +disagg: end compute reward +...``` + +`log_progress(message)` 直接输出一行裸文本:stdout 常开;设置 `VERL_TRACE_DIR` +时同步逐行落盘 `{trace_dir}/progress/{YYYYMMDD-HH}.log`(按小时轮转)。 +函数由 trace 工具模块提供(见 [trace_module.py](trace_module.py)),import 方式: + +```python +from verl.utils.tracing import log_progress +``` + +**零侵入规则同样适用**:只插入 `log_progress(...)` 调用(及 fetch 包装生成器), +不改动任何业务逻辑。 + +前端"训练阶段"面板的 11 个阶段与插桩对应关系(训推一体 v0 与训推分离 +separation 两条路径均已插桩;separation 的 `SeparateRayPPOTrainer` 继承 +v0 的 fetch 包装器): + +| # | 前端阶段 | 日志行(一体版无前缀 / 分离版 `disagg: ` 前缀) | 锚点 | +|---|---|---|---| +| 1 | 获取训练 Batch | `Trainer: try to/successfully load global batch@N, chunk[0/0] from replay_queue` | **仅分离版**:`_logged_train_dataloader()` 包装生成器(一体版用原始 for 循环,horizon 一体版无此行) | +| 1' | rollout 生成 | `start/end generation` | **仅一体版**:`marked_timer("gen"...)` 内(horizon 一体版单侧 `start generation`) | +| 2 | Reward 计算 | `start/end compute reward` | `marked_timer("reward"...)` 前后 | +| 3 | LogProbs 重算 | `start/end compute log prob` | `marked_timer("old_log_prob"...)` 前后(bypass 模式不进入→前端置灰) | +| 4 | Ref LogProbs | `start/end compute ref log prob` | `marked_timer(str(Role.RefPolicy)...)` 前后(无 ref policy 不进入→置灰) | +| 5 | Advantage 计算 | `start/end compute adv` | `marked_timer("adv"...)` 前后 | +| 5' | Critic values | `start/end compute values` | `marked_timer("values"...)` 前后(仅 `use_critic` 时进入;horizon 一体版单侧 `start compute values`) | +| 6 | Prompt 准确率更新 | `start/end update prompt2accs` | **两路径均无对应实现**(horizon 异步过采样特有),不插桩→前端置灰 | +| 7 | Actor 参数更新 | `start/end update actor` | `marked_timer("update_actor"...)` 前后 | +| 8 | 指标记录 | `start/end record metrics` | metrics 汇总块(v0:`# training metrics` 起至 `logger.log(...)`;separation:`_fit_collect_metrics`) | +| 9 | 参数同步 | `start/end sync params` | `marked_timer("update_weights"...)` 前后 | +| 10 | 验证 (Eval) | `start/end eval` | `marked_timer("testing"...)` 前后 | +| 11 | Checkpoint 保存 | `start/end save checkpoint` | `marked_timer("save_checkpoint"...)` 前后 | + +--- + +## A. fetch batch 包装(**仅分离版**;`SeparateRayPPOTrainer` 继承基类方法) + +在 `fit()` 的 epoch 循环内把 `for batch_dict in self.train_dataloader:` 改为: + +```python + for epoch in range(current_epoch, self.config.trainer.total_epochs): + self._train_dataloader_iter = iter(self.train_dataloader) + for batch_dict in self._logged_train_dataloader(): +``` + +并在类内新增(不改动循环体,零重缩进): + +```python + def _logged_train_dataloader(self): + """Yield train batches with horizon-style fetch progress logs.""" + while True: + log_progress( + f"Trainer: try to load global batch@{self.global_steps}, chunk[0/0] from replay_queue" + ) + try: + batch_dict = next(self._train_dataloader_iter) + except StopIteration: + return + log_progress( + f"Trainer: successfully load global batch@{self.global_steps}, chunk[0/0] from replay_queue" + ) + yield batch_dict +``` + +语义与原始 for 循环完全等价(每 epoch 新建 iterator,StopIteration 结束 epoch)。 + +--- + +## B. 阶段 start/end 对(两路径同名锚点,前缀按路径) + +每个阶段一对 start/end,start 插在 `with marked_timer(...)` 行**之前**, +end 插在块内最后锚点**之后**(缩进与锚点一致);内容即上表第 3 列字符串, +**一体版无前缀、分离版带 `disagg: ` 前缀**,例如: + +```python + log_progress('disagg: start compute reward') + with marked_timer("reward", timing_raw, color="yellow"): + ... + log_progress('disagg: end compute reward') +``` + +注意: +- 一体版额外输出 `start/end generation`(rollout 阶段);分离版 rollout 在 + 独立服务端,训练循环不输出 generation 行。 +- 每 step 开头输出 horizon 格式的 step banner:一体版 + `==================== Epoch {e}, consumed {k} batches, Step {n} ====================`, + 分离版 `==================== Step {n} ====================`; + heartbeat 不再输出。 +- fit() 开头输出初始化行:horizon 原文 `Total training steps: {N}` + + 逐行 dict 风格配置dump(前端解析格式,每行一条): + `'train_batch_size': N,` / `'rollout_n': N,` / `'ppo_mini_batch_size': N,` / + `'num_minibatches': N,` / `'total_training_steps': N` + (num_minibatches = ppo_epochs * train_batch_size * rollout_n // ppo_mini_batch_size)。 + +--- + +## C. record metrics 包裹 + +- v0:在 `# training metrics` 注释行后插 `log_progress('disagg: start record metrics')`, + 在 `logger.log(data=metrics, step=self.global_steps)` 后插 + `log_progress('disagg: end record metrics')`。 +- separation:在 `_fit_collect_metrics()` 首尾各插一行。 + +--- + +## 完成标准 + +- fetch 包装(仅分离版)+ 10 个可覆盖阶段(两路径各一套;一体版另含 + generation 对);prompt2accs 明确不插桩(无实现)。 +- 一体版行无前缀、分离版行带 `disagg: ` 前缀;无 heartbeat、无 step banner、 + 无 `[TRAIN_PROGRESS]` 结构化前缀残留: + `grep -rn "heartbeat\|TRAIN_PROGRESS" verl/ --include="*.py"` 应为 0。 +- `python -m py_compile` 通过。 +- `grep -rn "log_progress(" verl/ --include="*.py" | grep -v "utils/tracing"` 计数 = 50 + (v0 27 + separation 23;含初始化 2 行/路径、banner、fetch 包装 2 行仅分离版走到)。 +- 冒烟验证:不设置 VERL_TRACE_DIR 也能在 stdout 看到 `disagg: start compute reward` + 等裸行;设置后 `{trace_dir}/progress/{YYYYMMDD-HH}.log` 行数 > 0 且内容同 stdout。 + +--- + +## rl-tracing 进度增强(生成吞吐 + 已完成 rollout) + +训练进度面板(job_progress API)解析 master-0 原始日志,以下两类行由 verl +主动产出,不依赖 sglang 引擎统计日志(开源版 sglang Engine 的 +`log_level=error` + `disable_log_stats=True` 会屏蔽引擎自身的 +`gen throughput` 行): + +1. **生成吞吐** — `gen throughput (token/s): {X:.2f}` + - 解析正则:`gen throughput \(token/s\):\s*([\d.]+)`。 + - v0:`ray_trainer.py` gen 窗口结束、`gen_batch_output` slice 之后, + `log_gen_throughput(gen_batch_output, timing_raw.get("gen"))` + (token 数 = attention_mask 响应区求和,缺 mask 退化为 numel)。 + - separation:`separation/ray_trainer.py` 同位置同样调用。 + - v1:`v1/trainer_base.py` `_step_once` gen 窗口结束后,**仅非 disagg + 模式**(`if not _pp`),从 KVBatchMeta tags 的 `response_len` 求和后 + `log_gen_throughput_tokens(...)`;separate_async 的 gen 窗口是 + replay-buffer 取批等待,语义不是生成吞吐,故不打。 +2. **已完成 rollout** — `Task {uid}/{traj_uid} is completed, score={score} status=Finished` + - 计数条件:`like '%is completed, score=%' and like '%status=Finished%'`。 + - 唯一打点:`experimental/agent_loop/agent_loop.py` `_run_agent_loop` + 的 `_agent_loop_postprocess` 返回后(v0+v1 共用),逐轨迹一行, + 跳过 validation;仅在 VERL_TRACE_DIR 设置时输出(`log_rollout_completed` + 内部用 `_resolve_trace_dir()` 门控);**不带 `TONGYI:` 前缀**。 + - score 取 `AgentLoopOutput.reward_score`(`_compute_score` 原地写入); + 不能取 `_agent_loop_postprocess` 的返回值 —— v0 返回包装输出, + v1 TQ 子类返回 None。None 时输出 0.0。 + +注意:本增强后,`grep -rn "log_progress(" verl/ --include="*.py"` 的计数 +校验基线随之变化(v0/separation 各 +1 处调用,agent_loop 经由 +`log_rollout_completed` 间接输出)。 + +--- + +## step 级窗口行(rollout/train 起止) + +每个 global step 输出四行,格式固定 ``Step {N} {phase} {boundary}``(step 号后 +无冒号),由 `log_step_window(global_step, phase, boundary)`(trace.py)统一输出, +stdout 与 `{trace_dir}/progress/{YYYYMMDD-HH}.log` 同步落盘: + +``` +Step 3 rollout started +Step 3 rollout finished +Step 3 train started +Step 3 train finished +``` + +插桩位置(均不走 sample-flow CSV): +- **rollout started/finished**: + - v0 `ray_trainer.py` / separation:gen 窗口内 `start generation` 之前 / + gen 窗口结束 slice 之后(与 `log_gen_throughput` 同位置)。 + - v1 `v1/trainer_base.py` `_step_once`:`on_sample_begin()` 之后 / + gen 窗口结束之后(所有 trainer_mode 都输出;separate_async 下该窗口 + 语义是 replay-buffer 取批等待)。 +- **train started/finished**:`start update actor` 之前 / + `end update actor` 之后(v0、v1 `_step_once`、separation + `_fit_update_actor` 三处);critic warmup 跳过 actor 更新时不输出。 + +窗口语义:v0 为同步生成包络(driver 时钟两端同源);v1 rollout 窗口是 +"批次就绪(消费)窗口",真实生产时间见逐轨迹 `agent_task_started`。 diff --git a/.qoder/skills/verl-trace-instrument/trace_data_formats.md b/.qoder/skills/verl-trace-instrument/trace_data_formats.md new file mode 100644 index 00000000..87a2fe32 --- /dev/null +++ b/.qoder/skills/verl-trace-instrument/trace_data_formats.md @@ -0,0 +1,239 @@ +# verl v1 分支 Trace 数据格式说明(verl-tracing-v1) + +本文档描述 `verl-tracing-v1` 分支产出的三类 trace 数据的**格式定义、字段语义与示例**, +并逐项说明在 xfuse-controller-py 的 rl-tracing 页面展示时,当前后端/前端代码 +**能正常消费的部分**与**因格式变动/数据差异无法正确展示的部分**。 + +三类数据均由 `VERL_TRACE_DIR` 单一开关控制,落盘到 `{VERL_TRACE_DIR}` 下三个子目录: + +| 数据 | 生成模块 | 文件 | 粒度 | collector 摄入(`collectorConfigOverride.json`) | +|---|---|---|---|---| +| sample-flow | `verl/utils/tracing/trace.py` `trace_flow` | `sample/{hostname}-pid{pid}.csv` | 每轨迹/样本每事件一行 | `suffixes=[".csv"]`,`csv_headers` 定列名,**文件本身无表头行** | +| trajectory | `verl/utils/tracing/trajectory_store.py` | `trajectory/trajectory_{pid}_{ts}.jsonl` | 每完成一条轨迹一行 JSON | `suffixes=[".jsonl"]`,`labels.cat=trajectory`,`expandJsonl=false` | +| progress | `verl/utils/tracing/trace.py` `log_progress` | `progress/{YYYYMMDD-HH}.log`(按小时轮转) | step/阶段级裸文本行 | `labels.cat=progress`,`collectExisting=true` | + +三者统一发往 SLS project `post-train-metric` / logstore `test-rl-trace`; +discovery 文件:`/tmp/verl_trace_sample_path.txt`、`/tmp/verl_trace_trajectory_path.txt`、 +`/tmp/verl_progress_path.txt`(节点本地,多机训练需到各节点读取)。 + +--- + +## 1. Sample-flow CSV(样本/轨迹流转事件) + +### 1.1 行格式 + +每行 6 列(列名由 collector 的 `csv_headers` 注入,文件内不含表头): + +``` +job_id, cat, sample_id, sample_status, timestamp, content +``` + +| 列 | 说明 | +|---|---| +| `job_id` | 取环境变量 `DLC_JOB_ID`,未设置时为 `ENV_NOTFOUND_DLC_JOB_ID` | +| `cat` | 恒为 `sample`(开源版无 horizon 的 `cat=megatron` 行) | +| `sample_id` | `traj_uid` 为数字时是 `{uid}/{traj_uid}`(轨迹级行),否则是 `{uid}/sample`(样本级行) | +| `sample_status` | `stage` 为 `TRAIN_UPDATE`/`FILTER_TRAJECTORIES` 时 `COMPLETE`,否则 `RUNNING` | +| `timestamp` | 毫秒时间戳 | +| `content` | `from={模块}\|to={模块}\|stage={阶段}\|global_step={N}\|detail={细节}`,管道分隔,global_step/detail 可缺省 | + +v1 的 `traj_uid` 来源于 TransferQueue key `{uid}_{session_id}_{index}` 的解析 +(`_trace_ids_from_batch`),因此轨迹级行的 sample_id 形如 `{uid}/{session_id}`。 +validation 样本的 trace 被整体抑制,不会出现在 CSV 中。 + +### 1.2 v1 阶段清单(按流水线顺序) + +| # | stage | detail | 行粒度 | 打点位置 | +|---|---|---|---|---| +| 1 | DATA_PREPROCESS | uid_generated | uid/sample | trainer_base(数据集 uid 生成) | +| 2 | DATA_PREPROCESS | traj_uid_generated_by_repeat | uid/N | ray_trainer/v1 分发前 repeat 扩增 | +| 3 | DATA_DISPATCH | dispatched_to_agent_loop_manager | uid/N | trainer_base → AgentLoopManager | +| 4 | DATA_DISPATCH | dispatched_to_agent_loop_worker | uid/N | AgentLoopManager → Worker(带 global_step) | +| 5 | AGENT_LOOP | agent_task_started | uid/N | Worker 建 asyncio task(带 global_step) | +| 6 | LLM_GENERATE | sending_generate_request ↔ received_generate_response | uid/N | 每轮生成一对(多轮 = 多对) | +| 7 | AGENT_TOOL | tool_executed_name={name} | uid/N | 仅 tool_agent_loop | +| 8 | AGENT_LOOP | agent_run_completed | uid/N | agent 执行结束(带 global_step) | +| 9 | ROLLOUT_REWARD | reward_score_computed_score={X} | uid/N | reward loop 打分(streaming reward 模式) | +| 10 | TRAIN_SAMPLE | sampled_from_replay_buffer | uid/N | **v1 独有**,ReplayBuffer.sample 取批 | +| 11 | TRAIN_REWARD | reward_extracted | uid/N | colocate reward 提取 | +| 12 | TRAIN_LOGPROB | old_log_prob_computed | uid/N | | +| 13 | TRAIN_REF_LOGPROB | ref_log_prob_computed | uid/N | 可选 | +| 14 | TRAIN_ADVANTAGE | advantage_computed | uid/N | | +| 15 | TRAIN_UPDATE | actor_parameters_updated | uid/N | 训练终态,sample_status=COMPLETE | +| 16 | SYNC_PARAMS | weights_synced_global_step={N} | `*/*` | 权重同步到 rollout | + +### 1.3 示例 + +```csv +ENV_NOTFOUND_DLC_JOB_ID,sample,92ee95cb-xxxx/sample,RUNNING,1787215029000,from=DataLoader|to=PPOTrainerV1|stage=DATA_PREPROCESS|global_step=1|detail=uid_generated +ENV_NOTFOUND_DLC_JOB_ID,sample,92ee95cb-xxxx/3,RUNNING,1787215029010,from=PPOTrainerV1|to=AgentLoopManagerTQ|stage=DATA_DISPATCH|global_step=1|detail=dispatched_to_agent_loop_worker +ENV_NOTFOUND_DLC_JOB_ID,sample,92ee95cb-xxxx/3,RUNNING,1787215029349,from=AgentLoopWorkerTQ|to=AgentLoop(single_turn_agent)|stage=AGENT_LOOP|global_step=1|detail=agent_task_started +ENV_NOTFOUND_DLC_JOB_ID,sample,92ee95cb-xxxx/3,RUNNING,1787215029400,from=AgentLoop(single_turn_agent)|to=SGLang|stage=LLM_GENERATE|detail=sending_generate_request +ENV_NOTFOUND_DLC_JOB_ID,sample,92ee95cb-xxxx/3,RUNNING,1787215030500,from=SGLang|to=AgentLoop(single_turn_agent)|stage=LLM_GENERATE|detail=received_generate_response +ENV_NOTFOUND_DLC_JOB_ID,sample,92ee95cb-xxxx/3,RUNNING,1787215030600,from=AgentLoop(single_turn_agent)|to=AgentLoopWorker|stage=AGENT_LOOP|global_step=1|detail=agent_run_completed +ENV_NOTFOUND_DLC_JOB_ID,sample,92ee95cb-xxxx/3,RUNNING,1787215030700,from=RewardLoopWorker|to=AgentLoopWorker|stage=ROLLOUT_REWARD|detail=reward_score_computed_score=1.0 +ENV_NOTFOUND_DLC_JOB_ID,sample,92ee95cb-xxxx/3,RUNNING,1787215043421,from=ReplayBuffer|to=PPOTrainerV1|stage=TRAIN_SAMPLE|global_step=1|detail=sampled_from_replay_buffer +ENV_NOTFOUND_DLC_JOB_ID,sample,92ee95cb-xxxx/3,COMPLETE,1787215060000,from=ActorRolloutWorker|to=PPOTrainerV1|stage=TRAIN_UPDATE|global_step=1|detail=actor_parameters_updated +``` + +### 1.4 rl-tracing 页面展示现状与差异 + +后端(`src/routes/rl_tracing.py`)查询条件 `(cat="sample" or cat="megatron") and job_id="{dlc}"`, +列结构与 v1 完全一致,**样本列表、样本详情时间线、漏斗/阶段耗时的骨架可以正常展示**,但存在以下差异: + +| 展示项 | 现状 | 差异原因 | +|---|---|---| +| 样本列表/终态统计 | ✅ 正常(终态判定用 `actor_parameters_updated`) | — | +| 样本详情 stage/detail 说明 | ⚠️ 部分 | 前端 `STAGE_DOCS`/`DETAIL_DOCS` 是 horizon 词表:`AGENT_LOOP`、`LLM_GENERATE`、`AGENT_TOOL`、`TRAIN_SAMPLE`、`SYNC_PARAMS` 及 `dispatched_to_agent_loop_*`、`sampled_from_replay_buffer`、`reward_score_computed_score=` 等 v1 命名无图例文案(页面按原始名展示,不报错) | +| 流转分析-漏斗/Sankey | ⚠️ 部分为空 | SQL 里程碑依赖 horizon 的 `dispatched_to_instance`/`trajectory_pushed`/`trajectories_filtered_and_ready`/`FORM_BATCH`,v1 均无 → "下发到推理实例/Rollout 完成入缓冲/筛选入选"计数恒 0;"生成轨迹/开始执行/完成训练"正常 | +| 流转分析-阶段耗时表 | ⚠️ 部分为空 | `dispatch_wait`/`to_buffer`/`buffer_wait`/`train`(依赖 t_form)无数据;多轮生成、端到端(生成→训练完成)正常 | +| 流转分析-在途卡点 | ✅ 运行中有数据 | 终态仅"已训练"一种(无过采样取消/丢弃);任务结束后 inflight=0 显示"无在途轨迹"属预期 | +| 流转分析-Step 生产-消费视图 | ❌ 恒空 | 依赖 `FORM_BATCH` 的 global_step 锚点与 `cat=megatron` 批级训练行,v1 两者皆无。**step 窗口数据已改由 progress 行承载(见 3.4),后端需新增解析才能展示** | + +--- + +## 2. Trajectory JSONL(轨迹明细) + +### 2.1 行格式 + +文件名 `trajectory_{pid}_{ts}.jsonl`(pid 仅用于文件命名)。每行一个 JSON 对象, +**顶层只保留两个索引列,其余字段全部折叠进 `raw_message` JSON 字符串** +(SLS 摄入约定:扁平索引列 + 单一 raw payload 列): + +```json +{ "prompt_uid": "", "sample_index": , "raw_message": "" } +``` + +`raw_message` 解析后的字段(`Trajectory` dataclass,序列化时已裁剪): + +| 字段 | 说明 | +|---|---| +| request_id | traj_uid(v1 = TQ key 的 session_id) | +| n_samples | rollout.n | +| global_step | 轨迹所属训练步 | +| created_at_ts / completed_at_ts / total_time_ms | 轨迹起止时间戳与总耗时 | +| turns[] | 每轮一个 Turn(见下) | +| initial_prompt_text | 首轮 prompt 的 decode 文本 | +| agent_metrics | agent 侧计时(generate_sequences/tool_calls/compute_score/num_preempted) | + +Turn 字段:`turn_index`、`prompt_token_count`、`completion_text`、`completion_token_count`、 +`finish_reason`、`observation_text`、`start_ts`/`end_ts`、`observation_start_ts`/`observation_end_ts`、 +`turn_total_time_ms`(未采集时由 completion+observation 推导)、`completion_time_ms`、`observation_time_ms`、 +`sglang_time_ms`/`sglang_queue_time_ms`/`sglang_ttft_ms`/`sglang_prefill_ms`/`sglang_decode_time_ms`、 +`template_overhead_tokens`/`template_overhead_text`(turn_index≥1 时计算)。 + +序列化裁剪规则(`TrajectoryStore`): +- **删除大字段**:`initial_prompt_token_ids`、turn 级 `completion_token_ids`/`observation_token_ids`; +- **删除架构不可得字段**:`replica_id`、turn 级 `prompt_text`/`sglang_rid`; +- 单轮场景 `observation_*`/`template_overhead_*` 为 null 属正常。 + +validation 轨迹不落盘。 + +### 2.2 示例 + +```json +{"prompt_uid": "92ee95cb-f14d-4a51-92a5-9a2a0f88d1d5", "sample_index": 3, "raw_message": "{\"request_id\": \"3\", \"n_samples\": 4, \"global_step\": 3, \"created_at_ts\": 1787215072.088314, \"completed_at_ts\": 1787215073.420837, \"total_time_ms\": 1332.52, \"turns\": [{\"turn_index\": 0, \"prompt_token_count\": 87, \"completion_text\": \"To determine how much Joanie...\", \"completion_token_count\": 512, \"finish_reason\": \"length\", \"observation_text\": null, \"start_ts\": 1787215072.1198363, \"end_ts\": 1787215073.321284, \"turn_total_time_ms\": 1201.45, \"completion_time_ms\": 1201.45, \"observation_time_ms\": null, \"sglang_time_ms\": 1111.64, \"sglang_queue_time_ms\": null, \"sglang_ttft_ms\": null, \"sglang_prefill_ms\": null, \"sglang_decode_time_ms\": 1111.64, \"template_overhead_tokens\": null, \"template_overhead_text\": null}], \"initial_prompt_text\": \"<|im_start|>system\\nYou are Qwen...\", \"agent_metrics\": {\"generate_sequences\": 1.201, \"tool_calls\": 0.0, \"compute_score\": 0.0547, \"num_preempted\": -1}}"} +``` + +### 2.3 相对旧格式的变动 + +1. **顶层结构**:旧格式每行是完整扁平 JSON(request_id/prompt_uid/turns 等全在顶层); + 新格式顶层仅 `prompt_uid`/`sample_index`/`raw_message`,其余在 `raw_message` 字符串内; +2. **rank_id 字段移除**:旧格式中的 `rank_id`(值为 worker 进程 pid,语义不准确)已从 + schema 中删除,不再输出;消费端无需处理(`Trajectory.from_dict` 对未知字段自动忽略, + 旧数据仍可加载); +3. collector 侧 `expandJsonl: false`,每行作为**一条**日志进入 `test-rl-trace`(带 `cat=trajectory` 标签),不再展开成多行。 + +### 2.4 rl-tracing 页面展示现状与差异 + +**❌ 当前后端无法展示新格式轨迹**。差异: + +| 维度 | 后端现状 | 新格式 | +|---|---|---| +| 数据源 | 轨迹下载 API(`sample_trajs_download`/`job_trajs_download`)查询 **`sgl-metric` logstore**(horizon 侧 per-turn 展开摄入) | v1 轨迹在 **`test-rl-trace`**,`cat=trajectory`,每条轨迹**一行** | +| 索引列 | 依赖展开列 `prompt_uid`/`sample_index`/`turn_index` 做 `group by` 与逐行拉取 | `prompt_uid`/`sample_index` 仍在顶层可索引;`turn_index` 在 raw_message 内,不可直接 SQL | +| 还原逻辑 | `_traj_jsonl` 把 SLS 行剔除系统字段后直接拼回 jsonl | 需要改为:按 `cat=trajectory` 查 test-rl-trace → 解析每行 `raw_message`(JSON 字符串)→ 还原完整轨迹对象 | + +需要的后端改动:新增一个针对 `test-rl-trace` 的轨迹查询分支 +(`job_id="{dlc}" and cat="trajectory" | select prompt_uid, sample_index, ...`), +zip 导出时把 `raw_message` 解析为对象再写 jsonl;前端下载入口不用变。 + +--- + +## 3. Progress 日志(训练进度裸行) + +### 3.1 文件与行约定 + +- 文件:`{VERL_TRACE_DIR}/progress/{YYYYMMDD-HH}.log`,按小时轮转;同时无条件输出到 stdout; +- 每行一条完整日志,horizon 裸格式,**无结构化前缀**; +- 前缀规则:v1 `trainer_mode == 'separate_async'` 时阶段行带 `disagg: ` 前缀, + colocate/sync 模式无前缀;step 级窗口行/吞吐行/完成行**一律无前缀**。 + +### 3.2 行分类 + +| 类别 | 行格式 | 产出位置 | +|---|---|---| +| 初始化 | `Total training steps: {N}` | fit 开始 | +| 配置 dump | `'train_batch_size': {N},` / `'rollout_n': {N},` / `'ppo_mini_batch_size': {N},` / `'num_minibatches': {N},` / `'total_training_steps': {N}` | 逐行 dict 格式,供前端任务配置解析 | +| step banner | `==================== Epoch {e}, consumed {b} batches, Step {N} ====================`(colocate)/ `==================== Step {N} ====================`(disagg) | 每步开始 | +| 阶段对 | `[disagg: ]start generation` / `end generation`、`start/end compute reward`、`start/end compute log prob`、`start/end compute ref log prob`、`start/end compute values`、`start/end compute adv`、`start/end update actor`、`start/end save checkpoint`、`start/end eval`、`start/end record metrics`、`start/end sync params` | `_step_once` 各阶段首尾 | +| fetch 行(仅 disagg) | `disagg: try to load global batch@{N}, chunk[0/0] from replay_queue` / `disagg: successfully load global batch@{N}, ...` | replay_buffer.sample 前后 | +| 生成吞吐 | `gen throughput (token/s): {X:.2f}` | gen 窗口后;colocate 为真实生成吞吐,separate_async 不输出 | +| 已完成 rollout | `Task {uid}/{traj_uid} is completed, score={X} status=Finished` | agent_loop 每条训练轨迹完成时(跳过 validation,trace 开启时才输出) | +| step 级窗口 | `Step {N} rollout started` / `Step {N} rollout finished` / `Step {N} train started` / `Step {N} train finished` | gen 窗口与 update_actor 首尾,**每步 4 行** | + +### 3.3 单个 step 的示例序列(colocate 模式) + +```text +==================== Epoch 0, consumed 0 batches, Step 1 ==================== +Step 1 rollout started +start generation +end generation +Step 1 rollout finished +gen throughput (token/s): 4823.51 +Task 92ee95cb-f14d-4a51-92a5-9a2a0f88d1d5/0 is completed, score=1.0 status=Finished +Task 92ee95cb-f14d-4a51-92a5-9a2a0f88d1d5/1 is completed, score=0.0 status=Finished +start compute reward +end compute reward +start compute log prob +end compute log prob +start compute ref log prob +end compute ref log prob +start compute adv +end compute adv +Step 1 train started +start update actor +end update actor +Step 1 train finished +start sync params +end sync params +``` + +### 3.4 rl-tracing 页面展示现状与差异 + +训练进度 tab(`/api/job_progress`)解析的是 **job-log logstore 的 master-0 stdout** +(raw_message),而 progress 文件与 stdout 内容一致,因此: + +| 行 | 展示现状 | 说明 | +|---|---|---| +| 配置 dump | ✅ | 后端按 `train_batch_size` 等关键词 + 任务起点窗口解析 | +| step banner | ✅ | 匹配 `%= Step %` / `%, Step %` | +| 阶段对 | ✅ | 匹配 `start generation`/`start compute reward` 等枚举;`disagg: ` 前缀行走 trainer 标记解析 | +| fetch 行 | ✅ | 匹配 `%load global batch%` | +| gen throughput | ✅ | `misc` 查询匹配 `%gen throughput (token/s)%`,正则 `gen throughput \(token/s\):\s*([\d.]+)` —— **与旧 SGLang 引擎日志格式刻意对齐,前端"生成吞吐"卡片直接可用** | +| 已完成 rollout | ✅ | `finished` 计数匹配 `%is completed, score=%` + `%status=Finished%`,速率用最近 120 条估算 —— 与 horizon agentic 行格式对齐(仅去掉了 `TONGYI:` 前缀,计数条件不受影响) | +| **Step 窗口行** | ❌ 后端未解析 | `Step {N} rollout|train started|finished` 是新增行,`job_progress` 与流转分析都没有消费它。它是当前**唯一**承载"每 step rollout 起止 + 训练起止"的数据源(CSV 里已无对应锚点),页面若要展示 Step 生产-消费视图,需在后端新增解析:从 job-log(或 test-rl-trace `cat=progress`)按 `Step (\d+) (rollout|train) (started|finished)` 聚合出每步四个时间戳 | + +--- + +## 4. 差异汇总(给 xfuse-controller-py 的适配清单) + +| # | 数据 | 展示项 | 状态 | 需要的改动 | +|---|---|---|---|---| +| 1 | CSV | 样本列表/详情/漏斗骨架 | ✅ 可用 | — | +| 2 | CSV | stage/detail 图例 | ⚠️ 文案缺失 | 前端 `STAGE_DOCS`/`DETAIL_DOCS` 增补 v1 命名 | +| 3 | CSV | 漏斗"下发/入缓冲/筛选"列、部分阶段耗时 | ⚠️ 恒空 | horizon 专属里程碑,开源同步架构无对应事件,建议按 v1 里程碑重写或隐藏 | +| 4 | CSV | Step 生产-消费视图(trace 口径) | ❌ 恒空 | FORM_BATCH/cat=megatron 不存在;改用 progress 的 Step 窗口行(见 6) | +| 5 | JSONL | 轨迹浏览/下载 | ❌ 不可用 | 数据源从 sgl-metric 展开列改为 test-rl-trace `cat=trajectory` + `raw_message` 解析;`rank_id` 字段已移除 | +| 6 | progress | Step 窗口(每步 rollout/训练起止) | ❌ 未解析 | 后端新增 `Step {N} ...` 行解析,可作为 Step 视图新数据源 | +| 7 | progress | 生成吞吐 / 已完成 rollout | ✅ 可用 | 格式已对齐现有解析,无需改动 | diff --git a/.qoder/skills/verl-trace-instrument/trace_module.py b/.qoder/skills/verl-trace-instrument/trace_module.py new file mode 100644 index 00000000..a9b1ff9e --- /dev/null +++ b/.qoder/skills/verl-trace-instrument/trace_module.py @@ -0,0 +1,313 @@ +""" +Trace utility for tracking (uid, traj_uid) flow in RL training pipeline. +""" + +import atexit +import contextvars +import csv +import os +import queue +import socket +import threading +import time +from typing import Sequence, Union + +# --------------------------------------------------------------------------- +# Single switch +# --------------------------------------------------------------------------- +# VERL_TRACE_DIR enables file output (or call set_trace_dir() before the +# first trace event; set_trace_dir() takes priority over the env var). +# Under the trace dir, sample-flow / megatron CSVs go into a ``sample/`` +# subdirectory; trajectory JSONL goes into ``trajectory/``. +# Both stdout and file output for trace events are gated by trace_dir being +# configured via set_trace_dir() or the VERL_TRACE_DIR env var. +# --------------------------------------------------------------------------- + +_TRACE_ENV_VAR = "VERL_TRACE_DIR" + +_SHUTDOWN = object() + +_trace_queue = None +_writer_thread = None +_filepath = None +_initialized = False +_config_trace_dir = None +_lock = threading.Lock() + +# Column order of the sample-flow CSV rows. The header itself is NOT +# written to the file; consumers (e.g. the SLS collector csv_headers) supply it. +_CSV_HEADER = ["job_id", "cat", "sample_id", "sample_status", "timestamp", "content"] + + +def _writer_loop(filepath, q): + f = open(filepath, "a", newline="", encoding="utf-8") + writer = csv.writer(f) + try: + while True: + batch = [] + try: + item = q.get(timeout=5.0) + except queue.Empty: + continue + if item is _SHUTDOWN: + f.flush() + return + batch.append(item) + while not q.empty() and len(batch) < 64: + item = q.get_nowait() + if item is _SHUTDOWN: + for row in batch: + writer.writerow(row) + f.flush() + return + batch.append(item) + for row in batch: + writer.writerow(row) + # Flush after each batch so SLS Logtail always sees consistent + # bytes rather than stale I/O-buffer fragments. + f.flush() + finally: + f.close() + + +def set_trace_dir(path): + """Set trace dir from config. Must be called before the first trace event.""" + global _config_trace_dir + _config_trace_dir = path + + +def _resolve_trace_dir(): + """Return the base trace directory, or None if unset. + + Priority: set_trace_dir() > VERL_TRACE_DIR. + """ + return _config_trace_dir or os.environ.get(_TRACE_ENV_VAR) or None + + +def get_trace_dir(): + """Public accessor for the base trace directory (or None).""" + return _resolve_trace_dir() + + +def _ensure_initialized(): + """Lazy-initialise the per-process background writer. + + Returns True iff file output is enabled (VERL_TRACE_DIR is set). + Safe to call from any thread. + """ + global _trace_queue, _writer_thread, _filepath, _initialized + if _initialized: + return _trace_queue is not None + + with _lock: + if _initialized: + return _trace_queue is not None + + trace_dir = _resolve_trace_dir() + if not trace_dir: + # Stay uninitialized so a later set_trace_dir() call (or a + # VERL_TRACE_DIR set after import) can still enable file output. + return False + _initialized = True + + os.makedirs(trace_dir, exist_ok=True) + sample_dir = os.path.join(trace_dir, "sample") + os.makedirs(sample_dir, exist_ok=True) + + # Expose sample path so external tools can discover it. + try: + with open("/tmp/verl_trace_sample_path.txt", "w") as _f: + _f.write(sample_dir) + except OSError: + pass + + hostname = socket.gethostname() + pid = os.getpid() + _filepath = os.path.join(sample_dir, f"{hostname}-pid{pid}.csv") + + _trace_queue = queue.Queue(maxsize=10000) + _writer_thread = threading.Thread( + target=_writer_loop, + args=(_filepath, _trace_queue), + daemon=True, + name="flow-trace-writer", + ) + _writer_thread.start() + atexit.register(_shutdown) + return True + + +def _shutdown(): + global _trace_queue, _writer_thread + if _trace_queue is not None: + try: + _trace_queue.put(_SHUTDOWN, timeout=5) + except queue.Full: + return + if _writer_thread is not None: + _writer_thread.join(timeout=10) + + +def _enqueue(row): + try: + _trace_queue.put_nowait(row) + except queue.Full: + pass + + +# --------------------------------------------------------------------------- +# Sample-flow tracing +# --------------------------------------------------------------------------- + +# Task-local switch to silence sample-flow tracing for validation samples. +# Each agent-loop sample runs in its own asyncio.Task (which owns a copy of the +# contextvars context), so setting it at the task entry point silences every +# trace_flow call made deeper in that sample's call chain (LLM_GENERATE, +# ROLLOUT_REWARD, ...) without touching each call site. +_suppress_trace = contextvars.ContextVar("verl_suppress_sample_trace", default=False) + + +def suppress_sample_trace(flag: bool = True): + """Silence (or re-enable) trace_flow for the current asyncio task / thread. + + Used to skip validation-sample rollouts, which never enter the training + stages and would otherwise show up as forever-RUNNING samples. + """ + _suppress_trace.set(bool(flag)) + + +def trace_flow(uid: Union[str, Sequence[str]], + traj_uid: Union[str, Sequence[str]], + from_module: str, + to_module: str, + stage: str, + detail: str = "", + global_step: int = -1): + """ + 打印数据流转追踪信息。 + + ``uid`` 和 ``traj_uid`` 既可以是单个字符串，也可以是字符串序列（list / + numpy array 等）。当传入序列时，函数内部自动按元素逐条记录，调用侧无需 + 显式写 for 循环。若一个为序列、另一个为标量字符串，标量会被广播到与序列 + 等长。 + + Args: + uid: 样本唯一标识符，单个字符串或字符串序列 + traj_uid: trajectory 唯一标识符，单个字符串或字符串序列 + from_module: 当前模块名称 + to_module: 即将转入的模块名称 + stage: 流转阶段名称 + detail: 可选的额外信息 + global_step: 当前训练全局步数(-1 表示未知) + + Use example: + # single item + trace_flow(uid='xxx', traj_uid='0', ...) + + # batch – pass the whole array directly + trace_flow(uid=batch.non_tensor_batch['uid'], + traj_uid=batch.non_tensor_batch['traj_uid'], ...) + """ + if _suppress_trace.get(): + return + + log_enabled = _ensure_initialized() + + if not log_enabled: + return + + # ---- normalise to lists ---- + uids = [uid] if isinstance(uid, str) else uid + traj_uids = [traj_uid] if isinstance(traj_uid, str) else traj_uid + + n_uid, n_traj = len(uids), len(traj_uids) + if n_uid == 1 and n_traj > 1: + uids = list(uids) * n_traj + elif n_traj == 1 and n_uid > 1: + traj_uids = list(traj_uids) * n_uid + + # ---- fields shared across all items in this batch ---- + timestamp_iso = time.strftime("%Y-%m-%dT%H:%M:%S", time.localtime()) + timestamp_ms = int(time.time() * 1000) + detail_str = f" | detail={detail}" if detail else "" + step_str = f" | global_step={global_step}" if global_step != -1 else "" + job_id = os.environ.get("DLC_JOB_ID", "ENV_NOTFOUND_DLC_JOB_ID") + sample_status = "COMPLETE" if stage in {"TRAIN_UPDATE", "FILTER_TRAJECTORIES"} else "RUNNING" + + parts = [f"from={from_module}", f"to={to_module}", f"stage={stage}"] + if global_step != -1: + parts.append(f"global_step={global_step}") + if detail: + parts.append(f"detail={detail}") + content = "|".join(parts) + + for u, t in zip(uids, traj_uids): + print(f"[FLOW_TRACE] {timestamp_iso} | uid={u} | traj_uid={t} | " + f"from={from_module} | to={to_module} | stage={stage}{step_str}{detail_str}") + + sample_id = f"{u}/{t}" if str(t).isdigit() else f"{u}/sample" + _enqueue([job_id, "sample", sample_id, sample_status, timestamp_ms, content]) + + +# --------------------------------------------------------------------------- +# Training progress logging (step/stage level, stdout only) +# --------------------------------------------------------------------------- + +ENABLE_PROGRESS_LOG = True + +_PROGRESS_LOCK = threading.Lock() +_progress_file = None +_progress_hour_key = None + + +def _write_progress_file(line: str): + """Append a progress line to ``{trace_dir}/progress/{YYYYMMDD-HH}.log``. + + Files rotate hourly (name = year-month-day-hour, e.g. ``20260811-15.log``) + and output is gated by the trace dir, like the sample-flow CSVs. + Never raises. + """ + global _progress_file, _progress_hour_key + trace_dir = _resolve_trace_dir() + if not trace_dir: + return + hour_key = time.strftime("%Y%m%d-%H", time.localtime()) + with _PROGRESS_LOCK: + try: + if hour_key != _progress_hour_key: + if _progress_file is not None: + try: + _progress_file.close() + except OSError: + pass + _progress_file = None + progress_dir = os.path.join(trace_dir, "progress") + os.makedirs(progress_dir, exist_ok=True) + _progress_file = open( + os.path.join(progress_dir, f"{hour_key}.log"), "a", encoding="utf-8") + _progress_hour_key = hour_key + _progress_file.write(line + "\n") + _progress_file.flush() + except OSError: + pass + + +def log_progress(message: str): + """输出一行训练进度,内容与 verl-horizon 原始日志一致。 + + 行内容直接采用 horizon 的裸格式,例如: + ``disagg: start compute reward`` / ``disagg: end compute reward`` / + ``Trainer: try to load global batch@N, chunk[0/0] from replay_queue`` 等, + 前端训练阶段面板直接解析这些行。 + + stdout 无条件输出;设置了 VERL_TRACE_DIR 时同时逐行落盘到 + ``{trace_dir}/progress/{YYYYMMDD-HH}.log``(按小时轮转,文件名即年月日-时, + 如 ``20260811-15.log``)。 + + Args: + message: 完整的一行日志内容(不含换行)。 + """ + if not ENABLE_PROGRESS_LOG: + return + print(message, flush=True) + _write_progress_file(message) diff --git a/.qoder/skills/verl-trace-instrument/trajectory_points.md b/.qoder/skills/verl-trace-instrument/trajectory_points.md new file mode 100644 index 00000000..cae0d4b5 --- /dev/null +++ b/.qoder/skills/verl-trace-instrument/trajectory_points.md @@ -0,0 +1,293 @@ +# Trajectory 监控插桩点清单(对齐最新 verl-horizon) + +最新 verl-horizon 的 trajectory 机制:`Turn`/`Trajectory`/`TurnRecord` 数据模型 +(`trajectory_types.py`)+ 线程安全 `TrajectoryTracker`(`trajectory_tracker.py`, +register → append turns → build → unregister 生命周期)+ `TrajectoryStore` +(`trajectory_store.py`,JSONL 单线程异步写、100MB 轮转、裁剪 token_ids 大字段)。 +采集端在每次 LLM generate 调用包一个 `TurnRecord`;dump 端在 agent 完成后 +**中央点位**统一 build + write(horizon:`BaseAgent.safe_run` build → +`AgentWorker.run_with_scheduler_ctx` push 前 write)。 + +本清单把该机制落到 verl-latest 的 AgentLoop 架构: + +- `AgentLoopBase` ↔ horizon `BaseAgent`(持有 tracker) +- `AgentLoopWorker` ↔ horizon `AgentWorker`(持有 store、中央 dump) +- `server_manager.generate` ↔ horizon `run_generate`(TurnRecord 采集点) + +**前置条件**:trace 工具模块四个文件已就位(见 SKILL.md Step 2)。 +未设置 `VERL_TRACE_DIR` 时 `get_trace_dir()` 为 None,store 不创建、 +不 register、不 build,全链路静默跳过。 + +**零侵入规则同样适用**:只追加采集/dump 代码与 `trace_traj_key` kwarg 透传, +不改 output 构造与状态机逻辑。 + +--- + +## 记录 schema(horizon Trajectory dataclass + 序列化适配) + +每行一个 JSON。dataclass 保持 horizon 原样(兼容 from_dict),但 +`TrajectoryStore._to_compact_dict` 在序列化时做适配: +- **裁剪大字段**:`initial_prompt_token_ids`、`completion_token_ids`、`observation_token_ids`; +- **删除 verl-latest 架构不可获取的字段**:`replica_id`(agent loop 经 + LLMServerClient 访问推理,无 per-request replica 概念;`rank_id` 字段 + 已从 schema 移除)、 + Turn 级 `prompt_text`(采集路径不填,horizon 也为 None)、`sglang_rid` + (引擎请求 id 不经 TokenOutput 透出); +- **sglang_* 时序字段保留并可填值**:sglang(enable_metrics=True,verl 启动时 + 已强制开启)在每个响应的 meta_info 里返回 per-request 的 e2e_latency/ + queue_time/prefill_finished_ts/request_received_ts/prefill_launch_latency, + `async_sglang_server.py` 已将其透传到 `TokenOutput.extra_fields`,采集端用 + `TurnRecord.fill_sglang_response(meta=extra_fields)` 派生 sglang_time_ms/ + sglang_queue_time_ms/sglang_ttft_ms/sglang_prefill_ms/sglang_decode_time_ms + (vLLM 后端无此数据,保持 null); +- **派生字段**:`turn_total_time_ms` 未设置时取 completion_time_ms + + observation_time_ms(单轮即 completion 耗时); +- **条件字段保留**(单轮为 null 属正常,tool 多轮会填值):`observation_*`、 + `template_overhead_*`(turn_index ≥ 1 时计算)。 + +落盘后的行格式(SLS 摄入约定:顶层仅保留索引列,其余全部折叠进 +`raw_message` JSON 字符串;完整字段定义见 trajectory_types_module.py): + +```json +{ + "prompt_uid": "uuid4...", // 顶层索引列 1:uid,过采样组共享 + "sample_index": 1, // 顶层索引列 2:traj_uid 解析 + "raw_message": "{ // 其余全部字段的 JSON 字符串 + \"request_id\": \"1\", // traj_uid + \"n_samples\": 4, // rollout.n + \"global_step\": 1, + \"created_at_ts\": 0.0, \"completed_at_ts\": 0.0, \"total_time_ms\": 0.0, + \"turns\": [{ // 每个 Turn = 一次 generate(action)+ tool 响应(observation) + \"turn_index\": 0, \"prompt_token_count\": 104, + \"completion_text\": \"...\", \"completion_token_count\": 228, \"finish_reason\": \"stop\", + \"observation_text\": null, \"start_ts\": 0.0, \"end_ts\": 0.0, + \"completion_time_ms\": 0.0, \"turn_total_time_ms\": 0.0, + \"observation_start_ts\": null, \"observation_end_ts\": null, \"observation_time_ms\": null, + \"sglang_time_ms\": 11280.0, \"sglang_queue_time_ms\": 350.0, \"sglang_ttft_ms\": 300.0, + \"sglang_prefill_ms\": 50.0, \"sglang_decode_time_ms\": 10630.0, + \"template_overhead_tokens\": null, \"template_overhead_text\": null + }], + \"initial_prompt_text\": \"...\", // 首个 TurnRecord 的 prompt decode + \"agent_metrics\": {\"generate_sequences\": 0.0, \"tool_calls\": 0.0} + }" +} +``` + +--- + +## C0: tracker 与 store 挂载 + +**文件 1**: `verl/experimental/agent_loop/agent_loop.py` + +顶部 import: +```python +from verl.utils.tracing import ( + TrajectoryStore, + TrajectoryTracker, + get_trace_dir, + log_progress, + trace_flow, +) +``` + +**锚点 1**(`AgentLoopBase.__init__` 内): +```python + self.mm_processor_kwargs = self.data_config.get("mm_processor_kwargs", {}) +``` +锚点之后插入: +```python + # Per-agent-loop trajectory tracker (aligned with verl-horizon BaseAgent): + # each run() call registers/appends turn records under a unique traj_key. + self._traj_tracker = TrajectoryTracker() +``` + +**锚点 2**(`AgentLoopWorker.__init__` 末尾,`RolloutTraceConfig.init(...)` 块之后): +```python + # trajectory dump — always under VERL_TRACE_DIR/trajectory (aligned with + # verl-horizon AgentWorker): one JSONL record per completed trajectory. + self._trajectory_store = None + self._trajectory_instance_id = os.getpid() + _trace_dir = get_trace_dir() + if _trace_dir: + _dump_dir = os.path.join(_trace_dir, "trajectory") + self._trajectory_store = TrajectoryStore(_dump_dir, self._trajectory_instance_id) + logger.info(f"TrajectoryStore initialized: dir={_dump_dir}, instance={self._trajectory_instance_id}") + # Expose trajectory path so external tools can discover it. + try: + with open("/tmp/verl_trace_trajectory_path.txt", "w") as _f: + _f.write(_dump_dir) + except OSError: + pass +``` + +--- + +## C1: 中央 register/build/dump(AgentLoopWorker._run_agent_loop) + +**文件**: `verl/experimental/agent_loop/agent_loop.py` + +**锚点**: +```python + output: AgentLoopOutput = await agent_loop.run(sampling_params, **kwargs) +``` + +run() 调用**前** register(镜像 horizon `BaseAgent.safe_run`): +```python + # trajectory recording — aligned with verl-horizon BaseAgent.safe_run: + # register before run(), dump after completion; skip when trace is off. + _traj_key = None + _traj_uid = None + _traj_created_at = None + _trace_item = None + if self._trajectory_store is not None: + _trace_item = SimpleNamespace( + non_tensor_batch=kwargs, meta_info={"global_step": trajectory["step"]} + ) + _traj_key, _traj_uid = agent_loop._traj_tracker.register_trajectory(_trace_item) + _traj_created_at = time.time() + kwargs["trace_traj_key"] = _traj_key +``` +(`SimpleNamespace` 适配 tracker 期望的 item 接口:`non_tensor_batch.get(uid/traj_uid)`、 +`meta_info.get(global_step)`;需顶部 `from types import SimpleNamespace`、`import time`。) + +锚点之后保留样本点 5 的 trace_flow,然后把原 `return await self._agent_loop_postprocess(...)` +改为:先 `kwargs.pop("trace_traj_key", None)`(该键仅供 run() 内部采集, +防止泄漏进 `_compute_score` 的 reward DataProto),再取结果、dump、最后返回; +且外层用 `try/except Exception: unregister_trajectory; raise` 包裹 run()/postprocess, +保证失败路径也注销(否则异常样本的 TurnRecord 永久泄漏在 `_records`): +```python + # trace_traj_key is only consumed inside run(); drop it before postprocess + kwargs.pop("trace_traj_key", None) + result = await self._agent_loop_postprocess(output, trajectory["validate"], **kwargs) + + # dump trajectory before returning — training side doesn't need it + if _traj_key is not None: + try: + # agent_metrics may be a pydantic AgentLoopMetrics; coerce to a + # JSON-serializable dict so TrajectoryStore.write never fails. + _agent_metrics = getattr(output, "metrics", None) + if hasattr(_agent_metrics, "model_dump"): + _agent_metrics = _agent_metrics.model_dump() + elif not isinstance(_agent_metrics, (dict, type(None))): + try: + _agent_metrics = dict(_agent_metrics) + except Exception: + _agent_metrics = str(_agent_metrics) + _traj = agent_loop._traj_tracker.build_trajectory( + _trace_item, _traj_key, _traj_uid, _traj_created_at, + tokenizer=self.tokenizer, + instance_id=self._trajectory_instance_id, + config=self.config, + agent_metrics=_agent_metrics, + ) + self._trajectory_store.write(_traj) + except Exception as _e: + logger.warning(f"TrajectoryStore dump failed: {_e}") + finally: + agent_loop._traj_tracker.unregister_trajectory(_traj_key) + return result +``` +注意:`run()` 与 `_agent_loop_postprocess()` 需包在 `try/except Exception` 中, +异常分支同样调用 `unregister_trajectory(_traj_key)` 后 `raise`——成功与失败路径都不得遗留注册项。 + +--- + +## C2: single_turn TurnRecord 采集 + +**文件**: `verl/experimental/agent_loop/single_turn_agent_loop.py` + +顶部 import:`from verl.utils.tracing import TurnRecord, trace_flow` + +**锚点**:样本点 6/7 的 `with simple_timer("generate_sequences", metrics):` 块。 +sending trace 之后、`with` 块之前创建记录;received trace 之前 fill + 落账: +```python + # trajectory recording: one TurnRecord per generate call (aligned with + # verl-horizon BaseAgent.run_generate); no-op when tracing is off. + _traj_key = kwargs.get('trace_traj_key') + _rec = TurnRecord(prompt_ids) if _traj_key else None + with simple_timer("generate_sequences", metrics): + output: TokenOutput = await self.server_manager.generate(...) + if _rec is not None: + try: + try: + _completion_text = self.tokenizer.decode(output.token_ids, skip_special_tokens=False) + except Exception: + _completion_text = "" + # fill_sglang_response sets completion fields and derives sglang_* + # timing (e2e/queue/ttft/prefill/decode) from TokenOutput.extra_fields + # when the sglang backend forwards them (enable_metrics=True). + _rec.fill_sglang_response( + text=_completion_text, + response_ids=list(output.token_ids), + finish_reason=str(output.stop_reason) if output.stop_reason else "", + meta=dict(output.extra_fields or {}), + ) + # meta carries no 'completion_tokens' key here; fall back to ids length. + if _rec.turn.completion_token_count is None: + _rec.turn.completion_token_count = len(output.token_ids) + finally: + _rec.stamp_end() + self._traj_tracker.append_turn_record(_traj_key, _rec) +``` +(`self._traj_tracker` 由 C0 挂在 AgentLoopBase,子类继承。) + +--- + +## C3: tool_agent_loop TurnRecord 采集 + observation 补丁 + +**文件**: `verl/experimental/agent_loop/tool_agent_loop.py` + +顶部 import:`from verl.utils.tracing import TurnRecord, trace_flow` + +1. **步骤 A**(run() 内,样本点 8 挂载处)追加一行: +```python + agent_data.trace_traj_key = kwargs.get('trace_traj_key') +``` +2. **_handle_generating_state**:与 C2 相同模式(`TurnRecord(agent_data.prompt_ids)` + + `fill_sglang_response` + completion_token_count 兜底),fill 后 + `agent_data.trace_last_record = _rec`(供 observation 补丁定位最后一个 turn)。 +3. **_handle_processing_tools_state**:锚点 +```python + previous_messages = list(agent_data.messages) +``` +之后记录 `_trace_pre_msg_len = len(agent_data.messages)`;锚点 +```python + agent_data.messages.extend(add_messages) +``` +之后补丁 observation(镜像 horizon `TurnRecord.patch_observation`): +```python + # trajectory recording: attach tool observations to the last turn record + _trace_rec = getattr(agent_data, 'trace_last_record', None) + if _trace_rec is not None: + try: + _trace_rec.patch_observation( + agent_data.messages, _trace_pre_msg_len, len(agent_data.messages), tokenizer=self.tokenizer + ) + except Exception: + pass +``` + +--- + +## 自定义 agent loop 的扩展方式 + +自定义 AgentLoopBase 子类自动继承 `_traj_tracker`;只需在每次 +`server_manager.generate` 调用处按 C2 模式包 `TurnRecord`(traj_key 取 +`kwargs.get('trace_traj_key')`),中央 dump(C1)无需改动。 + +--- + +## 完成标准 + +- C0-C3 均插入,相关文件 `python -m py_compile` 通过。 +- 设置 `VERL_TRACE_DIR` 跑通后,`$VERL_TRACE_DIR/trajectory/` 下出现 + `trajectory_{instance}_{ts}.jsonl`,每行一个合法 JSON,顶层仅 + prompt_uid/sample_index/raw_message 三个键,raw_message 为 JSON 字符串, + 解析后含 request_id/global_step/turns/initial_prompt_text, + 且不含 `initial_prompt_token_ids` 与 `rank_id`(已裁剪/移除);**抽查时 head 只取单个文件** + (多文件会输出 `==> 文件名 <==` 头行)。 +- 轨迹条数 ≈ TRAIN_BATCH_SIZE × rollout.n × steps(+ validation)。 +- 不设置 `VERL_TRACE_DIR` 时无 trajectory 目录、无报错(静默跳过)。 +- 若 JSONL 全为 0 字节:到 `/tmp/ray/session_latest/logs/worker-*.err` + grep `TrajectoryStore dump failed` 定位被吞异常(常见:metrics 未做 + pydantic → dict 强转)。 diff --git a/.qoder/skills/verl-trace-instrument/trajectory_store_module.py b/.qoder/skills/verl-trace-instrument/trajectory_store_module.py new file mode 100644 index 00000000..cd36f9e9 --- /dev/null +++ b/.qoder/skills/verl-trace-instrument/trajectory_store_module.py @@ -0,0 +1,136 @@ +"""Writes Trajectory objects to JSONL files with automatic rotation.""" + +import json +import logging +import os +import time +from concurrent.futures import ThreadPoolExecutor +from typing import Optional + +from verl.utils.tracing.trajectory_types import Trajectory + +logger = logging.getLogger(__name__) + + +class TrajectoryStore: + """Writes Trajectory objects to JSONL files with automatic rotation. + + Each trajectory is serialized as a compact JSON line. + When the current file exceeds ``max_file_size_mb`` (default 100 MB), a new + file is created automatically. File I/O is offloaded to a single-thread + executor so it never blocks the inference event loop. + """ + + _DEFAULT_MAX_MB = 100 + + def __init__(self, dump_dir: str, instance_id: int, + max_file_size_mb: int = _DEFAULT_MAX_MB): + self.dump_dir = dump_dir + self.instance_id = instance_id + self.max_file_size_bytes = max_file_size_mb * 1024 * 1024 + self._file: Optional[object] = None + self._file_path: Optional[str] = None + self._executor = ThreadPoolExecutor(max_workers=1) + os.makedirs(dump_dir, exist_ok=True) + self._open_new_file() + + def _open_new_file(self): + if self._file is not None: + try: + self._file.close() + except Exception: + pass + self._file = None + ts = f"{time.strftime('%Y%m%d_%H%M%S')}_{int(time.time() * 1000) % 1000:03d}" + self._file_path = os.path.join( + self.dump_dir, + f'trajectory_{self.instance_id}_{ts}.jsonl', + ) + self._file = open(self._file_path, 'a') + + def _check_rotation(self): + try: + if self._file_path and os.path.getsize(self._file_path) >= self.max_file_size_bytes: + self._open_new_file() + except OSError: + pass + + def _write_sync(self, line: str): + try: + self._check_rotation() + self._file.write(line + '\n') + self._file.flush() + except Exception as e: + logger.error(f'TrajectoryStore write failed: {e}') + + # Fields stripped during serialization to save storage space. + _TRAJECTORY_DROP_FIELDS = frozenset({'initial_prompt_token_ids'}) + _TURN_DROP_FIELDS = frozenset({'completion_token_ids', 'observation_token_ids'}) + + # Fields unavailable in the verl-latest architecture (kept in the dataclass + # for horizon-schema compatibility, dropped at serialization time): + # - replica_id: the agent loop talks to an LLMServerClient, there is no + # per-request engine-replica notion; + # - prompt_text: never populated on the collection path (large, horizon + # leaves it None as well). + # NOTE: the Turn sglang_* timing fields are NOT dropped: sglang + # (enable_metrics=True) returns per-request e2e_latency/queue_time/prefill + # timestamps in the response meta_info, which async_sglang_server.py + # forwards via TokenOutput.extra_fields and the agent loops fill into the + # TurnRecord. sglang_rid stays dropped: the engine's request id is not + # exposed through TokenOutput. + _TRAJECTORY_UNAVAILABLE_FIELDS = frozenset({'replica_id'}) + _TURN_UNAVAILABLE_FIELDS = frozenset({ + 'prompt_text', + 'sglang_rid', + }) + + def _to_compact_dict(self, trajectory: Trajectory) -> dict: + """Convert trajectory to dict, dropping large token-id fields and + fields unavailable in the verl-latest architecture, and deriving + turn_total_time_ms when not set by the collector.""" + d = trajectory.to_dict() + for key in self._TRAJECTORY_DROP_FIELDS | self._TRAJECTORY_UNAVAILABLE_FIELDS: + d.pop(key, None) + for turn in d.get('turns', []): + for key in self._TURN_DROP_FIELDS | self._TURN_UNAVAILABLE_FIELDS: + turn.pop(key, None) + # Derive total turn time when the collector did not set it + # (patch_observation only sets it for tool turns). + if turn.get('turn_total_time_ms') is None: + turn['turn_total_time_ms'] = (turn.get('completion_time_ms') or 0.0) + ( + turn.get('observation_time_ms') or 0.0 + ) + return d + + # Top-level keys kept outside ``raw_message``; everything else is folded + # into the raw_message JSON string (SLS ingestion convention: flat index + # columns + one raw payload column). + _TOP_LEVEL_KEYS = ('prompt_uid', 'sample_index') + + def write(self, trajectory: Trajectory): + """Asynchronously write a trajectory to the current JSONL file. + + Line format: ``{"prompt_uid": ..., "sample_index": ..., + "raw_message": ""}``. + """ + d = self._to_compact_dict(trajectory) + line = json.dumps({ + key: d.get(key) + for key in self._TOP_LEVEL_KEYS + } | { + 'raw_message': json.dumps( + {k: v for k, v in d.items() if k not in self._TOP_LEVEL_KEYS}, + ensure_ascii=False, + ), + }, ensure_ascii=False) + self._executor.submit(self._write_sync, line) + + def close(self): + self._executor.shutdown(wait=True) + if self._file is not None: + try: + self._file.close() + except Exception: + pass + self._file = None diff --git a/.qoder/skills/verl-trace-instrument/trajectory_tracker_module.py b/.qoder/skills/verl-trace-instrument/trajectory_tracker_module.py new file mode 100644 index 00000000..38b5776e --- /dev/null +++ b/.qoder/skills/verl-trace-instrument/trajectory_tracker_module.py @@ -0,0 +1,226 @@ +"""Thread-safe tracker for accumulating TurnRecords during multi-turn agent execution. + +Manages the lifecycle of trajectory recording: + register → append turns → build → unregister + +Extracted from BaseAgent to keep trajectory-tracking logic in the tracing module. +""" + +import logging +import os +import threading +import time +from typing import Dict, List, Optional + +import numpy as np + +from verl.utils.tracing.trajectory_types import Trajectory, TurnRecord + +logger = logging.getLogger(__name__) +logger.setLevel(os.getenv('VERL_PPO_LOGGING_LEVEL', 'WARN')) + + +class TrajectoryTracker: + """Thread-safe tracker for accumulating TurnRecords during multi-turn agent execution. + + Each agent instance owns one TrajectoryTracker. FastAPI may invoke + callbacks on different uvicorn worker threads, so all access to the + internal dict is guarded by a Lock. + """ + + def __init__(self): + self._records: Dict[str, List[TurnRecord]] = {} + self._lock = threading.Lock() + + # ------------------------------------------------------------------ + # Static utilities + # ------------------------------------------------------------------ + + @staticmethod + def make_traj_key(uid, sample_idx) -> str: + """Build a trajectory map key: ``'prompt-{uid}/sample-{sample_idx}'``.""" + return f'prompt-{uid}/sample-{sample_idx}' + + @staticmethod + def extract_str(val) -> Optional[str]: + """Extract a plain string from *val* (str, np.ndarray, or None).""" + if val is None: + return None + if isinstance(val, np.ndarray): + return str(val.item()) if val.size == 1 else str(val) + return str(val) + + @staticmethod + def parse_sample_index(traj_uid) -> Optional[int]: + """Parse sample_index from traj_uid. + + In repeat_n(), traj_uid is assigned as str(idx) where idx is 0..n-1. + We try to parse it as an int; if it is not a simple integer string + (e.g. a UUID for non-over-sampling paths), return None. + """ + if traj_uid is None: + return None + s = traj_uid + if isinstance(s, np.ndarray): + s = str(s.item()) if s.size == 1 else str(s) + else: + s = str(s) + try: + return int(s) + except (ValueError, TypeError): + return None + + @staticmethod + def extract_n_samples(config) -> Optional[int]: + """Extract the over-sampling factor *n* from config. + + Source: ``config.actor_rollout_ref.rollout.over_sampling.n`` + Falls back to ``config.actor_rollout_ref.rollout.n`` (group size). + Returns None if neither is available. + """ + try: + rollout_cfg = config.actor_rollout_ref.rollout + over_sampling_n = rollout_cfg.get('over_sampling', {}).get('n', None) + if over_sampling_n is not None: + return int(over_sampling_n) + group_n = rollout_cfg.get('n', None) + if group_n is not None: + return int(group_n) + except (AttributeError, KeyError, TypeError): + pass + return None + + # ------------------------------------------------------------------ + # Lifecycle: register → append → get → build → unregister + # ------------------------------------------------------------------ + + def register_trajectory(self, item) -> tuple: + """Create an empty turn-records list for *item* and return ``(traj_key, traj_uid)``. + + ``traj_uid`` alone is just the over-sampling index (0..n-1) and can + collide across different prompts processed concurrently, so we combine + it with uid to form a unique key. + """ + traj_uid_raw = item.non_tensor_batch.get('traj_uid', '') + traj_uid = self.extract_str(traj_uid_raw) + if not traj_uid: + traj_uid = str(id(item)) + logger.warning( + "Missing traj_uid for item uid=%s, falling back to id(). " + "meta_info_keys=%s non_tensor_keys=%s validate=%s global_step=%s", + item.non_tensor_batch.get('uid', '?'), + list(item.meta_info.keys()) if item.meta_info else [], + list(item.non_tensor_batch.keys()) if item.non_tensor_batch else [], + item.meta_info.get('validate', '?'), + item.meta_info.get('global_step', '?'), + ) + uid_raw = item.non_tensor_batch.get('uid', '') + uid = self.extract_str(uid_raw) or '' + traj_key = self.make_traj_key(uid, traj_uid) + with self._lock: + if traj_key in self._records: + logger.warning( + "register_trajectory: key %s is already registered; overwriting " + "%d in-flight turn record(s)", traj_key, len(self._records[traj_key])) + self._records[traj_key] = [] + return traj_key, traj_uid + + def get_turn_records(self, traj_key: str) -> List[TurnRecord]: + """Return the turn-records list for *traj_key*. Thread-safe.""" + with self._lock: + return self._records.get(traj_key, []) + + def append_turn_record(self, traj_key: str, record: TurnRecord): + """Append *record* to the trajectory identified by *traj_key*. + + Logs a warning if *traj_key* is None (caller forgot to pass it). + """ + if traj_key is None: + logger.warning('append_turn_record called without traj_key, record dropped') + return + with self._lock: + records = self._records.get(traj_key) + if records is not None: + records.append(record) + + def unregister_trajectory(self, traj_key: str): + """Remove the turn-records list for *traj_key*. Called at ``safe_run()`` end.""" + with self._lock: + self._records.pop(traj_key, None) + + # ------------------------------------------------------------------ + # Build final Trajectory from accumulated TurnRecords + # ------------------------------------------------------------------ + + def build_trajectory(self, item, traj_key, traj_uid, created_at, + tokenizer, instance_id, config, + agent_metrics=None) -> Trajectory: + """Build a :class:`Trajectory` from the accumulated turn records. + + Parameters + ---------- + item : DataProtoItem + The original data item (used for uid / global_step metadata). + traj_key : str + Key used to look up turn records in this tracker. + traj_uid : str + Raw trajectory uid, used as ``Trajectory.request_id``. + created_at : float + Unix timestamp when the trajectory started. + tokenizer : PreTrainedTokenizer + Tokenizer for decoding prompts / overhead text. + instance_id : int + The agent worker instance id (process pid; kept for file naming + and horizon-schema ``replica_id`` compatibility). + config + Config object for extracting ``n_samples``. + agent_metrics : dict, optional + Per-trajectory metrics collected by the agent. + """ + completed_at = time.time() + turn_records = self.get_turn_records(traj_key) + + for i, rec in enumerate(turn_records): + rec.turn.turn_index = i + + # Compute per-turn template overhead (static chat-template tokens between turns) + for i in range(1, len(turn_records)): + prev = turn_records[i - 1].turn + cur = turn_records[i].turn + if (cur.prompt_token_count is not None + and prev.prompt_token_count is not None + and prev.completion_token_count is not None): + obs_tokens = len(prev.observation_token_ids) if prev.observation_token_ids else 0 + overhead = (cur.prompt_token_count + - prev.prompt_token_count + - prev.completion_token_count + - obs_tokens) + cur.template_overhead_tokens = overhead + if overhead > 0: + offset = prev.prompt_token_count + prev.completion_token_count + cur.template_overhead_text = tokenizer.decode( + turn_records[i]._prompt_ids[offset:offset + overhead], + skip_special_tokens=False, + ) + + initial_prompt = None + initial_prompt_ids = None + if turn_records: + initial_prompt_ids = list(turn_records[0]._prompt_ids) + initial_prompt = tokenizer.decode(initial_prompt_ids, skip_special_tokens=False) + + return Trajectory( + request_id=str(traj_uid), + replica_id=instance_id, + global_step=item.meta_info.get('global_step', None), + created_at_ts=created_at, + completed_at_ts=completed_at, + total_time_ms=(completed_at - created_at) * 1000, + turns=[rec.turn for rec in turn_records], + prompt_uid=self.extract_str(item.non_tensor_batch.get('uid', None)), + sample_index=self.parse_sample_index(item.non_tensor_batch.get('traj_uid', None)), + n_samples=self.extract_n_samples(config), + initial_prompt_text=initial_prompt, + initial_prompt_token_ids=initial_prompt_ids, + agent_metrics=agent_metrics, + ) diff --git a/.qoder/skills/verl-trace-instrument/trajectory_types_module.py b/.qoder/skills/verl-trace-instrument/trajectory_types_module.py new file mode 100644 index 00000000..81f1f652 --- /dev/null +++ b/.qoder/skills/verl-trace-instrument/trajectory_types_module.py @@ -0,0 +1,321 @@ +"""Core data types for trajectory recording. + +Turn — one complete RL step (action + observation) +Trajectory — full multi-turn trajectory for one inference request +TurnRecord — mutable builder that wraps a Turn during recording + +## Multi-turn data flow +## +## Each Turn = one complete RL step = action + observation. +## Corresponds to the (state, action, observation) tuple in agent RL. +## +## messages = [system_prompt, user_task] +## +## -- Turn 0 -------------------------------------------------- +## [action] send prompt_ids -> sglang generate -> model_output_0 = "call tool X" +## [observation] agent executes tool X -> tool_result_0 +## -- Turn 0 end -- +## messages += [{assistant: model_output_0}, {tool: tool_result_0}] +## +## -- Turn 1 -------------------------------------------------- +## [action] send prompt_ids (with history) -> sglang generate -> model_output_1 = "call tool Y" +## [observation] agent executes tool Y -> tool_result_1 +## -- Turn 1 end -- +## messages += [{assistant: model_output_1}, {tool: tool_result_1}] +## +## -- Turn 2 -------------------------------------------------- +## [action] send prompt_ids (with all history) -> sglang generate -> "final answer" +## [observation] None (no tool call in the last turn) +## -- Turn 2 end -- + +## Turn timing hierarchy +## +## turn_total_time_ms +## ├── completion_time_ms (client-side /generate call: start_ts → end_ts) +## │ └── sglang_time_ms (server-side total) +## │ ├── sglang_queue_time_ms +## │ ├── sglang_ttft_ms (prefill) +## │ └── sglang_decode_time_ms +## └── observation_time_ms (tool execution: observation_start_ts → observation_end_ts) +""" + +import json +import time +from dataclasses import asdict, dataclass, field +from typing import List, Optional + + +@dataclass +class Turn: + """One Turn = one complete RL step = action + observation. + + Corresponds to the (state, action, observation, reward) tuple in agent RL: + - state: the full context sent to sglang (prompt_ids); length recorded by prompt_token_count + - action: model output (completion_text / completion_token_ids), obtained via run_generate() + - observation: environment feedback (tool execution result), recorded in observation_* fields + - reward: typically only present in the last Turn, recorded at the Trajectory level + + The "role" is implicit: the action part is always the assistant's model output, + and the observation part (when present) is always the tool/environment response. + No explicit role field is needed. + """ + + # -- identity --------------------------------------------------------------- + turn_index: int + # 0-based index of this turn within the Trajectory. + + # -- prompt (input to sglang) ----------------------------------------------- + prompt_token_count: Optional[int] = None + # Number of input tokens sent to sglang for this call. + + prompt_text: Optional[str] = None + # Decoded text of the prompt sent to sglang. Optional (large). + + # -- completion (model output) ---------------------------------------------- + completion_text: str = "" + # Text content of the model output (decoded text from sglang). + + completion_token_ids: List[int] = field(default_factory=list) + # Output/completion token IDs only (does NOT include prompt tokens). + + completion_token_count: Optional[int] = None + # Number of new output tokens generated by sglang. + + finish_reason: str = "" + # Why sglang stopped generating: "stop", "length", "tool_call", or error strings. + + # -- observation (tool / environment feedback) ------------------------------ + observation_text: Optional[str] = None + # Text result returned by the tool/environment. None when no tool call. + + observation_token_ids: Optional[List[int]] = None + # Token IDs of the tool result. None when no tool call. + + # -- timing (see tree in module docstring) ---------------------------------- + start_ts: float = 0.0 + # Client-side timestamp before the HTTP POST to sglang /generate. + + end_ts: float = 0.0 + # Client-side timestamp after the sglang /generate response is received. + + observation_start_ts: Optional[float] = None + # Timestamp when tool execution starts (= end_ts). + + observation_end_ts: Optional[float] = None + # Timestamp when tool execution completes. + + turn_total_time_ms: Optional[float] = None + # Total turn duration: completion_time_ms + observation_time_ms. + + completion_time_ms: float = 0.0 + # Client-side wall-clock of the sglang /generate call (start_ts → end_ts). + + observation_time_ms: Optional[float] = None + # Tool execution duration in milliseconds. + + # -- sglang server metrics -------------------------------------------------- + sglang_rid: str = "" + # sglang request ID (uuid4 hex) for cross-referencing with sglang logs. + + sglang_time_ms: Optional[float] = None + # sglang server-side end-to-end latency (milliseconds). + + sglang_queue_time_ms: Optional[float] = None + # Time spent waiting in sglang's scheduler queue (milliseconds). + + sglang_ttft_ms: Optional[float] = None + # Time To First Token (milliseconds): prefill phase duration. + + sglang_prefill_ms: Optional[float] = None + # Pure GPU prefill duration (milliseconds) from scheduler-side timestamps. + + sglang_decode_time_ms: Optional[float] = None + # Decode phase duration (milliseconds): sglang_time - queue - ttft. + + template_overhead_tokens: Optional[int] = None + template_overhead_text: Optional[str] = None + # Static chat-template tokens added between turns (e.g. <|im_start|>user\n…<|im_end|>\n<|im_start|>assistant\n\n). + # Computed as: prompt_token_count[N] - prompt_token_count[N-1] - completion_token_count[N-1] - observation_token_count[N-1] + # Only set for turn_index >= 1. Typically ~13 for Qwen-style chat templates. + + +@dataclass +class Trajectory: + """Complete multi-turn trajectory for one inference request. + + Represents the full lifecycle of a single rollout: from safe_run() start + through all sglang /generate calls and tool executions until completion. + + In over-sampling scenarios (n > 1), the same prompt produces N + Trajectory objects. Use prompt_uid / sample_index / n_samples to group them. + """ + + # -- identity --------------------------------------------------------------- + request_id: str + # Unique identifier for this trajectory (from traj_uid). + + prompt_uid: Optional[str] = None + # Stable identifier for the original prompt. Shared across over-samples. + + sample_index: Optional[int] = None + # Index of this trajectory among the N samples for the same prompt (0..N-1). + + n_samples: Optional[int] = None + # Total number of samples requested for this prompt. + + # -- infrastructure --------------------------------------------------------- + replica_id: int = -1 + # Which sglang engine replica served this request (0..num_replicas-1). + + global_step: Optional[int] = None + # Model version at dispatch time — the weights used to generate this trajectory. + # In disaggregated mode, the model may advance further by the time this + # trajectory completes and is consumed for training. + + # -- timing ----------------------------------------------------------------- + created_at_ts: float = 0.0 + # Unix timestamp when safe_run() begins processing. + + completed_at_ts: float = 0.0 + # Unix timestamp when safe_run() finishes processing. + + total_time_ms: float = 0.0 + # Total wall-clock duration in milliseconds. + + # -- data ------------------------------------------------------------------- + turns: List[Turn] = field(default_factory=list) + # Ordered list of Turn objects, one per sglang /generate call. + + initial_prompt_text: Optional[str] = None + initial_prompt_token_ids: Optional[List[int]] = None + + # -- agent-level metrics ---------------------------------------------------- + agent_metrics: Optional[dict] = None + # Per-trajectory metrics collected by the agent (e.g. tool call counts, + # token usage breakdowns). Passed through from result.meta_info['metrics']. + + def to_dict(self) -> dict: + return asdict(self) + + def to_json(self) -> str: + return json.dumps(self.to_dict(), ensure_ascii=False, separators=(',', ':')) + + @classmethod + def from_dict(cls, d: dict) -> 'Trajectory': + turns_data = d.pop('turns', []) + turns = [] + for t in turns_data: + turn_fields = Turn.__dataclass_fields__ + filtered = {k: v for k, v in t.items() if k in turn_fields} + turns.append(Turn(**filtered)) + traj_fields = {f for f in cls.__dataclass_fields__} + filtered_d = {k: v for k, v in d.items() if k in traj_fields} + return cls(turns=turns, **filtered_d) + + @classmethod + def from_json(cls, s: str) -> 'Trajectory': + return cls.from_dict(json.loads(s)) + + +class TurnRecord: + """Mutable builder that wraps a Turn dataclass. + + Created at the start of each run_generate() call. Methods mutate the + inner Turn as sglang response data and tool observations arrive. + Once the trajectory is complete, callers read .turn directly. + """ + + __slots__ = ('turn', '_prompt_ids') + + def __init__(self, prompt_ids, rid=''): + self._prompt_ids = prompt_ids + self.turn = Turn( + turn_index=-1, + completion_text='', completion_token_ids=[], finish_reason='', + start_ts=time.time(), end_ts=0.0, completion_time_ms=0.0, + sglang_rid=rid, + prompt_token_count=len(prompt_ids) if prompt_ids else 0, + ) + + def stamp_end(self): + """Set end_ts to now and compute completion_time_ms. Called in finally block.""" + self.turn.end_ts = time.time() + self.turn.completion_time_ms = (self.turn.end_ts - self.turn.start_ts) * 1000 + + def mark_skipped(self, finish_reason='length'): + """Mark this turn as skipped (e.g. zero tokens left) and stamp end.""" + self.turn.finish_reason = finish_reason + self.stamp_end() + + def fill_sglang_response(self, text, response_ids, finish_reason, meta): + """Fill in fields from a successful sglang /generate response. + + NOTE: queue_time, prefill_finished_ts, request_received_ts, and + prefill_launch_latency are only present in meta_info when the sglang + server is launched with --enable-metrics. Without it, sglang_ttft_ms, + sglang_queue_time_ms, sglang_prefill_ms, and sglang_decode_time_ms + will all be None. + """ + t = self.turn + t.completion_text = text + t.completion_token_ids = response_ids + t.finish_reason = finish_reason + t.completion_token_count = meta.get("completion_tokens") + e2e = meta.get("e2e_latency") + t.sglang_time_ms = e2e * 1000 if e2e is not None else None + prefill_ts = meta.get("prefill_finished_ts") + recv_ts = meta.get("request_received_ts") + t.sglang_ttft_ms = (prefill_ts - recv_ts) * 1000 if (prefill_ts is not None and recv_ts is not None) else None + queue_t = meta.get("queue_time") + t.sglang_queue_time_ms = queue_t * 1000 if queue_t is not None else None + # GPU-side prefill duration (independent of IPC overhead) + pll = meta.get("prefill_launch_latency") + t.sglang_prefill_ms = pll * 1000 if pll is not None else None + # Derive decode time = server total - queue - prefill + if t.sglang_time_ms is not None: + q = t.sglang_queue_time_ms or 0 + p = t.sglang_ttft_ms or 0 + t.sglang_decode_time_ms = max(0, t.sglang_time_ms - q - p) + + def patch_observation(self, messages, pre_msg_length, msg_len, tokenizer=None): + """Attach tool/environment observation from new messages since last turn. + + Between consecutive LLM calls, the agent framework executes tools and + appends observation messages. This extracts those observations and + records them (with timing and token counts) on this turn. + """ + if pre_msg_length >= msg_len: + return + new_messages = messages[pre_msg_length:] + obs_parts = [] + for m in new_messages: + if m.get('role') not in ('tool', 'function'): + continue + content = m.get('content', '') or '' + if isinstance(content, list): + content = '\n'.join( + c.get('text', '') if isinstance(c, dict) else str(c) + for c in content + ) + obs_parts.append(str(content)) + if not obs_parts: + return + t = self.turn + t.observation_text = '\n'.join(obs_parts) + if tokenizer is not None: + try: + t.observation_token_ids = tokenizer.encode( + t.observation_text, add_special_tokens=False + ) if t.observation_text else None + except Exception: + t.observation_token_ids = None + callback_ts = time.time() + t.observation_start_ts = t.end_ts if t.end_ts else None + t.observation_end_ts = callback_ts + t.observation_time_ms = ( + (callback_ts - t.observation_start_ts) * 1000 + if t.observation_start_ts is not None else None + ) + # Derive turn total + obs = t.observation_time_ms or 0 + t.turn_total_time_ms = t.completion_time_ms + obs diff --git a/scripts/coding_tool.py b/scripts/coding_tool.py new file mode 100644 index 00000000..be837a0e --- /dev/null +++ b/scripts/coding_tool.py @@ -0,0 +1,50 @@ +# Copyright 2026 +# +# 多轮 coding 冒烟测试用的本地代码解释器工具。 +# 通过 verl.tools.function_tool 的 @function_tool 装饰器注册,schema 由 +# 函数签名 + Google 风格 docstring 自动推断,经 +# actor_rollout_ref.rollout.multi_turn.function_tool_path 加载。 +# +# 设计:每次调用在独立子进程中执行模型生成的 Python 代码,带超时保护, +# 返回 stdout/stderr/退出码。完全本地,不依赖 sandbox-fusion 等外部服务。 + +import subprocess +import sys + +from verl.tools.function_tool import function_tool + +_MAX_OUTPUT_CHARS = 1200 +_TIMEOUT_SEC = 10 + + +@function_tool +def python_code_interpreter(code: str) -> str: + """Execute a python code snippet and return its stdout and stderr. + + Use this tool to run python code to help solve the given problem. + The code runs in a fresh subprocess; print the values you want to + observe. If the code fails, fix it and call this tool again. + + Args: + code: The python code to execute. Use print() to surface results. + + Returns: + Execution result as a string: stdout, stderr and exit code. + """ + try: + proc = subprocess.run( + [sys.executable, "-I", "-c", code], + capture_output=True, + text=True, + timeout=_TIMEOUT_SEC, + ) + parts = [] + if proc.stdout: + parts.append(f"stdout:\n{proc.stdout[:_MAX_OUTPUT_CHARS]}") + if proc.stderr: + parts.append(f"stderr:\n{proc.stderr[:_MAX_OUTPUT_CHARS]}") + if proc.returncode != 0: + parts.append(f"[exit code: {proc.returncode}]") + return "\n".join(parts) if parts else "[executed successfully, no output]" + except subprocess.TimeoutExpired: + return f"[error: code execution timed out after {_TIMEOUT_SEC}s]" diff --git a/scripts/prep_gsm8k_coding.py b/scripts/prep_gsm8k_coding.py new file mode 100644 index 00000000..4478f0a7 --- /dev/null +++ b/scripts/prep_gsm8k_coding.py @@ -0,0 +1,59 @@ +# Copyright 2026 +# +# 由已预处理好的 gsm8k parquet 派生 "coding 多轮" 版本数据集: +# 仅在 prompt 前注入一条 system 消息,引导模型用 python_code_interpreter +# 工具写代码解题并多轮修正;data_source/reward_model 保持不变, +# 复用 verl.utils.reward_score.gsm8k 的规则奖励(解析 `#### `)。 +# +# 用法: +# python3 scripts/prep_gsm8k_coding.py \ +# --src_dir ~/data/gsm8k --save_dir ~/data/gsm8k_coding + +import argparse +import os + +import pandas as pd + +SYSTEM_PROMPT = ( + "You are a careful problem solver. You have access to the `python_code_interpreter` " + "tool, which executes Python code and returns its stdout/stderr. " + "Solve the problem by writing Python code and calling the tool to run it. " + "Inspect the execution result: if there is an error or the result looks wrong, " + "fix the code and call the tool again. After you are confident about the result, " + "give the final numeric answer in the format `#### `." +) + + +def convert(df: pd.DataFrame) -> pd.DataFrame: + def inject_system(prompt): + msgs = list(prompt) + system_msg = {"role": "system", "content": SYSTEM_PROMPT} + if msgs and msgs[0].get("role") == "system": + msgs[0] = system_msg + else: + msgs = [system_msg] + msgs + return msgs + + df = df.copy() + df["prompt"] = df["prompt"].map(inject_system) + return df + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--src_dir", default=os.path.expanduser("~/data/gsm8k"), + help="目录,需包含 gsm8k.py 预处理产出的 train/test.parquet") + parser.add_argument("--save_dir", default=os.path.expanduser("~/data/gsm8k_coding"), + help="coding 版数据集输出目录") + args = parser.parse_args() + + os.makedirs(args.save_dir, exist_ok=True) + for split in ("train", "test"): + src = os.path.join(args.src_dir, f"{split}.parquet") + if not os.path.isfile(src): + raise FileNotFoundError(f"缺少源数据 {src},请先运行 examples/data_preprocess/gsm8k.py") + df = pd.read_parquet(src) + out = convert(df) + dst = os.path.join(args.save_dir, f"{split}.parquet") + out.to_parquet(dst) + print(f"[prep_gsm8k_coding] {split}: {len(out)} 条 -> {dst}") diff --git a/scripts/run_trace_smoke_test.sh b/scripts/run_trace_smoke_test.sh new file mode 100644 index 00000000..082697e5 --- /dev/null +++ b/scripts/run_trace_smoke_test.sh @@ -0,0 +1,241 @@ +#!/usr/bin/env bash +# Trace 冒烟测试:8xA100 训推一体(colocated)GRPO 小模型训练 +# +# 目的:验证对齐最新 verl-horizon 实现后的三类追踪输出: +# 1. 样本流转 trace(trace_flow,支持批量 uid/traj_uid + global_step) +# -> [FLOW_TRACE] stdout + CSV 落盘到 $VERL_TRACE_DIR/sample/(cat=sample) +# 2. 训练进度跟踪(训推一体/训推分离) +# -> 与 verl-horizon 一致的裸格式行(stdout + $VERL_TRACE_DIR/progress/{YYYYMMDD-HH}.log +# 按小时轮转):训推一体无前缀(如 'start compute reward'), +# 训推分离带 'disagg: ' 前缀(如 'disagg: start compute reward')、 +# 另有 'Trainer: successfully load global batch@N ...' fetch 行 +# 3. trajectory store(Turn/Trajectory/TurnRecord/TrajectoryTracker/ +# TrajectoryStore):agent loop 完成后把完整多轮轨迹以 JSONL dump 到 +# $VERL_TRACE_DIR/trajectory/(trace 写入格式:顶层 prompt_uid/sample_index/ +# raw_message;raw_message JSON 内含 horizon schema:request_id/turns[]/ +# initial_prompt_text 等,自动去除 token_ids 大字段) +# +# 覆盖路径(3 个 step 内全部命中): +# v1(默认):DATA_PREPROCESS(uid) -> DATA_DISPATCH -> AGENT_LOOP(session 注入 +# traj_uid) -> LLM_GENERATE -> ROLLOUT_REWARD -> TRAIN_SAMPLE -> TRAIN_REWARD* +# -> TRAIN_LOGPROB -> TRAIN_REF_LOGPROB -> TRAIN_ADVANTAGE -> TRAIN_UPDATE +# -> SYNC_PARAMS(*TRAIN_REWARD 仅 colocated reward 时出现) +# v0(USE_V1=False):DATA_PREPROCESS(uid/traj_uid) -> DATA_DISPATCH -> AGENT_LOOP +# -> LLM_GENERATE -> TRAIN_REWARD -> TRAIN_LOGPROB -> TRAIN_REF_LOGPROB +# -> TRAIN_ADVANTAGE -> TRAIN_UPDATE -> SYNC_PARAMS;step 2 触发 testing(P10) +# +# 注意:对齐最新 horizon 后,[FLOW_TRACE] stdout 与 CSV 均由 VERL_TRACE_DIR +# 统一门控(未设置时两者都静默);进度行 stdout 常开、设置后同步落盘。 +# 发现文件:/tmp/verl_trace_sample_path.txt、/tmp/verl_trace_trajectory_path.txt +# +# 用法: +# bash scripts/run_trace_smoke_test.sh +# 可选环境变量: +# MODEL_PATH 本地或 HF 模型路径(默认:本地 ModelScope 快照,否则在线下载 +# Qwen/Qwen2.5-0.5B-Instruct 到 MODELSCOPE_CACHE,否则 HF 直连) +# TRACE_DIR trace 输出目录(默认 <仓库根>/test/trace) +# INFER_BACKEND vllm 或 sglang(默认 sglang) +# SGLANG_ATTENTION_BACKEND sglang attention 后端(默认 triton:A100 不支持 fa3, +# 且本环境未装 flashinfer;设为 flashinfer 需自行安装) +# USE_V1 是否用 V1 trainer(默认 True:TaskRunnerV1 + TransferQueue, +# v1 sync 模式即本分支插桩覆盖的主路径;依赖外部包 transfer_queue, +# pip install TransferQueue;设 False 回退 v0 RayPPOTrainer 路径) +# DATALOADER_NUM_WORKERS 默认 0:训练结束时 Ray 清理阶段会 SIGKILL DataLoader +# 多进程 worker,触发 torch signal handler 报 +# "DataLoader worker is killed by signal: Killed" 并使脚本 +# 在自检前中断;0 为进程内加载,无子进程即无退出期报错 +# (GSM8K 数据量极小,不影响冒烟性能) +# MODELSCOPE_CACHE ModelScope 缓存目录(默认 $HOME/modelscope_cache,HF 不可达时用于取模型/数据) + +set -xeuo pipefail + +cd "$(dirname "$0")/.." + +# ModelScope 离线兑底:本环境无法直连 huggingface.co 时,优先使用本地 ModelScope 快照 +MODELSCOPE_CACHE=${MODELSCOPE_CACHE:-${HOME}/modelscope_cache} +MS_QWEN_DIR="${MODELSCOPE_CACHE}/models/Qwen--Qwen2.5-0.5B-Instruct/snapshots/master" +MS_GSM8K_DIR="${MODELSCOPE_CACHE}/datasets/AI-ModelScope--gsm8k/snapshots/master" + +########################### 冒烟测试配置 ########################### +# ---- 0. 模型准备(本地快照 -> modelscope 在线下载 -> HF 直连) ---- +if [ -n "${MODEL_PATH:-}" ]; then + : # 用户显式指定,保持不变 +elif [ -d "${MS_QWEN_DIR}" ]; then + MODEL_PATH="${MS_QWEN_DIR}" # 本地 ModelScope 快照 +elif python3 -c "import modelscope" 2>/dev/null; then + echo "[INFO] 本地无 Qwen2.5-0.5B-Instruct 快照,从 ModelScope 下载(约 1GB)..." + python3 -c " +from modelscope.hub.snapshot_download import snapshot_download +snapshot_download('Qwen/Qwen2.5-0.5B-Instruct', cache_dir='${MODELSCOPE_CACHE}')" + if [ ! -d "${MS_QWEN_DIR}" ]; then + echo "[FAIL] ModelScope 下载完成但未找到快照目录: ${MS_QWEN_DIR}" >&2 + exit 1 + fi + MODEL_PATH="${MS_QWEN_DIR}" +else + MODEL_PATH=Qwen/Qwen2.5-0.5B-Instruct +fi +INFER_BACKEND=${INFER_BACKEND:-sglang} +TRACE_DIR=${TRACE_DIR:-$(pwd)/test/trace} +NPROC_PER_NODE=${NPROC_PER_NODE:-8} +USE_V1=${USE_V1:-True} + +# v1 trainer(TaskRunnerV1)依赖外部包 TransferQueue,启动前预检 +if [ "${USE_V1}" = "True" ] || [ "${USE_V1}" = "true" ]; then + if ! python3 -c "import transfer_queue" 2>/dev/null; then + echo "[FAIL] trainer.use_v1=True 需要 transfer_queue 包,请先安装:pip install TransferQueue" >&2 + echo " 或设 USE_V1=False 回退 v0 RayPPOTrainer 路径" >&2 + exit 1 + fi +fi + +# 小 batch:64 prompt x n=4 = 256 条轨迹/step,0.5B 模型单 step 分钟级 +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-64} +ROLLOUT_N=${ROLLOUT_N:-4} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-512} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-512} +# 只跑 3 个 step;step 2 做一次验证(覆盖 eval 进度点) +TOTAL_STEPS=${TOTAL_STEPS:-3} +TEST_FREQ=${TEST_FREQ:-2} + +DATA_DIR=${DATA_DIR:-$HOME/data/gsm8k} +TRAIN_FILE=${TRAIN_FILE:-$DATA_DIR/train.parquet} +TEST_FILE=${TEST_FILE:-$DATA_DIR/test.parquet} +########################### 配置结束 ########################### + +# ---- trace 开关 ---- +export VERL_TRACE_DIR=${TRACE_DIR} +mkdir -p "${TRACE_DIR}" + +# 确保仓库根目录在 sys.path 中,使 examples 脚本与 python -m verl.* 可导入 verl +export PYTHONPATH="$(pwd)${PYTHONPATH:+:$PYTHONPATH}" + +# ---- 环境预检(新机器缺依赖时快速失败) ---- +python3 -c "import torch; assert torch.cuda.is_available(), 'CUDA unavailable'" \ + || { echo "[FAIL] torch/CUDA 检查失败,请确认 CUDA 环境" >&2; exit 1; } +for _mod in ray transformers datasets "${INFER_BACKEND}"; do + python3 -c "import ${_mod}" 2>/dev/null \ + || { echo "[FAIL] 缺少依赖: ${_mod},请先安装" >&2; exit 1; } +done +python3 -c "import verl" 2>/dev/null \ + || { echo "[FAIL] 无法 import verl,需在仓库根目录运行或已安装 verl" >&2; exit 1; } + +# ---- 1. 数据准备(GSM8K parquet,已存在则跳过) ---- +# huggingface.co 不可达时依次兑底:本地 ModelScope 快照 -> modelscope 在线下载 -> HF 直连 +if [ ! -f "${TRAIN_FILE}" ] || [ ! -f "${TEST_FILE}" ]; then + mkdir -p "${DATA_DIR}" + if [ ! -d "${MS_GSM8K_DIR}" ] && python3 -c "import modelscope" 2>/dev/null; then + python3 -c " +from modelscope.hub.snapshot_download import snapshot_download +snapshot_download('AI-ModelScope/gsm8k', repo_type='dataset', cache_dir='${MODELSCOPE_CACHE}')" + fi + if [ -d "${MS_GSM8K_DIR}" ]; then + python3 examples/data_preprocess/gsm8k.py \ + --local_dataset_path "${MS_GSM8K_DIR}" --local_save_dir "${DATA_DIR}" + else + python3 examples/data_preprocess/gsm8k.py --local_save_dir "${DATA_DIR}" + fi +fi + +# ---- 2. 启动训推一体 GRPO 训练 ---- +TRAIN_EXIT=0 +python3 -m verl.trainer.main_ppo \ + algorithm.adv_estimator=grpo \ + data.train_files="${TRAIN_FILE}" \ + data.val_files="${TEST_FILE}" \ + data.train_batch_size=${TRAIN_BATCH_SIZE} \ + data.max_prompt_length=${MAX_PROMPT_LENGTH} \ + data.max_response_length=${MAX_RESPONSE_LENGTH} \ + data.filter_overlong_prompts=True \ + data.truncation='error' \ + data.dataloader_num_workers=${DATALOADER_NUM_WORKERS:-0} \ + algorithm.use_kl_in_reward=False \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.actor.ppo_mini_batch_size=32 \ + actor_rollout_ref.actor.use_dynamic_bsz=True \ + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=8192 \ + actor_rollout_ref.actor.use_kl_loss=True \ + actor_rollout_ref.actor.kl_loss_coef=0.001 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.actor.fsdp_config.param_offload=False \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \ + actor_rollout_ref.rollout.name=${INFER_BACKEND} \ + +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend=${SGLANG_ATTENTION_BACKEND:-triton} \ + actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \ + actor_rollout_ref.rollout.n=${ROLLOUT_N} \ + actor_rollout_ref.rollout.free_cache_engine=True \ + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=8192 \ + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=8192 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + trainer.balance_batch=True \ + trainer.val_before_train=False \ + trainer.logger='["console"]' \ + trainer.project_name=verl_trace_smoke \ + trainer.experiment_name=qwen2_5_0_5b_grpo_smoke \ + trainer.n_gpus_per_node=${NPROC_PER_NODE} \ + trainer.nnodes=1 \ + trainer.save_freq=-1 \ + trainer.test_freq=${TEST_FREQ} \ + trainer.total_training_steps=${TOTAL_STEPS} \ + trainer.use_v1=${USE_V1} \ + "$@" || TRAIN_EXIT=$? + +# ---- 3. 结果自检 ---- +echo "================ TRACE SMOKE TEST RESULT ================" +echo "CSV 文件:" +ls -l "${TRACE_DIR}/sample/" || echo "[FAIL] 未发现 CSV,检查 VERL_TRACE_DIR 注入与插桩点 1-2" +echo "" +echo "样本流转行数(期望 > 0):" +wc -l "${TRACE_DIR}/sample/"*.csv 2>/dev/null || true +echo "" +echo "各 stage 覆盖情况(期望出现 DATA_PREPROCESS/DATA_DISPATCH/AGENT_LOOP/LLM_GENERATE/TRAIN_* ;v1 另有 ROLLOUT_REWARD/TRAIN_SAMPLE):" +cut -d',' -f6 "${TRACE_DIR}/sample/"*.csv 2>/dev/null | grep -o 'stage=[A-Z_]*' | grep -v '^stage=$' | sort | uniq -c || true +echo "" +echo "训练进度 log(期望 progress/ 下按小时轮转的 {YYYYMMDD-HH}.log,行数 > 0):" +ls -l "${TRACE_DIR}/progress/" 2>/dev/null || echo "[FAIL] 未发现 progress 目录,检查 VERL_TRACE_DIR 与 log_progress 插桩" +wc -l "${TRACE_DIR}/progress/"*.log 2>/dev/null || true +echo "" +echo "训练阶段覆盖(训推一体无前缀 'start/end ...';训推分离带 'disagg: ' 前缀):" +cat "${TRACE_DIR}/progress/"*.log 2>/dev/null | grep -oE '(disagg: )?(start|end) [a-z0-9& ]*' | sort | uniq -c || true +grep -c 'successfully load global batch' "${TRACE_DIR}/progress/"*.log 2>/dev/null || true +echo "" +echo "trajectory dump(JSONL,期望 > 0 个文件,文件名为 trajectory_{instance}_{ts}.jsonl):" +ls -l "${TRACE_DIR}/trajectory/" 2>/dev/null || echo "[FAIL] 未发现 trajectory 目录,检查 AgentLoopWorker 插桩与 VERL_TRACE_DIR 是否到达 AgentLoopWorker" +echo "" +echo "轨迹条数(期望 ≈ TRAIN_BATCH_SIZE x ROLLOUT_N x TOTAL_STEPS):" +wc -l "${TRACE_DIR}/trajectory/"*.jsonl 2>/dev/null | tail -1 || true +echo "" +echo "轨迹 schema 抽查(期望顶层 prompt_uid/sample_index/raw_message,raw_message 内含 request_id/turns,不含 token-id 字段):" +# 注意:head 传入多文件会输出 "==> 文件名 <==" 头行导致 JSON 解析失败,故只取第一个文件 +_first_traj=$(ls "${TRACE_DIR}/trajectory/"*.jsonl 2>/dev/null | head -1) +if [ -n "${_first_traj}" ]; then + head -1 "${_first_traj}" | python3 -c " +import json, sys +try: + top = json.loads(sys.stdin.readline()) + top_keys = ('prompt_uid', 'sample_index', 'raw_message') + missing = [k for k in top_keys if k not in top] + assert not missing, f'missing top-level keys: {missing}' + d = json.loads(top['raw_message']) + keys = ('request_id', 'turns', 'initial_prompt_text') + missing = [k for k in keys if k not in d] + assert not missing, f'missing raw_message keys: {missing}' + assert 'initial_prompt_token_ids' not in d, 'token-id fields should be trimmed' + print(f'schema OK: turns={len(d[\"turns\"])}, request_id={d[\"request_id\"]}, prompt_uid={top[\"prompt_uid\"]}') +except Exception as e: + print(f'[FAIL] trajectory schema 检查失败: {e}') +" || true +else + echo "[FAIL] 未发现任何 trajectory JSONL 文件" +fi +echo "==========================================================" +if [ "${TRAIN_EXIT}" -ne 0 ]; then + echo "[WARN] 训练进程退出码 ${TRAIN_EXIT}(若非预期,自检结果见上),按原退出码返回" + exit "${TRAIN_EXIT}" +fi diff --git a/scripts/run_trace_smoke_test_7b.sh b/scripts/run_trace_smoke_test_7b.sh new file mode 100755 index 00000000..8276aa83 --- /dev/null +++ b/scripts/run_trace_smoke_test_7b.sh @@ -0,0 +1,230 @@ +#!/usr/bin/env bash +# Trace 冒烟测试(大模型版):8xA100 训推一体(colocated)GRPO Qwen3.8-27B 训练 +# +# 与 run_trace_smoke_test.sh(0.5B)结构一致,差异仅在更大模型与配套规模参数 +# (rollout TP=2 + actor 参数/优化器 offload), +# 同样验证对齐最新 verl-horizon 实现后的三类追踪输出: +# 1. 样本流转 trace(trace_flow,支持批量 uid/traj_uid + global_step) +# -> [FLOW_TRACE] stdout + CSV 落盘到 $VERL_TRACE_DIR/sample/(cat=sample) +# 2. 训练进度跟踪(训推一体/训推分离) +# -> 与 verl-horizon 一致的裸格式行(stdout + $VERL_TRACE_DIR/progress/{YYYYMMDD-HH}.log +# 按小时轮转):训推一体无前缀(如 'start compute reward'), +# 训推分离带 'disagg: ' 前缀 +# 3. trajectory store:agent loop 完成后把完整多轮轨迹以 JSONL dump 到 +# $VERL_TRACE_DIR/trajectory/(trace 写入格式:顶层 prompt_uid/sample_index/ +# raw_message,raw_message JSON 内为 horizon schema,自动去除 token_ids 大字段) +# +# 覆盖路径(3 个 step 内全部命中): +# DATA_PREPROCESS(uid/traj_uid) -> DATA_DISPATCH -> AGENT_LOOP -> LLM_GENERATE +# -> TRAIN_REWARD -> TRAIN_LOGPROB -> TRAIN_REF_LOGPROB -> TRAIN_ADVANTAGE +# -> TRAIN_UPDATE -> SYNC_PARAMS;step 2 触发 testing(P10) +# +# 用法: +# bash scripts/run_trace_smoke_test_7b.sh +# 可选环境变量: +# MODEL_PATH 本地或 HF 模型路径(默认:本地 ModelScope 快照,否则在线下载 +# Qwen/Qwen3.8-27B 到 MODELSCOPE_CACHE,约 54GB) +# TRACE_DIR trace 输出目录(默认 <仓库根>/test/trace_7b) +# INFER_BACKEND vllm 或 sglang(默认 sglang) +# SGLANG_ATTENTION_BACKEND sglang attention 后端(默认 triton:A100 不支持 fa3) +# USE_V1 是否用 V1 trainer(默认 True:TaskRunnerV1 + TransferQueue, +# v1 sync 模式即本分支插桩覆盖的主路径,样本 trace/进度/ +# trajectory 全支持;设 False 回退 v0 RayPPOTrainer 路径) +# DATALOADER_NUM_WORKERS 默认 0:训练结束时 Ray 清理阶段会 SIGKILL DataLoader +# 多进程 worker,使脚本在自检前中断;0 为进程内加载无此问题 +# MODELSCOPE_CACHE ModelScope 缓存目录(默认 $HOME/modelscope_cache) + +set -xeuo pipefail + +cd "$(dirname "$0")/.." + +MODELSCOPE_CACHE=${MODELSCOPE_CACHE:-${HOME}/modelscope_cache} +MS_QWEN_DIR="${MODELSCOPE_CACHE}/models/Qwen--Qwen3.8-27B/snapshots/master" +MS_GSM8K_DIR="${MODELSCOPE_CACHE}/datasets/AI-ModelScope--gsm8k/snapshots/master" + +########################### 冒烟测试配置 ########################### +# ---- 0. 模型准备(本地快照 -> modelscope 在线下载 -> HF 直连) ---- +if [ -n "${MODEL_PATH:-}" ]; then + : # 用户显式指定,保持不变 +elif [ -d "${MS_QWEN_DIR}" ]; then + MODEL_PATH="${MS_QWEN_DIR}" # 本地 ModelScope 快照 +elif python3 -c "import modelscope" 2>/dev/null; then + echo "[INFO] 本地无 Qwen3.8-27B 快照,从 ModelScope 下载(约 54GB)..." + python3 -c " +from modelscope.hub.snapshot_download import snapshot_download +snapshot_download('Qwen/Qwen3.8-27B', cache_dir='${MODELSCOPE_CACHE}')" + if [ ! -d "${MS_QWEN_DIR}" ]; then + echo "[FAIL] ModelScope 下载完成但未找到快照目录: ${MS_QWEN_DIR}" >&2 + exit 1 + fi + MODEL_PATH="${MS_QWEN_DIR}" +else + MODEL_PATH=Qwen/Qwen3.8-27B +fi +INFER_BACKEND=${INFER_BACKEND:-sglang} +TRACE_DIR=${TRACE_DIR:-$(pwd)/test/trace_7b} +NPROC_PER_NODE=${NPROC_PER_NODE:-8} +USE_V1=${USE_V1:-True} + +# v1 trainer(TaskRunnerV1)依赖外部包 TransferQueue,启动前预检 +if [ "${USE_V1}" = "True" ] || [ "${USE_V1}" = "true" ]; then + if ! python3 -c "import transfer_queue" 2>/dev/null; then + echo "[FAIL] trainer.use_v1=True 需要 transfer_queue 包,请先安装:pip install TransferQueue" >&2 + echo " 或设 USE_V1=False 回退 v0 RayPPOTrainer 路径" >&2 + exit 1 + fi +fi + +# 中等 batch:64 prompt x n=4 = 256 条轨迹/step;7B 模型单 step 数分钟级 +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-64} +ROLLOUT_N=${ROLLOUT_N:-4} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-512} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-512} +# 只跑 3 个 step;step 2 做一次验证(覆盖 testing 进度点) +TOTAL_STEPS=${TOTAL_STEPS:-3} +TEST_FREQ=${TEST_FREQ:-2} + +DATA_DIR=${DATA_DIR:-$HOME/data/gsm8k} +TRAIN_FILE=${TRAIN_FILE:-$DATA_DIR/train.parquet} +TEST_FILE=${TEST_FILE:-$DATA_DIR/test.parquet} +########################### 配置结束 ########################### + +# ---- trace 开关 ---- +export VERL_TRACE_DIR=${TRACE_DIR} +mkdir -p "${TRACE_DIR}" + +# 确保仓库根目录在 sys.path 中,使 examples 脚本与 python -m verl.* 可导入 verl +export PYTHONPATH="$(pwd)${PYTHONPATH:+:$PYTHONPATH}" + +# ---- 环境预检(新机器缺依赖时快速失败) ---- +python3 -c "import torch; assert torch.cuda.is_available(), 'CUDA unavailable'" \ + || { echo "[FAIL] torch/CUDA 检查失败,请确认 CUDA 环境" >&2; exit 1; } +for _mod in ray transformers datasets "${INFER_BACKEND}"; do + python3 -c "import ${_mod}" 2>/dev/null \ + || { echo "[FAIL] 缺少依赖: ${_mod},请先安装" >&2; exit 1; } +done +python3 -c "import verl" 2>/dev/null \ + || { echo "[FAIL] 无法 import verl,需在仓库根目录运行或已安装 verl" >&2; exit 1; } + +# ---- 1. 数据准备(GSM8K parquet,已存在则跳过) ---- +if [ ! -f "${TRAIN_FILE}" ] || [ ! -f "${TEST_FILE}" ]; then + mkdir -p "${DATA_DIR}" + if [ ! -d "${MS_GSM8K_DIR}" ] && python3 -c "import modelscope" 2>/dev/null; then + python3 -c " +from modelscope.hub.snapshot_download import snapshot_download +snapshot_download('AI-ModelScope/gsm8k', repo_type='dataset', cache_dir='${MODELSCOPE_CACHE}')" + fi + if [ -d "${MS_GSM8K_DIR}" ]; then + python3 examples/data_preprocess/gsm8k.py \ + --local_dataset_path "${MS_GSM8K_DIR}" --local_save_dir "${DATA_DIR}" + else + python3 examples/data_preprocess/gsm8k.py --local_save_dir "${DATA_DIR}" + fi +fi + +# ---- 2. 启动训推一体 GRPO 训练 ---- +# 7B 显存要点(8xA100-80GB,FSDP 分 8 片): +# - actor 参数+梯度+AdamW 状态切片后每卡约 15GB,无需 param/optimizer offload +# - rollout 与训练 colocate:sglang 显存占比降到 0.35,free_cache_engine=True +# - ref 模型不训练,仅算 logprob,param_offload=True 省显存 +# - dynamic bsz 按 token 数打包,max_token_len_per_gpu 控制单卡峰值 +TRAIN_EXIT=0 +python3 -m verl.trainer.main_ppo \ + algorithm.adv_estimator=grpo \ + data.train_files="${TRAIN_FILE}" \ + data.val_files="${TEST_FILE}" \ + data.train_batch_size=${TRAIN_BATCH_SIZE} \ + data.max_prompt_length=${MAX_PROMPT_LENGTH} \ + data.max_response_length=${MAX_RESPONSE_LENGTH} \ + data.filter_overlong_prompts=True \ + data.truncation='error' \ + data.dataloader_num_workers=${DATALOADER_NUM_WORKERS:-0} \ + algorithm.use_kl_in_reward=False \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.actor.ppo_mini_batch_size=32 \ + actor_rollout_ref.actor.use_dynamic_bsz=True \ + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=16384 \ + actor_rollout_ref.actor.use_kl_loss=True \ + actor_rollout_ref.actor.kl_loss_coef=0.001 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.actor.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + actor_rollout_ref.rollout.name=${INFER_BACKEND} \ + +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend=${SGLANG_ATTENTION_BACKEND:-triton} \ + actor_rollout_ref.rollout.tensor_model_parallel_size=2 \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \ + actor_rollout_ref.rollout.n=${ROLLOUT_N} \ + actor_rollout_ref.rollout.free_cache_engine=True \ + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=16384 \ + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=16384 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + trainer.balance_batch=True \ + trainer.val_before_train=False \ + trainer.logger='["console"]' \ + trainer.project_name=verl_trace_smoke \ + trainer.experiment_name=qwen3_8_27b_grpo_smoke \ + trainer.n_gpus_per_node=${NPROC_PER_NODE} \ + trainer.nnodes=1 \ + trainer.save_freq=-1 \ + trainer.test_freq=${TEST_FREQ} \ + trainer.total_training_steps=${TOTAL_STEPS} \ + trainer.use_v1=${USE_V1} \ + "$@" || TRAIN_EXIT=$? + +# ---- 3. 结果自检 ---- +echo "================ TRACE SMOKE TEST RESULT (7B) ================" +echo "CSV 文件:" +ls -l "${TRACE_DIR}/sample/" || echo "[FAIL] 未发现 CSV,检查 VERL_TRACE_DIR 注入与插桩点 1-2" +echo "" +echo "样本流转行数(期望 > 0):" +wc -l "${TRACE_DIR}/sample/"*.csv 2>/dev/null || true +echo "" +echo "各 stage 覆盖情况(期望出现 DATA_PREPROCESS/DATA_DISPATCH/AGENT_LOOP/LLM_GENERATE/TRAIN_* ;v1 另有 ROLLOUT_REWARD/TRAIN_SAMPLE):" +cut -d',' -f6 "${TRACE_DIR}/sample/"*.csv 2>/dev/null | grep -o 'stage=[A-Z_]*' | grep -v '^stage=$' | sort | uniq -c || true +echo "" +echo "训练进度 log(期望 progress/ 下按小时轮转的 {YYYYMMDD-HH}.log,行数 > 0):" +ls -l "${TRACE_DIR}/progress/" 2>/dev/null || echo "[FAIL] 未发现 progress 目录" +wc -l "${TRACE_DIR}/progress/"*.log 2>/dev/null || true +echo "" +echo "训练阶段覆盖(训推一体无前缀 'start/end ...';训推分离带 'disagg: ' 前缀):" +cat "${TRACE_DIR}/progress/"*.log 2>/dev/null | grep -oE '(disagg: )?(start|end) [a-z0-9& ]*' | sort | uniq -c || true +echo "" +echo "trajectory dump(JSONL,期望 > 0 个文件,文件名为 trajectory_{instance}_{ts}.jsonl):" +ls -l "${TRACE_DIR}/trajectory/" 2>/dev/null || echo "[FAIL] 未发现 trajectory 目录,检查 AgentLoopWorker 插桩与 VERL_TRACE_DIR 是否到达 AgentLoopWorker" +echo "" +echo "轨迹条数(期望 ≈ TRAIN_BATCH_SIZE x ROLLOUT_N x TOTAL_STEPS = 768):" +wc -l "${TRACE_DIR}/trajectory/"*.jsonl 2>/dev/null | tail -1 || true +echo "" +echo "轨迹 schema 抽查(期望顶层 prompt_uid/sample_index/raw_message,raw_message 内含 request_id/turns,不含 token-id 字段):" +# 注意:head 传入多文件会输出 "==> 文件名 <==" 头行导致 JSON 解析失败,故只取第一个文件 +_first_traj=$(ls "${TRACE_DIR}/trajectory/"*.jsonl 2>/dev/null | head -1) +if [ -n "${_first_traj}" ]; then + head -1 "${_first_traj}" | python3 -c " +import json, sys +try: + top = json.loads(sys.stdin.readline()) + top_keys = ('prompt_uid', 'sample_index', 'raw_message') + missing = [k for k in top_keys if k not in top] + assert not missing, f'missing top-level keys: {missing}' + d = json.loads(top['raw_message']) + keys = ('request_id', 'turns', 'initial_prompt_text') + missing = [k for k in keys if k not in d] + assert not missing, f'missing raw_message keys: {missing}' + assert 'initial_prompt_token_ids' not in d, 'token-id fields should be trimmed' + print(f'schema OK: turns={len(d[\"turns\"])}, request_id={d[\"request_id\"]}, prompt_uid={top[\"prompt_uid\"]}') +except Exception as e: + print(f'[FAIL] trajectory schema 检查失败: {e}') +" || true +else + echo "[FAIL] 未发现任何 trajectory JSONL 文件" +fi +echo "==============================================================" +if [ "${TRAIN_EXIT}" -ne 0 ]; then + echo "[WARN] 训练进程退出码 ${TRAIN_EXIT}(若非预期,自检结果见上),按原退出码返回" + exit "${TRAIN_EXIT}" +fi diff --git a/scripts/run_trace_smoke_test_7b_multinode.sh b/scripts/run_trace_smoke_test_7b_multinode.sh new file mode 100755 index 00000000..5ce633ff --- /dev/null +++ b/scripts/run_trace_smoke_test_7b_multinode.sh @@ -0,0 +1,313 @@ +#!/usr/bin/env bash +# Trace 冒烟测试(大模型多机版):N 机 x 8xA100 训推一体(colocated)GRPO Qwen3.8-27B 训练 +# +# 与 run_trace_smoke_test_7b.sh(单机 8 卡)结构一致,差异仅在多机 Ray 组网: +# - 每台机器都用同一脚本启动,RANK=0 为主节点(ray head + 训练提交), +# 其余节点以 RANK>0 + MASTER_ADDR 指向主节点 IP 加入集群 +# - 节点数与节点序号优先读容器默认注入的 WORLD_SIZE / RANK 环境变量 +# - 训练验证的三类追踪输出与单机版相同: +# 1. 样本流转 trace(trace_flow)-> [FLOW_TRACE] stdout + CSV 落盘 $VERL_TRACE_DIR/sample/ +# 2. 训练进度跟踪 -> horizon 裸格式行(训推一体无前缀,如 'start compute reward'; +# stdout + 按小时落盘 $VERL_TRACE_DIR/progress/{YYYYMMDD-HH}.log) +# 3. trajectory store -> $VERL_TRACE_DIR/trajectory/ JSONL(schema 对齐 horizon) +# +# 多机注意事项: +# - MODEL_PATH / TRAIN_FILE / TEST_FILE / TRACE_DIR 建议放在各节点同路径的共享存储 +# (NFS 等)上;否则需预先在每台机器准备好相同的模型与数据文件 +# - Ray actor 可能被调度到任意节点:共享存储下自检可直接看 head 的 TRACE_DIR; +# 非共享存储时,各节点的 sample/trajectory/progress 输出分散在本机磁盘, +# 需自行从各 worker 节点汇总后再自检 +# - VERL_TRACE_DIR/PYTHONPATH 由 main_ppo 的 runtime_env 从 driver 侧转发; +# 如需各节点 daemon 也直接可见,可在 `ray start` 之前自行 export +# +# 用法(每台机器各执行一次;容器调度器通常会注入 RANK/WORLD_SIZE/MASTER_ADDR, +# 直接运行即可,无需显式传参): +# # 主节点(rank 0) +# RANK=0 WORLD_SIZE=2 bash scripts/run_trace_smoke_test_7b_multinode.sh +# # 其余节点(rank 从 1 开始) +# RANK=1 WORLD_SIZE=2 MASTER_ADDR=<主节点IP> bash scripts/run_trace_smoke_test_7b_multinode.sh +# 可选环境变量: +# MODEL_PATH 本地或 HF 模型路径(默认:本地 ModelScope 快照,否则在线下载 +# Qwen/Qwen3.8-27B 到 MODELSCOPE_CACHE,约 54GB) +# TRACE_DIR trace 输出目录(默认 <仓库根>/test/trace_7b_multinode) +# INFER_BACKEND vllm 或 sglang(默认 sglang) +# SGLANG_ATTENTION_BACKEND sglang attention 后端(默认 triton:A100 不支持 fa3) +# USE_V1 是否用 V1 trainer(默认 True:TaskRunnerV1 + TransferQueue, +# 同单机版说明;所有节点都需安装 transfer_queue 包) +# DATALOADER_NUM_WORKERS 默认 0(同单机版:避免训练结束 Ray 清理 SIGKILL 多进程 worker) +# MODELSCOPE_CACHE ModelScope 缓存目录(默认 $HOME/modelscope_cache) +# WORLD_SIZE 总节点数(默认 2,容器环境变量同名优先) +# RANK 当前节点序号(默认 0,即主节点) +# MASTER_ADDR 主节点 IP(RANK>0 时必填,默认 127.0.0.1) +# RAY_PORT Ray GCS 端口(默认 6379) +# GLOO_SOCKET_IFNAME / NCCL_SOCKET_IFNAME 跨节点控制面网卡(默认 bond0) +# NCCL_IB_GID_INDEX RoCE GID(默认 3=RoCEv2,本网络唯一可用);NCCL_IB_HCA 默认 mlx5_bond +# RAY_JOIN_TIMEOUT_SEC worker 等待加入集群的超时(默认 600) +# STOP_RAY_AFTER_DONE 训练结束后是否 ray stop 关闭整集群(默认 1) + +set -xeuo pipefail + +cd "$(dirname "$0")/.." + +MODELSCOPE_CACHE=${MODELSCOPE_CACHE:-${HOME}/modelscope_cache} +MS_QWEN_DIR="${MODELSCOPE_CACHE}/models/Qwen--Qwen3.8-27B/snapshots/master" +MS_GSM8K_DIR="${MODELSCOPE_CACHE}/datasets/AI-ModelScope--gsm8k/snapshots/master" + +########################### 多机组网配置 ########################### +NNODES=${WORLD_SIZE:-2} +NODE_RANK=${RANK:-0} +MASTER_ADDR=${MASTER_ADDR:-127.0.0.1} +RAY_PORT=${RAY_PORT:-6379} +RAY_JOIN_TIMEOUT_SEC=${RAY_JOIN_TIMEOUT_SEC:-600} +STOP_RAY_AFTER_DONE=${STOP_RAY_AFTER_DONE:-1} +if [ "${NODE_RANK}" -gt 0 ] && [ "${MASTER_ADDR}" = "127.0.0.1" ]; then + echo "[FAIL] RANK=${NODE_RANK} 的 worker 节点必须通过 MASTER_ADDR 指定主节点 IP" >&2 + exit 1 +fi +# verl Worker 的 RANK/WORLD_SIZE 由 WorkerGroup 按 actor 显式注入(见 +# single_controller/ray/base.py);此处 unset 防止脚本环境经 `ray start` 泄漏进 +# Ray daemon,被未显式覆盖的裸 actor(如 AgentLoopWorker)错误继承 +unset RANK WORLD_SIZE + +# ---- 0. 模型准备(本地快照 -> modelscope 在线下载 -> HF 直连) ---- +if [ -n "${MODEL_PATH:-}" ]; then + : # 用户显式指定,保持不变 +elif [ -d "${MS_QWEN_DIR}" ]; then + MODEL_PATH="${MS_QWEN_DIR}" # 本地 ModelScope 快照 +elif python3 -c "import modelscope" 2>/dev/null; then + echo "[INFO] 本地无 Qwen3.8-27B 快照,从 ModelScope 下载(约 54GB)..." + python3 -c " +from modelscope.hub.snapshot_download import snapshot_download +snapshot_download('Qwen/Qwen3.8-27B', cache_dir='${MODELSCOPE_CACHE}')" + if [ ! -d "${MS_QWEN_DIR}" ]; then + echo "[FAIL] ModelScope 下载完成但未找到快照目录: ${MS_QWEN_DIR}" >&2 + exit 1 + fi + MODEL_PATH="${MS_QWEN_DIR}" +else + MODEL_PATH=Qwen/Qwen3.8-27B +fi +INFER_BACKEND=${INFER_BACKEND:-sglang} +TRACE_DIR=${TRACE_DIR:-$(pwd)/test/trace_7b_multinode} +NPROC_PER_NODE=${NPROC_PER_NODE:-8} +USE_V1=${USE_V1:-True} + +# v1 trainer(TaskRunnerV1)依赖外部包 TransferQueue,启动前预检(所有节点) +if [ "${USE_V1}" = "True" ] || [ "${USE_V1}" = "true" ]; then + if ! python3 -c "import transfer_queue" 2>/dev/null; then + echo "[FAIL] trainer.use_v1=True 需要 transfer_queue 包,请先安装:pip install TransferQueue" >&2 + echo " 或设 USE_V1=False 回退 v0 RayPPOTrainer 路径" >&2 + exit 1 + fi +fi + +# 多机默认把 batch 按节点数放大,保持每卡负载与单机版一致; +# 需满足可被 NNODES*NPROC_PER_NODE 相关的 DP 切分整除 +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-$((64 * NNODES))} +ROLLOUT_N=${ROLLOUT_N:-4} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-512} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-512} +# 只跑 3 个 step;step 2 做一次验证(覆盖 eval 进度点) +TOTAL_STEPS=${TOTAL_STEPS:-3} +TEST_FREQ=${TEST_FREQ:-2} + +DATA_DIR=${DATA_DIR:-$HOME/data/gsm8k} +TRAIN_FILE=${TRAIN_FILE:-$DATA_DIR/train.parquet} +TEST_FILE=${TEST_FILE:-$DATA_DIR/test.parquet} +########################### 配置结束 ########################### + +# ---- trace 开关 ---- +export VERL_TRACE_DIR=${TRACE_DIR} +mkdir -p "${TRACE_DIR}" + +# 确保仓库根目录在 sys.path 中,使 examples 脚本与 python -m verl.* 可导入 verl +export PYTHONPATH="$(pwd)${PYTHONPATH:+:$PYTHONPATH}" + +# ---- 环境预检(新机器缺依赖时快速失败,每台节点都执行) ---- +python3 -c "import torch; assert torch.cuda.is_available(), 'CUDA unavailable'" \ + || { echo "[FAIL] torch/CUDA 检查失败,请确认 CUDA 环境" >&2; exit 1; } +for _mod in ray transformers datasets "${INFER_BACKEND}"; do + python3 -c "import ${_mod}" 2>/dev/null \ + || { echo "[FAIL] 缺少依赖: ${_mod},请先安装" >&2; exit 1; } +done +python3 -c "import verl" 2>/dev/null \ + || { echo "[FAIL] 无法 import verl,需在仓库根目录运行或已安装 verl" >&2; exit 1; } + +# ---- 1. 数据准备(GSM8K parquet) ---- +# 主节点负责生成;worker 节点仅校验可见(共享存储场景)。非共享存储时请先在 +# 每台机器上准备好相同的 TRAIN_FILE/TEST_FILE 再启动 +if [ "${NODE_RANK}" -eq 0 ]; then + if [ ! -f "${TRAIN_FILE}" ] || [ ! -f "${TEST_FILE}" ]; then + mkdir -p "${DATA_DIR}" + if [ ! -d "${MS_GSM8K_DIR}" ] && python3 -c "import modelscope" 2>/dev/null; then + python3 -c " +from modelscope.hub.snapshot_download import snapshot_download +snapshot_download('AI-ModelScope/gsm8k', repo_type='dataset', cache_dir='${MODELSCOPE_CACHE}')" + fi + if [ -d "${MS_GSM8K_DIR}" ]; then + python3 examples/data_preprocess/gsm8k.py \ + --local_dataset_path "${MS_GSM8K_DIR}" --local_save_dir "${DATA_DIR}" + else + python3 examples/data_preprocess/gsm8k.py --local_save_dir "${DATA_DIR}" + fi + fi +else + if [ ! -f "${TRAIN_FILE}" ] || [ ! -f "${TEST_FILE}" ]; then + echo "[FAIL] worker 节点不可见训练数据: ${TRAIN_FILE} / ${TEST_FILE}" >&2 + echo " 请确认数据在共享存储上,或预先在每台节点准备相同数据" >&2 + exit 1 + fi +fi + +# ---- 2. 组 Ray 集群 ---- +# 跨节点集合通信配置(本机实测结论): +# 1. 网络只支持 RoCEv2(GID 3);RoCEv1(GID 2)丢包(ib_write_bw syndrom 0x81) +# 2. 镜像自带的 HPC-X NCCL 插件(libnccl-net.so, IBext v10)GID 选择不对, +# 会让首次跨节点集合通信永久挂死,必须 NCCL_NET_PLUGIN=none 绕开; +# 之后 NCCL 用内置 IB 传输,RDMA(RoCEv2)依然生效,实测带宽正常 +# 3. 多网卡(bond0-4)机器需把控制面网卡固定到 bond0 +# 必须在 ray start 之前导出,由 daemon 继承后传递给所有 actor +export GLOO_SOCKET_IFNAME=${GLOO_SOCKET_IFNAME:-bond0} +export NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME:-bond0} +export NCCL_NET_PLUGIN=${NCCL_NET_PLUGIN:-none} +export NCCL_IB_GID_INDEX=${NCCL_IB_GID_INDEX:-3} +export NCCL_IB_HCA=${NCCL_IB_HCA:-mlx5_bond} +if [ "${NODE_RANK}" -eq 0 ]; then + ray start --head --port="${RAY_PORT}" --dashboard-host=0.0.0.0 --disable-usage-stats + + # 等待全部节点加入(alive 节点数 >= NNODES) + # 注意:Ray 2.56+ 已移除 `ray nodes` 命令,改用 State API 的 `ray list nodes` + echo "[INFO] 等待 ${NNODES} 个节点加入 Ray 集群(超时 ${RAY_JOIN_TIMEOUT_SEC}s)..." + _deadline=$(( $(date +%s) + RAY_JOIN_TIMEOUT_SEC )) + while true; do + _alive=$(ray list nodes --filter "state=ALIVE" 2>/dev/null | grep -c ' ALIVE ' || true) + if [ "${_alive}" -ge "${NNODES}" ]; then + echo "[INFO] 集群就绪: ${_alive}/${NNODES} 节点 ALIVE" + break + fi + if [ "$(date +%s)" -ge "${_deadline}" ]; then + echo "[FAIL] 等待节点加入超时,当前 ${_alive}/${NNODES}。请确认各 worker 已执行:" >&2 + echo " RANK= MASTER_ADDR=<本机IP> bash $0" >&2 + ray stop >/dev/null 2>&1 || true + exit 1 + fi + sleep 5 + done +else + ray start --address="${MASTER_ADDR}:${RAY_PORT}" --disable-usage-stats + echo "[INFO] RANK=${NODE_RANK} 已加入 Ray 集群 ${MASTER_ADDR}:${RAY_PORT}" + echo "[INFO] worker 节点无需运行训练,本脚本到此等待主节点结束;Ctrl+C 或主节点 ray stop 后退出" + # 阻塞等待:head 侧 ray stop 会终止整集群 daemon,届时本节点 ray 心跳断开 + while ray status >/dev/null 2>&1; do + sleep 30 + done + echo "[INFO] Ray 集群已断开,worker 节点退出" + exit 0 +fi + +# ---- 3. 启动训推一体 GRPO 训练(仅主节点) ---- +# 7B 显存要点与单机版一致(FSDP 按 NNODES*NPROC_PER_NODE 切片,单卡占用更低); +# rollout 与训练 colocate:sglang 显存占比 0.35,free_cache_engine=True +TRAIN_EXIT=0 +python3 -m verl.trainer.main_ppo \ + algorithm.adv_estimator=grpo \ + data.train_files="${TRAIN_FILE}" \ + data.val_files="${TEST_FILE}" \ + data.train_batch_size=${TRAIN_BATCH_SIZE} \ + data.max_prompt_length=${MAX_PROMPT_LENGTH} \ + data.max_response_length=${MAX_RESPONSE_LENGTH} \ + data.filter_overlong_prompts=True \ + data.truncation='error' \ + data.dataloader_num_workers=${DATALOADER_NUM_WORKERS:-0} \ + algorithm.use_kl_in_reward=False \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.actor.ppo_mini_batch_size=32 \ + actor_rollout_ref.actor.use_dynamic_bsz=True \ + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=16384 \ + actor_rollout_ref.actor.use_kl_loss=True \ + actor_rollout_ref.actor.kl_loss_coef=0.001 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.actor.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + actor_rollout_ref.rollout.name=${INFER_BACKEND} \ + +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend=${SGLANG_ATTENTION_BACKEND:-triton} \ + actor_rollout_ref.rollout.tensor_model_parallel_size=2 \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \ + actor_rollout_ref.rollout.n=${ROLLOUT_N} \ + actor_rollout_ref.rollout.free_cache_engine=True \ + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=16384 \ + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=16384 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + trainer.balance_batch=True \ + trainer.val_before_train=False \ + trainer.logger='["console"]' \ + trainer.project_name=verl_trace_smoke \ + trainer.experiment_name=qwen3_8_27b_grpo_smoke_multinode \ + trainer.n_gpus_per_node=${NPROC_PER_NODE} \ + trainer.nnodes=${NNODES} \ + trainer.save_freq=-1 \ + trainer.test_freq=${TEST_FREQ} \ + trainer.total_training_steps=${TOTAL_STEPS} \ + trainer.use_v1=${USE_V1} \ + "$@" || TRAIN_EXIT=$? + +# ---- 4. 结果自检(仅主节点;共享存储下覆盖全集群输出) ---- +echo "================ TRACE SMOKE TEST RESULT (7B multinode, ${NNODES} nodes) ================" +echo "CSV 文件:" +ls -l "${TRACE_DIR}/sample/" || echo "[FAIL] 未发现 CSV,检查 VERL_TRACE_DIR 注入与插桩点 1-2" +echo "" +echo "样本流转行数(期望 > 0):" +wc -l "${TRACE_DIR}/sample/"*.csv 2>/dev/null || true +echo "" +echo "各 stage 覆盖情况(期望出现 DATA_PREPROCESS/DATA_DISPATCH/AGENT_LOOP/LLM_GENERATE/TRAIN_* ;v1 另有 ROLLOUT_REWARD/TRAIN_SAMPLE):" +cut -d',' -f6 "${TRACE_DIR}/sample/"*.csv 2>/dev/null | grep -o 'stage=[A-Z_]*' | grep -v '^stage=$' | sort | uniq -c || true +echo "" +echo "训练进度 log(期望 progress/ 下按小时轮转的 {YYYYMMDD-HH}.log,行数 > 0):" +ls -l "${TRACE_DIR}/progress/" 2>/dev/null || echo "[FAIL] 未发现 progress 目录" +wc -l "${TRACE_DIR}/progress/"*.log 2>/dev/null || true +echo "" +echo "训练阶段覆盖(训推一体无前缀 'start/end ...';训推分离带 'disagg: ' 前缀):" +cat "${TRACE_DIR}/progress/"*.log 2>/dev/null | grep -oE '(disagg: )?(start|end) [a-z0-9& ]*' | sort | uniq -c || true +echo "" +echo "trajectory dump(JSONL,期望 > 0 个文件,文件名为 trajectory_{instance}_{ts}.jsonl):" +ls -l "${TRACE_DIR}/trajectory/" 2>/dev/null || echo "[FAIL] 未发现 trajectory 目录,检查 AgentLoopWorker 插桩与 VERL_TRACE_DIR 是否到达 AgentLoopWorker" +echo "" +echo "轨迹条数(期望 ≈ TRAIN_BATCH_SIZE x ROLLOUT_N x TOTAL_STEPS = $((TRAIN_BATCH_SIZE * ROLLOUT_N * TOTAL_STEPS))):" +wc -l "${TRACE_DIR}/trajectory/"*.jsonl 2>/dev/null | tail -1 || true +echo "" +echo "轨迹 schema 抽查(期望含 request_id/prompt_uid/turns,不含 initial_prompt_token_ids):" +# 注意:head 传入多文件会输出 "==> 文件名 <==" 头行导致 JSON 解析失败,故只取第一个文件 +_first_traj=$(ls "${TRACE_DIR}/trajectory/"*.jsonl 2>/dev/null | head -1) +if [ -n "${_first_traj}" ]; then + head -1 "${_first_traj}" | python3 -c " +import json, sys +try: + d = json.loads(sys.stdin.readline()) + keys = ('request_id', 'prompt_uid', 'sample_index', 'turns', 'initial_prompt_text') + missing = [k for k in keys if k not in d] + assert not missing, f'missing keys: {missing}' + assert 'initial_prompt_token_ids' not in d, 'token-id fields should be trimmed' + print(f'schema OK: turns={len(d[\"turns\"])}, request_id={d[\"request_id\"]}, prompt_uid={d[\"prompt_uid\"]}') +except Exception as e: + print(f'[FAIL] trajectory schema 检查失败: {e}') +" || true +else + echo "[FAIL] 未发现任何 trajectory JSONL 文件" +fi +echo "==============================================================" + +# ---- 5. 清理 Ray 集群 ---- +if [ "${STOP_RAY_AFTER_DONE}" -eq 1 ]; then + echo "[INFO] ray stop 关闭整集群(worker 节点 daemon 一并终止)" + ray stop >/dev/null 2>&1 || true +fi + +if [ "${TRAIN_EXIT}" -ne 0 ]; then + echo "[WARN] 训练进程退出码 ${TRAIN_EXIT}(若非预期,自检结果见上),按原退出码返回" + exit "${TRAIN_EXIT}" +fi diff --git a/scripts/run_trace_smoke_test_multiturn.sh b/scripts/run_trace_smoke_test_multiturn.sh new file mode 100755 index 00000000..43934a70 --- /dev/null +++ b/scripts/run_trace_smoke_test_multiturn.sh @@ -0,0 +1,265 @@ +#!/usr/bin/env bash +# Trace 冒烟测试(coding 多轮版):8xA100 训推一体 GRPO + tool_agent 多轮工具调用 +# +# 任务形态:模型用 python_code_interpreter 工具写 Python 代码解题(本地子进程 +# 沙箱,超时保护,无外部服务依赖),观察执行结果后多轮修正,最后输出 +# `#### `;奖励复用 gsm8k 规则奖励。数据集由本地 gsm8k parquet 派生 +# (注入 coding 风格 system prompt),全程离线可跑。 +# +# 与单轮冒烟脚本相比,本脚本额外验证: +# - tool_agent 多轮 agent loop(hermes 工具格式,Qwen 系列原生支持) +# - trajectory dump 的 turns > 1 多轮轨迹(assistant/tool 交替) +# 默认使用 Qwen3.8-27B + 最多 8 轮 assistant 交互,让模型真实地 +# 多次调用工具、观察执行结果并修正代码,产生更复杂的多轮轨迹。 +# 三类追踪输出同单轮版: +# 1. 样本流转 trace -> $VERL_TRACE_DIR/sample/(cat=sample) CSV +# 2. 训练进度:与 verl-horizon 一致的裸格式行(stdout + progress/ 按小时轮转 log, +# 训推一体无前缀,如 'start compute reward') +# 3. trajectory JSONL -> $VERL_TRACE_DIR/trajectory/ +# +# 用法: +# bash scripts/run_trace_smoke_test_multiturn.sh +# 可选环境变量: +# MODEL_PATH 本地或 HF 模型路径(默认:本地 ModelScope 快照,否则在线下载 +# Qwen/Qwen3.8-27B,否则 HF 直连) +# TRACE_DIR trace 输出目录(默认 <仓库根>/test/trace_multiturn) +# INFER_BACKEND vllm 或 sglang(默认 sglang) +# SGLANG_ATTENTION_BACKEND sglang attention 后端(默认 triton) +# USE_V1 是否用 V1 trainer(默认 True:TaskRunnerV1 + TransferQueue, +# 同单轮冒烟脚本;设 False 回退 v0 RayPPOTrainer 路径) +# MAX_ASSISTANT_TURNS 最大 assistant 轮数(默认 8) +# DATALOADER_NUM_WORKERS 默认 0(原因同单轮版) +# MODELSCOPE_CACHE ModelScope 缓存目录(默认 $HOME/modelscope_cache) + +set -xeuo pipefail + +cd "$(dirname "$0")/.." + +MODELSCOPE_CACHE=${MODELSCOPE_CACHE:-${HOME}/modelscope_cache} +MS_QWEN_DIR="${MODELSCOPE_CACHE}/models/Qwen--Qwen3.8-27B/snapshots/master" +MS_GSM8K_DIR="${MODELSCOPE_CACHE}/datasets/AI-ModelScope--gsm8k/snapshots/master" + +########################### 冒烟测试配置 ########################### +# ---- 0. 模型准备(本地快照 -> modelscope 在线下载 -> HF 直连) ---- +if [ -n "${MODEL_PATH:-}" ]; then + : # 用户显式指定,保持不变 +elif [ -d "${MS_QWEN_DIR}" ]; then + MODEL_PATH="${MS_QWEN_DIR}" # 本地 ModelScope 快照 +elif python3 -c "import modelscope" 2>/dev/null; then + echo "[INFO] 本地无 Qwen3.8-27B 快照,从 ModelScope 下载(约 54GB)..." + python3 -c " +from modelscope.hub.snapshot_download import snapshot_download +snapshot_download('Qwen/Qwen3.8-27B', cache_dir='${MODELSCOPE_CACHE}')" + if [ ! -d "${MS_QWEN_DIR}" ]; then + echo "[FAIL] ModelScope 下载完成但未找到快照目录: ${MS_QWEN_DIR}" >&2 + exit 1 + fi + MODEL_PATH="${MS_QWEN_DIR}" +else + MODEL_PATH=Qwen/Qwen3.8-27B +fi +INFER_BACKEND=${INFER_BACKEND:-sglang} +TRACE_DIR=${TRACE_DIR:-$(pwd)/test/trace_multiturn} +NPROC_PER_NODE=${NPROC_PER_NODE:-8} +USE_V1=${USE_V1:-True} + +# v1 trainer(TaskRunnerV1)依赖外部包 TransferQueue,启动前预检 +if [ "${USE_V1}" = "True" ] || [ "${USE_V1}" = "true" ]; then + if ! python3 -c "import transfer_queue" 2>/dev/null; then + echo "[FAIL] trainer.use_v1=True 需要 transfer_queue 包,请先安装:pip install TransferQueue" >&2 + echo " 或设 USE_V1=False 回退 v0 RayPPOTrainer 路径" >&2 + exit 1 + fi +fi + +# 小 batch:64 prompt x n=4 = 256 条轨迹/step +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-64} +ROLLOUT_N=${ROLLOUT_N:-4} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-512} +# 多轮:总响应预算 4096(所有轮共享),最多 8 轮 assistant 交互, +# 工具回复上限 1024(代码执行输出可能较长) +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-4096} +MAX_ASSISTANT_TURNS=${MAX_ASSISTANT_TURNS:-8} +MAX_TOOL_RESPONSE_LENGTH=${MAX_TOOL_RESPONSE_LENGTH:-1024} +# 只跑 3 个 step;step 2 做一次验证(覆盖 testing 进度点) +TOTAL_STEPS=${TOTAL_STEPS:-3} +TEST_FREQ=${TEST_FREQ:-2} + +BASE_DATA_DIR=${BASE_DATA_DIR:-$HOME/data/gsm8k} +DATA_DIR=${DATA_DIR:-$HOME/data/gsm8k_coding} +TRAIN_FILE=${TRAIN_FILE:-$DATA_DIR/train.parquet} +TEST_FILE=${TEST_FILE:-$DATA_DIR/test.parquet} +########################### 配置结束 ########################### + +# ---- trace 开关 ---- +export VERL_TRACE_DIR=${TRACE_DIR} +mkdir -p "${TRACE_DIR}" + +# 确保仓库根目录在 sys.path 中,使 examples 脚本与 python -m verl.* 可导入 verl +export PYTHONPATH="$(pwd)${PYTHONPATH:+:$PYTHONPATH}" + +# ---- 环境预检(新机器缺依赖时快速失败) ---- +python3 -c "import torch; assert torch.cuda.is_available(), 'CUDA unavailable'" \ + || { echo "[FAIL] torch/CUDA 检查失败,请确认 CUDA 环境" >&2; exit 1; } +for _mod in ray transformers datasets "${INFER_BACKEND}"; do + python3 -c "import ${_mod}" 2>/dev/null \ + || { echo "[FAIL] 缺少依赖: ${_mod},请先安装" >&2; exit 1; } +done +python3 -c "import verl" 2>/dev/null \ + || { echo "[FAIL] 无法 import verl,需在仓库根目录运行或已安装 verl" >&2; exit 1; } + +# ---- 1. 数据准备 ---- +# 1a. 基础 gsm8k parquet(已存在则跳过;兑底链同单轮版) +if [ ! -f "${BASE_DATA_DIR}/train.parquet" ] || [ ! -f "${BASE_DATA_DIR}/test.parquet" ]; then + mkdir -p "${BASE_DATA_DIR}" + if [ ! -d "${MS_GSM8K_DIR}" ] && python3 -c "import modelscope" 2>/dev/null; then + python3 -c " +from modelscope.hub.snapshot_download import snapshot_download +snapshot_download('AI-ModelScope/gsm8k', repo_type='dataset', cache_dir='${MODELSCOPE_CACHE}')" + fi + if [ -d "${MS_GSM8K_DIR}" ]; then + python3 examples/data_preprocess/gsm8k.py \ + --local_dataset_path "${MS_GSM8K_DIR}" --local_save_dir "${BASE_DATA_DIR}" + else + python3 examples/data_preprocess/gsm8k.py --local_save_dir "${BASE_DATA_DIR}" + fi +fi +# 1b. 派生 coding 多轮版(注入 system prompt,已存在则跳过) +if [ ! -f "${TRAIN_FILE}" ] || [ ! -f "${TEST_FILE}" ]; then + python3 scripts/prep_gsm8k_coding.py \ + --src_dir "${BASE_DATA_DIR}" --save_dir "${DATA_DIR}" +fi + +# ---- 2. 启动训推一体 GRPO 多轮工具调用训练 ---- +TRAIN_EXIT=0 +python3 -m verl.trainer.main_ppo \ + algorithm.adv_estimator=grpo \ + data.train_files="${TRAIN_FILE}" \ + data.val_files="${TEST_FILE}" \ + data.train_batch_size=${TRAIN_BATCH_SIZE} \ + data.max_prompt_length=${MAX_PROMPT_LENGTH} \ + data.max_response_length=${MAX_RESPONSE_LENGTH} \ + data.filter_overlong_prompts=True \ + data.truncation='error' \ + data.dataloader_num_workers=${DATALOADER_NUM_WORKERS:-0} \ + algorithm.use_kl_in_reward=False \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.actor.ppo_mini_batch_size=32 \ + actor_rollout_ref.actor.use_dynamic_bsz=True \ + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=24576 \ + actor_rollout_ref.actor.use_kl_loss=True \ + actor_rollout_ref.actor.kl_loss_coef=0.001 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.actor.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + actor_rollout_ref.rollout.name=${INFER_BACKEND} \ + +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend=${SGLANG_ATTENTION_BACKEND:-triton} \ + actor_rollout_ref.rollout.tensor_model_parallel_size=2 \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \ + actor_rollout_ref.rollout.n=${ROLLOUT_N} \ + actor_rollout_ref.rollout.free_cache_engine=True \ + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=24576 \ + actor_rollout_ref.rollout.agent.default_agent_loop=tool_agent \ + actor_rollout_ref.rollout.multi_turn.enable=True \ + actor_rollout_ref.rollout.multi_turn.format=hermes \ + actor_rollout_ref.rollout.multi_turn.function_tool_path="$(pwd)/scripts/coding_tool.py" \ + actor_rollout_ref.rollout.multi_turn.max_assistant_turns=${MAX_ASSISTANT_TURNS} \ + actor_rollout_ref.rollout.multi_turn.max_tool_response_length=${MAX_TOOL_RESPONSE_LENGTH} \ + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=24576 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + trainer.balance_batch=True \ + trainer.val_before_train=False \ + trainer.logger='["console"]' \ + trainer.project_name=verl_trace_smoke \ + trainer.experiment_name=qwen3_8_27b_grpo_multiturn_tool_smoke \ + trainer.n_gpus_per_node=${NPROC_PER_NODE} \ + trainer.nnodes=1 \ + trainer.save_freq=-1 \ + trainer.test_freq=${TEST_FREQ} \ + trainer.total_training_steps=${TOTAL_STEPS} \ + trainer.use_v1=${USE_V1} \ + "$@" || TRAIN_EXIT=$? + +# ---- 3. 结果自检 ---- +echo "============ TRACE SMOKE TEST RESULT (MULTITURN CODING) ============" +echo "CSV 文件:" +ls -l "${TRACE_DIR}/sample/" || echo "[FAIL] 未发现 CSV,检查 VERL_TRACE_DIR 注入与插桩点 1-2" +echo "" +echo "样本流转行数(期望 > 0):" +wc -l "${TRACE_DIR}/sample/"*.csv 2>/dev/null || true +echo "" +echo "各 stage 覆盖情况(期望出现 DATA_PREPROCESS/DATA_DISPATCH/AGENT_LOOP/LLM_GENERATE/TRAIN_* ;v1 另有 ROLLOUT_REWARD/TRAIN_SAMPLE):" +cut -d',' -f6 "${TRACE_DIR}/sample/"*.csv 2>/dev/null | grep -o 'stage=[A-Z_]*' | grep -v '^stage=$' | sort | uniq -c || true +echo "" +echo "训练进度 log(期望 progress/ 下按小时轮转的 {YYYYMMDD-HH}.log,行数 > 0):" +ls -l "${TRACE_DIR}/progress/" 2>/dev/null || echo "[FAIL] 未发现 progress 目录" +wc -l "${TRACE_DIR}/progress/"*.log 2>/dev/null || true +echo "" +echo "训练阶段覆盖(训推一体无前缀 'start/end ...';训推分离带 'disagg: ' 前缀):" +cat "${TRACE_DIR}/progress/"*.log 2>/dev/null | grep -oE '(disagg: )?(start|end) [a-z0-9& ]*' | sort | uniq -c || true +echo "" +echo "trajectory dump(JSONL):" +ls -l "${TRACE_DIR}/trajectory/" 2>/dev/null || echo "[FAIL] 未发现 trajectory 目录" +echo "" +echo "轨迹条数(期望 ≈ TRAIN_BATCH_SIZE x ROLLOUT_N x TOTAL_STEPS):" +wc -l "${TRACE_DIR}/trajectory/"*.jsonl 2>/dev/null | tail -1 || true +echo "" +echo "多轮 turns 分布(本脚本关注点:期望出现 turns>1 的轨迹):" +python3 -c " +import glob, json +turns = [] +for f in sorted(glob.glob('${TRACE_DIR}/trajectory/*.jsonl')): + with open(f) as fp: + for line in fp: + try: + top = json.loads(line) + # 新格式:顶层 prompt_uid/sample_index/raw_message, + # turns 在 raw_message JSON 字符串内 + d = json.loads(top['raw_message']) if 'raw_message' in top else top + turns.append(len(d.get('turns') or [])) + except Exception: + pass +if not turns: + print('[FAIL] 未解析到任何轨迹') +else: + multi = sum(1 for t in turns if t > 1) + print(f'轨迹总数={len(turns)}, avg_turns={sum(turns)/len(turns):.2f}, ' + f'max_turns={max(turns)}, 多轮(>1)占比={multi/len(turns):.1%}') + if multi == 0: + print('[WARN] 没有多轮轨迹:模型可能未调用工具直接作答,' + '可调高温度/换更强模型,或检查 hermes tool 解析') +" || true +echo "" +echo "轨迹 schema 抽查(期望顶层 prompt_uid/sample_index/raw_message,raw_message 内含 request_id/turns,不含 token-id 字段):" +_first_traj=$(ls "${TRACE_DIR}/trajectory/"*.jsonl 2>/dev/null | head -1) +if [ -n "${_first_traj}" ]; then + head -1 "${_first_traj}" | python3 -c " +import json, sys +try: + top = json.loads(sys.stdin.readline()) + top_keys = ('prompt_uid', 'sample_index', 'raw_message') + missing = [k for k in top_keys if k not in top] + assert not missing, f'missing top-level keys: {missing}' + d = json.loads(top['raw_message']) + keys = ('request_id', 'turns', 'initial_prompt_text') + missing = [k for k in keys if k not in d] + assert not missing, f'missing raw_message keys: {missing}' + assert 'initial_prompt_token_ids' not in d, 'token-id fields should be trimmed' + assert 'rank_id' not in d, 'rank_id should be removed from schema' + print(f'schema OK: turns={len(d[\"turns\"])}, request_id={d[\"request_id\"]}, prompt_uid={top[\"prompt_uid\"]}') +except Exception as e: + print(f'[FAIL] trajectory schema 检查失败: {e}') +" || true +else + echo "[FAIL] 未发现任何 trajectory JSONL 文件" +fi +echo "======================================================================" +if [ "${TRAIN_EXIT}" -ne 0 ]; then + echo "[WARN] 训练进程退出码 ${TRAIN_EXIT}(若非预期,自检结果见上),按原退出码返回" + exit "${TRAIN_EXIT}" +fi diff --git a/verl/experimental/agent_loop/agent_loop.py b/verl/experimental/agent_loop/agent_loop.py index 4a1f8a94..59d462b1 100644 --- a/verl/experimental/agent_loop/agent_loop.py +++ b/verl/experimental/agent_loop/agent_loop.py @@ -31,7 +31,9 @@ import asyncio import logging import os import random +import time from abc import ABC, abstractmethod +from types import SimpleNamespace from typing import Any, Optional from uuid import uuid4 @@ -66,6 +68,14 @@ from verl.utils.tokenizer import ( ) from verl.utils.tokenizer.chat_template import apply_chat_template, initialize_system_prompt, initialize_turn_separator from verl.utils.tokenizer.continuous_token_wiring import create_continuous_token_builder +from verl.utils.tracing import ( + TrajectoryStore, + TrajectoryTracker, + get_trace_dir, + log_rollout_completed, + suppress_sample_trace, + trace_flow, +) from verl.workers.config import ( HFModelConfig, RolloutConfig, @@ -83,7 +93,7 @@ class AgentLoopMetrics(BaseModel): generate_sequences: float = 0.0 tool_calls: float = 0.0 - compute_score: float = 0.0 + reward_score: float = 0.0 num_preempted: int = -1 # -1 means not available @@ -235,6 +245,9 @@ class AgentLoopBase(ABC): self.data_config = data_config.config self.apply_chat_template_kwargs = self.data_config.get("apply_chat_template_kwargs", {}) self.mm_processor_kwargs = self.data_config.get("mm_processor_kwargs", {}) + # Per-agent-loop trajectory tracker (aligned with verl-horizon BaseAgent): + # each run() call registers/appends turn records under a unique traj_key. + self._traj_tracker = TrajectoryTracker() self.continuous_token_builder = None self.enable_continuous_token = False continuous_token_config = self.data_config.continuous_token @@ -569,6 +582,22 @@ class AgentLoopWorker: trace_config.get("max_samples_per_step_per_worker", None), ) + # trajectory dump — always under VERL_TRACE_DIR/trajectory (aligned with + # verl-horizon AgentWorker): one JSONL record per completed trajectory. + self._trajectory_store = None + self._trajectory_instance_id = os.getpid() + _trace_dir = get_trace_dir() + if _trace_dir: + _dump_dir = os.path.join(_trace_dir, "trajectory") + self._trajectory_store = TrajectoryStore(_dump_dir, self._trajectory_instance_id) + logger.info(f"TrajectoryStore initialized: dir={_dump_dir}, instance={self._trajectory_instance_id}") + # Expose trajectory path so external tools can discover it. + try: + with open("/tmp/verl_trace_trajectory_path.txt", "w") as _f: + _f.write(_dump_dir) + except OSError: + pass + def _get_mm_processor_kwargs(self, audio_data: Optional[list[Any]] = None) -> dict[str, Any]: """Return multimodal processor kwargs with audio sampling-rate defaults.""" mm_processor_kwargs = dict(self.mm_processor_kwargs or {}) @@ -657,6 +686,12 @@ class AgentLoopWorker: for i in range(len(batch)): trace_this_sample = i in traced_indices kwargs = {k: v[i] for k, v in batch.non_tensor_batch.items() if k != "__do_sample__"} + if not validate: + trace_flow(uid=kwargs.get('uid', 'unknown'), traj_uid=kwargs.get('traj_uid', 'N/A'), + from_module='AgentLoopWorker', + to_module=f"AgentLoop({kwargs.get('agent_name', 'single_turn')})", + stage='AGENT_LOOP', detail='agent_task_started', + global_step=batch.meta_info.get('global_steps', -1)) sample_sampling_params = dict(sampling_params) if not validate and per_sample_do_sample is not None and not bool(per_sample_do_sample[i]): apply_greedy_sampling_params(sample_sampling_params) @@ -681,6 +716,11 @@ class AgentLoopWorker: trace: bool = True, **kwargs, ) -> _InternalAgentLoopOutput: + # validation samples never enter the training stages; silence their + # sample-flow trace for this sample task's whole call chain + # (LLM_GENERATE / ROLLOUT_REWARD / ... — see suppress_sample_trace) + if trajectory.get("validate", False): + suppress_sample_trace() with rollout_trace_attr( step=trajectory["step"], sample_index=trajectory["sample_index"], @@ -704,8 +744,77 @@ class AgentLoopWorker: data_config=DictConfigWrap(self.config.data), tools=ToolListWrap(self.tools), ) - output: AgentLoopOutput = await agent_loop.run(sampling_params, **kwargs) - return await self._agent_loop_postprocess(output, trajectory["validate"], **kwargs) + # trajectory recording — aligned with verl-horizon BaseAgent.safe_run: + # register before run(), dump after completion; skip when trace is off + # and for validation samples (they are not training trajectories). + _traj_key = None + _traj_uid = None + _traj_created_at = None + _trace_item = None + if self._trajectory_store is not None and not trajectory.get("validate", False): + _trace_item = SimpleNamespace( + non_tensor_batch=kwargs, meta_info={"global_step": trajectory["step"]} + ) + _traj_key, _traj_uid = agent_loop._traj_tracker.register_trajectory(_trace_item) + _traj_created_at = time.time() + kwargs["trace_traj_key"] = _traj_key + + try: + output: AgentLoopOutput = await agent_loop.run(sampling_params, **kwargs) + # trace_traj_key is only consumed inside run() for per-turn record + # tracking; drop it before postprocess so it never leaks into + # _compute_score's kwargs (packed into the reward DataProto's + # non_tensor_batch and sent to RewardLoopWorker). + kwargs.pop("trace_traj_key", None) + trace_flow(uid=kwargs.get('uid', 'unknown'), traj_uid=kwargs.get('traj_uid', 'N/A'), + from_module=f'AgentLoop({agent_name})', to_module='AgentLoopWorker', + stage='AGENT_LOOP', detail='agent_run_completed', + global_step=trajectory.get('step', -1)) + result = await self._agent_loop_postprocess(output, trajectory["validate"], **kwargs) + except Exception: + # Release the tracker entry on failure too; otherwise the turn + # records of failed samples leak in _records forever. + if _traj_key is not None: + agent_loop._traj_tracker.unregister_trajectory(_traj_key) + raise + # per-trajectory rollout completion marker — the rl-tracing + # progress panel counts 'is completed, score=' + 'status=Finished' + # lines as finished rollouts (validation samples are skipped). + # Score comes from the AgentLoopOutput mutated in place by + # _compute_score: _agent_loop_postprocess return values differ by + # subclass (v0 returns the padded output, v1 TQ returns None). + if not trajectory.get("validate", False): + _score_src = output[-1] if isinstance(output, list) else output + log_rollout_completed(kwargs.get('uid', 'unknown'), + kwargs.get('traj_uid', 'N/A'), + getattr(_score_src, 'reward_score', None)) + + # dump trajectory before returning — training side doesn't need it + if _traj_key is not None: + try: + # agent_metrics may be a pydantic AgentLoopMetrics; coerce to a + # JSON-serializable dict so TrajectoryStore.write never fails. + _agent_metrics = getattr(output, "metrics", None) + if hasattr(_agent_metrics, "model_dump"): + _agent_metrics = _agent_metrics.model_dump() + elif not isinstance(_agent_metrics, (dict, type(None))): + try: + _agent_metrics = dict(_agent_metrics) + except Exception: + _agent_metrics = str(_agent_metrics) + _traj = agent_loop._traj_tracker.build_trajectory( + _trace_item, _traj_key, _traj_uid, _traj_created_at, + tokenizer=self.tokenizer, + instance_id=self._trajectory_instance_id, + config=self.config, + agent_metrics=_agent_metrics, + ) + self._trajectory_store.write(_traj) + except Exception as _e: + logger.warning(f"TrajectoryStore dump failed: {_e}") + finally: + agent_loop._traj_tracker.unregister_trajectory(_traj_key) + return result def _pad_token_ids( self, @@ -964,63 +1073,66 @@ class AgentLoopWorker: final_output = outputs[-1] if final_output.reward_score is None and enable_async_reward: - timing = {} - with simple_timer("compute_score", timing): - all_prompts, all_responses, all_input_ids, all_attention_mask, all_position_ids = [], [], [], [], [] - for output in outputs: - prompts = torch.tensor(output.prompt_ids, dtype=torch.int64) - responses = torch.tensor(output.response_ids, dtype=torch.int64) - input_ids = torch.cat([prompts, responses], dim=0) - attention_mask = torch.ones_like(input_ids, dtype=torch.int64) - multi_modal_inputs = self._compute_multi_modal_inputs(output, input_ids) - position_ids = self._compute_position_ids( - input_ids.unsqueeze(0), - attention_mask.unsqueeze(0), - multi_modal_inputs, - output.mm_processor_kwargs - if output.mm_processor_kwargs is not None - else self._get_mm_processor_kwargs( - output.multi_modal_data.get("audios") if output.multi_modal_data else None - ), - ).squeeze(0) - all_prompts.append(prompts) - all_responses.append(responses) - all_input_ids.append(input_ids) - all_attention_mask.append(attention_mask) - all_position_ids.append(position_ids) - - n = len(outputs) - batch = TensorDict( - { - "prompts": torch.nn.utils.rnn.pad_sequence(all_prompts, batch_first=True, padding_value=0), - "responses": torch.nn.utils.rnn.pad_sequence(all_responses, batch_first=True, padding_value=0), - "attention_mask": torch.nn.utils.rnn.pad_sequence( - all_attention_mask, batch_first=True, padding_value=0 - ), - "input_ids": torch.nn.utils.rnn.pad_sequence(all_input_ids, batch_first=True, padding_value=0), - "position_ids": torch.nn.utils.rnn.pad_sequence( - all_position_ids, batch_first=True, padding_value=0 - ), - }, - batch_size=n, - ) - non_tensor_batch = { - **{k: np.array([v] * n) for k, v in kwargs.items()}, - "__num_turns__": np.array([o.num_turns for o in outputs]), - "tool_extra_fields": np.array([o.extra_fields for o in outputs], dtype=object), - "prompt_len": np.array([len(o.prompt_ids) for o in outputs]), - "response_len": np.array([len(o.response_ids) for o in outputs]), - } - - data = DataProto( - batch=batch, - non_tensor_batch=non_tensor_batch, - ) - selected_reward_loop_worker_handle = random.choice(self.reward_loop_worker_handles) - result = await selected_reward_loop_worker_handle.compute_score.remote(data) - final_output.reward_score = result["reward_score"] - final_output.extra_fields["reward_extra_info"] = result["reward_extra_info"] - final_output.metrics.compute_score = timing["compute_score"] + all_prompts, all_responses, all_input_ids, all_attention_mask, all_position_ids = [], [], [], [], [] + for output in outputs: + prompts = torch.tensor(output.prompt_ids, dtype=torch.int64) + responses = torch.tensor(output.response_ids, dtype=torch.int64) + input_ids = torch.cat([prompts, responses], dim=0) + attention_mask = torch.ones_like(input_ids, dtype=torch.int64) + multi_modal_inputs = self._compute_multi_modal_inputs(output, input_ids) + position_ids = self._compute_position_ids( + input_ids.unsqueeze(0), + attention_mask.unsqueeze(0), + multi_modal_inputs, + output.mm_processor_kwargs + if output.mm_processor_kwargs is not None + else self._get_mm_processor_kwargs( + output.multi_modal_data.get("audios") if output.multi_modal_data else None + ), + ).squeeze(0) + all_prompts.append(prompts) + all_responses.append(responses) + all_input_ids.append(input_ids) + all_attention_mask.append(attention_mask) + all_position_ids.append(position_ids) + + n = len(outputs) + batch = TensorDict( + { + "prompts": torch.nn.utils.rnn.pad_sequence(all_prompts, batch_first=True, padding_value=0), + "responses": torch.nn.utils.rnn.pad_sequence(all_responses, batch_first=True, padding_value=0), + "attention_mask": torch.nn.utils.rnn.pad_sequence( + all_attention_mask, batch_first=True, padding_value=0 + ), + "input_ids": torch.nn.utils.rnn.pad_sequence(all_input_ids, batch_first=True, padding_value=0), + "position_ids": torch.nn.utils.rnn.pad_sequence( + all_position_ids, batch_first=True, padding_value=0 + ), + }, + batch_size=n, + ) + non_tensor_batch = { + **{k: np.array([v] * n) for k, v in kwargs.items()}, + "__num_turns__": np.array([o.num_turns for o in outputs]), + "tool_extra_fields": np.array([o.extra_fields for o in outputs], dtype=object), + "prompt_len": np.array([len(o.prompt_ids) for o in outputs]), + "response_len": np.array([len(o.response_ids) for o in outputs]), + } + + data = DataProto( + batch=batch, + non_tensor_batch=non_tensor_batch, + ) + selected_reward_loop_worker_handle = random.choice(self.reward_loop_worker_handles) + result = await selected_reward_loop_worker_handle.compute_score.remote(data) + final_output.reward_score = result["reward_score"] + final_output.extra_fields["reward_extra_info"] = result["reward_extra_info"] + trace_flow(uid=kwargs.get('uid', 'unknown'), traj_uid=kwargs.get('traj_uid', 'N/A'), + from_module='RewardLoopWorker', to_module='AgentLoopWorker', + stage='ROLLOUT_REWARD', + detail=f"reward_score_computed_score={result['reward_score']}") + # carry the trajectory's reward score into metrics (dumped as agent_metrics) + final_output.metrics.reward_score = final_output.reward_score async def _compute_teacher_logprobs( self, @@ -1238,6 +1350,13 @@ class AgentLoopManager: prompts.non_tensor_batch["priority"] = np.arange(len(prompts), dtype=np.int64) chunkes = prompts.chunk(len(self.agent_loop_workers)) + if not prompts.meta_info.get("validate", False): + for _worker_idx, _chunk in enumerate(chunkes): + _traj_uids = _chunk.non_tensor_batch.get("traj_uid", np.array(["N/A"] * len(_chunk))) + trace_flow(uid=_chunk.non_tensor_batch["uid"], traj_uid=_traj_uids, + from_module='AgentLoopManager', to_module=f'AgentLoopWorker_{_worker_idx}', + stage='DATA_DISPATCH', detail='dispatched_to_agent_loop_worker', + global_step=prompts.meta_info.get("global_steps", -1)) outputs = await asyncio.gather( *[ worker.generate_sequences.remote(chunk) @@ -1257,7 +1376,9 @@ class AgentLoopManager: timing = {} t_generate_sequences = np.array([metric["generate_sequences"] for chunk in metrics for metric in chunk]) t_tool_calls = np.array([metric["tool_calls"] for chunk in metrics for metric in chunk]) - t_compute_score = np.array([metric["compute_score"] for chunk in metrics for metric in chunk]) + reward_scores = np.array( + [metric["reward_score"] for chunk in metrics for metric in chunk], dtype=object + ) num_preempted = np.array([metric["num_preempted"] for chunk in metrics for metric in chunk]) timing["agent_loop/num_preempted/min"] = num_preempted.min() timing["agent_loop/num_preempted/max"] = num_preempted.max() @@ -1268,16 +1389,21 @@ class AgentLoopManager: timing["agent_loop/tool_calls/min"] = t_tool_calls.min() timing["agent_loop/tool_calls/max"] = t_tool_calls.max() timing["agent_loop/tool_calls/mean"] = t_tool_calls.mean() - timing["agent_loop/compute_score/min"] = t_compute_score.min() - timing["agent_loop/compute_score/max"] = t_compute_score.max() - timing["agent_loop/compute_score/mean"] = t_compute_score.mean() + valid_reward_scores = np.array([s for s in reward_scores if s is not None], dtype=float) + if valid_reward_scores.size > 0: + timing["agent_loop/reward_score/min"] = valid_reward_scores.min() + timing["agent_loop/reward_score/max"] = valid_reward_scores.max() + timing["agent_loop/reward_score/mean"] = valid_reward_scores.mean() # batch sequence generation is bounded by the slowest sample - slowest = np.argmax(t_generate_sequences + t_tool_calls + t_compute_score) + slowest = np.argmax(t_generate_sequences + t_tool_calls) prompt_length = output.batch["prompts"].shape[1] timing["agent_loop/slowest/generate_sequences"] = t_generate_sequences[slowest] timing["agent_loop/slowest/tool_calls"] = t_tool_calls[slowest] - timing["agent_loop/slowest/compute_score"] = t_compute_score[slowest] + # reward may not have been computed inside the agent loop (None); + # loggers like wandb reject None values, so skip the key in that case + if reward_scores[slowest] is not None: + timing["agent_loop/slowest/reward_score"] = reward_scores[slowest] timing["agent_loop/slowest/num_preempted"] = num_preempted[slowest] if "attention_mask" in output.batch: diff --git a/verl/experimental/agent_loop/single_turn_agent_loop.py b/verl/experimental/agent_loop/single_turn_agent_loop.py index dda94de5..c1d46020 100644 --- a/verl/experimental/agent_loop/single_turn_agent_loop.py +++ b/verl/experimental/agent_loop/single_turn_agent_loop.py @@ -19,6 +19,7 @@ from uuid import uuid4 from verl.experimental.agent_loop.agent_loop import AgentLoopBase, AgentLoopOutput, register from verl.utils.profiler import simple_timer from verl.utils.rollout_trace import rollout_trace_op +from verl.utils.tracing import TurnRecord, trace_flow from verl.workers.rollout.replica import TokenOutput logger = logging.getLogger(__file__) @@ -62,6 +63,13 @@ class SingleTurnAgentLoop(AgentLoopBase): # 3. generate sequences metrics = {} + trace_flow(uid=kwargs.get('uid', 'unknown'), traj_uid=kwargs.get('traj_uid', 'N/A'), + from_module='SingleTurnAgentLoop', to_module='RolloutReplica', + stage='LLM_GENERATE', detail='sending_generate_request') + # trajectory recording: one TurnRecord per generate call (aligned with + # verl-horizon BaseAgent.run_generate); no-op when tracing is off. + _traj_key = kwargs.get('trace_traj_key') + _rec = TurnRecord(prompt_ids) if _traj_key else None with simple_timer("generate_sequences", metrics): request_id = f"det-{priority}" if getattr(self.rollout_config, "full_determinism", False) else uuid4().hex output: TokenOutput = await self.server_manager.generate( @@ -74,6 +82,30 @@ class SingleTurnAgentLoop(AgentLoopBase): mm_processor_kwargs=mm_processor_kwargs, priority=priority, ) + if _rec is not None: + try: + try: + _completion_text = self.tokenizer.decode(output.token_ids, skip_special_tokens=False) + except Exception: + _completion_text = "" + # fill_sglang_response sets completion fields and derives sglang_* + # timing (e2e/queue/ttft/prefill/decode) from TokenOutput.extra_fields + # when the sglang backend forwards them (enable_metrics=True). + _rec.fill_sglang_response( + text=_completion_text, + response_ids=list(output.token_ids), + finish_reason=str(output.stop_reason) if output.stop_reason else "", + meta=dict(output.extra_fields or {}), + ) + # meta carries no 'completion_tokens' key here; fall back to ids length. + if _rec.turn.completion_token_count is None: + _rec.turn.completion_token_count = len(output.token_ids) + finally: + _rec.stamp_end() + self._traj_tracker.append_turn_record(_traj_key, _rec) + trace_flow(uid=kwargs.get('uid', 'unknown'), traj_uid=kwargs.get('traj_uid', 'N/A'), + from_module='RolloutReplica', to_module='SingleTurnAgentLoop', + stage='LLM_GENERATE', detail='received_generate_response') if metrics.get("num_preempted") is None: metrics["num_preempted"] = output.num_preempted if output.num_preempted is not None else -1 diff --git a/verl/experimental/agent_loop/tool_agent_loop.py b/verl/experimental/agent_loop/tool_agent_loop.py index 098db2b9..b6ed7247 100644 --- a/verl/experimental/agent_loop/tool_agent_loop.py +++ b/verl/experimental/agent_loop/tool_agent_loop.py @@ -34,6 +34,7 @@ from verl.tools.function_tool import FunctionTool, normalize_function_tool_retur from verl.tools.schemas import OpenAIFunctionCallSchema, OpenAIFunctionParsedSchema, ToolResponse from verl.utils.profiler import simple_timer from verl.utils.rollout_trace import rollout_trace_op +from verl.utils.tracing import TurnRecord, trace_flow from verl.workers.rollout.replica import TokenOutput logger = logging.getLogger(__file__) @@ -160,6 +161,10 @@ class ToolAgentLoop(AgentLoopBase): agent_data._active_tools = self.tools agent_data._active_tool_schemas = self.tool_schemas + agent_data.trace_uid = kwargs.get('uid', 'unknown') + agent_data.trace_traj_uid = kwargs.get('traj_uid', 'N/A') + agent_data.trace_traj_key = kwargs.get('trace_traj_key') + # State machine loop state = AgentState.PENDING while state != AgentState.TERMINATED: @@ -231,6 +236,15 @@ class ToolAgentLoop(AgentLoopBase): stop_token_ids = list(set((sampling_params.get("stop_token_ids") or []) + self.tool_parser.stop_token_ids)) sampling_params = {**sampling_params, "stop_token_ids": stop_token_ids} + # trajectory recording: one TurnRecord per generate call (aligned with + # verl-horizon BaseAgent.run_generate); no-op when tracing is off. + _traj_key = getattr(agent_data, 'trace_traj_key', None) + _rec = TurnRecord(agent_data.prompt_ids) if _traj_key else None + + trace_flow(uid=getattr(agent_data, 'trace_uid', 'unknown'), + traj_uid=getattr(agent_data, 'trace_traj_uid', 'N/A'), + from_module='ToolAgentLoop', to_module='RolloutReplica', + stage='LLM_GENERATE', detail='sending_generate_request') with simple_timer("generate_sequences", agent_data.metrics): output: TokenOutput = await self.server_manager.generate( request_id=agent_data.request_id, @@ -241,6 +255,32 @@ class ToolAgentLoop(AgentLoopBase): audio_data=agent_data.audio_data, mm_processor_kwargs=agent_data.mm_processor_kwargs, ) + trace_flow(uid=getattr(agent_data, 'trace_uid', 'unknown'), + traj_uid=getattr(agent_data, 'trace_traj_uid', 'N/A'), + from_module='RolloutReplica', to_module='ToolAgentLoop', + stage='LLM_GENERATE', detail='received_generate_response') + if _rec is not None: + try: + try: + _completion_text = self.tokenizer.decode(output.token_ids, skip_special_tokens=False) + except Exception: + _completion_text = "" + # fill_sglang_response sets completion fields and derives sglang_* + # timing (e2e/queue/ttft/prefill/decode) from TokenOutput.extra_fields + # when the sglang backend forwards them (enable_metrics=True). + _rec.fill_sglang_response( + text=_completion_text, + response_ids=list(output.token_ids), + finish_reason=str(output.stop_reason) if output.stop_reason else "", + meta=dict(output.extra_fields or {}), + ) + # meta carries no 'completion_tokens' key here; fall back to ids length. + if _rec.turn.completion_token_count is None: + _rec.turn.completion_token_count = len(output.token_ids) + finally: + _rec.stamp_end() + self._traj_tracker.append_turn_record(_traj_key, _rec) + agent_data.trace_last_record = _rec # first time to set num_preempted if agent_data.metrics.get("num_preempted") is None: agent_data.metrics["num_preempted"] = output.num_preempted if output.num_preempted is not None else -1 @@ -309,6 +349,7 @@ class ToolAgentLoop(AgentLoopBase): add_messages: list[dict[str, Any]] = [] new_images_this_turn: list[Any] = [] # Local variable instead of agent_data attribute previous_messages = list(agent_data.messages) + _trace_pre_msg_len = len(agent_data.messages) tasks = [] tool_call_names = [] @@ -374,6 +415,16 @@ class ToolAgentLoop(AgentLoopBase): agent_data.messages.extend(add_messages) + # trajectory recording: attach tool observations to the last turn record + _trace_rec = getattr(agent_data, 'trace_last_record', None) + if _trace_rec is not None: + try: + _trace_rec.patch_observation( + agent_data.messages, _trace_pre_msg_len, len(agent_data.messages), tokenizer=self.tokenizer + ) + except Exception: + pass + if self.enable_continuous_token and not new_images_this_turn: schemas = getattr(agent_data, "_active_tool_schemas", self.tool_schemas) merge_result, response_mask, response_logprobs = await self.ct_merge_non_assistant_msg( @@ -549,4 +600,8 @@ class ToolAgentLoop(AgentLoopBase): if attr_value is not None: tool_response_kwargs[attr_name] = attr_value + trace_flow(uid=getattr(agent_data, 'trace_uid', 'unknown'), + traj_uid=getattr(agent_data, 'trace_traj_uid', 'N/A'), + from_module='Tool', to_module='ToolAgentLoop', + stage='AGENT_TOOL', detail=f'tool_executed_name={tool_call.name}') return ToolResponse(**tool_response_kwargs), tool_reward, res diff --git a/verl/experimental/separation/ray_trainer.py b/verl/experimental/separation/ray_trainer.py index 42b6e8aa..d59a9167 100644 --- a/verl/experimental/separation/ray_trainer.py +++ b/verl/experimental/separation/ray_trainer.py @@ -47,6 +47,7 @@ from verl.utils.config import omega_conf_to_dataclass from verl.utils.debug import marked_timer from verl.utils.import_utils import load_class_from_fqn from verl.utils.metric import reduce_metrics +from verl.utils.tracing import log_gen_throughput, log_progress, log_step_window class SeparateRayPPOTrainer(RayPPOTrainer): @@ -331,8 +332,29 @@ class SeparateRayPPOTrainer(RayPPOTrainer): ) self.next_step_profile = False + # horizon-style init progress lines (stdout + progress log) + log_progress(f'Total training steps: {self.total_training_steps}') + _actor_cfg = self.config.actor_rollout_ref.actor + _num_minibatches = _actor_cfg.ppo_epochs * ( + self.config.data.train_batch_size + * self.config.actor_rollout_ref.rollout.n + // _actor_cfg.ppo_mini_batch_size + ) + # per-line dict-style config dump (the format parsed by the frontend) + for _cfg_line in ( + f"'train_batch_size': {self.config.data.train_batch_size},", + f"'rollout_n': {self.config.actor_rollout_ref.rollout.n},", + f"'ppo_mini_batch_size': {_actor_cfg.ppo_mini_batch_size},", + f"'num_minibatches': {_num_minibatches},", + f"'total_training_steps': {self.total_training_steps}", + ): + log_progress(_cfg_line) + for epoch in range(current_epoch, self.config.trainer.total_epochs): - for batch_dict in self.train_dataloader: + self._train_dataloader_iter = iter(self.train_dataloader) + for batch_dict in self._logged_train_dataloader(): + # per-step banner, horizon disagg format + log_progress('=' * 20 + f' Step {self.global_steps} ' + '=' * 20) self.epoch = epoch self.fit_step(batch_dict) if self.is_last_step: @@ -432,6 +454,9 @@ class SeparateRayPPOTrainer(RayPPOTrainer): with marked_timer("gen", timing_raw, color="red"): if self.curr_step_profile: self.async_rollout_manager.start_profile(global_step=self.global_steps) + # per-step rollout window boundary (progress log only, + # not part of the sample-flow CSV) + log_step_window(self.global_steps, 'rollout', 'started') combined_gen_output = self.async_rollout_manager.generate_sequences(combined_gen_batch) self.checkpoint_manager.sleep_replicas() if self.curr_step_profile: @@ -443,6 +468,10 @@ class SeparateRayPPOTrainer(RayPPOTrainer): gen_batch_output = combined_gen_output.slice(0, num_sampled_prompts) if "__do_sample__" in gen_batch_output.non_tensor_batch: gen_batch_output.pop(non_tensor_batch_keys=["__do_sample__"]) + # rollout window end (progress log) + generation throughput progress + # line (parsed by the rl-tracing progress panel) + log_step_window(self.global_steps, 'rollout', 'finished') + log_gen_throughput(gen_batch_output, timing_raw.get("gen")) if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX: gen_baseline_output = combined_gen_output.slice(num_sampled_prompts, None) @@ -487,6 +516,7 @@ class SeparateRayPPOTrainer(RayPPOTrainer): def _fit_compute_reward(self, batch: DataProto) -> DataProto: timing_raw = self.timing_raw + log_progress('disagg: start compute reward') with marked_timer("reward", timing_raw, color="yellow"): # compute reward model score if self.use_rm and "rm_scores" not in batch.batch.keys(): @@ -495,6 +525,7 @@ class SeparateRayPPOTrainer(RayPPOTrainer): # Compute or extract reward_tensor and reward_extra_infos_dict for training reward_tensor, reward_extra_infos_dict = extract_reward(batch) + log_progress('disagg: end compute reward') self.reward_tensor = reward_tensor self.reward_extra_infos_dict = reward_extra_infos_dict return batch @@ -517,6 +548,7 @@ class SeparateRayPPOTrainer(RayPPOTrainer): policy_loss_config=self.config.actor_rollout_ref.actor.policy_loss, ) else: # Recompute old_log_probs + log_progress('disagg: start compute log prob') with marked_timer("old_log_prob", timing_raw, color="blue"): old_log_prob, old_log_prob_mfu = self._compute_old_log_prob(batch) entropys = old_log_prob.batch["entropys"] @@ -541,6 +573,7 @@ class SeparateRayPPOTrainer(RayPPOTrainer): else: old_log_prob.batch.pop("routed_experts") batch = batch.union(old_log_prob) + log_progress('disagg: end compute log prob') if "rollout_log_probs" in batch.batch.keys(): # TODO: we may want to add diff of probs too. from verl.utils.debug.metrics import calculate_debug_metrics @@ -553,17 +586,21 @@ class SeparateRayPPOTrainer(RayPPOTrainer): def _fit_compute_ref_log_prob(self, batch: DataProto) -> DataProto: timing_raw = self.timing_raw if self.use_reference_policy: + log_progress('disagg: start compute ref log prob') with marked_timer(str(Role.RefPolicy), timing_raw, color="olive"): ref_log_prob = self._compute_ref_log_prob(batch) batch = batch.union(ref_log_prob) + log_progress('disagg: end compute ref log prob') return batch def _fit_compute_critic(self, batch: DataProto) -> DataProto: timing_raw = self.timing_raw if self.use_critic: + log_progress('disagg: start compute values') with marked_timer("values", timing_raw, color="cyan"): values = self._compute_values(batch) batch = batch.union(values) + log_progress('disagg: end compute values') return batch def _fit_compute_advantage(self, batch) -> DataProto: @@ -572,6 +609,7 @@ class SeparateRayPPOTrainer(RayPPOTrainer): reward_tensor = self.reward_tensor reward_extra_infos_dict = self.reward_extra_infos_dict + log_progress('disagg: start compute adv') with marked_timer("adv", timing_raw, color="brown"): # we combine with rule-based rm reward_extra_infos_dict: dict[str, list] @@ -620,6 +658,7 @@ class SeparateRayPPOTrainer(RayPPOTrainer): norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo, config=self.config.algorithm, ) + log_progress('disagg: end compute adv') return batch def _fit_update_critic(self, batch: DataProto) -> DataProto: @@ -638,8 +677,12 @@ class SeparateRayPPOTrainer(RayPPOTrainer): # implement critic warmup if self.config.trainer.critic_warmup <= self.global_steps: # update actor + log_step_window(self.global_steps, 'train', 'started') + log_progress('disagg: start update actor') with marked_timer("update_actor", timing_raw, color="red"): actor_output = self._update_actor(batch) + log_progress('disagg: end update actor') + log_step_window(self.global_steps, 'train', 'finished') actor_output_metrics = reduce_metrics(actor_output.meta_info["metrics"]) metrics.update(actor_output_metrics) @@ -649,8 +692,10 @@ class SeparateRayPPOTrainer(RayPPOTrainer): timing_raw = self.timing_raw if self.config.trainer.critic_warmup <= self.global_steps: # update weights from trainer to rollout + log_progress('disagg: start sync params') with marked_timer("update_weights", timing_raw, color="red"): sync_metrics = self.checkpoint_manager.update_weights(self.global_steps) + log_progress('disagg: end sync params') # Engines may report per-sync stats (e.g. the delta backends' changed # ratio / wire payload); surface them in this step's metrics. if sync_metrics: @@ -670,11 +715,13 @@ class SeparateRayPPOTrainer(RayPPOTrainer): if self.config.trainer.test_freq > 0 and ( self.is_last_step or self.global_steps % self.config.trainer.test_freq == 0 ): + log_progress('disagg: start eval') with marked_timer("testing", timing_raw, color="green"): val_metrics: dict = self._validate() if self.is_last_step: self.last_val_metrics = val_metrics metrics.update(val_metrics) + log_progress('disagg: end eval') def _fit_save_checkpoint(self): timing_raw = self.timing_raw @@ -695,6 +742,7 @@ class SeparateRayPPOTrainer(RayPPOTrainer): ): if esi_close_to_expiration: print("Force saving checkpoint: ESI instance expiration approaching.") + log_progress('disagg: start save checkpoint') with marked_timer("save_checkpoint", timing_raw, color="green"): # sleep replicas to avoid OOM during checkpoint saving # self.checkpoint_manager.sleep_replicas() @@ -702,6 +750,7 @@ class SeparateRayPPOTrainer(RayPPOTrainer): # wake replicas to avoid OOM during checkpoint saving # TODO: Check separation is needed. # self.checkpoint_manager.update_weights() + log_progress('disagg: end save checkpoint') def _fit_stop_profile(self, should_profiler=None): timing_raw = self.timing_raw @@ -728,6 +777,7 @@ class SeparateRayPPOTrainer(RayPPOTrainer): metrics = self.metrics timing_raw = self.timing_raw + log_progress('disagg: start record metrics') # collect metrics metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic)) metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw)) @@ -737,6 +787,7 @@ class SeparateRayPPOTrainer(RayPPOTrainer): # compute variance proxy metrics gradient_norm = metrics.get("actor/grad_norm", None) metrics.update(compute_variance_proxy_metrics(batch=batch, gradient_norm=gradient_norm)) + log_progress('disagg: end record metrics') def _fit_experimental(self, batch): # this is experimental and may be changed/removed in the future diff --git a/verl/single_controller/ray/base.py b/verl/single_controller/ray/base.py index bb6a6043..16bed8d2 100644 --- a/verl/single_controller/ray/base.py +++ b/verl/single_controller/ray/base.py @@ -648,6 +648,8 @@ class RayWorkerGroup(WorkerGroup): ) raise ValueError(f"Cannot override protected system env: {conflict_env_vars}") env_vars.update(worker_env) + if os.environ.get("VERL_TRACE_DIR"): + env_vars["VERL_TRACE_DIR"] = os.environ["VERL_TRACE_DIR"] import re cia_name = type(ray_cls_with_init.cls).__name__ diff --git a/verl/trainer/constants_ppo.py b/verl/trainer/constants_ppo.py index ec5b02a4..02734496 100644 --- a/verl/trainer/constants_ppo.py +++ b/verl/trainer/constants_ppo.py @@ -118,6 +118,11 @@ def get_ppo_ray_runtime_env(config=None): # Always forward these at call-time, not import-time. for key in ("PYTHONHASHSEED", "VERL_FULL_DETERMINISM", "VLLM_BATCH_INVARIANT", "VERL_RL_INSIGHT_ENABLE"): runtime_env["env_vars"][key] = os.environ.get(key, "0") + # Forward trace env vars (only when set) so bare @ray.remote actors such as + # AgentLoopWorker can see them regardless of how the Ray cluster was started. + for key in ("VERL_TRACE_DIR",): + if os.environ.get(key): + runtime_env["env_vars"][key] = os.environ[key] # Forward PYTHONPATH to Ray workers so packages exposed only via PYTHONPATH (e.g. the # Megatron-LM baked into the CI image at /workspace/Megatron-LM, which is not installed # into site-packages) stay importable. Workers do not inherit the driver's PYTHONPATH when diff --git a/verl/trainer/ppo/ray_trainer.py b/verl/trainer/ppo/ray_trainer.py index 908d36b2..e8c0a4e2 100644 --- a/verl/trainer/ppo/ray_trainer.py +++ b/verl/trainer/ppo/ray_trainer.py @@ -69,6 +69,7 @@ from verl.utils.py_functional import rename_dict from verl.utils.seqlen_balancing import calculate_workload, get_seqlen_balanced_partitions, log_seqlen_unbalance from verl.utils.skip.skip_manager import SkipManager from verl.utils.torch_functional import masked_mean +from verl.utils.tracing import log_gen_throughput, log_progress, log_step_window, trace_flow from verl.utils.tracking import ValidationGenerationsLogger from verl.workers.config import DistillationConfig, EngineConfig from verl.workers.rollout.llm_server import LLMServerManager @@ -1377,6 +1378,21 @@ class RayPPOTrainer: critic_output = DataProto.from_single_dict(data={}, meta_info={"metrics": output}) return critic_output + def _logged_train_dataloader(self): + """Yield train batches with horizon-style fetch progress logs.""" + while True: + log_progress( + f"Trainer: try to load global batch@{self.global_steps}, chunk[0/0] from replay_queue" + ) + try: + batch_dict = next(self._train_dataloader_iter) + except StopIteration: + return + log_progress( + f"Trainer: successfully load global batch@{self.global_steps}, chunk[0/0] from replay_queue" + ) + yield batch_dict + def fit(self): """ The training loop of PPO. @@ -1437,8 +1453,34 @@ class RayPPOTrainer: ) next_step_profile = False + # horizon-style init progress lines (stdout + progress log) + log_progress(f'Total training steps: {self.total_training_steps}') + _actor_cfg = self.config.actor_rollout_ref.actor + _num_minibatches = _actor_cfg.ppo_epochs * ( + self.config.data.train_batch_size + * self.config.actor_rollout_ref.rollout.n + // _actor_cfg.ppo_mini_batch_size + ) + # per-line dict-style config dump (the format parsed by the frontend) + for _cfg_line in ( + f"'train_batch_size': {self.config.data.train_batch_size},", + f"'rollout_n': {self.config.actor_rollout_ref.rollout.n},", + f"'ppo_mini_batch_size': {_actor_cfg.ppo_mini_batch_size},", + f"'num_minibatches': {_num_minibatches},", + f"'total_training_steps': {self.total_training_steps}", + ): + log_progress(_cfg_line) + for epoch in range(current_epoch, self.config.trainer.total_epochs): + batches_consumed = 0 for batch_dict in self.train_dataloader: + # per-step banner, horizon colocated format + log_progress( + '=' * 20 + + f' Epoch {epoch}, consumed {batches_consumed} batches, Step {self.global_steps} ' + + '=' * 20 + ) + batches_consumed += 1 if hasattr(self.actor_rollout_wg, "async_calls_finalize_fn_exec"): self.actor_rollout_wg.async_calls_finalize_fn_exec(blocking=False) metrics = {} @@ -1457,6 +1499,10 @@ class RayPPOTrainer: batch.non_tensor_batch["uid"] = np.array( [str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object ) + trace_flow(uid=batch.non_tensor_batch["uid"], traj_uid='N/A', + from_module='DataLoader', to_module='RayPPOTrainer', + stage='DATA_PREPROCESS', detail='uid_generated', + global_step=self.global_steps) gen_batch = self._get_gen_batch(batch) @@ -1464,6 +1510,14 @@ class RayPPOTrainer: gen_batch.meta_info["global_steps"] = self.global_steps rollout_n = self.config.actor_rollout_ref.rollout.n gen_batch_output = gen_batch.repeat(repeat_times=rollout_n, interleave=True) + gen_batch_output.non_tensor_batch["traj_uid"] = np.array( + [str(i % rollout_n) for i in range(len(gen_batch_output))], dtype=object + ) + trace_flow(uid=gen_batch_output.non_tensor_batch["uid"], + traj_uid=gen_batch_output.non_tensor_batch["traj_uid"], + from_module='RayPPOTrainer', to_module='RayPPOTrainer', + stage='DATA_PREPROCESS', detail='traj_uid_generated_by_repeat', + global_step=self.global_steps) if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX: # NOTE: REMAX needs one sampled rollout plus one greedy baseline per prompt. @@ -1485,8 +1539,13 @@ class RayPPOTrainer: with marked_timer("gen", timing_raw, color="red"): if curr_step_profile: self.llm_server_manager.start_profile() + # per-step rollout window boundary (progress log only, + # not part of the sample-flow CSV) + log_step_window(self.global_steps, 'rollout', 'started') + log_progress('start generation') combined_gen_output = self.async_rollout_manager.generate_sequences(combined_gen_batch) self.checkpoint_manager.sleep_replicas() + log_progress('end generation') if curr_step_profile: self.llm_server_manager.stop_profile() @@ -1496,6 +1555,17 @@ class RayPPOTrainer: gen_batch_output = combined_gen_output.slice(0, num_sampled_prompts) if "__do_sample__" in gen_batch_output.non_tensor_batch: gen_batch_output.pop(non_tensor_batch_keys=["__do_sample__"]) + # rollout window end (progress log) + generation throughput + # progress line (parsed by the rl-tracing progress panel) + log_step_window(self.global_steps, 'rollout', 'finished') + log_gen_throughput(gen_batch_output, timing_raw.get("gen")) + # traj_uid may be dropped by rollout postprocess (streaming reward mode); + # re-inject with the same interleave pattern as point 2 so downstream + # trace points can resolve it after batch.union(gen_batch_output). + if "traj_uid" not in gen_batch_output.non_tensor_batch: + gen_batch_output.non_tensor_batch["traj_uid"] = np.array( + [str(i % rollout_n) for i in range(len(gen_batch_output))], dtype=object + ) if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX: gen_baseline_output = combined_gen_output.slice(num_sampled_prompts, None) @@ -1533,6 +1603,7 @@ class RayPPOTrainer: continue images_seqlens_all.extend(multi_modal_input["images_seqlens"].tolist()) batch.meta_info["images_seqlens"] = images_seqlens_all + log_progress('start compute reward') with marked_timer("reward", timing_raw, color="yellow"): # compute reward model score if self.use_rm and "rm_scores" not in batch.batch.keys(): @@ -1541,6 +1612,12 @@ class RayPPOTrainer: # extract reward_tensor and reward_extra_infos_dict for training reward_tensor, reward_extra_infos_dict = extract_reward(batch) + _traj_uids = batch.non_tensor_batch.get("traj_uid", np.array(["N/A"] * len(batch))) + trace_flow(uid=batch.non_tensor_batch['uid'], traj_uid=_traj_uids, + from_module='RayPPOTrainer', to_module='RewardManager', + stage='TRAIN_REWARD', detail='reward_extracted', + global_step=self.global_steps) + log_progress('end compute reward') # Operating Mode Selection: # - Bypass mode: Sets old_log_probs = rollout_log_probs (2 policies: π_rollout, π_θ) @@ -1557,8 +1634,14 @@ class RayPPOTrainer: policy_loss_config=self.config.actor_rollout_ref.actor.policy_loss, ) else: # Recompute old_log_probs + log_progress('start compute log prob') with marked_timer("old_log_prob", timing_raw, color="blue"): old_log_prob, old_log_prob_mfu = self._compute_old_log_prob(batch) + _traj_uids = batch.non_tensor_batch.get("traj_uid", np.array(["N/A"] * len(batch))) + trace_flow(uid=batch.non_tensor_batch['uid'], traj_uid=_traj_uids, + from_module='ActorRolloutRefWorker', to_module='RayPPOTrainer', + stage='TRAIN_LOGPROB', detail='old_log_prob_computed', + global_step=self.global_steps) entropys = old_log_prob.batch["entropys"] response_masks = batch.batch["response_mask"] actor_config = self.config.actor_rollout_ref.actor @@ -1583,6 +1666,7 @@ class RayPPOTrainer: "it should not be set when using R2 mode." ) batch = batch.union(old_log_prob) + log_progress('end compute log prob') if "rollout_log_probs" in batch.batch.keys(): # TODO: we may want to add diff of probs too. from verl.utils.debug.metrics import calculate_debug_metrics @@ -1592,16 +1676,26 @@ class RayPPOTrainer: assert "old_log_probs" in batch.batch, f'"old_log_prob" not in {batch.batch.keys()=}' if self.use_reference_policy: # compute reference log_prob + log_progress('start compute ref log prob') with marked_timer(str(Role.RefPolicy), timing_raw, color="olive"): ref_log_prob = self._compute_ref_log_prob(batch) + _traj_uids = batch.non_tensor_batch.get("traj_uid", np.array(["N/A"] * len(batch))) + trace_flow(uid=batch.non_tensor_batch['uid'], traj_uid=_traj_uids, + from_module='RefPolicyWorker', to_module='RayPPOTrainer', + stage='TRAIN_REF_LOGPROB', detail='ref_log_prob_computed', + global_step=self.global_steps) batch = batch.union(ref_log_prob) + log_progress('end compute ref log prob') # compute values if self.use_critic: + log_progress('start compute values') with marked_timer("values", timing_raw, color="cyan"): values = self._compute_values(batch) batch = batch.union(values) + log_progress('end compute values') + log_progress('start compute adv') with marked_timer("adv", timing_raw, color="brown"): # we combine with rule-based rm reward_extra_infos_dict: dict[str, list] @@ -1648,6 +1742,12 @@ class RayPPOTrainer: norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo, config=self.config.algorithm, ) + _traj_uids = batch.non_tensor_batch.get("traj_uid", np.array(["N/A"] * len(batch))) + trace_flow(uid=batch.non_tensor_batch['uid'], traj_uid=_traj_uids, + from_module='RayPPOTrainer', to_module='RayPPOTrainer', + stage='TRAIN_ADVANTAGE', detail='advantage_computed', + global_step=self.global_steps) + log_progress('end compute adv') # update critic if self.use_critic: with marked_timer("update_critic", timing_raw, color="pink"): @@ -1661,8 +1761,17 @@ class RayPPOTrainer: self.checkpoint_manager.update_weights(self.global_steps) else: # update actor + log_step_window(self.global_steps, 'train', 'started') + log_progress('start update actor') with marked_timer("update_actor", timing_raw, color="red"): actor_output = self._update_actor(batch) + _traj_uids = batch.non_tensor_batch.get("traj_uid", np.array(["N/A"] * len(batch))) + trace_flow(uid=batch.non_tensor_batch['uid'], traj_uid=_traj_uids, + from_module='ActorRolloutRefWorker', to_module='RayPPOTrainer', + stage='TRAIN_UPDATE', detail='actor_parameters_updated', + global_step=self.global_steps) + log_progress('end update actor') + log_step_window(self.global_steps, 'train', 'finished') # Check if the ESI (Elastic Server Instance)/training plan is close to expiration. esi_close_to_expiration = should_save_ckpt_esi( @@ -1683,12 +1792,20 @@ class RayPPOTrainer: ): if esi_close_to_expiration: print("Force saving checkpoint: ESI instance expiration approaching.") + log_progress('start save checkpoint') with marked_timer("save_checkpoint", timing_raw, color="green"): self._save_checkpoint() + log_progress('end save checkpoint') # update weights from trainer to rollout + log_progress('start sync params') with marked_timer("update_weights", timing_raw, color="red"): self.checkpoint_manager.update_weights(self.global_steps) + trace_flow(uid='*', traj_uid='*', from_module='CheckpointEngineManager', + to_module='RolloutReplica', stage='SYNC_PARAMS', + detail=f'weights_synced_global_step={self.global_steps}', + global_step=self.global_steps) + log_progress('end sync params') actor_output_metrics = reduce_metrics(actor_output.meta_info["metrics"]) metrics.update(actor_output_metrics) @@ -1702,11 +1819,13 @@ class RayPPOTrainer: if self.config.trainer.test_freq > 0 and ( is_last_step or self.global_steps % self.config.trainer.test_freq == 0 ): + log_progress('start eval') with marked_timer("testing", timing_raw, color="green"): val_metrics: dict = self._validate() if is_last_step: last_val_metrics = val_metrics metrics.update(val_metrics) + log_progress('end eval') with marked_timer("stop_profile", timing_raw): next_step_profile = ( @@ -1726,6 +1845,7 @@ class RayPPOTrainer: self.max_steps_duration = max(self.max_steps_duration, steps_duration) # training metrics + log_progress('start record metrics') metrics.update( { "training/global_step": self.global_steps, @@ -1764,6 +1884,7 @@ class RayPPOTrainer: # TODO: make a canonical logger that supports various backend logger.log(data=metrics, step=self.global_steps) + log_progress('end record metrics') progress_bar.update(1) self.global_steps += 1 diff --git a/verl/trainer/ppo/v1/agent_loop_tq.py b/verl/trainer/ppo/v1/agent_loop_tq.py index 05f56043..e87a258c 100644 --- a/verl/trainer/ppo/v1/agent_loop_tq.py +++ b/verl/trainer/ppo/v1/agent_loop_tq.py @@ -21,6 +21,7 @@ import os from typing import Any import ray +import numpy as np import torch import transfer_queue as tq from tensordict import NonTensorData, NonTensorStack, TensorDict @@ -33,6 +34,7 @@ from verl.experimental.agent_loop import ( ) from verl.utils.ray_utils import auto_await from verl.utils.tensordict_utils import list_of_dict_to_tensordict +from verl.utils.tracing import suppress_sample_trace, trace_flow logger = logging.getLogger(__name__) logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "INFO")) @@ -98,6 +100,12 @@ class AgentLoopWorkerTQ(AgentLoopWorker): logger.exception(f"Unsupported type {type(v)} for key {k}") # “fire-and-forget” background tasks + if not validate: + trace_flow(uid=str(prompt.get('uid', 'unknown')), traj_uid='N/A', + from_module='AgentLoopWorkerTQ', + to_module=f"AgentLoop({prompt.get('agent_name', 'single_turn')})", + stage='AGENT_LOOP', detail='agent_task_started', + global_step=int(prompt.get('global_steps', -1))) task = asyncio.create_task( self._run_prompt(prompt, sampling_params, trajectory=trajectory_info[i], trace=trace_this_sample) ) @@ -106,6 +114,11 @@ class AgentLoopWorkerTQ(AgentLoopWorker): async def _run_prompt(self, prompt: dict, sampling_params: dict, trajectory: dict, trace: bool = False) -> None: """Spawn multiple agent loops in parallel according to rollout.n or rollout.val_kwargs.n.""" + # validation samples never enter the training stages; silence their + # sample-flow trace for this prompt task and all session sub-tasks + # (contextvars are copied into tasks created below) + if trajectory.get("validate", False): + suppress_sample_trace() uid, partition_id = prompt["uid"], "train" if not trajectory["validate"] else "val" await tq.async_kv_put(key=uid, partition_id=partition_id, tag={"status": "running"}) tasks = [] @@ -121,9 +134,23 @@ class AgentLoopWorkerTQ(AgentLoopWorker): tasks = [] for i in range(n): + trace_flow(uid=str(prompt.get('uid', 'unknown')), traj_uid=str(i), + from_module='AgentLoopWorkerTQ', to_module='AgentLoopWorkerTQ', + stage='AGENT_LOOP', + detail=f'session_started_session_id={i}_rollout_n={n}', + global_step=int(prompt.get('global_steps', -1))) + # traj_uid injection: v1 uses session_id as the per-trajectory id + # (TQ key format {uid}_{session_id}_{index}); merge into a copy so a + # pre-existing "traj_uid" in the prompt batch never causes a + # duplicate-keyword crash and wins over the derived value. + _session_prompt = {"traj_uid": str(i), **prompt} task = asyncio.create_task( self._run_agent_loop( - run_sampling_params, trajectory=trajectory, trace=trace, session_id=i, **prompt + run_sampling_params, + trajectory=trajectory, + trace=trace, + session_id=i, + **_session_prompt, ) ) tasks.append(task) @@ -249,6 +276,18 @@ class AgentLoopManagerTQ(AgentLoopManager): prompts (TensorDict): Input batch from train or validation dataset. """ chunkes = prompts.chunk(len(self.agent_loop_workers)) + _val = prompts.get("validate", False) + _val = bool(getattr(_val, "data", _val)) + for _worker_idx, _chunk in enumerate(chunkes): + _uids = _chunk.get("uid", None) + if _uids is not None and not _val: + _gs = _chunk.get("global_steps", -1) + _gs = getattr(_gs, "data", _gs) + trace_flow(uid=[str(getattr(_uids[i], "data", _uids[i])) for i in range(len(_chunk))], + traj_uid='N/A', + from_module='AgentLoopManagerTQ', to_module=f'AgentLoopWorkerTQ_{_worker_idx}', + stage='DATA_DISPATCH', detail='dispatched_to_agent_loop_worker', + global_step=int(_gs) if isinstance(_gs, (int, np.integer)) else -1) ray.get( [ worker.generate_sequences.remote(chunk) diff --git a/verl/trainer/ppo/v1/trainer_base.py b/verl/trainer/ppo/v1/trainer_base.py index 9a0d3759..707ecc6a 100644 --- a/verl/trainer/ppo/v1/trainer_base.py +++ b/verl/trainer/ppo/v1/trainer_base.py @@ -85,6 +85,7 @@ from verl.utils.py_functional import rename_dict from verl.utils.seqlen_balancing import calculate_workload, get_seqlen_balanced_partitions, log_seqlen_unbalance from verl.utils.skip import SkipManager from verl.utils.tracking import DapoFilteredRewardTableLogger, Tracking, ValidationGenerationsLogger +from verl.utils.tracing import log_gen_throughput_tokens, log_progress, log_step_window, trace_flow from verl.workers.config import CriticConfig, DistillationConfig, HFModelConfig from verl.workers.engine_workers import ActorRolloutRefWorker, TrainingWorker, TrainingWorkerConfig from verl.workers.rollout.llm_server import LLMServerClient, LLMServerManager @@ -102,6 +103,24 @@ logger = logging.getLogger(__name__) logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "INFO")) +def _trace_ids_from_batch(batch) -> tuple[list[str], list[str]]: + """Trace-only helper: parse (uid, traj_uid) pairs from TransferQueue + trajectory keys ``{uid}_{session_id}_{index}`` via rsplit with isdigit + validation; padding samples (``is_padding`` tag) are skipped.""" + _uids, _traj_uids = [], [] + for _k, _tag in zip(batch.keys, batch.tags): + if not isinstance(_tag, dict) or _tag.get("is_padding", False): + continue + _parts = str(_k).rsplit("_", 2) + if len(_parts) == 3 and _parts[1].isdigit() and _parts[2].isdigit(): + _uids.append(_parts[0]) + _traj_uids.append(_parts[1]) + else: + _uids.append(str(_k)) + _traj_uids.append("N/A") + return _uids, _traj_uids + + def _tq_supports_checkpoint() -> bool: """Whether the installed TransferQueue can snapshot/restore its state for checkpoint consistency.""" try: @@ -439,9 +458,39 @@ class PPOTrainer(ABC): self.next_step_profile = False self.on_train_begin() + + # horizon-style init progress lines (stdout + progress log) + _pp = 'disagg: ' if self.trainer_mode == 'separate_async' else '' + log_progress(f'Total training steps: {self.total_training_steps}') + _actor_cfg = self.config.actor_rollout_ref.actor + _num_minibatches = _actor_cfg.ppo_epochs * ( + self.config.data.train_batch_size + * self.config.actor_rollout_ref.rollout.n + // _actor_cfg.ppo_mini_batch_size + ) + # per-line dict-style config dump (the format parsed by the frontend) + for _cfg_line in ( + f"'train_batch_size': {self.config.data.train_batch_size},", + f"'rollout_n': {self.config.actor_rollout_ref.rollout.n},", + f"'ppo_mini_batch_size': {_actor_cfg.ppo_mini_batch_size},", + f"'num_minibatches': {_num_minibatches},", + f"'total_training_steps': {self.total_training_steps}", + ): + log_progress(_cfg_line) + last_val_metrics = None while current_epoch < self.config.trainer.total_epochs and self.global_steps <= self.total_training_steps: is_last_step = self.global_steps >= self.total_training_steps + # per-step banner, horizon format (colocated with Epoch/consumed batches, disagg Step only) + if _pp: + log_progress('=' * 20 + f' Step {self.global_steps} ' + '=' * 20) + else: + log_progress( + '=' * 20 + + f' Epoch {current_epoch}, consumed {(self.global_steps - 1) % self.steps_per_epoch} batches,' + + f' Step {self.global_steps} ' + + '=' * 20 + ) metrics = {} self.timing_raw = {} @@ -457,8 +506,10 @@ class PPOTrainer(ABC): if self.config.trainer.save_freq > 0 and ( is_last_step or self.global_steps % self.config.trainer.save_freq == 0 ): + log_progress(_pp + 'start save checkpoint') with marked_timer("save_checkpoint", self.timing_raw, color="green"): self._save_checkpoint() + log_progress(_pp + 'end save checkpoint') self.on_step_end() metrics.update(self._consume_sync_metrics()) @@ -467,6 +518,7 @@ class PPOTrainer(ABC): if self.config.trainer.test_freq > 0 and ( is_last_step or self.global_steps % self.config.trainer.test_freq == 0 ): + log_progress(_pp + 'start eval') with marked_timer("testing", self.timing_raw, color="green"): self.on_validate_begin() val_metrics: dict = self._validate() @@ -474,8 +526,10 @@ class PPOTrainer(ABC): if is_last_step: last_val_metrics = val_metrics metrics.update(val_metrics) + log_progress(_pp + 'end eval') # 5. record metrics + log_progress(_pp + 'start record metrics') self._compute_metrics(batch, metrics, self.timing_raw, global_steps=self.global_steps, epoch=current_epoch) # 6. dump rollout generations if enabled @@ -492,6 +546,7 @@ class PPOTrainer(ABC): self.dapo_filtered_reward_logger.log( self.config.trainer.logger, dapo_filtered_reward_counts, self.global_steps ) + log_progress(_pp + 'end record metrics') progress_bar.update(1) self.global_steps += 1 SkipManager.set_step(self.global_steps) @@ -535,43 +590,109 @@ class PPOTrainer(ABC): def _step_once(self, metrics: dict, timing_raw: dict, sample_batch_size: int) -> KVBatchMeta: """Run a single local update: sample one mini-batch and perform the full PPO pipeline once.""" + _pp = 'disagg: ' if self.trainer_mode == 'separate_async' else '' # 1. sample batch from replay buffer with marked_timer("gen", timing_raw, color="red"): self.on_sample_begin() + # per-step rollout window start boundary (progress log only, + # not part of the sample-flow CSV) + log_step_window(self.global_steps, 'rollout', 'started') + if _pp: + log_progress( + f'Trainer: try to load global batch@{self.global_steps}, chunk[0/0] from replay_queue' + ) + else: + log_progress('start generation') batch, off_policy_metrics = self.replay_buffer.sample( global_steps=self.global_steps, partition_id="train", batch_size=sample_batch_size, ) + if _pp: + log_progress( + f'Trainer: successfully load global batch@{self.global_steps}, chunk[0/0] from replay_queue' + ) + else: + log_progress('end generation') metrics.update(off_policy_metrics) + _t_uids, _t_trajs = _trace_ids_from_batch(batch) + trace_flow(uid=_t_uids, traj_uid=_t_trajs, + from_module='ReplayBuffer', to_module='PPOTrainerV1', + stage='TRAIN_SAMPLE', detail='sampled_from_replay_buffer', + global_step=self.global_steps) batch.extra_info["temperature"] = self.config.actor_rollout_ref.rollout.temperature self.on_sample_end() + # per-step rollout window end boundary (progress log only) + log_step_window(self.global_steps, 'rollout', 'finished') + # generation throughput progress line — colocated modes only; in + # separate_async the gen window is a replay-buffer fetch wait rather + # than rollout production, so the numbers would be misleading. + if not _pp: + # KVBatchMeta carries no tensors; token count comes from the + # per-trajectory response_len tag set by AgentLoopWorkerTQ. + _gen_tokens = sum( + int(_tag.get("response_len", 0) or 0) + for _tag in batch.tags + if isinstance(_tag, dict) and not _tag.get("is_padding", False) + ) + log_gen_throughput_tokens(_gen_tokens, timing_raw.get("gen")) + # 2. [OPTIONAL] compute reward score with colocated reward model if self.reward_loop_manager.reward_loop_worker_handles is None: + log_progress(_pp + 'start compute reward') with marked_timer("reward", timing_raw, color="yellow"): batch = self._compute_reward_colocate(batch, metrics=metrics) + _t_uids, _t_trajs = _trace_ids_from_batch(batch) + trace_flow(uid=_t_uids, traj_uid=_t_trajs, + from_module='PPOTrainerV1', to_module='RewardLoopManager', + stage='TRAIN_REWARD', detail='reward_extracted', + global_step=self.global_steps) + log_progress(_pp + 'end compute reward') # 3. balance batch across data parallel groups batch = self._balance_batch(batch, metrics=metrics) # 4. compute old_log_prob + log_progress(_pp + 'start compute log prob') with marked_timer("old_log_prob", timing_raw, color="blue"): batch = self._compute_old_log_prob(batch, metrics=metrics) + _t_uids, _t_trajs = _trace_ids_from_batch(batch) + trace_flow(uid=_t_uids, traj_uid=_t_trajs, + from_module='ActorRolloutWorker', to_module='PPOTrainerV1', + stage='TRAIN_LOGPROB', detail='old_log_prob_computed', + global_step=self.global_steps) + log_progress(_pp + 'end compute log prob') # 5. [OPTIONAL] compute ref_log_prob if self.use_reference_policy: + log_progress(_pp + 'start compute ref log prob') with marked_timer("ref", timing_raw, color="olive"): batch = self._compute_ref_log_prob(batch, metrics=metrics) + _t_uids, _t_trajs = _trace_ids_from_batch(batch) + trace_flow(uid=_t_uids, traj_uid=_t_trajs, + from_module='RefPolicyWorker', to_module='PPOTrainerV1', + stage='TRAIN_REF_LOGPROB', detail='ref_log_prob_computed', + global_step=self.global_steps) + log_progress(_pp + 'end compute ref log prob') # 6. [OPTIONAL] compute critic values if self.use_critic: + log_progress(_pp + 'start compute values') with marked_timer("values", timing_raw, color="cyan"): batch = self._compute_values(batch, metrics=metrics) + log_progress(_pp + 'end compute values') # 7. compute advantage and return + log_progress(_pp + 'start compute adv') with marked_timer("adv", timing_raw, color="brown"): batch = self._compute_advantage(batch, metrics=metrics) + _t_uids, _t_trajs = _trace_ids_from_batch(batch) + trace_flow(uid=_t_uids, traj_uid=_t_trajs, + from_module='PPOTrainerV1', to_module='PPOTrainerV1', + stage='TRAIN_ADVANTAGE', detail='advantage_computed', + global_step=self.global_steps) + log_progress(_pp + 'end compute adv') # 8. [OPTIONAL] update critic if self.use_critic: @@ -580,8 +701,17 @@ class PPOTrainer(ABC): # 9. update actor if self.config.trainer.critic_warmup <= self.global_steps: + log_step_window(self.global_steps, 'train', 'started') + log_progress(_pp + 'start update actor') with marked_timer("update_actor", timing_raw, color="red"): batch = self._update_actor(batch, metrics=metrics) + _t_uids, _t_trajs = _trace_ids_from_batch(batch) + trace_flow(uid=_t_uids, traj_uid=_t_trajs, + from_module='ActorRolloutWorker', to_module='PPOTrainerV1', + stage='TRAIN_UPDATE', detail='actor_parameters_updated', + global_step=self.global_steps) + log_progress(_pp + 'end update actor') + log_step_window(self.global_steps, 'train', 'finished') return batch @@ -1323,6 +1453,9 @@ class PPOTrainer(ABC): batch_dict = next(self.train_dataloader_it) batch_dict["uid"] = np.array([str(uuid.uuid4()) for _ in range(len(batch_dict["raw_prompt"]))], dtype=object) + trace_flow(uid=batch_dict["uid"], traj_uid='N/A', from_module='DataLoader', + to_module='PPOTrainerV1', stage='DATA_PREPROCESS', detail='uid_generated', + global_step=self.global_steps) return tu.get_tensordict(batch_dict) def _next_train_batch(self, num_prompts: int | None = None) -> TensorDict: @@ -1357,6 +1490,10 @@ class PPOTrainer(ABC): else: tq.kv_batch_put(keys=list(batch["uid"]), partition_id="train", tags=tags) + trace_flow(uid=[str(getattr(_uid, "data", _uid)) for _uid in batch["uid"]], traj_uid='N/A', + from_module='PPOTrainerV1', to_module='AgentLoopManagerTQ', + stage='DATA_DISPATCH', detail='dispatched_to_agent_loop_manager', + global_step=self.global_steps) self.agent_loop_manager.generate_sequences(batch) return len(batch) diff --git a/verl/trainer/ppo/v1/trainer_colocate_async.py b/verl/trainer/ppo/v1/trainer_colocate_async.py index 12e31a35..edc5c7df 100644 --- a/verl/trainer/ppo/v1/trainer_colocate_async.py +++ b/verl/trainer/ppo/v1/trainer_colocate_async.py @@ -16,6 +16,7 @@ import os from verl.trainer.ppo.v1.trainer_base import PPOTrainer, register_trainer from verl.utils.debug import marked_timer +from verl.utils.tracing import log_progress, trace_flow from verl.workers.rollout.llm_server import FullyAsyncLLMServerClient logger = logging.getLogger(__name__) @@ -46,11 +47,17 @@ class PPOTrainerColocateAsync(PPOTrainer): logger.info(f"Added {num_warmup_batches} warmup batches to the agent loop manager") def on_step_end(self): + log_progress('start sync params') with marked_timer("update_weights", self.timing_raw, color="red"): # wake up all replicas to update weights self.checkpoint_manager.update_weights(self.global_steps) + trace_flow(uid='*', traj_uid='*', from_module='CheckpointEngineManager', + to_module='RolloutReplica', stage='SYNC_PARAMS', + detail=f'weights_synced_global_step={self.global_steps}', + global_step=self.global_steps) # resume generation self.checkpoint_manager.resume_generation_replicas() + log_progress('end sync params') def on_sample_end(self): # abort all unfinished requests and pause generation diff --git a/verl/trainer/ppo/v1/trainer_separate_async.py b/verl/trainer/ppo/v1/trainer_separate_async.py index 18a06ee2..46eec0c3 100644 --- a/verl/trainer/ppo/v1/trainer_separate_async.py +++ b/verl/trainer/ppo/v1/trainer_separate_async.py @@ -25,6 +25,7 @@ from verl.trainer.ppo.utils import Role, need_reward_model from verl.trainer.ppo.v1.trainer_base import PPOTrainer, register_trainer from verl.utils.config import omega_conf_to_dataclass from verl.utils.debug import marked_timer +from verl.utils.tracing import log_progress, trace_flow from verl.workers.rollout.llm_server import FullyAsyncLLMServerClient, LLMServerManager logger = logging.getLogger(__name__) @@ -159,10 +160,16 @@ class PPOTrainerSeparateAsync(PPOTrainer): self.switch_to_trainer() def on_step_end(self): + log_progress('disagg: start sync params') with marked_timer("update_weights", self.timing_raw, color="red"): # wake up all replicas to update weights; the manager returns the # engines' per-sync metrics (empty for backends that track none) self._pending_sync_metrics = self.standalone_checkpoint_manager.update_weights(self.global_steps) + trace_flow(uid='*', traj_uid='*', from_module='CheckpointEngineManager', + to_module='RolloutReplica', stage='SYNC_PARAMS', + detail=f'weights_synced_global_step={self.global_steps}', + global_step=self.global_steps) + log_progress('disagg: end sync params') def _get_n_gpus_for_throughput(self) -> int: """Include standalone rollout GPUs in the throughput denominator. diff --git a/verl/trainer/ppo/v1/trainer_sync.py b/verl/trainer/ppo/v1/trainer_sync.py index e0cc4c59..5678c1eb 100644 --- a/verl/trainer/ppo/v1/trainer_sync.py +++ b/verl/trainer/ppo/v1/trainer_sync.py @@ -16,6 +16,7 @@ import os from verl.trainer.ppo.v1.trainer_base import PPOTrainer, register_trainer from verl.utils.debug import marked_timer +from verl.utils.tracing import log_progress, trace_flow logger = logging.getLogger(__name__) logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "INFO")) @@ -33,9 +34,15 @@ class PPOTrainerSync(PPOTrainer): self.checkpoint_manager.update_weights(self.global_steps) def on_step_end(self): + log_progress('start sync params') with marked_timer("update_weights", self.timing_raw, color="red"): # wake up all replicas to update weights self.checkpoint_manager.update_weights(self.global_steps) + trace_flow(uid='*', traj_uid='*', from_module='CheckpointEngineManager', + to_module='RolloutReplica', stage='SYNC_PARAMS', + detail=f'weights_synced_global_step={self.global_steps}', + global_step=self.global_steps) + log_progress('end sync params') def on_sample_end(self): # sleep all replicas to discard weights and kv cache diff --git a/verl/utils/tracing/__init__.py b/verl/utils/tracing/__init__.py new file mode 100644 index 00000000..b7c6061a --- /dev/null +++ b/verl/utils/tracing/__init__.py @@ -0,0 +1,43 @@ +"""Unified tracing module for verl. + +Consolidates three tracing facilities (aligned with the latest +verl-horizon implementation): +- Sample-flow tracing: trace_flow (per-sample CSV under {trace_dir}/sample/) +- Training progress tracing: log_progress ([TRAIN_PROGRESS] stdout + hourly + {trace_dir}/progress/{YYYYMMDD-HH}.log, step/stage level for both colocated + and separation trainers) +- Trajectory JSONL recording: Turn, Trajectory, TurnRecord, TrajectoryStore, + TrajectoryTracker +""" + +from verl.utils.tracing.trace import ( + trace_flow, + log_progress, + log_gen_throughput, + log_gen_throughput_tokens, + log_rollout_completed, + log_step_window, + get_trace_dir, + set_trace_dir, + suppress_sample_trace, +) +from verl.utils.tracing.trajectory_types import Turn, Trajectory, TurnRecord +from verl.utils.tracing.trajectory_store import TrajectoryStore +from verl.utils.tracing.trajectory_tracker import TrajectoryTracker + +__all__ = [ + 'trace_flow', + 'log_progress', + 'log_gen_throughput', + 'log_gen_throughput_tokens', + 'log_rollout_completed', + 'log_step_window', + 'get_trace_dir', + 'set_trace_dir', + 'suppress_sample_trace', + 'Turn', + 'Trajectory', + 'TurnRecord', + 'TrajectoryStore', + 'TrajectoryTracker', +] diff --git a/verl/utils/tracing/trace.py b/verl/utils/tracing/trace.py new file mode 100644 index 00000000..ea304d82 --- /dev/null +++ b/verl/utils/tracing/trace.py @@ -0,0 +1,427 @@ +""" +Trace utility for tracking (uid, traj_uid) flow in RL training pipeline. +""" + +import atexit +import contextvars +import csv +import logging +import os +import queue +import socket +import threading +import time +from typing import Sequence, Union + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Single switch +# --------------------------------------------------------------------------- +# VERL_TRACE_DIR enables file output (or call set_trace_dir() before the +# first trace event; set_trace_dir() takes priority over the env var). +# Under the trace dir, sample-flow / megatron CSVs go into a ``sample/`` +# subdirectory; trajectory JSONL goes into ``trajectory/``. +# Both stdout and file output for trace events are gated by trace_dir being +# configured via set_trace_dir() or the VERL_TRACE_DIR env var. +# --------------------------------------------------------------------------- + +_TRACE_ENV_VAR = "VERL_TRACE_DIR" + +_SHUTDOWN = object() + +_trace_queue = None +_writer_thread = None +_filepath = None +_initialized = False +_config_trace_dir = None +_lock = threading.Lock() + +# Column order of the sample-flow CSV rows. The header itself is NOT +# written to the file; consumers (e.g. the SLS collector csv_headers) supply it. +_CSV_HEADER = ["job_id", "cat", "sample_id", "sample_status", "timestamp", "content"] + + +def _writer_loop(filepath, q): + f = open(filepath, "a", newline="", encoding="utf-8") + writer = csv.writer(f) + try: + while True: + batch = [] + try: + item = q.get(timeout=5.0) + except queue.Empty: + continue + if item is _SHUTDOWN: + f.flush() + return + batch.append(item) + while not q.empty() and len(batch) < 64: + item = q.get_nowait() + if item is _SHUTDOWN: + for row in batch: + writer.writerow(row) + f.flush() + return + batch.append(item) + for row in batch: + writer.writerow(row) + # Flush after each batch so SLS Logtail always sees consistent + # bytes rather than stale I/O-buffer fragments. + f.flush() + finally: + f.close() + + +def set_trace_dir(path): + """Set trace dir from config. Must be called before the first trace event.""" + global _config_trace_dir + _config_trace_dir = path + + +def _resolve_trace_dir(): + """Return the base trace directory, or None if unset. + + Priority: set_trace_dir() > VERL_TRACE_DIR. + """ + return _config_trace_dir or os.environ.get(_TRACE_ENV_VAR) or None + + +def get_trace_dir(): + """Public accessor for the base trace directory (or None).""" + return _resolve_trace_dir() + + +def _ensure_initialized(): + """Lazy-initialise the per-process background writer. + + Returns True iff file output is enabled (VERL_TRACE_DIR is set). + Safe to call from any thread. + """ + global _trace_queue, _writer_thread, _filepath, _initialized + if _initialized: + return _trace_queue is not None + + with _lock: + if _initialized: + return _trace_queue is not None + + trace_dir = _resolve_trace_dir() + if not trace_dir: + # Stay uninitialized so a later set_trace_dir() call (or a + # VERL_TRACE_DIR set after import) can still enable file output. + return False + _initialized = True + + os.makedirs(trace_dir, exist_ok=True) + sample_dir = os.path.join(trace_dir, "sample") + os.makedirs(sample_dir, exist_ok=True) + + # Expose sample path so external tools can discover it. + try: + with open("/tmp/verl_trace_sample_path.txt", "w") as _f: + _f.write(sample_dir) + except OSError: + pass + + hostname = socket.gethostname() + pid = os.getpid() + _filepath = os.path.join(sample_dir, f"{hostname}-pid{pid}.csv") + + _trace_queue = queue.Queue(maxsize=10000) + _writer_thread = threading.Thread( + target=_writer_loop, + args=(_filepath, _trace_queue), + daemon=True, + name="flow-trace-writer", + ) + _writer_thread.start() + atexit.register(_shutdown) + return True + + +def _shutdown(): + global _trace_queue, _writer_thread + if _trace_queue is not None: + try: + _trace_queue.put(_SHUTDOWN, timeout=5) + except queue.Full: + return + if _writer_thread is not None: + _writer_thread.join(timeout=10) + + +def _enqueue(row): + try: + _trace_queue.put_nowait(row) + except queue.Full: + pass + + +# --------------------------------------------------------------------------- +# Sample-flow tracing +# --------------------------------------------------------------------------- + +# Task-local switch to silence sample-flow tracing for validation samples. +# Each agent-loop sample runs in its own asyncio.Task (which owns a copy of the +# contextvars context), so setting it at the task entry point silences every +# trace_flow call made deeper in that sample's call chain (LLM_GENERATE, +# ROLLOUT_REWARD, ...) without touching each call site. +_suppress_trace = contextvars.ContextVar("verl_suppress_sample_trace", default=False) + + +def suppress_sample_trace(flag: bool = True): + """Silence (or re-enable) trace_flow for the current asyncio task / thread. + + Used to skip validation-sample rollouts, which never enter the training + stages and would otherwise show up as forever-RUNNING samples. + """ + _suppress_trace.set(bool(flag)) + + +def trace_flow(uid: Union[str, Sequence[str]], + traj_uid: Union[str, Sequence[str]], + from_module: str, + to_module: str, + stage: str, + detail: str = "", + global_step: int = -1): + """ + 打印数据流转追踪信息。 + + ``uid`` 和 ``traj_uid`` 既可以是单个字符串，也可以是字符串序列（list / + numpy array 等）。当传入序列时，函数内部自动按元素逐条记录，调用侧无需 + 显式写 for 循环。若一个为序列、另一个为标量字符串，标量会被广播到与序列 + 等长。 + + Args: + uid: 样本唯一标识符，单个字符串或字符串序列 + traj_uid: trajectory 唯一标识符，单个字符串或字符串序列 + from_module: 当前模块名称 + to_module: 即将转入的模块名称 + stage: 流转阶段名称 + detail: 可选的额外信息 + global_step: 当前训练全局步数(-1 表示未知) + + Use example: + # single item + trace_flow(uid='xxx', traj_uid='0', ...) + + # batch – pass the whole array directly + trace_flow(uid=batch.non_tensor_batch['uid'], + traj_uid=batch.non_tensor_batch['traj_uid'], ...) + """ + if _suppress_trace.get(): + return + + log_enabled = _ensure_initialized() + + if not log_enabled: + return + + # ---- normalise to lists ---- + uids = [uid] if isinstance(uid, str) else uid + traj_uids = [traj_uid] if isinstance(traj_uid, str) else traj_uid + + n_uid, n_traj = len(uids), len(traj_uids) + if n_uid == 1 and n_traj > 1: + uids = list(uids) * n_traj + elif n_traj == 1 and n_uid > 1: + traj_uids = list(traj_uids) * n_uid + elif n_uid != n_traj: + # zip() below silently truncates to the shorter side; surface caller bugs + logger.warning( + "trace_flow: uid/traj_uid length mismatch (%d vs %d) and neither side " + "is a broadcastable scalar; extra entries will be dropped (stage=%s)", + n_uid, n_traj, stage) + + # ---- fields shared across all items in this batch ---- + timestamp_iso = time.strftime("%Y-%m-%dT%H:%M:%S", time.localtime()) + timestamp_ms = int(time.time() * 1000) + detail_str = f" | detail={detail}" if detail else "" + step_str = f" | global_step={global_step}" if global_step != -1 else "" + job_id = os.environ.get("DLC_JOB_ID", "ENV_NOTFOUND_DLC_JOB_ID") + sample_status = "COMPLETE" if stage in {"TRAIN_UPDATE", "FILTER_TRAJECTORIES"} else "RUNNING" + + parts = [f"from={from_module}", f"to={to_module}", f"stage={stage}"] + if global_step != -1: + parts.append(f"global_step={global_step}") + if detail: + parts.append(f"detail={detail}") + content = "|".join(parts) + + for u, t in zip(uids, traj_uids): + print(f"[FLOW_TRACE] {timestamp_iso} | uid={u} | traj_uid={t} | " + f"from={from_module} | to={to_module} | stage={stage}{step_str}{detail_str}") + + sample_id = f"{u}/{t}" if str(t).isdigit() else f"{u}/sample" + _enqueue([job_id, "sample", sample_id, sample_status, timestamp_ms, content]) + + +# --------------------------------------------------------------------------- +# Training progress logging (step/stage level, stdout only) +# --------------------------------------------------------------------------- + +ENABLE_PROGRESS_LOG = True + +_PROGRESS_LOCK = threading.Lock() +_progress_file = None +_progress_hour_key = None +_progress_path_exposed = False + + +def _write_progress_file(line: str): + """Append a progress line to ``{trace_dir}/progress/{YYYYMMDD-HH}.log``. + + Files rotate hourly (name = year-month-day-hour, e.g. ``20260811-15.log``) + and output is gated by the trace dir, like the sample-flow CSVs. + The progress directory is exposed via ``/tmp/verl_progress_path.txt`` so + external tools can discover it. + Never raises. + """ + global _progress_file, _progress_hour_key, _progress_path_exposed + trace_dir = _resolve_trace_dir() + if not trace_dir: + return + hour_key = time.strftime("%Y%m%d-%H", time.localtime()) + with _PROGRESS_LOCK: + try: + if hour_key != _progress_hour_key: + if _progress_file is not None: + try: + _progress_file.close() + except OSError: + pass + _progress_file = None + progress_dir = os.path.join(trace_dir, "progress") + os.makedirs(progress_dir, exist_ok=True) + if not _progress_path_exposed: + # Expose progress path so external tools can discover it. + try: + with open("/tmp/verl_progress_path.txt", "w") as _f: + _f.write(progress_dir) + _progress_path_exposed = True + except OSError: + pass + _progress_file = open( + os.path.join(progress_dir, f"{hour_key}.log"), "a", encoding="utf-8") + _progress_hour_key = hour_key + _progress_file.write(line + "\n") + _progress_file.flush() + except OSError: + pass + + +def log_progress(message: str): + """输出一行训练进度,内容与 verl-horizon 原始日志一致。 + + 行内容直接采用 horizon 的裸格式,例如: + ``disagg: start compute reward`` / ``disagg: end compute reward`` / + ``Trainer: try to load global batch@N, chunk[0/0] from replay_queue`` 等, + 前端训练阶段面板直接解析这些行。 + + stdout 无条件输出;设置了 VERL_TRACE_DIR 时同时逐行落盘到 + ``{trace_dir}/progress/{YYYYMMDD-HH}.log``(按小时轮转,文件名即年月日-时, + 如 ``20260811-15.log``)。 + + Args: + message: 完整的一行日志内容(不含换行)。 + """ + if not ENABLE_PROGRESS_LOG: + return + print(message, flush=True) + _write_progress_file(message) + + +def log_gen_throughput_tokens(n_tokens: int, gen_seconds: float): + r"""输出一行生成吞吐进度行,格式与前端解析正则对齐。 + + rl-tracing 训练进度面板用 ``gen throughput \(token/s\):\s*([\d.]+)`` 解析 + master-0 原始日志,因此行内容必须严格保持 + ``gen throughput (token/s): {X:.2f}``。不依赖 sglang 引擎统计日志 + (开源版 sglang Engine 的 log_level=error 会屏蔽引擎自身的统计行)。 + 参数非法时静默跳过,绝不抛出。 + + Args: + n_tokens: 本步生成的 token 总数(非 pad 响应 token)。 + gen_seconds: 生成窗口耗时(秒),即 marked_timer("gen") 的计时结果。 + """ + try: + if n_tokens is None or gen_seconds is None or gen_seconds <= 0 or n_tokens < 0: + return + log_progress(f"gen throughput (token/s): {int(n_tokens) / gen_seconds:.2f}") + except Exception: + pass + + +def log_gen_throughput(gen_output, gen_seconds: float): + """从 rollout 输出批次计算生成吞吐并输出进度行(v0 路径使用)。 + + token 数 = attention_mask 在响应区域的求和(即非 pad 的响应 token); + 无 attention_mask 时退化为 ``batch_size * response_length``。v1 路径的 + 批次是 KVBatchMeta(无 tensor),请自行从 tags 的 ``response_len`` 求和 + 后调用 :func:`log_gen_throughput_tokens`。 + + Args: + gen_output: rollout 输出(DataProto),需包含 ``responses``,可选 ``attention_mask``。 + gen_seconds: 生成窗口耗时(秒)。 + """ + try: + batch = getattr(gen_output, "batch", None) + if batch is None or "responses" not in batch.keys(): + return + responses = batch["responses"] + mask = batch.get("attention_mask", None) + if mask is not None and mask.dim() == 2 and mask.shape[-1] >= responses.shape[-1]: + prompt_len = mask.shape[-1] - responses.shape[-1] + n_tokens = int(mask[:, prompt_len:].sum()) + else: + n_tokens = int(responses.numel()) + log_gen_throughput_tokens(n_tokens, gen_seconds) + except Exception: + pass + + +def log_rollout_completed(uid, traj_uid, score=None): + """输出一行 rollout 完成进度行,格式与前端计数条件对齐。 + + rl-tracing 训练进度面板用 + ``raw_message like '%is completed, score=%' and raw_message like '%status=Finished%'`` + 统计已完成 rollout 数并用最近 120 条估算完成速率,因此行内容必须同时 + 包含 ``is completed, score=`` 与 ``status=Finished``。horizon 原始行带 + ``TONGYI:`` 前缀,此处按约定不带该前缀。仅在 trace 开启(设置了 + VERL_TRACE_DIR)时输出。绝不抛出。 + + Args: + uid: 样本唯一标识符。 + traj_uid: trajectory 唯一标识符。 + score: 该轨迹的 reward 分数(可为 None,输出 0.0)。 + """ + try: + if not _resolve_trace_dir(): + return + if score is None: + score = 0.0 + log_progress(f"Task {uid}/{traj_uid} is completed, score={score} status=Finished") + except Exception: + pass + + +def log_step_window(global_step: int, phase: str, boundary: str): + """输出一行 step 级窗口进度行(rollout/train 的起止)。 + + 格式固定为 ``Step {N} {phase} {boundary}``(step 号后无冒号),例如 + ``Step 3 rollout started`` / ``Step 3 train finished``, + 每个 global step 各四行,供外部工具从 progress 日志重建 + 每步的 rollout 窗口与训练窗口。不走 sample-flow CSV。 + 绝不抛出。 + + Args: + global_step: 当前训练全局步数。 + phase: 'rollout' 或 'train'。 + boundary: 'started' 或 'finished'。 + """ + try: + log_progress(f"Step {global_step} {phase} {boundary}") + except Exception: + pass diff --git a/verl/utils/tracing/trajectory_store.py b/verl/utils/tracing/trajectory_store.py new file mode 100644 index 00000000..cd36f9e9 --- /dev/null +++ b/verl/utils/tracing/trajectory_store.py @@ -0,0 +1,136 @@ +"""Writes Trajectory objects to JSONL files with automatic rotation.""" + +import json +import logging +import os +import time +from concurrent.futures import ThreadPoolExecutor +from typing import Optional + +from verl.utils.tracing.trajectory_types import Trajectory + +logger = logging.getLogger(__name__) + + +class TrajectoryStore: + """Writes Trajectory objects to JSONL files with automatic rotation. + + Each trajectory is serialized as a compact JSON line. + When the current file exceeds ``max_file_size_mb`` (default 100 MB), a new + file is created automatically. File I/O is offloaded to a single-thread + executor so it never blocks the inference event loop. + """ + + _DEFAULT_MAX_MB = 100 + + def __init__(self, dump_dir: str, instance_id: int, + max_file_size_mb: int = _DEFAULT_MAX_MB): + self.dump_dir = dump_dir + self.instance_id = instance_id + self.max_file_size_bytes = max_file_size_mb * 1024 * 1024 + self._file: Optional[object] = None + self._file_path: Optional[str] = None + self._executor = ThreadPoolExecutor(max_workers=1) + os.makedirs(dump_dir, exist_ok=True) + self._open_new_file() + + def _open_new_file(self): + if self._file is not None: + try: + self._file.close() + except Exception: + pass + self._file = None + ts = f"{time.strftime('%Y%m%d_%H%M%S')}_{int(time.time() * 1000) % 1000:03d}" + self._file_path = os.path.join( + self.dump_dir, + f'trajectory_{self.instance_id}_{ts}.jsonl', + ) + self._file = open(self._file_path, 'a') + + def _check_rotation(self): + try: + if self._file_path and os.path.getsize(self._file_path) >= self.max_file_size_bytes: + self._open_new_file() + except OSError: + pass + + def _write_sync(self, line: str): + try: + self._check_rotation() + self._file.write(line + '\n') + self._file.flush() + except Exception as e: + logger.error(f'TrajectoryStore write failed: {e}') + + # Fields stripped during serialization to save storage space. + _TRAJECTORY_DROP_FIELDS = frozenset({'initial_prompt_token_ids'}) + _TURN_DROP_FIELDS = frozenset({'completion_token_ids', 'observation_token_ids'}) + + # Fields unavailable in the verl-latest architecture (kept in the dataclass + # for horizon-schema compatibility, dropped at serialization time): + # - replica_id: the agent loop talks to an LLMServerClient, there is no + # per-request engine-replica notion; + # - prompt_text: never populated on the collection path (large, horizon + # leaves it None as well). + # NOTE: the Turn sglang_* timing fields are NOT dropped: sglang + # (enable_metrics=True) returns per-request e2e_latency/queue_time/prefill + # timestamps in the response meta_info, which async_sglang_server.py + # forwards via TokenOutput.extra_fields and the agent loops fill into the + # TurnRecord. sglang_rid stays dropped: the engine's request id is not + # exposed through TokenOutput. + _TRAJECTORY_UNAVAILABLE_FIELDS = frozenset({'replica_id'}) + _TURN_UNAVAILABLE_FIELDS = frozenset({ + 'prompt_text', + 'sglang_rid', + }) + + def _to_compact_dict(self, trajectory: Trajectory) -> dict: + """Convert trajectory to dict, dropping large token-id fields and + fields unavailable in the verl-latest architecture, and deriving + turn_total_time_ms when not set by the collector.""" + d = trajectory.to_dict() + for key in self._TRAJECTORY_DROP_FIELDS | self._TRAJECTORY_UNAVAILABLE_FIELDS: + d.pop(key, None) + for turn in d.get('turns', []): + for key in self._TURN_DROP_FIELDS | self._TURN_UNAVAILABLE_FIELDS: + turn.pop(key, None) + # Derive total turn time when the collector did not set it + # (patch_observation only sets it for tool turns). + if turn.get('turn_total_time_ms') is None: + turn['turn_total_time_ms'] = (turn.get('completion_time_ms') or 0.0) + ( + turn.get('observation_time_ms') or 0.0 + ) + return d + + # Top-level keys kept outside ``raw_message``; everything else is folded + # into the raw_message JSON string (SLS ingestion convention: flat index + # columns + one raw payload column). + _TOP_LEVEL_KEYS = ('prompt_uid', 'sample_index') + + def write(self, trajectory: Trajectory): + """Asynchronously write a trajectory to the current JSONL file. + + Line format: ``{"prompt_uid": ..., "sample_index": ..., + "raw_message": ""}``. + """ + d = self._to_compact_dict(trajectory) + line = json.dumps({ + key: d.get(key) + for key in self._TOP_LEVEL_KEYS + } | { + 'raw_message': json.dumps( + {k: v for k, v in d.items() if k not in self._TOP_LEVEL_KEYS}, + ensure_ascii=False, + ), + }, ensure_ascii=False) + self._executor.submit(self._write_sync, line) + + def close(self): + self._executor.shutdown(wait=True) + if self._file is not None: + try: + self._file.close() + except Exception: + pass + self._file = None diff --git a/verl/utils/tracing/trajectory_tracker.py b/verl/utils/tracing/trajectory_tracker.py new file mode 100644 index 00000000..38b5776e --- /dev/null +++ b/verl/utils/tracing/trajectory_tracker.py @@ -0,0 +1,226 @@ +"""Thread-safe tracker for accumulating TurnRecords during multi-turn agent execution. + +Manages the lifecycle of trajectory recording: + register → append turns → build → unregister + +Extracted from BaseAgent to keep trajectory-tracking logic in the tracing module. +""" + +import logging +import os +import threading +import time +from typing import Dict, List, Optional + +import numpy as np + +from verl.utils.tracing.trajectory_types import Trajectory, TurnRecord + +logger = logging.getLogger(__name__) +logger.setLevel(os.getenv('VERL_PPO_LOGGING_LEVEL', 'WARN')) + + +class TrajectoryTracker: + """Thread-safe tracker for accumulating TurnRecords during multi-turn agent execution. + + Each agent instance owns one TrajectoryTracker. FastAPI may invoke + callbacks on different uvicorn worker threads, so all access to the + internal dict is guarded by a Lock. + """ + + def __init__(self): + self._records: Dict[str, List[TurnRecord]] = {} + self._lock = threading.Lock() + + # ------------------------------------------------------------------ + # Static utilities + # ------------------------------------------------------------------ + + @staticmethod + def make_traj_key(uid, sample_idx) -> str: + """Build a trajectory map key: ``'prompt-{uid}/sample-{sample_idx}'``.""" + return f'prompt-{uid}/sample-{sample_idx}' + + @staticmethod + def extract_str(val) -> Optional[str]: + """Extract a plain string from *val* (str, np.ndarray, or None).""" + if val is None: + return None + if isinstance(val, np.ndarray): + return str(val.item()) if val.size == 1 else str(val) + return str(val) + + @staticmethod + def parse_sample_index(traj_uid) -> Optional[int]: + """Parse sample_index from traj_uid. + + In repeat_n(), traj_uid is assigned as str(idx) where idx is 0..n-1. + We try to parse it as an int; if it is not a simple integer string + (e.g. a UUID for non-over-sampling paths), return None. + """ + if traj_uid is None: + return None + s = traj_uid + if isinstance(s, np.ndarray): + s = str(s.item()) if s.size == 1 else str(s) + else: + s = str(s) + try: + return int(s) + except (ValueError, TypeError): + return None + + @staticmethod + def extract_n_samples(config) -> Optional[int]: + """Extract the over-sampling factor *n* from config. + + Source: ``config.actor_rollout_ref.rollout.over_sampling.n`` + Falls back to ``config.actor_rollout_ref.rollout.n`` (group size). + Returns None if neither is available. + """ + try: + rollout_cfg = config.actor_rollout_ref.rollout + over_sampling_n = rollout_cfg.get('over_sampling', {}).get('n', None) + if over_sampling_n is not None: + return int(over_sampling_n) + group_n = rollout_cfg.get('n', None) + if group_n is not None: + return int(group_n) + except (AttributeError, KeyError, TypeError): + pass + return None + + # ------------------------------------------------------------------ + # Lifecycle: register → append → get → build → unregister + # ------------------------------------------------------------------ + + def register_trajectory(self, item) -> tuple: + """Create an empty turn-records list for *item* and return ``(traj_key, traj_uid)``. + + ``traj_uid`` alone is just the over-sampling index (0..n-1) and can + collide across different prompts processed concurrently, so we combine + it with uid to form a unique key. + """ + traj_uid_raw = item.non_tensor_batch.get('traj_uid', '') + traj_uid = self.extract_str(traj_uid_raw) + if not traj_uid: + traj_uid = str(id(item)) + logger.warning( + "Missing traj_uid for item uid=%s, falling back to id(). " + "meta_info_keys=%s non_tensor_keys=%s validate=%s global_step=%s", + item.non_tensor_batch.get('uid', '?'), + list(item.meta_info.keys()) if item.meta_info else [], + list(item.non_tensor_batch.keys()) if item.non_tensor_batch else [], + item.meta_info.get('validate', '?'), + item.meta_info.get('global_step', '?'), + ) + uid_raw = item.non_tensor_batch.get('uid', '') + uid = self.extract_str(uid_raw) or '' + traj_key = self.make_traj_key(uid, traj_uid) + with self._lock: + if traj_key in self._records: + logger.warning( + "register_trajectory: key %s is already registered; overwriting " + "%d in-flight turn record(s)", traj_key, len(self._records[traj_key])) + self._records[traj_key] = [] + return traj_key, traj_uid + + def get_turn_records(self, traj_key: str) -> List[TurnRecord]: + """Return the turn-records list for *traj_key*. Thread-safe.""" + with self._lock: + return self._records.get(traj_key, []) + + def append_turn_record(self, traj_key: str, record: TurnRecord): + """Append *record* to the trajectory identified by *traj_key*. + + Logs a warning if *traj_key* is None (caller forgot to pass it). + """ + if traj_key is None: + logger.warning('append_turn_record called without traj_key, record dropped') + return + with self._lock: + records = self._records.get(traj_key) + if records is not None: + records.append(record) + + def unregister_trajectory(self, traj_key: str): + """Remove the turn-records list for *traj_key*. Called at ``safe_run()`` end.""" + with self._lock: + self._records.pop(traj_key, None) + + # ------------------------------------------------------------------ + # Build final Trajectory from accumulated TurnRecords + # ------------------------------------------------------------------ + + def build_trajectory(self, item, traj_key, traj_uid, created_at, + tokenizer, instance_id, config, + agent_metrics=None) -> Trajectory: + """Build a :class:`Trajectory` from the accumulated turn records. + + Parameters + ---------- + item : DataProtoItem + The original data item (used for uid / global_step metadata). + traj_key : str + Key used to look up turn records in this tracker. + traj_uid : str + Raw trajectory uid, used as ``Trajectory.request_id``. + created_at : float + Unix timestamp when the trajectory started. + tokenizer : PreTrainedTokenizer + Tokenizer for decoding prompts / overhead text. + instance_id : int + The agent worker instance id (process pid; kept for file naming + and horizon-schema ``replica_id`` compatibility). + config + Config object for extracting ``n_samples``. + agent_metrics : dict, optional + Per-trajectory metrics collected by the agent. + """ + completed_at = time.time() + turn_records = self.get_turn_records(traj_key) + + for i, rec in enumerate(turn_records): + rec.turn.turn_index = i + + # Compute per-turn template overhead (static chat-template tokens between turns) + for i in range(1, len(turn_records)): + prev = turn_records[i - 1].turn + cur = turn_records[i].turn + if (cur.prompt_token_count is not None + and prev.prompt_token_count is not None + and prev.completion_token_count is not None): + obs_tokens = len(prev.observation_token_ids) if prev.observation_token_ids else 0 + overhead = (cur.prompt_token_count + - prev.prompt_token_count + - prev.completion_token_count + - obs_tokens) + cur.template_overhead_tokens = overhead + if overhead > 0: + offset = prev.prompt_token_count + prev.completion_token_count + cur.template_overhead_text = tokenizer.decode( + turn_records[i]._prompt_ids[offset:offset + overhead], + skip_special_tokens=False, + ) + + initial_prompt = None + initial_prompt_ids = None + if turn_records: + initial_prompt_ids = list(turn_records[0]._prompt_ids) + initial_prompt = tokenizer.decode(initial_prompt_ids, skip_special_tokens=False) + + return Trajectory( + request_id=str(traj_uid), + replica_id=instance_id, + global_step=item.meta_info.get('global_step', None), + created_at_ts=created_at, + completed_at_ts=completed_at, + total_time_ms=(completed_at - created_at) * 1000, + turns=[rec.turn for rec in turn_records], + prompt_uid=self.extract_str(item.non_tensor_batch.get('uid', None)), + sample_index=self.parse_sample_index(item.non_tensor_batch.get('traj_uid', None)), + n_samples=self.extract_n_samples(config), + initial_prompt_text=initial_prompt, + initial_prompt_token_ids=initial_prompt_ids, + agent_metrics=agent_metrics, + ) diff --git a/verl/utils/tracing/trajectory_types.py b/verl/utils/tracing/trajectory_types.py new file mode 100644 index 00000000..81f1f652 --- /dev/null +++ b/verl/utils/tracing/trajectory_types.py @@ -0,0 +1,321 @@ +"""Core data types for trajectory recording. + +Turn — one complete RL step (action + observation) +Trajectory — full multi-turn trajectory for one inference request +TurnRecord — mutable builder that wraps a Turn during recording + +## Multi-turn data flow +## +## Each Turn = one complete RL step = action + observation. +## Corresponds to the (state, action, observation) tuple in agent RL. +## +## messages = [system_prompt, user_task] +## +## -- Turn 0 -------------------------------------------------- +## [action] send prompt_ids -> sglang generate -> model_output_0 = "call tool X" +## [observation] agent executes tool X -> tool_result_0 +## -- Turn 0 end -- +## messages += [{assistant: model_output_0}, {tool: tool_result_0}] +## +## -- Turn 1 -------------------------------------------------- +## [action] send prompt_ids (with history) -> sglang generate -> model_output_1 = "call tool Y" +## [observation] agent executes tool Y -> tool_result_1 +## -- Turn 1 end -- +## messages += [{assistant: model_output_1}, {tool: tool_result_1}] +## +## -- Turn 2 -------------------------------------------------- +## [action] send prompt_ids (with all history) -> sglang generate -> "final answer" +## [observation] None (no tool call in the last turn) +## -- Turn 2 end -- + +## Turn timing hierarchy +## +## turn_total_time_ms +## ├── completion_time_ms (client-side /generate call: start_ts → end_ts) +## │ └── sglang_time_ms (server-side total) +## │ ├── sglang_queue_time_ms +## │ ├── sglang_ttft_ms (prefill) +## │ └── sglang_decode_time_ms +## └── observation_time_ms (tool execution: observation_start_ts → observation_end_ts) +""" + +import json +import time +from dataclasses import asdict, dataclass, field +from typing import List, Optional + + +@dataclass +class Turn: + """One Turn = one complete RL step = action + observation. + + Corresponds to the (state, action, observation, reward) tuple in agent RL: + - state: the full context sent to sglang (prompt_ids); length recorded by prompt_token_count + - action: model output (completion_text / completion_token_ids), obtained via run_generate() + - observation: environment feedback (tool execution result), recorded in observation_* fields + - reward: typically only present in the last Turn, recorded at the Trajectory level + + The "role" is implicit: the action part is always the assistant's model output, + and the observation part (when present) is always the tool/environment response. + No explicit role field is needed. + """ + + # -- identity --------------------------------------------------------------- + turn_index: int + # 0-based index of this turn within the Trajectory. + + # -- prompt (input to sglang) ----------------------------------------------- + prompt_token_count: Optional[int] = None + # Number of input tokens sent to sglang for this call. + + prompt_text: Optional[str] = None + # Decoded text of the prompt sent to sglang. Optional (large). + + # -- completion (model output) ---------------------------------------------- + completion_text: str = "" + # Text content of the model output (decoded text from sglang). + + completion_token_ids: List[int] = field(default_factory=list) + # Output/completion token IDs only (does NOT include prompt tokens). + + completion_token_count: Optional[int] = None + # Number of new output tokens generated by sglang. + + finish_reason: str = "" + # Why sglang stopped generating: "stop", "length", "tool_call", or error strings. + + # -- observation (tool / environment feedback) ------------------------------ + observation_text: Optional[str] = None + # Text result returned by the tool/environment. None when no tool call. + + observation_token_ids: Optional[List[int]] = None + # Token IDs of the tool result. None when no tool call. + + # -- timing (see tree in module docstring) ---------------------------------- + start_ts: float = 0.0 + # Client-side timestamp before the HTTP POST to sglang /generate. + + end_ts: float = 0.0 + # Client-side timestamp after the sglang /generate response is received. + + observation_start_ts: Optional[float] = None + # Timestamp when tool execution starts (= end_ts). + + observation_end_ts: Optional[float] = None + # Timestamp when tool execution completes. + + turn_total_time_ms: Optional[float] = None + # Total turn duration: completion_time_ms + observation_time_ms. + + completion_time_ms: float = 0.0 + # Client-side wall-clock of the sglang /generate call (start_ts → end_ts). + + observation_time_ms: Optional[float] = None + # Tool execution duration in milliseconds. + + # -- sglang server metrics -------------------------------------------------- + sglang_rid: str = "" + # sglang request ID (uuid4 hex) for cross-referencing with sglang logs. + + sglang_time_ms: Optional[float] = None + # sglang server-side end-to-end latency (milliseconds). + + sglang_queue_time_ms: Optional[float] = None + # Time spent waiting in sglang's scheduler queue (milliseconds). + + sglang_ttft_ms: Optional[float] = None + # Time To First Token (milliseconds): prefill phase duration. + + sglang_prefill_ms: Optional[float] = None + # Pure GPU prefill duration (milliseconds) from scheduler-side timestamps. + + sglang_decode_time_ms: Optional[float] = None + # Decode phase duration (milliseconds): sglang_time - queue - ttft. + + template_overhead_tokens: Optional[int] = None + template_overhead_text: Optional[str] = None + # Static chat-template tokens added between turns (e.g. <|im_start|>user\n…<|im_end|>\n<|im_start|>assistant\n\n). + # Computed as: prompt_token_count[N] - prompt_token_count[N-1] - completion_token_count[N-1] - observation_token_count[N-1] + # Only set for turn_index >= 1. Typically ~13 for Qwen-style chat templates. + + +@dataclass +class Trajectory: + """Complete multi-turn trajectory for one inference request. + + Represents the full lifecycle of a single rollout: from safe_run() start + through all sglang /generate calls and tool executions until completion. + + In over-sampling scenarios (n > 1), the same prompt produces N + Trajectory objects. Use prompt_uid / sample_index / n_samples to group them. + """ + + # -- identity --------------------------------------------------------------- + request_id: str + # Unique identifier for this trajectory (from traj_uid). + + prompt_uid: Optional[str] = None + # Stable identifier for the original prompt. Shared across over-samples. + + sample_index: Optional[int] = None + # Index of this trajectory among the N samples for the same prompt (0..N-1). + + n_samples: Optional[int] = None + # Total number of samples requested for this prompt. + + # -- infrastructure --------------------------------------------------------- + replica_id: int = -1 + # Which sglang engine replica served this request (0..num_replicas-1). + + global_step: Optional[int] = None + # Model version at dispatch time — the weights used to generate this trajectory. + # In disaggregated mode, the model may advance further by the time this + # trajectory completes and is consumed for training. + + # -- timing ----------------------------------------------------------------- + created_at_ts: float = 0.0 + # Unix timestamp when safe_run() begins processing. + + completed_at_ts: float = 0.0 + # Unix timestamp when safe_run() finishes processing. + + total_time_ms: float = 0.0 + # Total wall-clock duration in milliseconds. + + # -- data ------------------------------------------------------------------- + turns: List[Turn] = field(default_factory=list) + # Ordered list of Turn objects, one per sglang /generate call. + + initial_prompt_text: Optional[str] = None + initial_prompt_token_ids: Optional[List[int]] = None + + # -- agent-level metrics ---------------------------------------------------- + agent_metrics: Optional[dict] = None + # Per-trajectory metrics collected by the agent (e.g. tool call counts, + # token usage breakdowns). Passed through from result.meta_info['metrics']. + + def to_dict(self) -> dict: + return asdict(self) + + def to_json(self) -> str: + return json.dumps(self.to_dict(), ensure_ascii=False, separators=(',', ':')) + + @classmethod + def from_dict(cls, d: dict) -> 'Trajectory': + turns_data = d.pop('turns', []) + turns = [] + for t in turns_data: + turn_fields = Turn.__dataclass_fields__ + filtered = {k: v for k, v in t.items() if k in turn_fields} + turns.append(Turn(**filtered)) + traj_fields = {f for f in cls.__dataclass_fields__} + filtered_d = {k: v for k, v in d.items() if k in traj_fields} + return cls(turns=turns, **filtered_d) + + @classmethod + def from_json(cls, s: str) -> 'Trajectory': + return cls.from_dict(json.loads(s)) + + +class TurnRecord: + """Mutable builder that wraps a Turn dataclass. + + Created at the start of each run_generate() call. Methods mutate the + inner Turn as sglang response data and tool observations arrive. + Once the trajectory is complete, callers read .turn directly. + """ + + __slots__ = ('turn', '_prompt_ids') + + def __init__(self, prompt_ids, rid=''): + self._prompt_ids = prompt_ids + self.turn = Turn( + turn_index=-1, + completion_text='', completion_token_ids=[], finish_reason='', + start_ts=time.time(), end_ts=0.0, completion_time_ms=0.0, + sglang_rid=rid, + prompt_token_count=len(prompt_ids) if prompt_ids else 0, + ) + + def stamp_end(self): + """Set end_ts to now and compute completion_time_ms. Called in finally block.""" + self.turn.end_ts = time.time() + self.turn.completion_time_ms = (self.turn.end_ts - self.turn.start_ts) * 1000 + + def mark_skipped(self, finish_reason='length'): + """Mark this turn as skipped (e.g. zero tokens left) and stamp end.""" + self.turn.finish_reason = finish_reason + self.stamp_end() + + def fill_sglang_response(self, text, response_ids, finish_reason, meta): + """Fill in fields from a successful sglang /generate response. + + NOTE: queue_time, prefill_finished_ts, request_received_ts, and + prefill_launch_latency are only present in meta_info when the sglang + server is launched with --enable-metrics. Without it, sglang_ttft_ms, + sglang_queue_time_ms, sglang_prefill_ms, and sglang_decode_time_ms + will all be None. + """ + t = self.turn + t.completion_text = text + t.completion_token_ids = response_ids + t.finish_reason = finish_reason + t.completion_token_count = meta.get("completion_tokens") + e2e = meta.get("e2e_latency") + t.sglang_time_ms = e2e * 1000 if e2e is not None else None + prefill_ts = meta.get("prefill_finished_ts") + recv_ts = meta.get("request_received_ts") + t.sglang_ttft_ms = (prefill_ts - recv_ts) * 1000 if (prefill_ts is not None and recv_ts is not None) else None + queue_t = meta.get("queue_time") + t.sglang_queue_time_ms = queue_t * 1000 if queue_t is not None else None + # GPU-side prefill duration (independent of IPC overhead) + pll = meta.get("prefill_launch_latency") + t.sglang_prefill_ms = pll * 1000 if pll is not None else None + # Derive decode time = server total - queue - prefill + if t.sglang_time_ms is not None: + q = t.sglang_queue_time_ms or 0 + p = t.sglang_ttft_ms or 0 + t.sglang_decode_time_ms = max(0, t.sglang_time_ms - q - p) + + def patch_observation(self, messages, pre_msg_length, msg_len, tokenizer=None): + """Attach tool/environment observation from new messages since last turn. + + Between consecutive LLM calls, the agent framework executes tools and + appends observation messages. This extracts those observations and + records them (with timing and token counts) on this turn. + """ + if pre_msg_length >= msg_len: + return + new_messages = messages[pre_msg_length:] + obs_parts = [] + for m in new_messages: + if m.get('role') not in ('tool', 'function'): + continue + content = m.get('content', '') or '' + if isinstance(content, list): + content = '\n'.join( + c.get('text', '') if isinstance(c, dict) else str(c) + for c in content + ) + obs_parts.append(str(content)) + if not obs_parts: + return + t = self.turn + t.observation_text = '\n'.join(obs_parts) + if tokenizer is not None: + try: + t.observation_token_ids = tokenizer.encode( + t.observation_text, add_special_tokens=False + ) if t.observation_text else None + except Exception: + t.observation_token_ids = None + callback_ts = time.time() + t.observation_start_ts = t.end_ts if t.end_ts else None + t.observation_end_ts = callback_ts + t.observation_time_ms = ( + (callback_ts - t.observation_start_ts) * 1000 + if t.observation_start_ts is not None else None + ) + # Derive turn total + obs = t.observation_time_ms or 0 + t.turn_total_time_ms = t.completion_time_ms + obs diff --git a/verl/workers/rollout/sglang_rollout/async_sglang_server.py b/verl/workers/rollout/sglang_rollout/async_sglang_server.py index b3329a82..4526129b 100644 --- a/verl/workers/rollout/sglang_rollout/async_sglang_server.py +++ b/verl/workers/rollout/sglang_rollout/async_sglang_server.py @@ -693,6 +693,19 @@ class SGLangHttpServer: extra_fields["spec_num_accepted_tokens"] = int(meta_info.get("spec_accept_token_num", 0)) extra_fields["spec_num_verify_steps"] = int(meta_info.get("spec_verify_ct", 0)) + # Forward per-request timing metrics (present when enable_metrics=True) so + # tracing can fill Turn sglang_* fields (e2e/queue/ttft/prefill/decode). + for _k in ( + "e2e_latency", + "queue_time", + "prefill_finished_ts", + "request_received_ts", + "prefill_launch_latency", + ): + _v = meta_info.get(_k) + if isinstance(_v, (int, float)): + extra_fields[_k] = _v + return TokenOutput( token_ids=token_ids, log_probs=log_probs, -- 2.43.0