Skip to content

Commit efcad5d

Browse files
Teq2412claude
andcommitted
feat(slime): add SlimeRunner as the Python entry point for training
Wraps train.sh + config.yaml + env vars behind a single Python class. Users instantiate SlimeRunner with a handful of per-experiment fields (exp_id, agent_runtime_arn, s3_bucket, model_dir, data_path, model_type) and call .train(num_rollout=N); everything else has a sensible default and an extra_flags escape hatch for slime / Megatron / SGLang CLI passthrough. Internally, .train() shells out to reproduce train.sh step-by-step: stop stale sglang/ray, ray start --head, source the slime model script, submit the slime training job via ray job submit. The toolkit config (agent_runtime_arn, s3_bucket, exp_id, etc.) is written to a temp YAML and passed via --custom-config-path so SlimeArtConfig.from_args continues to read it. train.sh stays in the repo as the low-level escape hatch and as the reference for what the class replicates. SETUP.md 3.4 now leads with the Python recipe; the bash recipe is kept as "advanced / debugging." Verified end-to-end on Qwen2.5-3B-Instruct + GSM8K with .train(num_rollout=10): reward climbs 0.25 -> 0.68 over 10 rollouts, all 10 train steps healthy, no failures. Plan: docs/roadmap/committed/slime-runner-entrypoint.md. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
1 parent 051d7df commit efcad5d

5 files changed

Lines changed: 513 additions & 21 deletions

File tree

src/agentcore_rl_toolkit/backends/slime/SETUP.md

Lines changed: 67 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -149,38 +149,84 @@ cp .wandb.env.example .wandb.env # optional; skip to disable wandb
149149

150150
### 3.4 Run training
151151

