Skip to content
Merged
Show file tree
Hide file tree
Changes from 10 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 90 additions & 0 deletions examples/experimental/openenv/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,90 @@
# OpenEnv Terminal-Bench-2 GRPO (GLM-4.7-Flash, single node)

Train GLM-4.7-Flash with GRPO on the HuggingFace [OpenEnv](https://github.com/huggingface/openenv)
**Terminal-Bench-2 (tbench2)** environment. A miles-side adapter runs the multi-turn
agentic loop (`reset(task_id)` → { policy emits one shell command → `step(exec)` →
feed output back } → `evaluate`) against an unmodified OpenEnv env server; the reward
is the binary pytest result (1.0 = all tests pass, else 0.0).

This guide targets a **single H200 node with 8 GPUs**. The run is colocated
(training + rollout on the same 8 GPUs): TP=4, EP=2, one SGLang engine per GPU.

## Prerequisites

- The node has Docker available (the env server launches one container per task).
- miles is installed and GLM-4.7-Flash weights are reachable (the launcher pulls
`zai-org/GLM-4.7-Flash` from HF and converts it to `torch_dist` on first run).
- Install the OpenEnv tbench2 env client (isolate it if its deps clash with the
miles image):

```bash
pip install -e <OpenEnv>/envs/tbench2_env
```

## 1. Build the prompt data

Clone the TB2 suite and emit one prompt row per `task_id`:

```bash
git clone --depth 1 https://github.com/laude-institute/terminal-bench-2.git /workspace/terminal-bench-2
python make_tbench2_data.py --tasks_dir /workspace/terminal-bench-2 --output /root/tbench2_train.jsonl
# add --n 8 for a small smoke subset
```

## 2. Start the env server

Run it in a separate shell (or off-node — see note). Docker mode gives real TB2
fidelity; it needs the Docker socket and pulls the per-task images on first use:

```bash
# Raise the open-file limit first (see Notes): the WebSocket env server holds an
# FD per live session + Docker connection and leaks sockets on unclean
# disconnects, so the default 1024 soft limit is exhausted on a long run.
ulimit -n 1048576
TB2_MODE=docker TB2_TASKS_DIR=/workspace/terminal-bench-2 MAX_CONCURRENT_ENVS=32 \
python -m tbench2_env.server.app --port 8003
```

`MAX_CONCURRENT_ENVS` caps live sandboxes; keep it at or below the rollout batch
concurrency. Per-task containers are heavy on disk — if you'd rather not colocate
them with the GPU workload, run the env server on a separate Docker host and point
the launcher at it via `--openenv-env-url http://<env-host>:8003`.

## 3. Launch training

```bash
python run-openenv-tbench2.py --openenv-env-url http://localhost:8003
```

Common overrides:

| Flag / env var | Default | Purpose |
| --- | --- | --- |
| `--openenv-env-url` | `http://localhost:8003` | Env server URL |
| `--prompt-data` | `/root/tbench2_train.jsonl` | Prompt set from step 1 |
| `--num-rollout` | (launcher) | Number of GRPO steps |
| `OPENENV_MAX_TURNS` | `30` | Max agent turns per episode |
| `OPENENV_MAX_ROLLOUT_TIME_SECONDS` | `3600` | Per-episode wall-clock cap; a straggler that exceeds it is terminated and scored 0 |
| `--dump-details <dir>` | off | Dump per-episode tokens/logprobs/masks/reward for inspection |
| `WANDB_KEY`, `--wandb-project`, `--wandb-team` | — | W&B logging |

## Notes

- **Reward signal.** The binary sparse reward needs a task subset where the base
policy *sometimes* succeeds (advantage variance). On the full TB2 suite,
GLM-4.7-Flash's low base solve-rate yields a near-flat GRPO signal — use a
variance-band subset (or a stronger base) to see a learning climb.
- **`_step` vs. rollout.** W&B `_step` is an internal log-call index that advances
several times per rollout; it is **not** the training step. Read the driver log's
`rollout N:` counter for true progress.
- **Sandbox leakage.** Upstream OpenEnv creates task containers with `remove=False`
and only tears them down on a clean session close (the idle reaper is off by
default), so an unclean disconnect (trainer crash) can orphan containers. Sweep
stale TB2 containers between runs, e.g. `docker rm -f` of any older than the
episode wall-cap.
- **Open-file limit.** The same unclean disconnects also leak socket FDs in the
env server process. On a long run under the default 1024 soft limit the accept
loop eventually fails every connection with `OSError: [Errno 24] Too many open
files`, silently throttling rollouts. Start the server with a raised limit
(`ulimit -n 1048576`, as in step 2); if a running server is already saturated,
restart it with the higher limit.
78 changes: 78 additions & 0 deletions examples/experimental/openenv/make_tbench2_data.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
"""Generate a prompt dataset for the OpenEnv Terminal-Bench-2 (tbench2) run.

The *tasks* are not hand-written: they are the
Terminal-Bench-2 suite shipped in the laude-institute/terminal-bench-2 repo. The
env serves the per-task instruction at reset(), so each prompt-data row only
needs the system prompt (how the agent should behave) plus the ``task_id`` in
metadata. ``openenv_agent_function`` drives the multi-turn loop and reads
metadata["task_id"].

task_ids are the top-level task directory names in the TB2 repo checkout
(e.g. "chess-best-move"); a valid task dir contains a ``task.toml``.

# all 89 tasks
python make_tbench2_data.py --tasks_dir /workspace/terminal-bench-2 \
--output /root/tbench2_train.jsonl
# a small smoke subset
python make_tbench2_data.py --tasks_dir /workspace/terminal-bench-2 \
--output /root/tbench2_smoke.jsonl --n 8
# an explicit subset
python make_tbench2_data.py --tasks_dir /workspace/terminal-bench-2 \
--tasks chess-best-move,circuit-fibsqrt
"""

import json
from pathlib import Path

from tap import Tap

# The agent contract must match openenv_agent_function._multi_turn: one shell
# command per turn inside a single ```bash block; TASK_COMPLETE to stop.
_SYSTEM = (
"You are an autonomous terminal agent solving a Terminal-Bench task. You will "
"be given the task instruction, then interact with a real Linux shell. On each "
"turn respond with EXACTLY ONE shell command inside a single ```bash code block "
"and nothing else. Inspect the environment, make the required changes, and "
"verify your work. When you are confident the task is fully complete, reply with "
"TASK_COMPLETE (with no code block)."
)


class Args(Tap):
tasks_dir: str = "/workspace/terminal-bench-2" # TB2 repo checkout
output: str = "/root/tbench2_train.jsonl"
n: int = 0 # 0 = all discovered tasks
tasks: str = "" # optional comma-separated explicit task_ids


def _discover_task_ids(tasks_dir: Path) -> list[str]:
return sorted(p.name for p in tasks_dir.iterdir() if (p / "task.toml").is_file())


def main() -> None:
args = Args().parse_args()
tasks_dir = Path(args.tasks_dir).expanduser().resolve()

if args.tasks:
task_ids = [t.strip() for t in args.tasks.split(",") if t.strip()]
else:
task_ids = _discover_task_ids(tasks_dir)
if args.n > 0:
task_ids = task_ids[: args.n]

missing = [t for t in task_ids if not (tasks_dir / t / "task.toml").is_file()]
if missing:
raise SystemExit(f"task_ids without a task.toml in {tasks_dir}: {missing}")

with open(args.output, "w") as f:
for tid in task_ids:
row = {
"prompt": [{"role": "system", "content": _SYSTEM}],
"metadata": {"task_id": tid},
}
f.write(json.dumps(row) + "\n")
print(f"Wrote {len(task_ids)} TB2 tasks to {args.output}")


if __name__ == "__main__":
main()
Loading
Loading