Skip to content

Commit 354633e

Browse files
HaoKang-Timmyergt10Hao Kangclaudehkang
authored
[example] ThunderAgent integration and R2Egym code agent rl training receipe (#1645)
## Summary This PR adds a ThunderAgent + SkyRL training recipe for fully async Harbor/R2EGym training with program-aware inference scheduling. The example lives under: - `examples/train/thunder_agent/` It includes: - a Harbor/R2EGym Qwen3-32B training entrypoint; - ThunderAgent-aware SkyRL inference routing; - rollout program release hooks; - weight-sync coordination around inference pause/resume; - Slurm scripts for the 5-node R2EGym 32B recipe; - dataset preparation and operational helpers; - router and scheduler regression tests; - documentation for the benchmark recipe and reproducible data split. ## Motivation ThunderAgent improves agent RL rollout throughput by scheduling inference at the program level rather than treating every request independently. This is especially useful for multi-turn agent workloads where trajectories have heterogeneous context lengths and execution times. This example demonstrates the integration on Harbor/R2EGym and documents the benchmark configuration used for the reported 40-step run. ## Key Files - `examples/train/thunder_agent/README.md` - End-to-end recipe documentation, hardware layout, data prep, launch flow, validation, and tunable parameters. - `examples/train/thunder_agent/main_harbor_thunder_agent.py` - Main fully async Harbor training entrypoint. - `examples/train/thunder_agent/main_thunder_agent.py` - Shared ThunderAgent experiment base that swaps SkyRL’s standard HTTP inference router for `ThunderAgentRouter`. - `examples/train/thunder_agent/skyrl_integration/router.py` - SkyRL integration router that composes ThunderAgent routes with SkyRL’s `/inference/v1/generate` and `/servers` endpoints. - `examples/train/thunder_agent/skyrl_integration/remote_inference_client.py` - Remote inference client wrapper that brackets pause/resume with ThunderAgent weight-sync begin/end calls. - `examples/train/thunder_agent/skyrl_integration/generator.py` - Harbor generator adapter that attaches routing IDs, releases ThunderAgent programs after trajectories, and preserves rollout logprobs when requested. - `examples/train/thunder_agent/scripts/r2egym_32b/` - Slurm and operational scripts for the R2EGym 32B recipe. - `tests/backends/skyrl_train/inference_servers/test_thunder_agent_router.py` - Router, program tracking, weight-sync, metadata stripping, and scheduler regression tests. ## Notes For Reviewers The recipe uses the new SkyRL HTTP inference path and requires: ```bash _SKYRL_USE_NEW_INFERENCE=1 The default documented benchmark layout is: - 1 merged head/rollout node with 8 H100s - Ray head - training driver - ThunderAgent proxy - 4 vLLM rollout servers at TP=2 - 4 trainer nodes with 8 H100s each - FSDP2 policy/ref training Total: 5 nodes, 40 H100 GPUs. The integration supports both: - embedded ThunderAgent router over external rollout servers; - externally provided ThunderAgent proxy via EXTERNAL_PROXY_URL / THUNDERAGENT_URL. --------- Co-authored-by: ergt10 <hshebswbd@gmail.com> Co-authored-by: Hao Kang <hkang@research-secure-hn.cloud.together.ai> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com> Co-authored-by: hkang <hkang@hkangs-MacBook-Air.local>
1 parent 59a8d4d commit 354633e

22 files changed

Lines changed: 4251 additions & 0 deletions

‎examples/train/README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ Welcome to the SkyRL-Train examples! In this folder you can find the following e
2929
- `flash_rl/`: Integration with [FlashRL’s](https://fengyao.notion.site/flash-rl) patched vLLM inference engine for high-throughput RL training.
3030
- `harbor/`: Custom [Harbor](https://harborframework.com/) Generator for training agents to solve TerminalBench tasks.
3131
- `mini_swe_agent/`: Integration with [Mini-SWE-Agent](https://github.com/SWE-agent/mini-swe-agent) to train coding agents on SWE-Bench via SkyRL.
32+
- `thunder_agent/`: ThunderAgent-accelerated Harbor/R2EGym training recipe with program-aware inference scheduling.
3233
- `../integrations/verifiers/`: Integration with PrimeIntellect's [Verifiers Library](https://github.com/PrimeIntellect-ai/verifiers) + [Environments Hub](https://app.primeintellect.ai/dashboard/environments?_gl=1*1vogwn8*_gcl_au*NjA1ODI2MTMxLjE3NjczOTkwMTM)
3334
- `../integrations/openenv/`: Integration with HuggingFace/Meta [OpenEnv](https://github.com/meta-pytorch/OpenEnv)
3435

‎examples/train/thunder_agent/README.md‎

Lines changed: 252 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,252 @@
1+
# ThunderAgent + SkyRL: R2EGym 32B Training Recipe
2+
3+
Train Qwen3-32B on R2EGym with ThunderAgent-accelerated rollout scheduling.
4+
5+
![TA vs no-TA stitched 40-step Harbor timeline](./docs/speedup.png)
6+
7+
**3.01× wall-clock speedup** over the no-TA baseline (8.84 h vs 26.58 h for 40 training steps / 10 epochs).
8+
9+
> **Scope.** This example uses SkyRL + Harbor integration but with MiniSWEAgent
10+
> rather than Terminus 2, and therefore is not token-in-token-out. The same
11+
> ThunderAgent setup should extend to Terminus 2 under Harbor, or to other
12+
> token-level agent harnesses, with minimal changes.
13+
14+
---
15+
16+
## Hardware
17+
18+
| Role | Nodes | GPUs | Notes |
19+
|---|---|---|---|
20+
| Head + rollout | 1 | 8 × H100 | Ray head, ThunderAgent proxy, training driver, 4 vLLM servers at TP=2 |
21+
| Trainer | 4 | 8 × H100 each | FSDP2 policy + ref model |
22+
23+
Total: 5 SLURM nodes, 40 H100 GPUs.
24+
25+
---
26+
27+
## Quick Start
28+
29+
Use the `sbatch` entrypoint for a fresh 5-node run. The lower-level stage
30+
wrapper is still available for debugging or for attaching to an existing
31+
allocation.
32+
33+
### Layout
34+
35+
The top level contains only the Python entrypoints, shared integration code,
36+
configs, docs, and script groups. The R2EGym 32B operational files are
37+
all under `scripts/r2egym_32b/`.
38+
39+
User-facing entrypoints:
40+
41+
| File | Purpose |
42+
|---|---|
43+
| `run_sbatch.sh` | One-command 5-node Slurm entrypoint |
44+
| `run_stages.sh` | Stage wrapper for existing allocations and retries |
45+
| `setup_env.sh` | Creates or validates the vLLM 0.20.1 CUDA 12.9 training venv |
46+
| `prepare_dataset.py` | Creates curated train/eval Harbor task directories |
47+
48+
Internal helpers called by `run_stages.sh`:
49+
50+
| File | Called by | Purpose |
51+
|---|---|---|
52+
| `start_rollout_servers.sh` | `run_stages.sh rollout` | Starts the four vLLM rollout servers on the rollout node |
53+
| `run_trainer.sh` | `run_stages.sh driver` | Runs the SkyRL training driver on the head node after Ray and rollout are ready |
54+
| `prepull_images.py` | `run_stages.sh head` | Lists or pre-pulls Docker images required by the curated tasks |
55+
| `cleanup_docker.sh` | `run_stages.sh cleanup-stage harbor_docker` | Cleans Docker containers and networks for this run |
56+
57+
The Harbor trial defaults live in `configs/harbor_trial/default.yaml`; reusable
58+
SkyRL integration code stays in `skyrl_integration/`.
59+
60+
### 1. Prepare Data
61+
62+
Download the four base difficulty-bucket datasets from HuggingFace:
63+
64+
```bash
65+
for BUCKET in trivial easy medium hard; do
66+
python examples/train_integrations/harbor/prepare_harbor_dataset.py \
67+
--dataset NovaSky-AI/r2egym-${BUCKET} \
68+
--output_dir ~/data/harbor/r2egym-${BUCKET}
69+
done
70+
```
71+
72+
Then generate the curated train/eval task directories used by this recipe:
73+
74+
```bash
75+
python examples/train/thunder_agent/scripts/r2egym_32b/prepare_dataset.py \
76+
--data-root ~/data/harbor
77+
```
78+
79+
This creates:
80+
- `~/data/harbor/r2egym-train256-medium-hard-v1/` with 256 task symlinks and `MANIFEST.json`
81+
- `~/data/harbor/r2egym-eval64-medium-hard-v1/` with 64 task symlinks and `MANIFEST.json`
82+
83+
### 2. Submit Training
84+
85+
```bash
86+
export REPO_ROOT=/path/to/skyrl-ta-pr-core
87+
cd "$REPO_ROOT"
88+
89+
# Required paths.
90+
export MODEL_PATH=/path/to/Qwen3-32B
91+
export DATA_ROOT=$HOME/data/harbor
92+
93+
# Keep the venv and wheel cache off small root filesystems when possible.
94+
export RECIPE_HOME=${SCRATCH:-$HOME/.cache}/skyrl-thunder-agent
95+
96+
# Optional run controls.
97+
export DOCKER_MODE=rootful
98+
export PREPULL_R2EGYM_IMAGES=true
99+
export ROLLOUT_ENFORCE_EAGER=true
100+
101+
sbatch \
102+
--partition=<partition> \
103+
--account=<account> \
104+
examples/train/thunder_agent/scripts/r2egym_32b/run_sbatch.sh
105+
```
106+
107+
The `sbatch` script requests 5 GPU nodes by default. In the default merged
108+
layout, node 0 runs Ray head, Docker setup, the training driver, and the 4
109+
rollout vLLM servers; nodes 1-4 are trainer workers. It then runs:
110+
111+
```text
112+
setup_env -> cleanup -> prepare -> head -> ray -> rollout -> status -> driver
113+
```
114+
115+
The wrapper defaults match the benchmark variant:
116+
117+
- `TRAIN_DATA=['$DATA_ROOT/r2egym-train256-medium-hard-v1']`
118+
- `EVAL_DATA=['$DATA_ROOT/r2egym-eval64-medium-hard-v1']`
119+
- `MAX_TRAIN_TASKS=256`, `MAX_EVAL_TASKS=64`
120+
- `FULL_EPOCHS=10`, `EVAL_INTERVAL_STEPS=4`, `CKPT_INTERVAL=4`
121+
- `USE_KL_LOSS=false`, `KL_LOSS_COEF=0.0`
122+
- `RUN_PREFLIGHT_CHECKS=false`, `AGENT_RUNTIME_PREFLIGHT=false`
123+
124+
The rollout launcher uses PR-core's native vLLM server module with
125+
`skyrl.backends.skyrl_train.inference_servers.vllm_worker.WorkerWrap` directly.
126+
127+
For CUDA 12.9 drivers, `scripts/r2egym_32b/setup_env.sh` installs the official
128+
vLLM release asset `vllm-0.20.1+cu129`; installing plain `vllm==0.20.1` from
129+
PyPI can select the CUDA 13 wheel.
130+
131+
`scripts/r2egym_32b/setup_env.sh` uses `uv` and installs an editable copy of
132+
the local repo into the training venv. Set `SETUP_ENV=check` when reusing an
133+
existing environment, or `SETUP_ENV=skip` when `PYTHON_BIN` and `RAY_BIN` are
134+
already exported.
135+
136+
`head` uses system Docker by default. If Docker Hub
137+
rate-limits anonymous pulls, run `docker login` on the head/rollout node before
138+
`head`, or keep `PREPULL_R2EGYM_IMAGES=true` so the wrapper fails early while
139+
pre-pulling the 320 curated train/eval task images instead of silently masking
140+
many Harbor trials during training. To inspect the exact image set:
141+
142+
```bash
143+
python examples/train/thunder_agent/scripts/r2egym_32b/prepull_images.py \
144+
--train-data "$TRAIN_DATA" \
145+
--eval-data "$EVAL_DATA" \
146+
--mode list
147+
```
148+
149+
For a temporary escape hatch while debugging non-Docker code paths, set
150+
`PREPULL_R2EGYM_IMAGES=false`; do not use that for a production recipe run.
151+
152+
### Existing Allocation
153+
154+
If you already have nodes allocated, use the stage wrapper directly. This is
155+
mainly useful for debugging failed stages without waiting for a new allocation.
156+
157+
```bash
158+
export WRAPPER="$REPO_ROOT/examples/train/thunder_agent/scripts/r2egym_32b/run_stages.sh"
159+
160+
export MERGED_JOB_ID=<slurm_job_id>
161+
export MERGED_NODE=<head_and_rollout_node>
162+
export TRAINER_NODE_SPECS="<trainer_node_1>:<job_id>,<trainer_node_2>:<job_id>,<trainer_node_3>:<job_id>,<trainer_node_4>:<job_id>"
163+
164+
bash "$WRAPPER" prepare
165+
bash "$WRAPPER" head
166+
bash "$WRAPPER" ray
167+
bash "$WRAPPER" rollout
168+
bash "$WRAPPER" status
169+
bash "$WRAPPER" driver
170+
```
171+
172+
For a single Slurm allocation, use the same `<job_id>` for all entries. The
173+
`sbatch` script constructs this format automatically.
174+
175+
### Validation
176+
177+
Before a full allocation, run the CPU-only checks:
178+
179+
```bash
180+
bash examples/train/thunder_agent/scripts/r2egym_32b/setup_env.sh check
181+
bash -n examples/train/thunder_agent/scripts/r2egym_32b/*.sh
182+
"$PYTHON_BIN" -m py_compile \
183+
examples/train/thunder_agent/main_harbor_thunder_agent.py \
184+
examples/train/thunder_agent/skyrl_integration/runtime_setup.py \
185+
examples/train/thunder_agent/scripts/r2egym_32b/prepull_images.py \
186+
examples/train/thunder_agent/scripts/r2egym_32b/prepare_dataset.py
187+
```
188+
189+
---
190+
191+
## Key Parameters
192+
193+
| Parameter | Default | Description |
194+
|---|---|---|
195+
| `MAX_TRAIN_TASKS` | 256 | Tasks per epoch |
196+
| `MAX_EVAL_TASKS` | 64 | Eval tasks |
197+
| `FULL_EPOCHS` | 10 | Training epochs |
198+
| `EVAL_INTERVAL_STEPS` | 4 | Steps between evals |
199+
| `HARBOR_AGENT_MAX_TURNS` | 25 | Max agent turns per trial |
200+
| `AGENT_TIMEOUT_SEC` | 9000 | Hard timeout per Harbor trial |
201+
| `MINI_SWE_MODEL_TIMEOUT_SEC` | 1200 | Per-LLM-call timeout |
202+
| `ROLLOUT_ENGINES` | 4 | Number of vLLM rollout servers |
203+
| `ROLLOUT_TP_SIZE` | 2 | Tensor-parallel size per server |
204+
| `THUNDER_AGENT_MODE` | `tr` | TA scheduler mode (`tr` = token-rate) |
205+
| `PREPULL_R2EGYM_IMAGES` | `true` | Pull all train/eval Harbor images during `head` |
206+
207+
All training hyperparameters can still be overridden through environment
208+
variables before invoking `run_sbatch.sh` or `run_stages.sh`.
209+
210+
For debugging only, `run_trainer.sh` can be called directly after Ray and the
211+
rollout servers are already up. Normal runs should go through `run_stages.sh`
212+
so node placement, rollout startup, readiness checks, and cleanup stay in one
213+
place:
214+
215+
```bash
216+
bash examples/train/thunder_agent/scripts/r2egym_32b/run_trainer.sh full \
217+
trainer.policy.optimizer_config.lr=5e-7 \
218+
harbor_trial_config.agent.kwargs.temperature=0.6
219+
```
220+
221+
---
222+
223+
## Environment Variables (URL auto-construction)
224+
225+
The run script accepts two input conventions for rollout server endpoints:
226+
227+
| Input | How to set |
228+
|---|---|
229+
| Direct URL list | `ROLLOUT_SERVER_URLS='["http://1.2.3.4:18000","http://1.2.3.4:18001",...]'` |
230+
| Host + port CSV | `ROLLOUT_HOST_IP=1.2.3.4` + `ROLLOUT_SERVER_PORTS_CSV=18000,18001,18002,18003` |
231+
232+
`EXTERNAL_PROXY_URL` (TA proxy) is resolved in this order:
233+
1. `EXTERNAL_PROXY_URL` if set explicitly
234+
2. `THUNDERAGENT_URL` if set (matches `run_harbor_benchmark.sh` convention)
235+
3. `http://$RAY_HEAD_IP:$SKYRL_INFERENCE_ROUTER_PORT` only when `USE_EXTERNAL_THUNDERAGENT_PROXY=1`
236+
4. Not set -> an embedded `ThunderAgentRouter` is started inside the trainer process
237+
238+
---
239+
240+
## Dataset Subsets (reproducible selection)
241+
242+
`scripts/r2egym_32b/prepare_dataset.py` reproduces the exact train/eval splits using a deterministic
243+
SHA-256 selection:
244+
245+
```
246+
seed: r2egym-medium-hard-v1-20260325
247+
train256: trivial=4 easy=16 medium=120 hard=116 (total 256)
248+
eval64: trivial=1 easy=4 medium=30 hard=29 (total 64)
249+
```
250+
251+
Train tasks do not overlap with eval tasks. Eval preserves `r2egym-eval32-medium-major-v1`
252+
as a prefix subset; train preserves `r2egym-train128-medium-major-v1` as a prefix.

‎examples/train/thunder_agent/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+

‎examples/train/thunder_agent/configs/harbor_trial/default.yaml‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
agent:
2+
name: mini-swe-agent
3+
model_name: hosted_vllm/Qwen3-32B
4+
kwargs:
5+
api_base: http://localhost:8080/v1
6+
max_turns: 20
7+
enable_summarize: false
8+
record_terminal_session: false
9+
store_all_messages: true
10+
temperature: 0.3
11+
environment:
12+
type: docker
13+
override_cpus: 2
14+
override_memory_mb: 4096
15+
override_storage_mb: 4096
16+
kwargs:
17+
auto_stop_interval_mins: null
18+
verifier:
19+
disable: false

‎examples/train/thunder_agent/docs/speedup.png‎

40.4 KB
Loading

0 commit comments

Comments
 (0)