152-
train.sh defaults target **8 × H100** (NUM_GPUS=8, TP_SIZE=2,
153-
ROLLOUT_GPUS_PER_ENGINE=2). For smaller clusters override via env
154-
(e.g. `NUM_GPUS=1 TP_SIZE=1 ROLLOUT_GPUS_PER_ENGINE=1` for a single
155-
GPU). Defaults also set `NUM_ROLLOUT=1` for smoke testing — bump to
156-
`NUM_ROLLOUT=100` (slime's production value) for a real run.
157-
158-
`SLIME_DIR` /
159-
`MEGATRON_DIR` need to point at the slime + Megatron-LM source trees
160-
(inside the `slimerl/slime:latest` container these are `/root/slime`
161-
and `/root/Megatron-LM`).
152+
Two entry points, same job under the hood:
153+
154+
- **`SlimeRunner` (Python)** — recommended. Write a short `train.py`;
155+
the class handles Ray/SGLang plumbing and slime CLI flags.
156+
- **`train.sh` (bash)** — escape hatch. Use when you need to override
157+
something the class doesn't surface, or to debug the raw slime CLI.
158+
159+
#### Python (`SlimeRunner`)
160+
161+
Defaults target **8 × H100** (`num_gpus=8`, `tp_size=2`,
162+
`rollout_gpus_per_engine=2`). Override kwargs for other cluster sizes.
163+
`.train(num_rollout=…)` defaults to 1 rollout for smoke testing — bump
164+
to 100 for a real run.
165+
166+
`slime_dir` / `megatron_dir` default to the in-container paths
167+
(`/root/slime` and `/root/Megatron-LM`); override for bare-metal
168+
installs.
169+
170+
```python
171+
# train.py — minimal 3B smoke test
172+
from agentcore_rl_toolkit.backends.slime import SlimeRunner
173+
174+
SlimeRunner(
175+
exp_id="gsm8k-3b-smoke",
176+
agent_runtime_arn="arn:aws:bedrock-agentcore:...",
177+
s3_bucket="your-bucket",
178+
model_dir="/path/to/Qwen2.5-3B-Instruct",
179+
data_path="/path/to/gsm8k_tiny.jsonl",
180+
model_type="qwen2.5-3B",
181+
).train(num_rollout=1)
182+
```
162183

163184
```bash
164185
cd /path/to/agentcore-rl-toolkit
186+
python src/agentcore_rl_toolkit/backends/slime/examples/math_agent/train.py
187+
```
188+
189+
32B on 8 GPUs:
190+
191+
```python
192+
SlimeRunner(
193+
exp_id="gsm8k-32b-run",
194+
agent_runtime_arn="arn:aws:bedrock-agentcore:...",
195+
s3_bucket="your-bucket",
196+
model_dir="/path/to/Qwen2.5-32B-Instruct",
197+
data_path="/path/to/gsm8k_tiny.jsonl",
198+
model_type="qwen2.5-32B",
199+
tp_size=8,
200+
rollout_gpus_per_engine=8,
201+
).train(num_rollout=5)
202+
```
203+
204+
Any slime/Megatron-LM/SGLang CLI flag that isn't surfaced as a named
205+
kwarg can be passed through `extra_flags`:
206+
207+
```python
208+
SlimeRunner(..., extra_flags=["--num-epoch", "3"]).train(num_rollout=50)
209+
```
165210

166-
# Qwen2.5-3B, 8 GPUs, 1 rollout (smoke test — train.sh defaults)
211+
If you prefer a YAML config, `SlimeRunner.from_yaml("config.yaml").train()`
212+
accepts the same keys.
213+
214+
#### Bash (`train.sh`) — escape hatch
215+
216+
`train.sh` takes the same knobs via env vars. It's kept as the low-level
217+
reference for what the Python class replicates, and as a debugging path
218+
for slime flag experiments.
219+
220+
```bash
221+
# 3B smoke test
167222
export SLIME_DIR=/root/slime \
168223
MEGATRON_DIR=/root/Megatron-LM \
169224
MODEL_DIR=/path/to/Qwen2.5-3B-Instruct \
170225
DATA_PATH=/path/to/gsm8k_tiny.jsonl
171226
bash src/agentcore_rl_toolkit/backends/slime/examples/math_agent/train.sh
172227
```
173228

174-
For 32B on 8 GPUs:
175-
176-
```bash
177-
export MODEL_DIR=/path/to/Qwen2.5-32B-Instruct \
178-
MODEL_TYPE=qwen2.5-32B \
179-
TP_SIZE=8 \
180-
ROLLOUT_GPUS_PER_ENGINE=8 \
181-
NUM_ROLLOUT=5
182-
bash src/agentcore_rl_toolkit/backends/slime/examples/math_agent/train.sh
183-
```
229+
For 32B on 8 GPUs, add `MODEL_TYPE=qwen2.5-32B TP_SIZE=8 ROLLOUT_GPUS_PER_ENGINE=8 NUM_ROLLOUT=5`.
184230

185231
### 3.5 Run evaluation
186232

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
"""Slime training backend for agentcore-rl-toolkit.
2+
3+
Primary entry point is :class:`SlimeRunner`. The ``integration/`` subpackage
4+
and ``patches/`` module are implementation detail — users shouldn't import
5+
them directly, but they are load-bearing (slime loads
6+
``agentcore_rl_toolkit.backends.slime.integration.rollout.generate_rollout``
7+
via ``--rollout-function-path`` at job-submit time).
8+
"""
9+
10+
from .runner import SlimeRunner
11+
12+
__all__ = ["SlimeRunner"]
Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
"""Train the strands math agent on GSM8K via slime — Python entry point.
2+
3+
Prerequisites (see ``SETUP.md`` for full instructions):
4+
- Inside a working slime environment (``slimerl/slime:latest`` container or equivalent).
5+
- Agent deployed to ACR; runtime ARN + S3 bucket handy.
6+
- Model checkpoint and training JSONL downloaded locally.
7+
8+
Run:
9+
python train.py
10+
"""
11+
12+
from agentcore_rl_toolkit.backends.slime import SlimeRunner
13+
14+
if __name__ == "__main__":
15+
SlimeRunner(
16+
exp_id="gsm8k-3b-smoke",
17+
agent_runtime_arn="arn:aws:bedrock-agentcore:<region>:<account>:runtime/<runtime-id>",
18+
s3_bucket="your-bucket-name",
19+
model_dir="/workspace/slime_workdir/models/Qwen2.5-3B-Instruct",
20+
data_path="/workspace/slime_workdir/data/gsm8k_tiny.jsonl",
21+
model_type="qwen2.5-3B",
22+
).train(num_rollout=1)

0 commit comments

Comments
 (0)