Source: llm/slime
Scale Agentic RL on your own Infrastructure — slime on SkyPilot Job Groups#
An end-to-end RL example: train a coding agent (Qwen3-14B by default) with slime on SkyPilot Job Groups, scaling the components of the RL stack independently.
A Job Group gives us the pieces to scale slime to larger workloads:
Gang-scheduled jobs — the trainer cluster (Megatron, RolloutManager) and N inference clusters (SGLang engines) run as independent jobs that start, run, and stop together.
Discovery — jobs reach each other by stable hostname (
sglang-0.<group>), so the trainer cluster’s router sends rollout traffic and synchronization calls to each engine.Heterogeneous placement — this example uses H100s throughout, but jobs can sit on different GPUs (e.g. cheaper hardware for inference) via
resources.accelerators.
[!NOTE] This example also uses SkyPilot Sandboxes, to run the agent’s untrusted code and grade it for reward.
Architecture#
The trainer cluster runs the RolloutManager (agent rollouts + router) on CPU and
the trainer (Megatron, GRPO) on GPU; each inference cluster serves one SGLang
engine that the rollouts generate against.
A single agent loop looks like:
Claim a warm sandbox and check out the buggy commit.
Act — the agent (mini-swe-agent) loops: each turn generates through slime’s
OpenAIAdapter→ router → an engine (tokens + logprobs captured), then runs its tool call viasandbox.exec.Grade — reapply the agent’s diff over the hidden tests, run the repo’s test command, parse with SWE-smith; resolved ⇒ reward 1.
Emit —
finish_sessiondrains the trajectory into loss-masked training samples.
Requirements#
SkyPilot with Job Groups and Sandboxes enabled. Jobs authenticate back to the API server automatically, so no token setup is needed for Sandbox claims.
SkyPilot Sandboxes — fast, isolated rollout/eval pods on your own Kubernetes. Connecting to a Sandboxes-enabled API server ships the
sky.sandboxSDK wheel to~/.sky/bin/wheels/, which the trainer mounts and installs. Check withls ~/.sky/bin/wheels/— it should listskypilot_sandbox_sdk-*.whl.A Kubernetes cluster with 5–7 free H100s for the 1, 2, and 3 engine runs. SkyPilot places the jobs for you, whether that capacity is on one 8×H100 node or fragmented across several.
A ReadWriteMany (RWX) storage class.
Running it#
# One-time: the RWX volume the trainer + engine(s) share for disk weight-sync.
sky volumes apply policy-volume.yaml
# Optional: wandb logging. --secret picks the value up from your environment
# (and keeps it redacted in logs / dashboard / the stored request).
export WANDB_API_KEY=<key>
# Launch the disaggregated coding-agent RL Job Group.
sky jobs launch -n coding-agent coding-agent.yaml --secret WANDB_API_KEY
Scale the inference fleet to two or three engines — same Job Group, more SGLang jobs:
sky jobs launch -n coding-agent-2eng coding-agent-2engine.yaml --secret WANDB_API_KEY
sky jobs launch -n coding-agent-3eng coding-agent-3engine.yaml --secret WANDB_API_KEY
Each additional engine is an identical SGLang job in the YAML, plus its name in the trainer’s SGLANG_MEMBERS list.
Layout#
The rollout logic is code/generate.py (rollout fn),
code/sandbox_env.py (sandbox pool helper), and code/dataset_swesmith.py (dataset
loader).
Each job’s setup/run lives in scripts/ — engine-{setup,run}.sh
(SGLang server) and trainer-{setup,run}.sh (deps + weight-sync volume, engine discovery, sandbox pools).
trainer-run.sh calls run-slime.sh, which builds the slime command from env vars.
Configuring#
Every axis is an env var in the trainer’s YAML:
Env |
Meaning |
Default |
|---|---|---|
|
|
async |
|
|
disk |
|
|
delta |
|
concurrent rollouts (= sandbox claims); global batch is derived from these |
23 × 8 |
|
warm sandboxes per repo pool (total ≈ repos × this); ≥ concurrent rollouts for all-warm claims |
184 |
|
agent turns per instance |
20 |
|
repos to train on (one warm pool per repo image); empty |
markupsafe |
|
shuffle instances each step (keep on for real training); |
1 |
|
|
0 |
Warm pools vs. on-demand sandboxes. With pools on, the run pre-warms SANDBOX_POOL_REPLICAS
boxes per repo image (total ≈ repos × replicas) so claims hit a warm box. SANDBOX_NO_POOL=1 skips
pre-warming and creates an on-demand sandbox per claim (a cold image pull) instead. Either way, a
claim that can’t get a warm box falls back to on-demand, so a partial or slow warm never stalls the run.
Scaling the inference fleet on our own H100s#
We ran this example on our own H100s with a single inference cluster. slime’s built-in metrics make the bottleneck obvious: each step is mostly rollout (train ~400 s, rollout >1200 s), and the rollout is itself inference-bound. Generation requests pile up in a growing queue:
1 engine · #queue-req over time over 2 steps (peak 122): ▁▆███▇▆▅▄▃▂▁▁▂▇██▇▆▅▄▂▁▁
We used Job Groups to scale inference engines, and tracked the effect on end-to-end latency and queued requests. The following tables show results across engine counts and async vs sync mode.
Results
sync mode — rollout then train, so step ≈ rollout + train:
engines |
rollout (s/step) |
train (s/step) |
step (s) |
#queue-req (peak / median) |
|---|---|---|---|---|
1 |
1268 |
~400 |
1697 |
122 / 62 |
2 |
643 |
~400 |
1047 |
39 / 0 |
3 |
620 |
~400 |
1032 |
4 / 0 |
async mode — train overlaps the rollout, so step ≈ rollout:
engines |
rollout (s/step) |
train (s/step) |
step (s) |
#queue-req (peak / median) |
|---|---|---|---|---|
1 |
1215 |
~400 |
1200 |
125 / 68 |
2 |
805 |
~400 |
840 |
40 / 0 |
3 |
647 |
~400 |
661 |
9 / 0 |
Step time falls drastically with increased inference throughput by scaling engines from 1 → 3 (1200 → 661 s (≈1.8×) in async mode). Sync mode runs rollouts and training sequentially on every step, so total step time speedups are less pronounced than the rollout throughput gains. At 3 engines, the queue is empty for most of the run, and peak queue depth is far lower (4 sync / 9 async).
Numbers vary per run. Extracted metrics from all six runs are in benchmark/metrics.json.
Components#
Job Group (SkyPilot) — gang-schedules the trainer + inference clusters as one unit with stable inter-job hostnames.
slime, external mode — Megatron trainer + N external SGLang engines.
Sandboxes (SkyPilot) — warm pools of CPU pods that run the agent’s untrusted code and tests.
SWE-smith — real-repo bug-fixing dataset, a gold standard for agentic coding RL.
mini-swe-agent — a small, popular agent harness; drives the read → edit → run-tests loop.
Weight sync — after each optimizer step the trainer publishes weights and the engines reload, over a shared RWX volume (
disk) or NCCL;deltamode ships only the changed bytes.
Further work#
Autoscaling inference fleets — the fleet is fixed per run today. For long-horizon rollouts with variable latency, inference engines should scale up and down with demand.
Pre-cached sandbox images — pre-pull the repo images onto cluster nodes so many-image runs (the full multi-repo dataset, or
SANDBOX_NO_POOL=1) don’t pay a cold pull per claim.Cross-cluster jobs — place inference and training on separate clusters. Weight sync then needs an object-store carrier instead of a shared volume.
Included files#
.skyignore
.venv
__pycache__
*.pyc
*.jsonl
*.log
.git
*.egg-info
.pytest_cache
policy-volume*.yaml
benchmark/README.md
Benchmark data
metrics.json holds the load-bearing numbers behind the Scaling the inference
fleet section of the top-level README — the raw metrics we
extracted from the run logs, so the tables there are auditable.
Six runs: 1 / 2 / 3 SGLang engines × sync / async, Qwen3-14B, GRPO,
disaggregated (external SGLang, TP1 per engine) on a SkyPilot Job Group.
Batch 184 (ROLLOUT_BATCH_SIZE 23 × N_SAMPLES_PER_PROMPT 8), 2 rollout steps each.
Step 2 is steady state — step 1 warms sandbox pools + the prefix cache.
Keyed by {mode}_{engines}eng. Per run:
steps.{1,2}— slime’s nativeperf/*per-step timing (step_time,rollout_time,train_time,train_wait_time,update_weights_time).rollouts_completed/rollouts_aborted— from the example’s own rollout logs.sglang— per-engine saturation from the SGLang logs: peak/median#running-reqand#queue-req, peak KV-cache fraction, and 10-bucket (“decile”) traces across the run. The decile traces are a readability downsample of the raw per-decode-batch stream (thousands of timestamped samples/engine); finer granularity is available from the same logs.
Single 2-step runs, so per-step variance is ~±15% — read these for the scaling
shape, not exact multipliers. Every number is parsed from logs the run itself
produces — slime’s perf/* lines and the example’s rollout logs (trainer job
log) plus SGLang’s decode-batch lines (engine job logs) — so a rerun of the
README commands yields the same metrics from sky jobs logs (raw logs not shipped).
benchmark/metrics.json
{
"_meta": {
"what": "Load-bearing metrics behind the README 'Scaling the inference fleet' tables. 1/2/3 engines x sync/async, Qwen3-14B, GRPO, disaggregated SGLang (TP1/engine) on SkyPilot Job Groups.",
"batch": 184,
"batch_note": "ROLLOUT_BATCH_SIZE 23 x N_SAMPLES_PER_PROMPT 8 = 184 concurrent rollouts.",
"steps_per_run": 2,
"steady_state_step": 2,
"steady_state_note": "Step 2 = steady state; step 1 warms sandbox pools + prefix cache.",
"variance_note": "Single 2-step runs; per-step variance ~\u00b115%. Read the scaling shape, not exact multipliers.",
"source": "Extracted from the run's own logs: slime perf/* + the example's rollout logs (trainer job log) and SGLang's decode-batch lines (engine job logs)."
},
"runs": {
"sync_1eng": {
"mode": "sync",
"engines": 1,
"rollouts_completed": 366,
"rollouts_aborted": 2,
"steps": {
"1": {
"step_time_sec": 1869.0,
"rollout_time_sec": 1401.6,
"train_time_sec": 447.5,
"train_wait_time_sec": 1421.5,
"update_weights_time_sec": 18.7,
"wait_time_ratio": 0.8
},
"2": {
"step_time_sec": 1697.5,
"rollout_time_sec": 1267.9,
"train_time_sec": 396.2,
"train_wait_time_sec": 1301.3,
"update_weights_time_sec": 31.6,
"wait_time_ratio": 0.8
}
},
"sglang": {
"_note": "per-engine saturation. deciles = 10 buckets across the run (spans both steps), a readability downsample of the raw per-decode-batch stream (~1.6-3.3k timestamped samples/engine); finer granularity is available from the same logs if needed.",
"n_engines": 1,
"running_req_peak_per_eng": [
168
],
"running_req_median_per_eng": [
39
],
"queue_req_peak_per_eng": [
122
],
"queue_req_median_per_eng": [
62
],
"kv_cache_frac_peak_per_eng": [
0.98
],
"eng0_running_req_deciles": [
95.7,
48.28,
38.57,
34.19,
25.35,
84.11,
53.19,
40.93,
34.46,
21.05
],
"eng0_queue_req_deciles": [
55.67,
112.01,
95.68,
50.02,
19.44,
25.21,
106.42,
96.49,
47.8,
10.34
],
"eng0_kv_cache_frac_deciles": [
0.78,
0.94,
0.95,
0.95,
0.88,
0.59,
0.94,
0.94,
0.95,
0.74
]
}
},
"sync_2eng": {
"mode": "sync",
"engines": 2,
"rollouts_completed": 368,
"rollouts_aborted": 0,
"steps": {
"1": {
"step_time_sec": 1276.9,
"rollout_time_sec": 820.0,
"train_time_sec": 436.1,
"train_wait_time_sec": 840.8,
"update_weights_time_sec": 18.8,
"wait_time_ratio": 0.7
},
"2": {
"step_time_sec": 1047.1,
"rollout_time_sec": 643.3,
"train_time_sec": 371.1,
"train_wait_time_sec": 676.0,
"update_weights_time_sec": 31.0,
"wait_time_ratio": 0.6
}
},
"sglang": {
"_note": "per-engine saturation. deciles = 10 buckets across the run (spans both steps), a readability downsample of the raw per-decode-batch stream (~1.6-3.3k timestamped samples/engine); finer granularity is available from the same logs if needed.",
"n_engines": 2,
"running_req_peak_per_eng": [
85,
92
],
"running_req_median_per_eng": [
38,
37
],
"queue_req_peak_per_eng": [
39,
36
],
"queue_req_median_per_eng": [
0,
0
],
"kv_cache_frac_peak_per_eng": [
0.98,
0.98
],
"eng0_running_req_deciles": [
61.54,
50.24,
40.19,
29.96,
7.89,
22.47,
57.63,
39.96,
31.35,
7.81
],
"eng0_queue_req_deciles": [
3.9,
28.22,
24.01,
0.95,
0.0,
0.0,
12.8,
28.02,
5.67,
0.0
],
"eng0_kv_cache_frac_deciles": [
0.5,
0.94,
0.94,
0.8,
0.29,
0.13,
0.78,
0.95,
0.84,
0.29
]
}
},
"sync_3eng": {
"mode": "sync",
"engines": 3,
"rollouts_completed": 366,
"rollouts_aborted": 2,
"steps": {
"1": {
"step_time_sec": 1039.2,
"rollout_time_sec": 662.1,
"train_time_sec": 357.4,
"train_wait_time_sec": 681.8,
"update_weights_time_sec": 18.7,
"wait_time_ratio": 0.7
},
"2": {
"step_time_sec": 1031.6,
"rollout_time_sec": 620.0,
"train_time_sec": 377.7,
"train_wait_time_sec": 653.9,
"update_weights_time_sec": 32.2,
"wait_time_ratio": 0.6
}
},
"sglang": {
"_note": "per-engine saturation. deciles = 10 buckets across the run (spans both steps), a readability downsample of the raw per-decode-batch stream (~1.6-3.3k timestamped samples/engine); finer granularity is available from the same logs if needed.",
"n_engines": 3,
"running_req_peak_per_eng": [
54,
55,
51
],
"running_req_median_per_eng": [
25,
24,
25
],
"queue_req_peak_per_eng": [
3,
3,
4
],
"queue_req_median_per_eng": [
0,
0,
0
],
"kv_cache_frac_peak_per_eng": [
0.97,
0.97,
0.99
],
"eng0_running_req_deciles": [
24.06,
35.73,
30.81,
22.59,
6.92,
5.52,
34.45,
38.94,
25.85,
8.08
],
"eng0_queue_req_deciles": [
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.02,
0.1,
0.0
],
"eng0_kv_cache_frac_deciles": [
0.14,
0.46,
0.69,
0.56,
0.18,
0.05,
0.33,
0.79,
0.77,
0.3
]
}
},
"async_1eng": {
"mode": "async",
"engines": 1,
"rollouts_completed": 484,
"rollouts_aborted": 89,
"steps": {
"1": {
"step_time_sec": 1827.9,
"rollout_time_sec": 1388.8,
"train_time_sec": 421.1,
"train_wait_time_sec": 1406.8,
"update_weights_time_sec": 16.8,
"wait_time_ratio": 0.8
},
"2": {
"step_time_sec": 1200.4,
"rollout_time_sec": 1215.0,
"train_time_sec": 373.3,
"train_wait_time_sec": 827.1,
"update_weights_time_sec": 32.0,
"wait_time_ratio": 0.7
}
},
"sglang": {
"_note": "per-engine saturation. deciles = 10 buckets across the run (spans both steps), a readability downsample of the raw per-decode-batch stream (~1.6-3.3k timestamped samples/engine); finer granularity is available from the same logs if needed.",
"n_engines": 1,
"running_req_peak_per_eng": [
167
],
"running_req_median_per_eng": [
41
],
"queue_req_peak_per_eng": [
125
],
"queue_req_median_per_eng": [
68
],
"kv_cache_frac_peak_per_eng": [
0.99
],
"eng0_running_req_deciles": [
79.49,
41.87,
34.41,
35.72,
74.05,
43.52,
32.07,
71.73,
48.53,
37.19
],
"eng0_queue_req_deciles": [
70.93,
104.64,
49.68,
6.45,
78.77,
93.33,
30.03,
22.02,
105.98,
74.78
],
"eng0_kv_cache_frac_deciles": [
0.82,
0.94,
0.95,
0.57,
0.92,
0.94,
0.94,
0.6,
0.94,
0.95
]
}
},
"async_2eng": {
"mode": "async",
"engines": 2,
"rollouts_completed": 367,
"rollouts_aborted": 1,
"steps": {
"1": {
"step_time_sec": 1189.0,
"rollout_time_sec": 749.1,
"train_time_sec": 420.0,
"train_wait_time_sec": 769.0,
"update_weights_time_sec": 18.9,
"wait_time_ratio": 0.6
},
"2": {
"step_time_sec": 840.3,
"rollout_time_sec": 804.8,
"train_time_sec": 421.9,
"train_wait_time_sec": 418.4,
"update_weights_time_sec": 32.3,
"wait_time_ratio": 0.5
}
},
"sglang": {
"_note": "per-engine saturation. deciles = 10 buckets across the run (spans both steps), a readability downsample of the raw per-decode-batch stream (~1.6-3.3k timestamped samples/engine); finer granularity is available from the same logs if needed.",
"n_engines": 2,
"running_req_peak_per_eng": [
82,
76
],
"running_req_median_per_eng": [
40,
36
],
"queue_req_peak_per_eng": [
31,
40
],
"queue_req_median_per_eng": [
0,
0
],
"kv_cache_frac_peak_per_eng": [
0.98,
0.98
],
"eng0_running_req_deciles": [
65.56,
52.17,
42.2,
24.59,
3.99,
57.92,
47.99,
40.25,
26.63,
5.94
],
"eng0_queue_req_deciles": [
1.3,
18.96,
11.11,
0.08,
0.0,
1.32,
17.3,
11.71,
0.62,
0.0
],
"eng0_kv_cache_frac_deciles": [
0.48,
0.94,
0.95,
0.71,
0.15,
0.48,
0.94,
0.95,
0.74,
0.19
]
}
},
"async_3eng": {
"mode": "async",
"engines": 3,
"rollouts_completed": 367,
"rollouts_aborted": 1,
"steps": {
"1": {
"step_time_sec": 1049.5,
"rollout_time_sec": 610.9,
"train_time_sec": 418.1,
"train_wait_time_sec": 631.5,
"update_weights_time_sec": 19.6,
"wait_time_ratio": 0.6
},
"2": {
"step_time_sec": 660.8,
"rollout_time_sec": 646.7,
"train_time_sec": 393.7,
"train_wait_time_sec": 267.1,
"update_weights_time_sec": 37.2,
"wait_time_ratio": 0.4
}
},
"sglang": {
"_note": "per-engine saturation. deciles = 10 buckets across the run (spans both steps), a readability downsample of the raw per-decode-batch stream (~1.6-3.3k timestamped samples/engine); finer granularity is available from the same logs if needed.",
"n_engines": 3,
"running_req_peak_per_eng": [
53,
59,
55
],
"running_req_median_per_eng": [
28,
27,
25
],
"queue_req_peak_per_eng": [
9,
5,
4
],
"queue_req_median_per_eng": [
0,
0,
0
],
"kv_cache_frac_peak_per_eng": [
0.97,
0.97,
0.97
],
"eng0_running_req_deciles": [
33.35,
41.84,
31.96,
15.17,
4.34,
29.63,
41.06,
32.24,
16.7,
3.21
],
"eng0_queue_req_deciles": [
0.0,
1.92,
0.49,
0.0,
0.0,
0.0,
0.38,
0.56,
0.0,
0.0
],
"eng0_kv_cache_frac_deciles": [
0.22,
0.91,
0.78,
0.47,
0.16,
0.22,
0.86,
0.86,
0.51,
0.1
]
}
}
}
}
code/dataset_swesmith.py
"""SWE-smith -> a dataset-agnostic sample schema (the headline workload).
SWE-smith (SWE-bench/SWE-smith on HF, 50k instances / 128 repos) is real-repo
bug-fixing with executable per-repo environments. This loader maps its instances
onto the SAME schema generate.py already consumes for MBPP (see generate.py
_get_metadata), so the rollout fn / adapter / masking are UNCHANGED — only the
dataset, the images, and the grader differ.
Verified instance fields (huggingface.co/datasets/SWE-bench/SWE-smith):
instance_id, repo, patch (GOLD fix), FAIL_TO_PASS, PASS_TO_PASS,
base_commit, image_name (per-REPO env image), problem_statement, created_at
Mapping to sample metadata:
task <- problem_statement
image <- image_name (per-repo; drives one warm pool per repo)
pool <- sanitized repo name (pool-per-repo)
workdir <- /testbed (swesmith ENV_NAME; SWE-bench lineage)
setup_cmd <- establish the BUGGY state in /testbed (see OPEN QUESTION)
eval_files <- __eval__.py grader + spec.json (F2P/P2P/repo)
eval_cmd <- python __eval__.py (exit 0 == resolved)
gold_patch <- patch (for the golden-patch harness self-check)
Three things differ from MBPP:
1. PER-REPO POOLS/IMAGES: emit_pools() lists the distinct (pool, image) pairs
so the launch YAML can create one warm pool per repo image.
2. BUGGY-STATE SETUP: the per-repo image is shared across many bugs, so the
specific bug must be applied per rollout BEFORE the agent runs, via the new
metadata["setup_cmd"] hook generate.py runs post-claim, pre-agent.
RESOLVED (swesmith/profiles/base.py): the bug is a git BRANCH named exactly
`instance_id` in the mirror baked into the per-repo image →
`git -C /testbed checkout -f {instance_id}`.
3. GRADER via swesmith's OWN registry (don't reinvent per-repo test commands).
RESOLVED (base.py): registry.get_from_inst(inst) → get_test_cmd(inst) →
log_parser(output); pass == swebench TestStatus.PASSED. The staged
__eval__.py uses exactly this (needs swesmith+swebench importable in-sandbox).
"""
from __future__ import annotations
import argparse
import json
import re
WORKDIR = "/testbed"
def _pool_name(repo: str, suffix: str = "") -> str:
"""One pool per repo image. Sanitize 'owner/name' -> a k8s-safe pool id.
``suffix`` (a run-unique token, e.g. the Job Group name) makes the pool name
per-RUN so each run OWNS its own pool — creates it, uses it, deletes it at end
(self-cleanup, no cross-run sharing/leak). Omit for the shared-by-repo behavior.
"""
base = "swesmith-" + re.sub(r"[^a-z0-9]+", "-", repo.lower()).strip("-")
if suffix:
suf = re.sub(r"[^a-z0-9]+", "-", suffix.lower()).strip("-")
# Server caps the pool/template name at 58 chars (RFC1123 63 - len("-pool"),
# see sandbox/core.py) — NOT 63. Reserve room for the suffix so the
# run-unique part is NEVER truncated (else two runs' pools could collide).
base = base[:58 - len(suf) - 1].strip("-") + "-" + suf
return base[:58].strip("-")
# Reward is graded HOST-side in generate._evaluate_swesmith (swesmith/swebench
# are NOT in the env images — proven by the golden-patch smoke). So the dataset
# precomputes the per-repo test command here (get_test_cmd, deterministic per
# instance) and carries the instance fields log_parser needs; generate.py runs
# the cmd in the sandbox and parses in-process. API from swesmith/profiles/base.py:
# cmd, _ = registry.get_from_inst(inst).get_test_cmd(inst)
def _test_cmd_for(inst: dict) -> str:
"""Precompute the repo test command via swesmith's registry (host-side, at
dataset-gen — the trainer has swesmith installed). Deterministic per instance."""
from swesmith.profiles import registry
cmd, _ = registry.get_from_inst(inst).get_test_cmd(inst)
return cmd
def row_to_sample(row: dict, pool_suffix: str = "") -> dict:
repo = row["repo"]
instance_id = row["instance_id"]
# The instance fields generate._evaluate_swesmith needs for host-side parsing.
inst = {
"instance_id": instance_id,
"repo": repo,
"image_name": row.get("image_name", ""),
"base_commit": row.get("base_commit", ""),
"FAIL_TO_PASS": list(row.get("FAIL_TO_PASS") or []),
"PASS_TO_PASS": list(row.get("PASS_TO_PASS") or []),
}
return {
"text": row["problem_statement"],
"label": instance_id,
"metadata": {
"instance_id": instance_id,
"task": row["problem_statement"],
"repo": repo,
"base_commit": row.get("base_commit", ""),
"image": row["image_name"],
"pool": _pool_name(repo, pool_suffix),
"workdir": WORKDIR,
# O1 RESOLVED (swesmith/profiles/base.py): the bug lives on a git
# BRANCH named exactly instance_id in the repo mirror baked into the
# per-repo image. `git checkout <instance_id>` reaches the buggy state
# (bug present, hidden tests removed — agent works blind to tests).
"setup_cmd": f"git -C {WORKDIR} checkout -f {instance_id}",
"files": {
}, # repo pre-baked in the image; nothing to stage pre-loop
# Host-parsed grading: generate.py reapplies the
# agent's fix onto HEAD~1 (tests present), runs test_cmd, parses via
# swesmith log_parser. No in-sandbox grader.
"eval_kind": "swesmith",
"test_cmd": _test_cmd_for(inst),
"swesmith_inst": inst,
"gold_patch": row.get(
"patch", ""), # golden-patch harness self-check (smoke)
},
}
def emit_pools(rows: list[dict]) -> list[dict]:
"""Distinct (pool, image) pairs for warm-pool creation (one per repo image)."""
seen: dict[str, str] = {}
for r in rows:
m = r["metadata"]
seen.setdefault(m["pool"], m["image"])
return [{"pool": p, "image": img} for p, img in seen.items()]
def main() -> None:
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--out", required=True, help="output JSONL path")
ap.add_argument("--pools-out",
help="optional: write distinct (pool,image) JSON here")
ap.add_argument("--dataset", default="SWE-bench/SWE-smith")
ap.add_argument("--split", default="train")
ap.add_argument("--repos",
nargs="*",
default=None,
help="subset to these repos (owner/name); default = all")
ap.add_argument("--per-repo",
type=int,
default=0,
help="cap instances per repo (0 = all)")
ap.add_argument("--limit", type=int, default=0, help="global cap (0 = all)")
ap.add_argument(
"--pool-suffix",
default="",
help=
"run-unique token appended to pool names → each run owns+cleans its own pool"
)
args = ap.parse_args()
from datasets import load_dataset # deferred: heavy import
ds = load_dataset(args.dataset, split=args.split)
# SUBSTRING match (case-insensitive): pass bare tokens ("markupsafe"); the
# dataset `repo` is a mirror string, not "owner/name".
tokens = [t.lower() for t in (args.repos or [])]
per_repo_count: dict[str, int] = {}
rows: list[dict] = []
for row in ds:
row = dict(row)
if tokens and not any(t in row["repo"].lower() for t in tokens):
continue
if args.per_repo:
c = per_repo_count.get(row["repo"], 0)
if c >= args.per_repo:
continue
per_repo_count[row["repo"]] = c + 1
rows.append(row_to_sample(row, args.pool_suffix))
if args.limit and len(rows) >= args.limit:
break
with open(args.out, "w", encoding="utf-8") as f:
for r in rows:
f.write(json.dumps(r, ensure_ascii=False) + "\n")
pools = emit_pools(rows)
if args.pools_out:
with open(args.pools_out, "w", encoding="utf-8") as f:
json.dump(pools, f)
print(f"wrote {len(rows)} samples to {args.out}; {len(pools)} pools: " +
", ".join(p["pool"] for p in pools))
if __name__ == "__main__":
main()
code/generate.py
"""Coding-agent RL: per-sample generate() for slime over SkyPilot sandboxes.
--custom-generate-function-path generate.generate (code/ on PYTHONPATH)
Adapted from slime/examples/coding_agent_rl/generate.py with one relocation:
the agent runs TRAINER-SIDE (this process), not as a CLI inside the sandbox.
The agent IS upstream mini-swe-agent (DefaultAgent + LitellmTextbasedModel +
the stock mini_textbased.yaml templates); we implement only its two protocol
seams: Environment = sky.sandbox exec (sandbox_env.SandboxEnvironment) and
Model = litellm pointed at slime's OpenAIAdapter on localhost with
api_key=session_id (litellm sends it as Authorization: Bearer, which is
exactly the adapter's session routing). Per sample:
claim sandbox from warm pool -> stage task files -> DefaultAgent.run in a
thread (model turns: adapter; actions: sandbox exec) -> stage eval files
-> run eval_cmd (exit 0 => reward 1.0) -> finish_session drains the
adapter trajectory into loss-masked Sample(s)
DefaultAgent is synchronous, so each in-flight rollout runs its loop via
asyncio.to_thread (thread-per-inflight-rollout; fine at our batch sizes).
The sandbox executes ONLY untrusted code; token capture and observation
loss-masking live entirely in the adapter/TrajectoryManager. Crashes and
timeouts return the abort-sample shape (mirrors coding_agent_rl's
_abort_result) and the sandbox is always terminated in finally.
Dataset-agnostic sample schema (see _get_metadata; produced by e.g.
dataset_mbpp.py): metadata carries task, files, eval_files, eval_cmd, pool,
workdir, instance_id.
"""
from __future__ import annotations
import asyncio
import concurrent.futures
import dataclasses
import logging
import os
import secrets
import time
import traceback
from typing import Any
os.environ.setdefault("MSWEA_SILENT_STARTUP",
"1") # upstream knob: no banner per Ray worker
from minisweagent import package_dir as _miniswe_package_dir
from minisweagent.agents.default import DefaultAgent
from minisweagent.environments.local import LocalEnvironment
from minisweagent.exceptions import Submitted
from minisweagent.models.litellm_textbased_model import LitellmTextbasedModel
import sandbox_env
from slime.agent.adapters import OpenAIAdapter
from slime.agent.aiohttp_threaded import FilteredAccessLogger
from slime.agent.aiohttp_threaded import run_app_in_thread
from slime.utils.misc import SingletonMeta
from slime.utils.processing_utils import load_tokenizer
from slime.utils.types import Sample
import yaml
logger = logging.getLogger(__name__)
# --- rollout thread pool ------------------------------------------------------
# Each rollout runs its whole synchronous agent loop (staging, agent.run, eval)
# via asyncio.to_thread, which dispatches to the event loop's DEFAULT
# ThreadPoolExecutor. That default is min(32, cpu+4) workers -> at a rollout batch
# > 32, only ~32 rollouts run at once no matter how many are launched, so
# concurrent generation pins at ~32 and the inference engine starves (queue empty,
# KV low) regardless of pool size / batch / engine count. Size the pool to the
# batch so every in-flight rollout runs concurrently. Set once, on first generate().
_ROLLOUT_EXECUTOR: concurrent.futures.ThreadPoolExecutor | None = None
def _ensure_rollout_executor() -> None:
global _ROLLOUT_EXECUTOR
if _ROLLOUT_EXECUTOR is not None:
return
def _int_env(name, default): # robust to unset / empty-string / malformed
try:
return int(os.environ.get(name) or default)
except (TypeError, ValueError):
return default
batch = _int_env("GLOBAL_BATCH_SIZE",
0) or (_int_env("ROLLOUT_BATCH_SIZE", 8) *
_int_env("N_SAMPLES_PER_PROMPT", 8))
max_workers = max(64,
batch + 16) # one thread per in-flight rollout + headroom
_ROLLOUT_EXECUTOR = concurrent.futures.ThreadPoolExecutor(
max_workers=max_workers, thread_name_prefix="rollout")
asyncio.get_running_loop().set_default_executor(_ROLLOUT_EXECUTOR)
logger.info(
"[agent] rollout thread pool: max_workers=%d (default was min(32,cpu+4)); batch=%d",
max_workers,
batch,
)
# Upstream config, verbatim (system/instance templates, mswea_bash_command
# fence, observation template with 10k elision, format-error template).
# Changing any of it later is a config override, not a fork.
_MINI_CONFIG = yaml.safe_load(
(_miniswe_package_dir / "config" / "mini_textbased.yaml").read_text())
@dataclasses.dataclass(frozen=True)
class AgentConfig:
litellm_model_name: str # "openai/..." = litellm's generic OpenAI-compatible provider
max_turns: int # failure-tail ceiling only
exec_timeout_sec: int
eval_timeout_sec: int
rollout_guard_sec: int
sandbox_lifetime_sec: int # server-side auto-reap; set well above the rollout guard so a slow tail is never reaped mid-rollout (it's only a leak backstop)
adapter_port: int # 0 = ephemeral; localhost-only, so no fixed port needed
sglang_url: str | None # override for the adapter's upstream; see _AdapterService
default_pool: str
default_workdir: str
local_exec: bool # dev-only: upstream LocalEnvironment instead of sandboxes
@classmethod
def from_env(cls) -> AgentConfig:
max_turns = int(os.environ.get("AGENT_MAX_TURNS", "5"))
exec_timeout = int(os.environ.get("AGENT_EXEC_TIMEOUT_SEC", "60"))
eval_timeout = int(os.environ.get("AGENT_EVAL_TIMEOUT_SEC", "120"))
# guard covers worst case: every turn hits the exec timeout, plus eval,
# plus slack for model latency and sandbox claim.
guard = int(os.environ.get("AGENT_ROLLOUT_GUARD_SEC", "0") or
0) or (max_turns * exec_timeout + eval_timeout + 300)
return cls(
litellm_model_name=os.environ.get("AGENT_LITELLM_MODEL",
"openai/slime-actor"),
max_turns=max_turns,
exec_timeout_sec=exec_timeout,
eval_timeout_sec=eval_timeout,
rollout_guard_sec=guard,
# Leak backstop only (a crashed trainer's boxes self-destruct after this).
# Sit it an hour past the worst-case rollout guard so a slow-tail rollout is
# never reaped mid-flight — but bounded, so a hard trainer crash doesn't leak
# boxes for hours and starve the next run's pool.
sandbox_lifetime_sec=int(
os.environ.get("AGENT_SANDBOX_LIFETIME_SEC", "0") or 0) or
(guard + 3600),
adapter_port=int(os.environ.get("ADAPTER_PORT", "0")),
sglang_url=os.environ.get("AGENT_SGLANG_URL") or None,
default_pool=os.environ.get("AGENT_POOL", "swesmith"),
default_workdir=os.environ.get("AGENT_WORKDIR", "/workspace"),
local_exec=os.environ.get("AGENT_LOCAL_EXEC", "") == "1",
)
CONFIG = AgentConfig.from_env()
class _AdapterService(metaclass=SingletonMeta):
"""One OpenAIAdapter per rollout-actor process, bound to localhost.
Unlike coding_agent_rl's service, no ADAPTER_PUBLIC_HOST: the only client
is the in-process litellm model, so 127.0.0.1 with an ephemeral port
suffices (and nothing outside the pod can reach the policy endpoint).
"""
def __init__(self, args) -> None:
self.tokenizer = load_tokenizer(args.hf_checkpoint,
trust_remote_code=True)
self.max_context_len = int(
getattr(args, "rollout_max_context_len", 0) or 0)
# args.sglang_router_ip/port DOES resolve under
# --rollout-external-engine-addrs: slime still starts a trainer-local
# sglang_router and writes its ip/port back onto args
# (slime/backends/sglang_utils/external.py:185-187, dispatched from
# slime/ray/rollout.py:1103-1104 inside RolloutManager.__init__,
# rollout.py:436 — i.e. before any rollout runs in this process), and
# the external engine registers itself as a router worker
# (sglang_engine.py:190-199). slime's own rollout path posts to the
# same URL (sglang_rollout.py:81,159). AGENT_SGLANG_URL is a debug
# escape hatch to bypass the router and hit the engine directly
# (e.g. http://sglang-0.$SKYPILOT_JOBGROUP_NAME:30000); leave unset
# in normal runs.
sglang_url = CONFIG.sglang_url or f"http://{args.sglang_router_ip}:{args.sglang_router_port}"
self.adapter = OpenAIAdapter(
tokenizer=self.tokenizer,
sglang_url=sglang_url,
tool_parser=getattr(args, "sglang_tool_call_parser", None) or None,
reasoning_parser=getattr(args, "sglang_reasoning_parser", None) or
None,
)
self.app_handle = run_app_in_thread(
self.adapter.app,
host="127.0.0.1",
port=CONFIG.adapter_port,
thread_name="openai-adapter",
# client disconnect must cancel the handler so the adapter's
# fire-and-forget /abort_request frees the sglang slot (see the
# comment in coding_agent_rl/generate.py:151-154).
runner_kwargs={
"handler_cancellation": True,
"access_log_class": FilteredAccessLogger
},
)
self.adapter_url = f"http://127.0.0.1:{self.app_handle.port}"
logger.info(
"[agent] tokenizer=%s adapter=%s sglang_upstream=%s max_context_len=%s",
args.hf_checkpoint,
self.adapter_url,
sglang_url,
self.max_context_len,
)
def _make_model(adapter_url: str, session_id: str) -> LitellmTextbasedModel:
"""Upstream litellm text-based model pointed at the adapter.
api_key=session_id is the whole session-routing story: litellm's OpenAI
provider sends it as ``Authorization: Bearer <sid>``, which the adapter
resolves via sid_from_bearer (slime/agent/adapters/common.py:367).
Sampling params (temperature, max_new_tokens, ...) come from the
session's defaults registered in open_session, not from litellm kwargs.
"""
model_cfg = dict(_MINI_CONFIG.get("model") or {})
model_kwargs = {
**(model_cfg.pop("model_kwargs", None) or {}), # upstream: drop_params etc.
"api_base": f"{adapter_url}/v1",
"api_key": session_id,
"timeout": 900, # generous per-turn read timeout; the rollout guard bounds the total
}
return LitellmTextbasedModel(
model_name=CONFIG.litellm_model_name,
model_kwargs=model_kwargs,
# "slime-actor" has no litellm pricing entry; cost is meaningless here.
cost_tracking="ignore_errors",
**model_cfg, # upstream observation_template + format_error_template
)
def _make_env(sb, workdir: str):
"""Environment seam: sandbox-backed, or upstream LocalEnvironment (dev-only)."""
env_vars = dict((_MINI_CONFIG.get("environment") or {}).get("env") or
{}) # PAGER=cat etc.
if sb is None:
return LocalEnvironment(cwd=workdir,
env=env_vars,
timeout=CONFIG.exec_timeout_sec)
return sandbox_env.SandboxEnvironment(sb,
cwd=workdir,
env=env_vars,
timeout=CONFIG.exec_timeout_sec)
def _make_agent(model, env) -> DefaultAgent:
agent_cfg = dict(_MINI_CONFIG.get("agent") or
{}) # upstream templates, verbatim
agent_cfg.pop("mode",
None) # interactive-CLI knob; not an AgentConfig field
agent_cfg["step_limit"] = CONFIG.max_turns
agent_cfg[
"cost_limit"] = 0 # disabled: adapter turns cost $0; step limit + guard cap the run
return DefaultAgent(model, env, **agent_cfg)
def _get_metadata(sample: Sample) -> dict[str, Any]:
"""Normalize the dataset row schema (see module docstring)."""
m = sample.metadata or {}
return {
"instance_id": str(
m.get("instance_id") or sample.label or sample.index or "unknown"),
"pool": m.get("pool") or CONFIG.default_pool,
"image":
m.get("image") or
None, # per-repo env image: on-demand creates (no-pool mode / pool-claim fallback)
"workdir": m.get("workdir") or CONFIG.default_workdir,
"task": m.get("task") or _coerce_prompt(sample.prompt),
"files": m.get("files") or {}, # staged before the agent loop
"setup_cmd": m.get(
"setup_cmd"
), # optional: run post-claim, pre-agent (SWE-smith buggy-state)
"eval_kind": m.get("eval_kind") or
"cmd", # "cmd" (MBPP eval_cmd) | "swesmith" (host-parse)
"eval_files": m.get("eval_files")
or {}, # staged AFTER the loop (hidden graders)
"eval_cmd": m.get("eval_cmd"), # exit 0 == solved (cmd kind)
"test_cmd": m.get("test_cmd"), # repo test command (swesmith kind)
"swesmith_inst":
m.get("swesmith_inst"), # F2P/P2P/repo for host-side log_parser
}
def _coerce_prompt(prompt) -> str:
# Same fallback as coding_agent_rl/swe.py:_coerce_prompt: plain string or
# first user message of a conversation.
if isinstance(prompt, str):
return prompt
if isinstance(prompt, list):
for m in prompt:
if isinstance(m, dict) and m.get("role") == "user" and isinstance(
m.get("content"), str):
return m["content"]
return ""
def _stage_files(env, files: dict[str, str], workdir: str) -> None:
"""Sync staging helper (runs inside asyncio.to_thread)."""
if isinstance(env, sandbox_env.SandboxEnvironment):
env.write_files(files)
else: # upstream LocalEnvironment (AGENT_LOCAL_EXEC)
sandbox_env.stage_files_locally(workdir, files)
class _TimedEnv:
"""Transparent wrapper that accumulates per-action exec wall-time.
Wraps EITHER our SandboxEnvironment or upstream LocalEnvironment so total
sandbox-exec time is captured uniformly (a counter inside SandboxEnvironment
would miss the local-exec smoke). mini-swe's DefaultAgent only calls
``execute`` / ``get_template_vars`` / ``serialize`` and reads ``.config``;
everything but ``execute`` delegates to the inner env via ``__getattr__``.
"""
def __init__(self, inner) -> None:
self.inner = inner
self.exec_seconds = 0.0
self.exec_count = 0
self.exec_durations: list[float] = [
] # per-call wall-times (for p50/p95/p99)
def execute(self, *args, **kwargs):
t = time.perf_counter()
try:
return self.inner.execute(*args, **kwargs)
finally:
# ``finally`` so a Submitted-raising action still counts its time
dt = time.perf_counter() - t
self.exec_seconds += dt
self.exec_count += 1
self.exec_durations.append(round(dt, 3))
def __getattr__(self, name): # config, get_template_vars, serialize, ...
return getattr(self.inner, name)
def _evaluate(env, eval_cmd: str, timeout_sec: int) -> float:
"""Binary reward: eval_cmd exit 0 => 1.0; crash/timeout/nonzero => 0.0.
Sync (runs inside asyncio.to_thread). If the test suite's stdout happens
to start with the submit sentinel, env.execute raises Submitted; that
requires returncode 0, so it still counts as a pass.
"""
try:
output = env.execute({"command": eval_cmd}, timeout=timeout_sec)
except Submitted:
return 1.0
except Exception as e: # noqa: BLE001 - a bad rollout must never raise out of eval
logger.warning("[agent] eval_cmd failed to run: %s", e)
return 0.0
return 1.0 if output.get("returncode") == 0 else 0.0
# SWE-smith prep: capture the agent's fix as a patch, restore the hidden tests
# (the bug commit is HEAD~1: instance_id tip = bug + tests REMOVED), reapply the
# fix on top, so the repo test command grades the agent's edits against tests it
# never saw. Mirrors swesmith/harness/utils.py. Source
# files are identical between the two commits (only test files differ), so the
# agent's source diff applies cleanly onto HEAD~1.
_SWESMITH_PREP = (
"set -e; cd {workdir}; "
"git add -A; git diff --cached HEAD > /tmp/agent.patch || true; "
"git checkout -f HEAD~1; "
"git apply /tmp/agent.patch || git apply --3way /tmp/agent.patch || true")
def _evaluate_swesmith(env, md: dict, timeout_sec: int) -> float:
"""SWE-smith reward: reapply the agent's fix onto the tests-present commit,
run the repo test command in the sandbox, parse the output HOST-side via
swesmith's registry (log_parser). Reward 1.0 iff every FAIL_TO_PASS and
PASS_TO_PASS test is PASSED. swesmith/swebench are imported here (trainer
process), NOT in the sandbox — the env images don't ship them.
"""
inst = md.get("swesmith_inst") or {}
test_cmd = md.get("test_cmd")
if not test_cmd or not inst:
logger.warning("[agent] swesmith eval missing test_cmd/inst for %s",
md.get("instance_id"))
return 0.0
f2p = inst.get("FAIL_TO_PASS") or []
p2p = inst.get("PASS_TO_PASS") or []
try:
prep = env.execute(
{"command": _SWESMITH_PREP.format(workdir=md["workdir"])},
timeout=180)
if prep.get("returncode") != 0:
logger.warning("[agent] swesmith prep rc=%s for %s",
prep.get("returncode"), md.get("instance_id"))
return 0.0
out = env.execute({"command": test_cmd}, timeout=timeout_sec)
# Host-side parse (lazy import: only the trainer has swesmith installed).
from swebench.harness.constants import TestStatus
from swesmith.profiles import registry
passed = TestStatus.PASSED.value
status = registry.get_from_inst(inst).log_parser(out.get("output", ""))
ok = all(status.get(t) == passed for t in f2p) and all(
status.get(t) == passed for t in p2p)
return 1.0 if ok else 0.0
except Exception as e: # noqa: BLE001 - a bad rollout must never raise out of eval
logger.warning("[agent] swesmith eval failed for %s: %s",
md.get("instance_id"), e)
return 0.0
async def generate(args, sample: Sample,
sampling_params: dict[str, Any]) -> list[Sample]:
"""Per-sample agent rollout with wall-clock guard (rollout_guard_sec)."""
_ensure_rollout_executor(
) # size the to_thread pool to the batch (see above)
state = _AdapterService(args)
md = _get_metadata(sample)
instance_id = md["instance_id"]
if md["eval_kind"] == "cmd" and not md["eval_cmd"]:
return _abort_result(sample, "missing_eval_cmd", instance_id)
if md["eval_kind"] == "swesmith" and not md.get("test_cmd"):
return _abort_result(sample, "missing_test_cmd", instance_id)
session_id = sample.session_id = _session_id(sample, instance_id)
state.adapter.open_session(
session_id,
sampling_defaults=sampling_params,
max_context_tokens=state.max_context_len,
)
t0 = time.time()
sb = None
try:
async with asyncio.timeout(CONFIG.rollout_guard_sec):
# (a) sandbox claim latency
t_claim = 0.0
if not CONFIG.local_exec:
_c = time.perf_counter()
sb = await sandbox_env.claim_sandbox(
pool=md["pool"],
lifetime_sec=CONFIG.sandbox_lifetime_sec,
workdir=md["workdir"],
image=md.get(
"image"
), # on-demand create: no-pool mode, or pool-claim fallback
)
t_claim = time.perf_counter() - _c
env_raw = _make_env(sb, md["workdir"])
env = _TimedEnv(env_raw) # (d) accumulate per-action exec time
agent = _make_agent(_make_model(state.adapter_url, session_id), env)
await asyncio.to_thread(_stage_files, env_raw, md["files"],
md["workdir"])
# Optional per-task setup run post-claim, pre-agent (e.g. SWE-smith
# establishing the buggy state in a shared per-repo image). No-op for
# MBPP (no setup_cmd). Runs via env_raw so it isn't counted as an
# agent action in the exec timer; a nonzero rc aborts the rollout.
if md.get("setup_cmd"):
# The buggy-state `git checkout -f` fails transiently under load
# (exec hiccup / box not fully settled / slow disk with 100s of
# concurrent boxes). Retry with backoff; log the real rc+output if
# it ultimately fails. (Single-shot aborted ~11% of rollouts.)
setup_out: dict = {}
for _attempt in range(3):
setup_out = await asyncio.to_thread(
env_raw.execute, {"command": md["setup_cmd"]},
timeout=CONFIG.exec_timeout_sec)
if setup_out.get("returncode") == 0:
break
if _attempt < 2:
await asyncio.sleep(2 * (_attempt + 1))
if setup_out.get("returncode") != 0:
logger.warning(
"[agent] setup_cmd rc=%s for %s (3 tries): %s",
setup_out.get("returncode"),
instance_id,
str(setup_out.get("output", ""))[-300:],
)
return _abort_result(sample, "setup_cmd_failed",
instance_id)
# (b) agent-loop wall time. DefaultAgent.run is sync (model.query +
# env.execute block), so each in-flight rollout gets a worker thread.
_a = time.perf_counter()
exit_info = await asyncio.to_thread(agent.run, md["task"])
t_agent = time.perf_counter() - _a
# capture agent-loop exec time BEFORE eval so it isn't polluted by it
t_exec = env.exec_seconds
n_exec = env.exec_count
# (c) eval time -- run via env_raw to keep the exec counter loop-only
_e = time.perf_counter()
if md["eval_kind"] == "swesmith":
# SWE-smith: reapply agent fix onto tests-present commit, run repo
# test cmd, parse host-side. No eval_files staged (grading is not
# an in-sandbox script). See _evaluate_swesmith.
reward = await asyncio.to_thread(_evaluate_swesmith, env_raw,
md, CONFIG.eval_timeout_sec)
else:
# MBPP/cmd: stage hidden graders AFTER the loop (policy can't read
# or edit them), then run eval_cmd (exit 0 == solved).
await asyncio.to_thread(_stage_files, env_raw, md["eval_files"],
md["workdir"])
reward = await asyncio.to_thread(_evaluate, env_raw,
md["eval_cmd"],
CONFIG.eval_timeout_sec)
t_eval = time.perf_counter() - _e
samples = await state.adapter.finish_session(
session_id,
base_sample=sample,
reward=
reward, # split evenly across fan-out samples by the manager
)
if not samples:
return _abort_result(sample, "adapter_session_empty",
instance_id)
# Per-rollout timing rides sample.metadata; rollout_metrics.py dedups
# by session_id (a fork emits >1 sample, all sharing these numbers).
timing = {
"claim_sec": round(
t_claim, 3), # sandbox CREATE/claim time (one per rollout)
"agent_sec": round(t_agent, 3),
"eval_sec": round(t_eval, 3),
"exec_sec": round(t_exec, 3),
"exec_count": n_exec,
"exec_calls": list(env.exec_durations
), # per-call exec wall-times (percentiles)
}
for s in samples:
# slime's TrajectoryManager.to_sample does NOT propagate
# session_id onto emitted samples, and at multi-turn a rollout
# fans out into >1 per-turn sample that all share base_sample.index
# (the PROMPT index, not rollout-unique). Stamp our unique sid so
# rollout_metrics can group per-rollout (else pass@1/timing sum
# across a whole group -> pass@1 > 1). See test_masking.py.
s.session_id = session_id
s.metadata = {
**(s.metadata or {}),
"exit_status": exit_info.get("exit_status", ""),
"n_turns": agent.n_calls,
"agent_timing": timing,
# Authoritative binary rollout reward. Metrics must read THIS,
# not sum s.reward: at multi-turn a rollout fans out into
# per-turn samples and slime re-weights their per-sample
# reward (summing them inflates pass@1 by ~turns). See
# rollout_metrics.compute_metrics.
"agent_reward": float(reward),
}
logger.info(
"[agent] %s: reward=%.2f exit=%s turns=%d elapsed=%.1fs segments=%d "
"claim=%.2fs agent=%.2fs eval=%.2fs exec=%.2fs(n=%d)",
instance_id,
reward,
exit_info.get("exit_status", ""),
agent.n_calls,
time.time() - t0,
len(samples),
t_claim,
t_agent,
t_eval,
t_exec,
n_exec,
)
return samples
except asyncio.TimeoutError:
logger.warning(
"[agent] %s: wall_clock_timeout after %.1fs (guard=%ds)",
instance_id,
time.time() - t0,
CONFIG.rollout_guard_sec,
)
return _abort_result(sample, "wall_clock_timeout", instance_id)
except Exception as e:
logger.warning("[agent] %s: rollout failed: %s\n%s", instance_id, e,
traceback.format_exc())
return _abort_result(sample, f"exception:{type(e).__name__}",
instance_id)
finally:
if sb is not None: # never leak a claimed sandbox
await sandbox_env.terminate_sandbox(sb)
await state.adapter.drop_session(session_id) # cleanup only, idempotent
def _session_id(sample: Sample, instance_id: str) -> str:
# Same scheme as coding_agent_rl/generate.py:_session_id; sids must be
# unique per agent run (open_session raises on duplicates).
if sample.session_id:
return sample.session_id
if sample.index is not None and sample.group_index is not None:
return f"agent-{instance_id}-{sample.index}-{sample.group_index}"
return f"agent-{instance_id}-{secrets.token_hex(8)}"
def _abort_result(sample: Sample, reason: str,
instance_id: str) -> list[Sample]:
"""Mark ``sample`` aborted in place; verbatim shape from
slime/examples/coding_agent_rl/generate.py:_abort_result."""
sample.tokens = [0, 0]
sample.response = ""
sample.response_length = 1
sample.loss_mask = [0]
sample.rollout_log_probs = [0.0]
sample.reward = 0.0
sample.remove_sample = True
sample.status = Sample.Status.ABORTED
sample.metadata = {**(sample.metadata or {}), "abort_reason": reason}
logger.warning("[agent] %s aborted: %s", instance_id, reason)
return [sample]
code/pool_cleanup.py
"""Self-cleanup: terminate a run's sandboxes, then delete its pools.
Called from the trainer's end-of-run trap so each run OWNS and tears down its own
pool(s) — no leaks, no cross-run sharing. Scoped strictly to the pools this run
created (read from the pools-json emitted by dataset_swesmith). Best-effort:
never raises, so it can't fail the trap.
Order matters: delete_pool 500s while any sandbox references the pool (verified),
so we terminate the pool's sandboxes FIRST, then delete the pool.
python code/pool_cleanup.py <pools-json> # e.g. ~/sky_workdir/swesmith-pools.json
python code/pool_cleanup.py --pools poolA poolB # explicit names
"""
import json
import sys
def _load_pool_names(argv) -> list[str]:
if len(argv) >= 2 and argv[1] == "--pools":
return argv[2:]
if len(argv) >= 2:
try:
with open(argv[1]) as f:
return [p["pool"] for p in json.load(f)]
except Exception as e:
print(f"[pool-cleanup] could not read {argv[1]}: {e}")
return []
return []
def main() -> None:
names = set(_load_pool_names(sys.argv))
if not names:
print("[pool-cleanup] no pools to clean")
return
try:
try:
import sky.sandbox as s
except ImportError:
import sky_sandbox as s
except Exception as e:
print(f"[pool-cleanup] sandbox SDK unavailable: {e}")
return
# 1) terminate every sandbox belonging to our pools (any status).
try:
boxes = s.ls()
except Exception as e:
print(f"[pool-cleanup] ls() failed: {e}")
boxes = []
ok = 0
for b in boxes:
rec = getattr(b, "_record", {}) if not isinstance(b, dict) else b
if rec.get("template_name") in names:
try:
b.terminate()
ok += 1
except Exception:
pass
print(
f"[pool-cleanup] terminated {ok} sandboxes across {len(names)} pool(s)")
# 2) delete the pools (now unblocked).
for name in names:
try:
s.delete_pool(name)
print(f"[pool-cleanup] deleted pool: {name}")
except Exception as e:
print(
f"[pool-cleanup] delete {name} failed (may retry manually): {str(e)[:80]}"
)
if __name__ == "__main__":
main()
code/rollout_metrics.py
"""Custom rollout-logging hook: pass@k + sandbox-timing aggregates.
Wired via ``--custom-rollout-log-function-path rollout_metrics.log_rollout_metrics``.
slime calls this from ``_log_rollout_data`` (slime/ray/rollout.py:1291-1295)
with the FLAT list of every sample in the step (already flattened at
rollout.py:662-663). Returning a falsy value lets slime's own default logging
(rollout/*, perf/*) still run afterward -- we only ADD metrics, never replace.
Two things generate() can't compute alone (it sees one rollout at a time):
* pass@k -- needs the whole n_samples group. Group identity is
``Sample.group_index``, set by the data source (slime/rollout/data_source.py:112)
and preserved through deepcopy + our finish_session (trajectory.py:238
copies group_index onto every emitted Sample). pass@1 = group mean reward,
pass@k = 1 if any rollout in the group solved.
* timing aggregates -- generate() stamps per-rollout timings on
metadata["agent_timing"]; we dedup by session_id (a fork emits >1 sample
with identical timing) and reduce to mean/p50/max.
Reward-splitting note: finish_session splits a rollout's reward evenly across
its fan-out samples (trajectory.py:323-326), so we FIRST sum reward back per
session_id to recover the rollout reward, then group sessions by group_index.
That is correct at single-turn (1 sample/session) and stays correct once
multi-turn forking appears.
"""
from __future__ import annotations
from collections import defaultdict
import logging
logger = logging.getLogger(__name__)
def _pctl(values: list[float], q: float) -> float:
if not values:
return 0.0
s = sorted(values)
i = min(len(s) - 1, max(0, int(round(q * (len(s) - 1)))))
return s[i]
def _reduce(values: list[float]) -> dict[str, float]:
if not values:
return {"mean": 0.0, "p50": 0.0, "p95": 0.0, "p99": 0.0, "max": 0.0}
return {
"mean": sum(values) / len(values),
"p50": _pctl(values, 0.5),
"p95": _pctl(values, 0.95),
"p99": _pctl(values, 0.99),
"max": max(values),
}
def compute_metrics(samples) -> dict[str, float]:
"""Pure function (unit-testable): flat sample list -> metric dict."""
# --- per-session reduction (one authoritative reward + timing per rollout) -
# Read the binary rollout reward generate() stamped on metadata["agent_reward"]
# (set once per rollout). We must NOT sum s.reward: at multi-turn a rollout
# fans out into per-turn samples and slime re-weights their per-sample reward,
# so summing inflates pass@1 by ~turns (observed 1.56 vs a true 0.83). Fall
# back to summing s.reward only for older data without the stamp.
session_reward: dict[str, float] = {}
session_reward_fallback: dict[str, float] = defaultdict(float)
session_group: dict[str, int] = {}
session_timing: dict[str, dict] = {}
for s in samples:
sid = s.session_id or f"idx-{s.index}"
session_group[sid] = s.group_index
md = s.metadata or {}
if "agent_reward" in md:
session_reward[sid] = float(md["agent_reward"])
else:
session_reward_fallback[sid] += float(s.reward or 0.0)
if "agent_timing" in md and sid not in session_timing:
session_timing[sid] = md["agent_timing"]
for sid, r in session_reward_fallback.items():
session_reward.setdefault(sid, r)
# --- pass@k over groups of rollouts ---
group_rewards: dict[int, list[float]] = defaultdict(list)
for sid, r in session_reward.items():
group_rewards[session_group.get(sid)].append(r)
pass_at_1 = [sum(rs) / len(rs) for rs in group_rewards.values() if rs
] # per-group mean
pass_at_k = [
1.0 if any(r >= 1.0
for r in rs) else 0.0
for rs in group_rewards.values()
if rs
]
metrics: dict[str, float] = {}
if pass_at_1:
metrics["pass@1"] = sum(pass_at_1) / len(pass_at_1)
if pass_at_k:
metrics["pass@k"] = sum(pass_at_k) / len(pass_at_k)
metrics["num_groups"] = float(len(pass_at_k))
# --- multi-turn shape (turns) + packed-vs-forked ratio (masked_frac) ------
# turns = metadata["n_turns"] (= agent.n_calls): >1 confirms the agent iterates.
# masked_frac is NOT a masking-correctness check (that is guaranteed by
# slime's TrajectoryManager: observations never enter a trained response
# region -- proven both ways in test_masking.py). It is a packed-vs-forked
# SIGNAL: a linear chain packs into one sample only while each turn extends
# the prior as an exact token-prefix; mini-swe's chat-template re-tokenization
# drifts at turn boundaries, so turns FORK into per-turn samples whose response
# is one turn (all trainable -> masked_frac 0) with observations in the
# stripped leading prompt. masked_frac > 0 therefore means turns PACKED (obs
# as loss_mask=0 spans); ~0 means they forked. Both are correct; the number
# just tells us which path slime took. Skip aborted samples (loss_mask=[0]).
turns: list[float] = []
masked_fracs: list[float] = []
for s in samples:
if getattr(s, "remove_sample", False):
continue
md = s.metadata or {}
if "n_turns" in md:
turns.append(float(md["n_turns"]))
lm = getattr(s, "loss_mask", None)
if lm: # fraction of response tokens that are masked (observations)
masked_fracs.append(1.0 - (sum(lm) / len(lm)))
if turns:
metrics["turns/mean"] = sum(turns) / len(turns)
metrics["turns/max"] = max(turns)
if masked_fracs: # >0 confirms multi-turn observation masking is engaged
metrics["masked_frac/mean"] = sum(masked_fracs) / len(masked_fracs)
metrics["masked_frac/max"] = max(masked_fracs)
# --- sandbox timing aggregates (deduped per session) ---
# claim_sec = sandbox CREATE time, one per rollout -> percentiles are per-rollout.
# (p50/p95/p99 for all; the create-time distribution is the warm-pool/cold-create
# signal we want when tuning image caching on R2E-like per-repo images.)
for key in ("claim_sec", "agent_sec", "eval_sec", "exec_sec"):
vals = [float(t[key]) for t in session_timing.values() if key in t]
for stat, v in _reduce(vals).items():
metrics[f"sandbox_{key}/{stat}"] = v
# exec_call = per-CALL exec wall-time, flattened across every action in the step
# -> per-call percentiles (distinct from exec_sec which is per-rollout total).
exec_calls = [
d for t in session_timing.values() for d in (t.get("exec_calls") or [])
]
for stat, v in _reduce([float(d) for d in exec_calls]).items():
metrics[f"sandbox_exec_call/{stat}"] = v
metrics["sandbox_exec_call/count"] = float(len(exec_calls))
exec_counts = [
float(t["exec_count"])
for t in session_timing.values()
if "exec_count" in t
]
if exec_counts:
metrics["sandbox_exec_count/mean"] = sum(exec_counts) / len(exec_counts)
metrics["num_rollouts"] = float(len(session_reward))
return metrics
def log_rollout_metrics(rollout_id, args, samples, rollout_extra_metrics,
rollout_time) -> bool:
"""slime custom-rollout-log hook. Returns False so slime's default logging
still runs (we only add agent/* keys)."""
try:
metrics = compute_metrics(samples)
except Exception as e: # never break a training step over a metric
logger.warning("[metrics] compute failed: %s", e)
return False
# Greppable stdout line (works even when wandb is disabled).
logger.info("[metrics] step=%s %s", rollout_id, metrics)
print(f"[metrics] step={rollout_id} " +
" ".join(f"{k}={v:.4g}" for k, v in metrics.items()),
flush=True)
# Emit to wandb via slime's own logger if enabled (same path as the built-in
# rollout/perf metrics: logging_utils.log -> wandb.log).
try:
from slime.utils import logging_utils
from slime.utils.metric_utils import compute_rollout_step
from slime.utils.metric_utils import dict_add_prefix
if getattr(args, "use_wandb", False):
log_dict = dict_add_prefix(metrics, "agent/")
log_dict["rollout/step"] = compute_rollout_step(args, rollout_id)
logging_utils.log(args, log_dict, step_key="rollout/step")
except Exception as e:
logger.warning("[metrics] wandb emit skipped: %s", e)
return False # let slime log its default rollout/* + perf/* metrics too
code/sandbox_env.py
"""SkyPilot-sandbox Environment for mini-swe-agent, plus pool/claim helpers.
SandboxEnvironment satisfies the ``minisweagent.Environment`` protocol
(minisweagent/__init__.py:61-70: execute / get_template_vars / serialize) and
is structured after the installed ``environments/docker.py`` -- pydantic
config, sync ``execute(action, cwd, timeout)``, ``Submitted`` sentinel check
-- with docker-exec swapped for the sky.sandbox SDK's sync exec surface
(``Sandbox.exec(*argv, workdir=..., timeout_seconds=...) -> ExecHandle``,
``handle.wait()`` + ``handle.stdout.read()``).
Sync on purpose: the DefaultAgent loop runs inside ``asyncio.to_thread`` (one
thread per in-flight rollout), so everything here uses the SDK's plain sync
entrypoints. Output truncation is NOT done here: upstream's
mini_textbased.yaml observation template already elides >10k-char outputs,
and the SDK's ExecHandle bounds its own client-side buffer.
The sandbox executes ONLY untrusted code (agent actions, eval commands); the
agent loop and model client never run inside it.
"""
from __future__ import annotations
import os
import platform
import secrets
import time
from typing import Any
from minisweagent.exceptions import Submitted
from minisweagent.utils.serialize import recursive_merge
from pydantic import BaseModel
try:
import sky.sandbox as _sandbox # standalone skypilot_sandbox_sdk wheel
except ImportError:
try:
import sky_sandbox as _sandbox # top-level shim from the same wheel
except ImportError:
_sandbox = None # tolerated for AGENT_LOCAL_EXEC-only hosts
# Escape hatch for the FULL multi-repo dataset. This example warms ONE pool per repo
# image, which is fast and reliable for a pinned repo subset (the default). Across the
# whole dataset, though, that's many pools competing for cluster sandbox capacity and
# many independent warm-ups — a reliability cliff. Set SANDBOX_NO_POOL=1 to skip pools
# entirely and create an on-demand ad-hoc sandbox per claim from the instance's own
# image (create(image=...)). Trade-off: a cold image pull on each claim instead of a
# warm-pool hit, so per-claim latency is higher. Rule of thumb: pools for a pinned
# subset, no-pool for the full dataset.
_NO_POOL = os.environ.get("SANDBOX_NO_POOL") == "1"
class SandboxEnvironmentConfig(BaseModel):
cwd: str = "/workspace"
"""Working directory in which to execute commands."""
env: dict[str, str] = {}
"""Environment variables to set for each exec."""
timeout: int = 60
"""Per-action timeout in seconds (server-side exec timeout)."""
class SandboxEnvironment:
"""Executes each agent action as one stateless exec in a claimed sandbox.
mini-swe's execution model (independent subprocess per action,
environments/local.py) is semantically identical to sandbox exec, so the
port is one method.
"""
def __init__(self,
sb,
*,
config_class: type = SandboxEnvironmentConfig,
**kwargs) -> None:
self.sb = sb
self.config = config_class(**kwargs)
def execute(self,
action: dict,
cwd: str = "",
*,
timeout: int | None = None) -> dict[str, Any]:
"""Run one bash action; mirrors DockerEnvironment.execute's output dict."""
command = action.get("command", "")
cwd = cwd or self.config.cwd
timeout = timeout or self.config.timeout
try:
handle = self.sb.exec(
"bash",
"-c",
command,
workdir=cwd,
timeout_seconds=float(
timeout), # server-side kill: no orphaned commands
env=self.config.env or None,
)
# client wait bounded slightly above the server kill so we always
# collect the exit code the server assigns
returncode = handle.wait(timeout=timeout + 30)
# mini-swe merges stderr into stdout (local.py uses stderr=STDOUT);
# the SDK keeps separate streams, so concatenate. Decode defensively:
# agent commands run arbitrary code and can emit non-UTF-8 bytes.
out_b, err_b = handle.stdout.read(), handle.stderr.read()
output_text = ((out_b.decode("utf-8", "replace") if isinstance(
out_b, bytes) else out_b) + (err_b.decode(
"utf-8", "replace") if isinstance(err_b, bytes) else err_b))
output = {
"output": output_text,
"returncode": returncode,
"exception_info": ""
}
except Exception as e: # timeout / transport -> observation, not crash
output = {
"output": "",
"returncode": -1,
"exception_info": f"An error occurred while executing the command: {e}",
"extra": {
"exception_type": type(e).__name__,
"exception": str(e)
},
}
self._check_finished(output)
return output
def _check_finished(self, output: dict):
"""Raises Submitted if the output indicates task completion.
Verbatim from environments/docker.py::_check_finished (the sentinel
check is the Environment's job in upstream, not the agent's).
"""
lines = output.get("output", "").lstrip().splitlines(keepends=True)
if lines and lines[0].strip(
) == "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT" and output[
"returncode"] == 0:
submission = "".join(lines[1:])
raise Submitted({
"role": "exit",
"content": submission,
"extra": {
"exit_status": "Submitted",
"submission": submission
},
})
def get_template_vars(self, **kwargs) -> dict[str, Any]:
# Host uname, mirroring upstream DockerEnvironment.get_template_vars.
# The instance template renders it as <system_information>; trainer
# pods and sandbox images are both linux, so this is accurate enough.
return recursive_merge(self.config.model_dump(),
platform.uname()._asdict(), kwargs)
def serialize(self) -> dict:
return {
"info": {
"config": {
"environment": self.config.model_dump(mode="json"),
"environment_type": f"{self.__class__.__module__}.{self.__class__.__name__}",
}
}
}
def write_files(self, files: dict[str, str] | None) -> None:
"""Stage text files (path -> content; relative paths land under cwd).
Sync (called via asyncio.to_thread); raises on failure -- staging
errors are rollout aborts, not model observations.
"""
for path, content in (files or {}).items():
dst = path if path.startswith(
"/") else f"{self.config.cwd.rstrip('/')}/{path}"
parent = os.path.dirname(dst) or "/"
rc = self.sb.exec("mkdir", "-p", parent).wait()
if rc != 0:
raise RuntimeError(f"failed to mkdir {parent} (exit {rc})")
self.sb.write_text(content, dst)
def stage_files_locally(cwd: str, files: dict[str, str] | None) -> None:
"""Local-filesystem twin of SandboxEnvironment.write_files, for the
AGENT_LOCAL_EXEC path (which uses upstream LocalEnvironment)."""
for path, content in (files or {}).items():
dst = path if os.path.isabs(path) else os.path.join(cwd, path)
os.makedirs(os.path.dirname(dst) or ".", exist_ok=True)
with open(dst, "w", encoding="utf-8") as f:
f.write(content)
# --- pool / sandbox lifecycle helpers (used by generate.py and trainer setup) ---
async def claim_sandbox(
pool: str,
*,
lifetime_sec: float | None = None,
name: str | None = None,
workdir: str | None = None,
image: str | None = None,
):
"""Claim one sandbox for a rollout. Caller must terminate it.
Warm-pool mode (default): claim a ready sandbox from ``pool``; if the pool has
no ready box yet (slow/partial warm), fall back to an on-demand sandbox from
``image`` so the rollout still runs.
No-pool mode (``SANDBOX_NO_POOL=1``): always create an on-demand ad-hoc sandbox
from ``image`` — no pool needed, at the cost of a cold image pull per claim. Use
it for the full multi-repo dataset, where one-pool-per-image doesn't scale.
``lifetime_sec`` sets the server-side auto-reap (``create(timeout=...)``),
so a crashed trainer cannot leak sandboxes past that lifetime.
``workdir`` is created if given — generic pool images (python:*-slim)
don't ship /workspace, and every action cd's into it (exec workdir= is a
plain `cd`, rc=127 if missing; verified live 2026-07-14).
"""
if _sandbox is None:
raise RuntimeError(
"sky.sandbox SDK not installed (skypilot_sandbox_sdk wheel)")
name = name or f"r3-{secrets.token_hex(4)}" # unique to avoid name reuse across rollouts
if _NO_POOL:
# On-demand ad-hoc sandbox from the instance's own image (bypasses pools).
sb = await _sandbox.create.aio(name=name,
image=image,
cpus=1,
memory_gb=2,
timeout=lifetime_sec)
else:
try:
sb = await _sandbox.create.aio(name=name,
pool=pool,
timeout=lifetime_sec)
except Exception as e: # noqa: BLE001
# Pool claim failed — usually the pool hasn't warmed a ready box yet
# (slow / partial warm; see ensure_pool). Fall back to an on-demand
# sandbox from the repo image so the rollout still runs (cold image
# pull instead of a warm hit). Re-raise only if we have no image.
if image is None:
raise
print(
f"[sandbox] pool {pool!r} claim failed ({e}); falling back to "
f"on-demand sandbox {name!r}",
flush=True)
sb = await _sandbox.create.aio(name=name,
image=image,
cpus=1,
memory_gb=2,
timeout=lifetime_sec)
if workdir:
try:
handle = await sb.exec.aio("mkdir", "-p", workdir)
await handle.wait()
except Exception:
await terminate_sandbox(sb)
raise
return sb
async def terminate_sandbox(sb) -> None:
"""Best-effort terminate; never raises (used from finally blocks)."""
try:
await sb.terminate.aio()
except Exception:
pass
def ensure_pool(name: str,
*,
image: str,
replicas: int,
cpus: float = 1,
memory_gb: float = 2) -> None:
"""Create the warm pool, tolerating one that already exists.
Same pattern as sandbox_reward_server.py:196-220: create_pool, and on
failure (already exists) make sure it is at least the requested size.
Called once from trainer setup, not per rollout.
"""
if _sandbox is None:
raise RuntimeError(
"sky.sandbox SDK not installed (skypilot_sandbox_sdk wheel)")
# Time the create_pool call (greppable: "[pool-timing]"). This is the pool
# REGISTRATION cost; the image-pull / replica-warm tail then shows up as the
# cold-claim latency in sandbox_claim_sec p95/p99 (first claims wait for a
# replica incl. pull). Both matter for the R2E per-repo-image caching work:
# a per-image pool's first warm-up is dominated by image pull, so this is the
# baseline we'll optimize (pre-pull / image locality).
t0 = time.perf_counter()
try:
_sandbox.create_pool(name=name,
image=image,
cpus=cpus,
memory_gb=memory_gb,
replicas=replicas)
dt = time.perf_counter() - t0
print(
f"[pool-timing] created {name!r} image={image} replicas={replicas} create_call={dt:.2f}s",
flush=True)
return
except Exception as e: # noqa: BLE001
dt = time.perf_counter() - t0
msg = str(e)
# A 400 Bad Request means the create was REJECTED (e.g. name too long or
# invalid) — not recoverable, and set_pool_size can't fix it. Fail loudly.
if "400" in msg or "Bad Request" in msg:
raise RuntimeError(f"create_pool({name!r}) rejected: {msg}") from e
# Otherwise the pool either already exists OR didn't fully warm within the
# SDK timeout (msg carries "last seen N/replicas"). NEITHER is fatal: the
# pool exists and keeps warming, and any claim that can't get a ready box
# falls back to an on-demand sandbox (see claim_sandbox). Log how far it
# got and proceed rather than crashing the whole run on a slow/partial warm.
print(
f"[pool-timing] {name!r} not confirmed at full size ({msg}) in {dt:.2f}s; "
f"ensuring size {replicas} and proceeding (overflow claims fall back on-demand)",
flush=True)
try:
_sandbox.set_pool_size(name, replicas=replicas)
except Exception as e2: # noqa: BLE001
print(
f"[pool-timing] set_pool_size({name!r}) did not confirm full warm ({e2}); "
f"proceeding on the partial pool",
flush=True)
async def aclose() -> None:
"""Release the shared sandbox SDK session (process shutdown only)."""
if _sandbox is None:
return
try:
await _sandbox.aclose()
except Exception:
pass
coding-agent-2engine.yaml
# Coding-agent RL — 2-engine scale-out variant of coding-agent.yaml.
#
# Same Job Group, but the inference fleet is TWO external SGLang engines instead of
# one. Scaling out is fully declarative: add an `sglang2` job below and add
# its name to SGLANG_MEMBERS — trainer-run.sh health-gates every listed engine
# and passes all their addresses to slime (--rollout-external-engine-addrs), which
# registers both with the trainer-local router. This is the natural knob to
# relieve a rollout-bound run, and the direction that motivates elastic fleets.
#
# Footprint: trainer H100:4 + 2 x engine H100:1 = 6 GPU on one cluster.
#
# Prereq (once): sky volumes apply policy-volume.yaml
# Launch:
# sky jobs launch -n coding-agent-2eng coding-agent-2engine.yaml --secret WANDB_API_KEY # secret optional (wandb)
---
name: coding-agent-rl-2eng
execution: parallel
primary_tasks: [trainer]
termination_delay: 60s
---
name: sglang
resources:
infra: kubernetes # your H100 Kubernetes context, e.g. kubernetes/my-cluster
accelerators: H100:1
image_id: docker:slimerl/slime:nightly-dev-20260707a
workdir: .
volumes:
/shared/policy: slime-policy
setup: bash scripts/engine-setup.sh
run: bash scripts/engine-run.sh
---
name: sglang2
resources:
infra: kubernetes
accelerators: H100:1
image_id: docker:slimerl/slime:nightly-dev-20260707a
workdir: .
volumes:
/shared/policy: slime-policy
setup: bash scripts/engine-setup.sh
run: bash scripts/engine-run.sh
---
name: trainer
resources:
infra: kubernetes # SAME context as the engines — jobs co-schedule on one cluster
accelerators: H100:4
image_id: docker:slimerl/slime:nightly-dev-20260707a
workdir: .
volumes:
/shared/policy: slime-policy
file_mounts:
# Sandboxes SDK wheel: a sandbox-enabled API server ships it to the client at
# ~/.sky/bin/wheels; mount it in so trainer-setup.sh can pip install it (not on PyPI).
/wheels: ~/.sky/bin/wheels
envs:
SGLANG_MEMBERS: "sglang sglang2" # two engines -> both health-gated + discovered
MODEL_SCRIPT: "qwen3-14B.sh"
MODEL_DIR: "Qwen3-14B"
TP: "4"
ROLLOUT_GPUS_PER_ENGINE: "1"
SWESMITH_REPOS: "markupsafe" # pinned subset -> one warm pool per repo image
SWESMITH_PER_REPO: "16"
SWESMITH_POOL_SUFFIX: ""
# SANDBOX_NO_POOL=1 -> no warm pools; on-demand ad-hoc sandbox per claim. Use for
# the full dataset (drop SWESMITH_REPOS). See coding-agent.yaml / sandbox_env.py.
SANDBOX_NO_POOL: "0"
AGENT_WORKDIR: "/testbed"
AGENT_MAX_TURNS: "20"
AGENT_EVAL_TIMEOUT_SEC: "900"
# SANDBOX_POOL_REPLICAS must be >= ROLLOUT_BATCH_SIZE x N_SAMPLES_PER_PROMPT and
# <= your cluster's sandbox (CPU) capacity (see coding-agent.yaml). 184 = our
# benchmarked sizing; scale down if your sandbox cluster has less CPU headroom.
SANDBOX_POOL_REPLICAS: "184"
ROLLOUT_BATCH_SIZE: "23" # prompts/step
N_SAMPLES_PER_PROMPT: "8" # rollouts/prompt -> 23x8 = 184 concurrent rollouts (global batch derived)
NUM_ROLLOUT: "8"
ROLLOUT_SHUFFLE: "1" # shuffle instance order across steps; 0 = deterministic (reproducible)
ROLLOUT_MAX_CONTEXT_LEN: "32768"
MSWEA_SILENT_STARTUP: "1"
SLIME_MODE: "async" # async rollout/train overlap | "sync" serializes them
WEIGHT_TRANSPORT: "disk"
WEIGHT_MODE: "delta"
WANDB_PROJECT: "slime-coding-agent"
WANDB_GROUP: "coding-agent-2eng"
secrets:
# Optional wandb logging; provide with `--secret WANDB_API_KEY` (picked up
# from your environment, redacted in logs/dashboard). Unset = wandb off.
# Sandbox-API auth needs no secret: managed jobs get API server credentials
# injected automatically (api_server_access, on by default).
WANDB_API_KEY: ""
setup: bash scripts/trainer-setup.sh
run: bash scripts/trainer-run.sh
coding-agent-3engine.yaml
# Coding-agent RL — 3-engine scale-out variant of coding-agent.yaml.
#
# Same Job Group, THREE external SGLang engines. Note how little changes vs the
# 1- and 2-engine files: one more `sglang3` job (identical to the others) and
# one more name in SGLANG_MEMBERS. trainer-run.sh and run-slime.sh are untouched —
# scaling the inference fleet is purely declarative.
#
# In our benchmark this is roughly where co-scheduled inference stops paying off
# for this workload (the queue is drained and the engine GPUs begin to idle) —
# see the README's "Scaling the inference fleet" section.
#
# Footprint: trainer H100:4 + 3 x engine H100:1 = 7 GPU on one cluster.
#
# Prereq (once): sky volumes apply policy-volume.yaml
# Launch:
# sky jobs launch -n coding-agent-3eng coding-agent-3engine.yaml --secret WANDB_API_KEY # secret optional (wandb)
---
name: coding-agent-rl-3eng
execution: parallel
primary_tasks: [trainer]
termination_delay: 60s
---
name: sglang
resources:
infra: kubernetes # your H100 Kubernetes context, e.g. kubernetes/my-cluster
accelerators: H100:1
image_id: docker:slimerl/slime:nightly-dev-20260707a
workdir: .
volumes:
/shared/policy: slime-policy
setup: bash scripts/engine-setup.sh
run: bash scripts/engine-run.sh
---
name: sglang2
resources:
infra: kubernetes
accelerators: H100:1
image_id: docker:slimerl/slime:nightly-dev-20260707a
workdir: .
volumes:
/shared/policy: slime-policy
setup: bash scripts/engine-setup.sh
run: bash scripts/engine-run.sh
---
name: sglang3
resources:
infra: kubernetes
accelerators: H100:1
image_id: docker:slimerl/slime:nightly-dev-20260707a
workdir: .
volumes:
/shared/policy: slime-policy
setup: bash scripts/engine-setup.sh
run: bash scripts/engine-run.sh
---
name: trainer
resources:
infra: kubernetes # SAME context as the engines — jobs co-schedule on one cluster
accelerators: H100:4
image_id: docker:slimerl/slime:nightly-dev-20260707a
workdir: .
volumes:
/shared/policy: slime-policy
file_mounts:
# Sandboxes SDK wheel: a sandbox-enabled API server ships it to the client at
# ~/.sky/bin/wheels; mount it in so trainer-setup.sh can pip install it (not on PyPI).
/wheels: ~/.sky/bin/wheels
envs:
SGLANG_MEMBERS: "sglang sglang2 sglang3" # three engines -> all health-gated + discovered
MODEL_SCRIPT: "qwen3-14B.sh"
MODEL_DIR: "Qwen3-14B"
TP: "4"
ROLLOUT_GPUS_PER_ENGINE: "1"
SWESMITH_REPOS: "markupsafe" # pinned subset -> one warm pool per repo image
SWESMITH_PER_REPO: "16"
SWESMITH_POOL_SUFFIX: ""
# SANDBOX_NO_POOL=1 -> no warm pools; on-demand ad-hoc sandbox per claim. Use for
# the full dataset (drop SWESMITH_REPOS). See coding-agent.yaml / sandbox_env.py.
SANDBOX_NO_POOL: "0"
AGENT_WORKDIR: "/testbed"
AGENT_MAX_TURNS: "20"
AGENT_EVAL_TIMEOUT_SEC: "900"
# SANDBOX_POOL_REPLICAS must be >= ROLLOUT_BATCH_SIZE x N_SAMPLES_PER_PROMPT and
# <= your cluster's sandbox (CPU) capacity (see coding-agent.yaml). 184 = our
# benchmarked sizing; scale down if your sandbox cluster has less CPU headroom.
SANDBOX_POOL_REPLICAS: "184"
ROLLOUT_BATCH_SIZE: "23" # prompts/step
N_SAMPLES_PER_PROMPT: "8" # rollouts/prompt -> 23x8 = 184 concurrent rollouts (global batch derived)
NUM_ROLLOUT: "8"
ROLLOUT_SHUFFLE: "1" # shuffle instance order across steps; 0 = deterministic (reproducible)
ROLLOUT_MAX_CONTEXT_LEN: "32768"
MSWEA_SILENT_STARTUP: "1"
SLIME_MODE: "async" # async rollout/train overlap | "sync" serializes them
WEIGHT_TRANSPORT: "disk"
WEIGHT_MODE: "delta"
WANDB_PROJECT: "slime-coding-agent"
WANDB_GROUP: "coding-agent-3eng"
secrets:
# Optional wandb logging; provide with `--secret WANDB_API_KEY` (picked up
# from your environment, redacted in logs/dashboard). Unset = wandb off.
# Sandbox-API auth needs no secret: managed jobs get API server credentials
# injected automatically (api_server_access, on by default).
WANDB_API_KEY: ""
setup: bash scripts/trainer-setup.sh
run: bash scripts/trainer-run.sh
coding-agent.yaml
# Coding-agent RL on SkyPilot Job Groups + Sandboxes — one disaggregated Job Group.
#
# A slime trainer (Megatron) and an external SGLang inference engine are
# co-scheduled as ONE Job Group. The agentic rollout (mini-swe-agent) runs each
# turn's untrusted code in a SkyPilot Sandbox and grades it for reward; the
# trainer publishes updated weights over a shared RWX volume after each step.
#
# Each job's setup/run logic lives in scripts/ (engine-*.sh, trainer-*.sh) so
# this file stays a declarative manifest of the jobs + their resources + env.
#
# Footprint: trainer H100:4 + engine H100:1 = 5 GPU on one cluster.
# Scale the inference fleet with coding-agent-2engine.yaml / coding-agent-3engine.yaml.
#
# Prereq (once): sky volumes apply policy-volume.yaml
# Launch:
# sky jobs launch -n coding-agent coding-agent.yaml --secret WANDB_API_KEY # secret optional (wandb)
---
name: coding-agent-rl
execution: parallel
primary_tasks: [trainer]
termination_delay: 60s
---
name: sglang
resources:
infra: kubernetes # your H100 Kubernetes context, e.g. kubernetes/my-cluster
accelerators: H100:1 # engine tensor-parallel = 1 (1 H100/engine; validated repro sizing)
image_id: docker:slimerl/slime:nightly-dev-20260707a
workdir: . # so the job can see scripts/
volumes:
/shared/policy: slime-policy
setup: bash scripts/engine-setup.sh
run: bash scripts/engine-run.sh
---
name: trainer
resources:
infra: kubernetes # SAME context as the engine — jobs co-schedule on one cluster
accelerators: H100:4
image_id: docker:slimerl/slime:nightly-dev-20260707a
workdir: .
volumes:
/shared/policy: slime-policy
file_mounts:
# Sandboxes SDK wheel: a sandbox-enabled API server ships it to the client at
# ~/.sky/bin/wheels; mount it in so trainer-setup.sh can pip install it (not on PyPI).
/wheels: ~/.sky/bin/wheels
envs:
# --- inference fleet: the SGLang Job Group job names (one per engine) ---
SGLANG_MEMBERS: "sglang"
# --- model ---
MODEL_SCRIPT: "qwen3-14B.sh"
MODEL_DIR: "Qwen3-14B"
TP: "4"
ROLLOUT_GPUS_PER_ENGINE: "1"
# --- SWE-smith subset (bare repo tokens; one warm sandbox pool per repo image) ---
SWESMITH_REPOS: "markupsafe" # pinned subset -> one warm pool per repo image (fast, reliable)
SWESMITH_PER_REPO: "16"
SWESMITH_POOL_SUFFIX: ""
# Set SANDBOX_NO_POOL=1 to run WITHOUT warm pools — an on-demand ad-hoc sandbox is
# created per claim from the instance's own image. Prefer this for the FULL dataset
# (drop SWESMITH_REPOS): one pool per repo doesn't scale. See sandbox_env.py.
SANDBOX_NO_POOL: "0"
# --- agent loop ---
AGENT_WORKDIR: "/testbed"
AGENT_MAX_TURNS: "20"
AGENT_EVAL_TIMEOUT_SEC: "900"
# --- sizing: concurrent rollouts = ROLLOUT_BATCH_SIZE x N_SAMPLES_PER_PROMPT ---
# SANDBOX_POOL_REPLICAS must be >= that product (one sandbox claim per rollout)
# AND <= your cluster's sandbox (CPU) capacity. An oversized pool can't fully warm,
# so claims time out and rollouts fail — size it to the smaller of the two. 184 is
# our benchmarked sizing (warmed cleanly + held all ~184 concurrent claims); scale
# down if your sandbox cluster has less CPU headroom.
SANDBOX_POOL_REPLICAS: "184"
ROLLOUT_BATCH_SIZE: "23" # prompts/step
N_SAMPLES_PER_PROMPT: "8" # rollouts/prompt -> 23x8 = 184 concurrent rollouts (global batch derived)
NUM_ROLLOUT: "8"
ROLLOUT_SHUFFLE: "1" # shuffle instance order across steps; 0 = deterministic (reproducible)
ROLLOUT_MAX_CONTEXT_LEN: "32768"
MSWEA_SILENT_STARTUP: "1"
# --- infra axes ---
SLIME_MODE: "async" # async rollout/train overlap | "sync" serializes them
WEIGHT_TRANSPORT: "disk" # disk (shared RWX volume) | nccl
WEIGHT_MODE: "delta" # delta (changed bytes only) | full
# --- wandb (optional) ---
WANDB_PROJECT: "slime-coding-agent"
WANDB_GROUP: "coding-agent"
secrets:
# Optional wandb logging; provide with `--secret WANDB_API_KEY` (picked up
# from your environment, redacted in logs/dashboard). Unset = wandb off.
# Sandbox-API auth needs no secret: managed jobs get API server credentials
# injected automatically (api_server_access, on by default).
WANDB_API_KEY: ""
setup: bash scripts/trainer-setup.sh
run: bash scripts/trainer-run.sh
images/architecture.svg
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 1240 860" font-family="-apple-system, BlinkMacSystemFont, 'Segoe UI', Helvetica, Arial, sans-serif">
<defs>
<marker id="ab" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="7" markerHeight="7" orient="auto-start-reverse">
<path d="M0,0 L10,5 L0,10 z" fill="#2b5fd9"/>
</marker>
<marker id="ar" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="7" markerHeight="7" orient="auto-start-reverse">
<path d="M0,0 L10,5 L0,10 z" fill="#b23a20"/>
</marker>
<marker id="ag" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="7" markerHeight="7" orient="auto-start-reverse">
<path d="M0,0 L10,5 L0,10 z" fill="#2e7d5b"/>
</marker>
</defs>
<!-- paper -->
<rect x="0" y="0" width="1240" height="860" rx="18" fill="#f7f6f3"/>
<!-- ============ JobGroup boundary ============ -->
<rect x="48" y="78" width="1160" height="440" rx="14" fill="none" stroke="#2b5fd9" stroke-width="1.6" stroke-dasharray="2 6"/>
<rect x="72" y="62" width="226" height="32" rx="16" fill="#f7f6f3" stroke="#2b5fd9" stroke-width="1.4"/>
<text x="185" y="83" text-anchor="middle" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="12.5" letter-spacing="2" fill="#2b5fd9">SKYPILOT JOB GROUP</text>
<!-- ============ Trainer member ============ -->
<rect x="88" y="128" width="640" height="350" rx="10" fill="#edf1fa" stroke="#2b4f9e" stroke-width="1.3"/>
<text x="108" y="156" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="11" letter-spacing="1.5" fill="#2b4f9e">JOB 1</text>
<text x="108" y="182" font-size="18" font-weight="700" fill="#1f2430">Trainer cluster</text>
<rect x="636" y="142" width="76" height="24" rx="12" fill="#ffffff" stroke="#2b4f9e" stroke-opacity="0.4"/>
<text x="674" y="158" text-anchor="middle" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="11.5" fill="#2b4f9e">H100×4</text>
<!-- Megatron -->
<rect x="528" y="200" width="184" height="92" rx="8" fill="#ffffff" stroke="#7f97c4" stroke-width="1.2"/>
<text x="620" y="230" text-anchor="middle" font-size="13" font-weight="600" fill="#1f2430">Trainer</text>
<text x="620" y="248" text-anchor="middle" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="10.5" fill="#6b7280">Megatron · GPU</text>
<text x="620" y="272" text-anchor="middle" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="11.5" fill="#1f2430">④ GRPO + optimizer step</text>
<!-- RolloutManager -->
<rect x="108" y="240" width="360" height="200" rx="8" fill="#ffffff" stroke="#7f97c4" stroke-width="1.2"/>
<text x="124" y="268" font-size="14" font-weight="600" fill="#1f2430">RolloutManager</text>
<text x="124" y="286" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="10.5" fill="#6b7280">CPU · mini-swe-agent loops · router</text>
<rect x="124" y="340" width="104" height="52" rx="6" fill="#f3f6fc" stroke="#9db6ea" stroke-width="1.2"/>
<text x="176" y="369" text-anchor="middle" font-size="12.5" font-weight="600" fill="#1f2430">agent loop 1</text>
<rect x="240" y="340" width="104" height="52" rx="6" fill="#f3f6fc" stroke="#9db6ea" stroke-width="1.2"/>
<text x="292" y="369" text-anchor="middle" font-size="12.5" font-weight="600" fill="#1f2430">agent loop 2</text>
<rect x="356" y="340" width="96" height="52" rx="6" fill="none" stroke="#9db6ea" stroke-width="1.2" stroke-dasharray="4 4"/>
<text x="404" y="369" text-anchor="middle" font-size="12.5" fill="#6b7280">… loop N</text>
<!-- ============ Inference member(s) ============ -->
<rect x="790" y="128" width="400" height="350" rx="10" fill="#fbefe9" stroke="#a8402c" stroke-width="1.3"/>
<text x="810" y="156" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="11" letter-spacing="1.5" fill="#a8402c">JOBS 2…N+1</text>
<text x="810" y="182" font-size="18" font-weight="700" fill="#1f2430">Inference clusters</text>
<rect x="1078" y="142" width="96" height="24" rx="12" fill="#ffffff" stroke="#a8402c" stroke-opacity="0.4"/>
<text x="1126" y="158" text-anchor="middle" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="10.5" fill="#a8402c">H100×1 each</text>
<rect x="814" y="210" width="352" height="62" rx="8" fill="#ffffff" stroke="#c98872" stroke-width="1.2"/>
<text x="990" y="246" text-anchor="middle" font-size="13.5" font-weight="600" fill="#1f2430">SGLang engine</text>
<rect x="814" y="286" width="352" height="62" rx="8" fill="#ffffff" stroke="#c98872" stroke-width="1.2"/>
<text x="990" y="322" text-anchor="middle" font-size="13.5" font-weight="600" fill="#1f2430">SGLang engine</text>
<rect x="814" y="362" width="352" height="54" rx="8" fill="none" stroke="#c98872" stroke-width="1.2" stroke-dasharray="4 4"/>
<text x="990" y="394" text-anchor="middle" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="12" fill="#6b7280">… ×N engines</text>
<!-- ============ Sandboxes ============ -->
<rect x="100" y="612" width="424" height="160" rx="10" fill="#f2f1ed" stroke="#9aa1ad" stroke-width="1.2"/>
<text x="116" y="638" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="11" letter-spacing="1.5" fill="#5c636e">SKYPILOT SANDBOXES</text>
<text x="116" y="655" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="10.5" fill="#6b7280">CPU</text>
<rect x="124" y="664" width="104" height="66" rx="8" fill="#ffffff" stroke="#9aa1ad" stroke-width="1.2"/>
<text x="176" y="701" text-anchor="middle" font-size="12" font-weight="600" fill="#1f2430">sandbox 1</text>
<rect x="240" y="664" width="104" height="66" rx="8" fill="#ffffff" stroke="#9aa1ad" stroke-width="1.2"/>
<text x="292" y="701" text-anchor="middle" font-size="12" font-weight="600" fill="#1f2430">sandbox 2</text>
<rect x="356" y="664" width="96" height="66" rx="8" fill="none" stroke="#9aa1ad" stroke-width="1.2" stroke-dasharray="4 4"/>
<text x="404" y="701" text-anchor="middle" font-size="12" fill="#6b7280">… sandbox N</text>
<!-- ============ RWX volume ============ -->
<path d="M880,632 a100,12 0 0 0 200,0 v66 a100,12 0 0 1 -200,0 z" fill="#ffffff" stroke="#9aa1ad" stroke-width="1.2"/>
<ellipse cx="980" cy="632" rx="100" ry="12" fill="#ffffff" stroke="#9aa1ad" stroke-width="1.2"/>
<text x="980" y="668" text-anchor="middle" font-size="13.5" font-weight="600" fill="#1f2430">RWX volume</text>
<text x="980" y="686" text-anchor="middle" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="10.5" fill="#6b7280">/shared/policy</text>
<!-- ============ flows (drawn last, over the boxes) ============ -->
<!-- (1) RolloutManager <-> engines -->
<path d="M300,240 V104 H920 V128" fill="none" stroke="#2b5fd9" stroke-width="1.6" stroke-dasharray="6 5" marker-start="url(#ab)" marker-end="url(#ab)"/>
<text x="640" y="97" text-anchor="middle" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="11.5" fill="#2b5fd9">① each turn: messages ⇄ tokens + logprobs</text>
<!-- (2) each agent loop <-> its own sandbox -->
<path d="M176,392 V660" fill="none" stroke="#2e7d5b" stroke-width="1.6" stroke-dasharray="6 5" marker-start="url(#ag)" marker-end="url(#ag)"/>
<path d="M292,392 V660" fill="none" stroke="#2e7d5b" stroke-width="1.6" stroke-dasharray="6 5" marker-start="url(#ag)" marker-end="url(#ag)"/>
<path d="M404,392 V660" fill="none" stroke="#2e7d5b" stroke-width="1.6" stroke-dasharray="6 5" marker-start="url(#ag)" marker-end="url(#ag)"/>
<rect x="172" y="576" width="240" height="20" fill="#f7f6f3"/>
<text x="292" y="590" text-anchor="middle" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="11.5" fill="#2e7d5b">② checkout · exec · grade</text>
<!-- (3) trajectories to trainer -->
<path d="M468,330 H498 V246 H522" fill="none" stroke="#2b5fd9" stroke-width="1.6" stroke-dasharray="6 5" marker-end="url(#ab)"/>
<text x="478" y="352" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="11" fill="#2b5fd9">③ trajectories + logprobs + rewards</text>
<!-- (5) weights write: Megatron -> volume -->
<path d="M712,246 H760 V665 H874" fill="none" stroke="#b23a20" stroke-width="1.6" stroke-dasharray="6 5" marker-end="url(#ar)"/>
<rect x="766" y="573" width="112" height="20" fill="#f7f6f3"/>
<text x="772" y="587" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="11" fill="#b23a20">⑤ write weights</text>
<!-- (6) engines pull: volume -> inference -->
<path d="M980,618 V484" fill="none" stroke="#b23a20" stroke-width="1.6" stroke-dasharray="6 5" marker-end="url(#ar)"/>
<rect x="892" y="545" width="176" height="20" fill="#f7f6f3"/>
<text x="980" y="559" text-anchor="middle" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="11" fill="#b23a20">⑥ engines pull weights</text>
<!-- ============ legend ============ -->
<path d="M100,812 H140" stroke="#2b5fd9" stroke-width="1.6" stroke-dasharray="6 5"/>
<text x="148" y="816" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="11" fill="#5c636e">rollout data</text>
<path d="M260,812 H300" stroke="#2e7d5b" stroke-width="1.6" stroke-dasharray="6 5"/>
<text x="308" y="816" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="11" fill="#5c636e">sandbox actions</text>
<path d="M440,812 H480" stroke="#b23a20" stroke-width="1.6" stroke-dasharray="6 5"/>
<text x="488" y="816" font-family="ui-monospace, SFMono-Regular, Menlo, monospace" font-size="11" fill="#5c636e">policy weights</text>
</svg>
policy-volume.yaml
# ReadWriteMany volume the trainer and inference engine(s) share for disk-based
# weight synchronization (the trainer writes updated weights; engines read them).
# Apply once before launching: sky volumes apply policy-volume.yaml
name: slime-policy
type: k8s-pvc
size: 100Gi
infra: kubernetes # your H100 Kubernetes context, e.g. kubernetes/my-cluster
config:
access_mode: ReadWriteMany
pyproject.toml
[project]
name = "coding-agent-rl"
version = "0.1.0"
description = "Local dev dependencies for the coding-agent-rl example (SkyPilot JobGroups + Sandboxes)"
requires-python = ">=3.12"
dependencies = [
# Agent harness — used directly (unmodified loop); we subclass its
# Environment/Model seams. Brings litellm.
"mini-swe-agent>=2.4.5",
# slime CPU-path runtime deps (slime itself ships in the slimerl/slime image;
# its full requirements.txt pulls CUDA-only packages like ring_flash_attn, so
# it isn't installed here).
"aiohttp>=3.9",
"httpx[http2]",
"transformers>=4.44",
"torch",
"pyyaml",
"omegaconf",
"numpy",
"pillow", # slime.utils.processing_utils imports PIL
# 3b dataset loader
"datasets",
]
# The sky.sandbox SDK (SkyPilot Sandboxes) isn't on PyPI — a Sandboxes-enabled API
# server ships it to ~/.sky/bin/wheels/, and scripts/trainer-setup.sh installs it
# from the mounted /wheels dir on the cluster. See README.
#!/bin/bash
# slime launcher for the coding-agent RL loop. Builds the train_async.py / train.py
# argument set from env vars and submits it as a Ray job to the trainer's local Ray
# cluster. Everything is env-driven so the Job Group YAML is the single source of
# config; the notable knobs:
#
# Rollout --custom-generate-function-path generate.generate (agentic loop;
# code/ is on PYTHONPATH via the Ray runtime env below). Reward comes
# from the sandbox eval inside generate() — no reward model.
# Metrics --custom-rollout-log-function-path rollout_metrics.log_rollout_metrics
# (pass@k + sandbox timing; wandb is opt-in via WANDB_API_KEY).
# Sizing ROLLOUT_BATCH_SIZE x N_SAMPLES_PER_PROMPT = concurrent rollouts =
# concurrent sandbox claims; keep SANDBOX_POOL_REPLICAS >= that product.
# Topology TRAIN_SCRIPT=train_async.py -> async rollout/train overlap (default).
# WEIGHT_TRANSPORT=disk|nccl -> weight-sync carrier (default disk).
# WEIGHT_MODE=delta|full -> ship changed bytes only, or all.
#
# The rollout-actor env (sandbox SDK auth, AGENT_*/SANDBOX_* knobs) is passed through
# the Ray job runtime env — Ray actors don't reliably inherit the shell env.
set -ex
# will prevent ray from buffering stdout/stderr
export PYTHONUNBUFFERED=1
NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
if [ "$NVLINK_COUNT" -gt 0 ]; then
HAS_NVLINK=1
else
HAS_NVLINK=0
fi
echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"
if command -v nvidia-smi >/dev/null 2>&1; then
DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l | tr -d ' ')
else
DETECTED_GPUS=0
fi
NUM_GPUS=${NUM_GPUS:-${DETECTED_GPUS}}
if [ -z "$NUM_GPUS" ] || [ "$NUM_GPUS" -le 0 ]; then
NUM_GPUS=8
fi
echo "NUM_GPUS: $NUM_GPUS"
# Model axis. Defaults match this example (Qwen3-14B); override via the trainer
# YAML env to train a different model.
# MODEL_SCRIPT: slime arch args (scripts/models/*.sh)
# MODEL_DIR: checkpoint dir under /root (hf download target in the YAML)
# TP: tensor-model-parallel size (= trainer GPU count; 14B -> 4)
MODEL_SCRIPT="${MODEL_SCRIPT:-qwen3-14B.sh}"
MODEL_DIR="${MODEL_DIR:-Qwen3-14B}"
TP="${TP:-4}"
# Topology: SLIME_MODE=async (default; overlap rollout with training) | sync
# (serialize). Maps to slime's train_async.py / train.py.
SLIME_MODE="${SLIME_MODE:-async}"
TRAIN_SCRIPT=$([ "$SLIME_MODE" = "sync" ] && echo train.py || echo train_async.py)
# Vendored script lives outside the slime tree; point at the in-image copy.
SLIME_DIR="${SLIME_DIR:-/root/slime}"
source "${SLIME_DIR}/scripts/models/${MODEL_SCRIPT}"
# Our rollout code (generate.py, sandbox_env.py) synced via workdir.
AGENT_CODE_DIR="${AGENT_CODE_DIR:-$HOME/sky_workdir/code}"
PROMPT_DATA="${PROMPT_DATA:-$HOME/sky_workdir/mbppplus.jsonl}"
CKPT_ARGS=(
--hf-checkpoint /root/${MODEL_DIR}
--ref-load /root/${MODEL_DIR}_torch_dist
--load /root/${MODEL_DIR}_slime/
--save /root/${MODEL_DIR}_slime/
--save-interval ${SAVE_INTERVAL:-5}
)
# ROLLOUT_SHUFFLE=0 -> deterministic (same instances every step) for clean
# benchmark timing; default on (real training wants variety). Empty-array idiom
# so "off" contributes zero args (no empty token to argparse).
SHUFFLE_ARG=(); [ "${ROLLOUT_SHUFFLE:-1}" != "0" ] && SHUFFLE_ARG=(--rollout-shuffle)
ROLLOUT_ARGS=(
--prompt-data ${PROMPT_DATA}
--input-key text
--label-key label
--metadata-key metadata
# NO --apply-chat-template: generate.py's harness renders its own messages
"${SHUFFLE_ARG[@]}"
--num-rollout ${NUM_ROLLOUT:-3}
--rollout-batch-size ${ROLLOUT_BATCH_SIZE:-4}
--n-samples-per-prompt ${N_SAMPLES_PER_PROMPT:-4}
# Per-turn generation cap (feeds sampling_params -> adapter session
# defaults); MBPP solutions are short, 2048 bounds a runaway turn.
--rollout-max-response-len ${ROLLOUT_MAX_RESPONSE_LEN:-2048}
# Whole-session context cap enforced by the adapter (finish_reason=length
# past this).
--rollout-max-context-len ${ROLLOUT_MAX_CONTEXT_LEN:-8192}
--rollout-temperature ${ROLLOUT_TEMPERATURE:-1}
# Global (training) batch = train on exactly what we generated this step, so
# it's DERIVED from the two knobs above (not a separate config var).
--global-batch-size $(( ${ROLLOUT_BATCH_SIZE:-4} * ${N_SAMPLES_PER_PROMPT:-4} ))
--balance-data
--custom-generate-function-path generate.generate
# Adds agent/pass@1, agent/pass@k, agent/sandbox_*_sec/* to the metrics
# slime already logs; returns falsy so slime's default logging still runs.
--custom-rollout-log-function-path rollout_metrics.log_rollout_metrics
)
PERF_ARGS=(
--tensor-model-parallel-size ${TP:-2}
--sequence-parallel
--pipeline-model-parallel-size 1
--context-parallel-size 1
--expert-model-parallel-size 1
--expert-tensor-parallel-size 1
--recompute-granularity full
--recompute-method uniform
--recompute-num-layers 1
--use-dynamic-batch-size
--max-tokens-per-gpu ${MAX_TOKENS_PER_GPU:-9216}
)
GRPO_ARGS=(
--advantage-estimator grpo
--use-kl-loss
--kl-loss-coef 0.00
--kl-loss-type low_var_kl
--entropy-coef 0.00
--eps-clip 0.2
--eps-clip-high 0.28
)
OPTIMIZER_ARGS=(
--optimizer adam
--lr 1e-6
--lr-decay-style constant
--weight-decay 0.1
--adam-beta1 0.9
--adam-beta2 0.98
)
# wandb is OPT-IN via WANDB_API_KEY: unset => empty WANDB_ARGS => run stays
# fully disabled (never breaks a run for a missing key). slime derives the
# run name from --wandb-group (no --wandb-run-name arg exists; wandb_utils.py
# uses group as the name), so we encode the config in the group.
WANDB_ARGS=()
# SECURITY: this test EXPANDS WANDB_API_KEY, so `set -x` (on since the top of the
# script) would echo the key into the job log. Trace off around just the test;
# WANDB_ARGS itself holds only project/group/team, so re-enable right after.
set +x
if [ -n "${WANDB_API_KEY:-}" ]; then
WANDB_ARGS=(
--use-wandb
--wandb-project "${WANDB_PROJECT:-skypilot-rl-workloads}"
--wandb-group "${WANDB_GROUP:-coding-agent}"
${WANDB_TEAM:+--wandb-team "${WANDB_TEAM}"}
)
fi
set -x
SGLANG_ARGS=(
# per-engine GPU count (= each external sglang job's TP). Multi-engine
# scale-out passes several addrs in EXTERNAL_ENGINE_ADDR; slime auto-detects
# each engine's GPU count at discovery, but keep this consistent.
--rollout-num-gpus-per-engine ${ROLLOUT_GPUS_PER_ENGINE:-2}
--sglang-mem-fraction-static 0.7
)
MISC_ARGS=(
# default dropout in megatron is 0.1
--attention-dropout 0.0
--hidden-dropout 0.0
# should be good for model performance
--accumulate-allreduce-grads-in-fp32
--attention-softmax-in-fp32
# need to comment this when using model with MLA
--attention-backend flash
)
# --- Topology / weight-transport hook ---------------------------------------
# COLOCATE=1 -> single-job IPC (no external engine, no disk); actor.py
# forces UpdateWeightFromTensor when --colocate (actor.py:139-143).
# WEIGHT_TRANSPORT -> disaggregated carrier: disk (default, proven 3b) | nccl.
# nccl: slime opens a cross-job NCCL group (trainer rank0 + engine GPUs;
# update_weight_from_distributed.py:293-321) — the external job dials the
# TRAINER pod back on a random port. Drops the --update-weight-disk-* args.
# WEIGHT_MODE = full | delta (arguments.py:137). 'delta' bytewise-diffs weights
# vs a pinned-CPU snapshot of the last broadcast and ships only changed
# positions+values -- LOSSLESS apply (no arithmetic), so reward should overlay
# full; the payoff is bandwidth/latency (perf/update_weights_time), config-
# dependent (fewer changed bytes at low LR / bigger win at scale). Delta is
# NOT supported with --colocate and only rides nccl|disk (arguments.py:1734-1742).
WEIGHT_MODE="${WEIGHT_MODE:-full}"
if [ "${COLOCATE:-0}" = "1" ]; then
if [ "$WEIGHT_MODE" = "delta" ]; then echo "WARN: delta unsupported with --colocate; forcing full"; fi
TOPO_ARGS=( --colocate --update-weight-mode full )
echo "TOPOLOGY: colocated (IPC / update_weights_from_tensor)"
else
TOPO_ARGS=( --rollout-external-engine-addrs ${EXTERNAL_ENGINE_ADDR} --update-weight-mode ${WEIGHT_MODE} )
if [ "$WEIGHT_MODE" = "delta" ]; then
# delta REQUIRES a rollout-host-local (NVMe) checkpoint dir: each engine host
# pulls the published full/delta from --update-weight-disk-dir into here and
# reloads from it (arguments.py:224, engine.pull_weights). Container-local on
# the engine pod (per-host, NOT the shared volume). arguments.py:2016 enforces it.
TOPO_ARGS+=( --update-weight-encoding ${WEIGHT_ENCODING:-indices}
--update-weight-local-checkpoint-dir ${WEIGHT_LOCAL_CKPT_DIR:-/root/slime-local-ckpt} )
fi
if [ "${WEIGHT_TRANSPORT:-disk}" = "nccl" ]; then
TOPO_ARGS+=( --update-weight-transport nccl )
echo "TOPOLOGY: disaggregated, mode=${WEIGHT_MODE}, transport=nccl (cross-job NCCL group)"
else
TOPO_ARGS+=( --update-weight-transport disk --update-weight-disk-dir /shared/policy --update-weight-disk-keep-files )
echo "TOPOLOGY: disaggregated, mode=${WEIGHT_MODE}, transport=disk"
fi
fi
# launch the master node of ray in container
export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
# Idempotent: a prior run's Ray head survives job failure (daemonized).
# Probe slime's dashboard (8265; SkyPilot's own runtime uses 8266) and reuse.
if curl -sf http://127.0.0.1:8265/api/version >/dev/null 2>&1; then
echo "Ray head already running on 8265; reusing."
else
ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 --dashboard-agent-listen-port 52366 --metrics-export-port 8091
fi
# Runtime env for the ray job: PYTHONPATH gets Megatron + our rollout code
# (so load_function can import "generate.generate"); the AGENT_*/sandbox-SDK
# vars must be here because the RolloutManager actor (where generate() runs)
# gets its env from the ray runtime env, not from this shell.
# NCCL_DEBUG=INFO on the nccl variant so the engine/trainer logs print the
# chosen net transport -- grep `NET/IB` vs `NET/Socket` at reload to report
# whether the cross-job group rode RDMA or fell back to TCP.
NCCL_DEBUG_VAL="${NCCL_DEBUG:-}"
if [ "${WEIGHT_TRANSPORT:-disk}" = "nccl" ] && [ "${COLOCATE:-0}" != "1" ]; then
NCCL_DEBUG_VAL="${NCCL_DEBUG:-INFO}"
fi
# SECURITY: from here on we handle secrets (WANDB_API_KEY, the SA token) inside
# RUNTIME_ENV_JSON and the ray submit command. Turn OFF shell tracing so `set -x`
# never echoes their values into the job log. wandb picks up WANDB_API_KEY from the
# env below (no --wandb-key arg needed — that would re-leak it on the command line).
set +x
RUNTIME_ENV_JSON="{
\"env_vars\": {
\"PYTHONPATH\": \"/root/Megatron-LM/:${AGENT_CODE_DIR}\",
\"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
\"PYTORCH_CUDA_ALLOC_CONF\": \"${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}\",
\"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\",
\"NCCL_DEBUG\": \"${NCCL_DEBUG_VAL}\",
\"SKYPILOT_API_SERVER_ENDPOINT\": \"${SKYPILOT_API_SERVER_ENDPOINT:-}\",
\"SKYPILOT_SERVICE_ACCOUNT_TOKEN\": \"${SKYPILOT_SERVICE_ACCOUNT_TOKEN:-}\",
\"AGENT_POOL\": \"${AGENT_POOL:-swesmith}\",
\"AGENT_MAX_TURNS\": \"${AGENT_MAX_TURNS:-1}\",
\"AGENT_EXEC_TIMEOUT_SEC\": \"${AGENT_EXEC_TIMEOUT_SEC:-60}\",
\"AGENT_EVAL_TIMEOUT_SEC\": \"${AGENT_EVAL_TIMEOUT_SEC:-120}\",
\"AGENT_SANDBOX_LIFETIME_SEC\": \"${AGENT_SANDBOX_LIFETIME_SEC:-0}\",
\"AGENT_ROLLOUT_GUARD_SEC\": \"${AGENT_ROLLOUT_GUARD_SEC:-0}\",
\"AGENT_SGLANG_URL\": \"${AGENT_SGLANG_URL:-}\",
\"AGENT_LITELLM_MODEL\": \"${AGENT_LITELLM_MODEL:-openai/slime-actor}\",
\"AGENT_WORKDIR\": \"${AGENT_WORKDIR:-/workspace}\",
\"AGENT_LOCAL_EXEC\": \"${AGENT_LOCAL_EXEC:-}\",
\"SANDBOX_NO_POOL\": \"${SANDBOX_NO_POOL:-0}\",
\"ROLLOUT_BATCH_SIZE\": \"${ROLLOUT_BATCH_SIZE:-}\",
\"N_SAMPLES_PER_PROMPT\": \"${N_SAMPLES_PER_PROMPT:-}\",
\"WANDB_API_KEY\": \"${WANDB_API_KEY:-}\",
\"WANDB_GROUP\": \"${WANDB_GROUP:-coding-agent}\",
\"WANDB_PROJECT\": \"${WANDB_PROJECT:-skypilot-rl-workloads}\",
\"WANDB_TEAM\": \"${WANDB_TEAM:-}\",
\"MSWEA_SILENT_STARTUP\": \"1\"
}
}"
# Ray's dashboard job agent registers async after `ray start`; submitting too
# early yields "No available agent to submit job". We retry ONLY that
# submission-REJECTION case. A job that actually RAN and then failed must NOT be
# re-run: `ray job submit` blocks+tails and returns the JOB's exit code, so a
# training crash would otherwise loop forever, re-burning GPU and hiding the real
# error (observed live: an nccl weight-sync 400 at step 2 crash-looped 5x). So we
# distinguish the two by grepping the submit output and propagate real failures.
submit_log="/tmp/raysubmit-$$.log"
submit_ok=0
for attempt in 1 2 3 4 5 6; do
set +e
ray job submit --address="http://127.0.0.1:8265" \
--runtime-env-json="${RUNTIME_ENV_JSON}" \
-- python3 ${TRAIN_SCRIPT} \
--actor-num-nodes 1 \
--actor-num-gpus-per-node ${NUM_GPUS} \
${TOPO_ARGS[@]} \
${MODEL_ARGS[@]} \
${CKPT_ARGS[@]} \
${ROLLOUT_ARGS[@]} \
${OPTIMIZER_ARGS[@]} \
${GRPO_ARGS[@]} \
${WANDB_ARGS[@]} \
${PERF_ARGS[@]} \
${SGLANG_ARGS[@]} \
${MISC_ARGS[@]} 2>&1 | tee "${submit_log}"
rc=${PIPESTATUS[0]}
set -e
if [ "$rc" -eq 0 ]; then submit_ok=1; break; fi
# wandb TEARDOWN RACE: training can finish + save the final checkpoint, then wandb's
# atexit service-teardown hits an already-closed connection -> ConnectionResetError ->
# spurious rc=1 (surfaces AFTER `exit 0`). If the run demonstrably completed (checkpoint
# saved) and that teardown is the failure, treat it as success. Still fails on real
# errors (OOM / NCCL / pool-not-found / warm-pool timeout).
if grep -qiE 'successfully saved checkpoint from iteration' "${submit_log}" \
&& grep -qiE 'teardown_atexit|ConnectionResetError' "${submit_log}" \
&& ! grep -qiE 'CUDA out of memory|NCCL error|not found on any context|Timed out after [0-9]+s waiting for warm pool' "${submit_log}"; then
echo "run completed (final checkpoint saved); rc=$rc is the known wandb atexit teardown race — treating as SUCCESS."
submit_ok=1; break
fi
# Retry ONLY on submission-rejection (agent not up yet), never on a job crash.
if grep -qiE "No available agent|Failed to submit job|Connection refused|ConnectionError|500 Internal" "${submit_log}"; then
echo "ray agent not ready (attempt $attempt); retrying in 10s..."
sleep 10
else
echo "ray job RAN and FAILED (rc=$rc) — not retrying; surfacing the error above."
exit "$rc"
fi
done
[ "$submit_ok" -eq 1 ]
scripts/engine-run.sh
#!/bin/bash
# SGLang engine job — run. One inference server per Job Group engine job,
# reachable by the trainer at <job>-0.<jobgroup>:30000 (stable hostname).
# Env (all optional; defaults match this example):
# MODEL_DIR local checkpoint dir under /root (default Qwen3-14B)
# ENGINE_TP tensor-parallel size per engine (default 1 = 1 H100)
# SGLANG_MEM_FRACTION static KV/mem fraction (default 0.7)
set -ex
MODEL_DIR="${MODEL_DIR:-Qwen3-14B}"
ENGINE_TP="${ENGINE_TP:-1}"
python -m sglang.launch_server --model-path /root/${MODEL_DIR} --tp ${ENGINE_TP} \
--host 0.0.0.0 --port 30000 --mem-fraction-static ${SGLANG_MEM_FRACTION:-0.7}
scripts/engine-setup.sh
#!/bin/bash
# SGLang engine job — setup. Runs on each inference job of the Job Group.
# Just the model checkpoint: pull it once into the engine pod's local disk.
# Env (all optional; defaults match this example's Qwen3-14B):
# MODEL_DIR local dir under /root (default Qwen3-14B)
# MODEL_HF_REPO HF repo to download (default Qwen/Qwen3-14B)
set -ex
MODEL_DIR="${MODEL_DIR:-Qwen3-14B}"
MODEL_HF_REPO="${MODEL_HF_REPO:-Qwen/Qwen3-14B}"
pip install -q -U "huggingface_hub[cli]"
[ -d /root/${MODEL_DIR} ] || hf download ${MODEL_HF_REPO} --local-dir /root/${MODEL_DIR}
scripts/trainer-run.sh
#!/bin/bash
# Trainer job — run. The rollout-side orchestration around the slime launch:
# 1. wait for EVERY engine job to serve (all-or-nothing gang start),
# 2. build the engine-address list slime discovers,
# 3. warm one sandbox pool per repo image (unless SANDBOX_NO_POOL=1),
# 4. hand off to run-slime.sh (which builds + submits the slime Ray job).
#
# Scaling the inference fleet is entirely declarative: add an SGLang job to
# the YAML and add its name to SGLANG_MEMBERS below — this script is unchanged
# for 1, 2, or N engines.
# SGLANG_MEMBERS space-separated Job Group engine job names (default "sglang").
# Each is reached at <job>-0.<jobgroup>:30000.
set -ex
# Self-clean the run's sandbox pool on exit (success OR failure).
trap 'python ~/sky_workdir/code/pool_cleanup.py ~/sky_workdir/swesmith-pools.json 2>&1 || true' EXIT
# 1 + 2: health-gate every engine, then assemble the discovery address list.
SGLANG_MEMBERS="${SGLANG_MEMBERS:-sglang}"
ENGINE_ADDRS=""
for m in ${SGLANG_MEMBERS}; do
URL="http://${m}-0.${SKYPILOT_JOBGROUP_NAME}:30000"
healthy=0
for i in $(seq 1 180); do curl -sf "${URL}/health" >/dev/null 2>&1 && { echo "engine ${m} healthy"; healthy=1; break; }; sleep 10; done
[ "$healthy" = 1 ] || { echo "ERROR: engine ${m} never became healthy after 30m" >&2; exit 1; }
ENGINE_ADDRS="${ENGINE_ADDRS}${ENGINE_ADDRS:+ }${m}-0.${SKYPILOT_JOBGROUP_NAME}:30000"
done
echo "engine addrs: ${ENGINE_ADDRS}"
# 3: one warm sandbox pool per repo image — skipped in no-pool mode, where each
# rollout instead creates an on-demand ad-hoc sandbox from its own image.
if [ "${SANDBOX_NO_POOL:-0}" != "1" ]; then
PYTHONPATH=~/sky_workdir/code python - <<'PYEOF'
import json, os
import sandbox_env
with open(os.path.expanduser('~/sky_workdir/swesmith-pools.json')) as f:
pools = json.load(f)
reps = int(os.environ.get('SANDBOX_POOL_REPLICAS', '184'))
for p in pools:
sandbox_env.ensure_pool(p['pool'], image=p['image'], replicas=reps)
PYEOF
else echo "SANDBOX_NO_POOL=1 -> no warm pools; using on-demand ad-hoc sandboxes"; fi
# 4: hand off to slime.
cd /root/slime
set -o pipefail
EXTERNAL_ENGINE_ADDR="${ENGINE_ADDRS}" \
PROMPT_DATA="$HOME/sky_workdir/swesmith.jsonl" \
bash ~/sky_workdir/run-slime.sh
scripts/trainer-setup.sh
#!/bin/bash
# Trainer job — setup. Installs the rollout deps (agent harness, sandbox SDK,
# SWE-smith grading), pre-converts the HF checkpoint to Megatron torch_dist, and
# generates the SWE-smith dataset + warm-pool manifest. Idempotent — each guarded
# step is skipped if its output already exists.
# Env (all optional; defaults match this example's Qwen3-14B):
# MODEL_DIR / MODEL_HF_REPO / MODEL_SCRIPT model axis (see run-slime.sh)
# SWESMITH_REPOS / SWESMITH_PER_REPO / SWESMITH_POOL_SUFFIX dataset subset
set -ex
MODEL_DIR="${MODEL_DIR:-Qwen3-14B}"
MODEL_HF_REPO="${MODEL_HF_REPO:-Qwen/Qwen3-14B}"
MODEL_SCRIPT="${MODEL_SCRIPT:-qwen3-14B.sh}"
pip install -q -U "huggingface_hub[cli]"
# SkyPilot Sandboxes SDK (platform feature; provides sky.sandbox used by sandbox_env.py).
# It's a standalone wheel the sandbox-enabled API server ships to the client at
# ~/.sky/bin/wheels, mounted into this pod at /wheels by the trainer's file_mounts.
# --force-reinstall --no-deps: the dev0 version can otherwise silently no-op; its
# deps are light and satisfied by the base image + the installs below.
pip install --force-reinstall --no-deps /wheels/skypilot_sandbox_sdk-*.whl
pip install -q mini-swe-agent==2.4.5 litellm pillow datasets
# Host-side grading: SWE-smith's test command + log parser run on the TRAINER
# (not in the sandbox image), so swesmith/swebench are installed here.
pip install -q swebench
pip install -q "swesmith @ git+https://github.com/SWE-bench/SWE-smith.git" || pip install -q swesmith
[ -d /root/${MODEL_DIR} ] || hf download ${MODEL_HF_REPO} --local-dir /root/${MODEL_DIR}
if [ ! -d /root/${MODEL_DIR}_torch_dist ]; then
cd /root/slime
source scripts/models/${MODEL_SCRIPT}
PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
${MODEL_ARGS[@]} --hf-checkpoint /root/${MODEL_DIR} --save /root/${MODEL_DIR}_torch_dist
fi
cd ~/sky_workdir
[ -f swesmith.jsonl ] || python code/dataset_swesmith.py \
--out swesmith.jsonl --pools-out swesmith-pools.json \
${SWESMITH_REPOS:+--repos ${SWESMITH_REPOS}} --per-repo ${SWESMITH_PER_REPO} \
${SWESMITH_POOL_SUFFIX:+--pool-suffix ${SWESMITH_POOL_SUFFIX}}