Metadata-Version: 2.4
Name: collimate-rl
Version: 0.5.0
Summary: Use a Collimate microVM control plane as the sandboxed reward grader for RL post-training — one interface over the node-local daemon or the GA gateway (any col_* key tier); SkyRL, OpenEnv and slime adapters included.
Author-email: "Omer Gero (Gero Consulting)" <hello@collimate.ai>
License: Apache-2.0
Project-URL: Homepage, https://collimate.ai
Keywords: rl,rlhf,grpo,sandbox,microvm,skyrl,openenv,slime,reward
Requires-Python: >=3.8
Description-Content-Type: text/markdown
License-File: LICENSE
Provides-Extra: skyrl
Requires-Dist: skyrl-gym; extra == "skyrl"
Provides-Extra: test
Requires-Dist: pytest>=7; extra == "test"
Provides-Extra: demo
Requires-Dist: rich>=13; extra == "demo"
Dynamic: license-file

# collimate-rl — Collimate for RL post-training

Run RL rollout execution and reward grading inside forked
[Collimate](https://collimate.ai) microVMs: each rollout gets an isolated,
reproducible copy-on-write fork of a warm template — isolation for untrusted
model output, density for the GRPO group fan-out — through the RL frameworks
you already use (SkyRL, OpenEnv, slime).

There are two ways to wire Collimate into an RL stack; this package ships both.

| Mode | What Collimate does | Needs | Status |
|---|---|---|---|
| **B — exec-grader** (this package) | runs each rollout in a forked microVM via `/v1/exec` to produce an isolated, reproducible **reward** | only today's control plane | ✅ done; tested vs real microVMs |
| **A — OpenEnv runtime** | runs an OpenEnv env-server **image** as a forked microVM, reached over a per-session endpoint | the control plane's `expose_ports` (per-clone DNAT) + an app-warm template (`WAIT_HTTP`) | ✅ end-to-end on KVM: `from_docker_image(provider=CollimateProvider)` drives `reset/step` over **WebSocket** into a fork that resumes a live server baked in the snapshot; 24 concurrent persistent-WS sessions / 0 leaked VMs. Provider in `collimate_rl.openenv` |

Mode B is itself a legitimate RL-environment integration: the environment logic
lives in the trainer, and the model's code is graded inside a Collimate microVM
instead of in-process — isolation for untrusted output, density for the GRPO
group fan-out, and an optional per-rollout seed for common random numbers.

## One interface, both transports, any key tier

Everything in this package — `grade`/`grade_batch`, `Sandbox`, and the three
framework adapters (SkyRL, OpenEnv, slime) — runs over **either** wire with the
**same interface**:

| transport | client | auth |
|---|---|---|
| node-local `serve` daemon (dev box, BYOC node, in-pod SESSION socket) | `CollimateClient` | none (node-local) |
| GA gateway (the managed cloud) | `GatewayClient` | Bearer `col_*` key — demo / pro / enterprise |

`GatewayClient` implements the daemon client's session surface 1:1
(`create_session` / `session_exec` / `exec_once` / `get_session` /
`list_sessions` / `delete_session`), and `collimate_rl.connect()` picks the
transport from what you name (explicit args first, then the environment):

```python
from collimate_rl import connect, grade, PropertyReward

client = connect()                          # env-resolved: in-pod socket >
                                            # COLLIMATE_API_KEY (gateway) >
                                            # COLLIMATE_SERVE_URL > localhost daemon
client = connect("http://127.0.0.1:8080")   # a daemon, explicitly
client = connect(api_key="col_demo_...")    # the gateway, any key tier

client.images()                             # the ready-to-run image catalog
                                            # (gateway; a daemon client raises
                                            # a clear "gateway only" error)

res = grade(client, "dkr-python311", candidate, spec)   # identical either way
```

On the managed cloud, start from the catalog: `images()` lists the images the
platform keeps ready (id / name / tier / description), and a catalog id goes
straight into `create_sandbox` — pools are platform-managed, so there is
nothing to size. Baking your own image (`Template.bake`) is the pro path.

## Environments & sessions (pro/ent)

The concepts: an **environment** is the thing you run in (a catalog image, or
one you baked); a **session** is the bounded period you use it at width; a
**sandbox** is one instance. Pools exist, but they are the platform's job —
you only ever declare *how wide the next session is*:

```python
with client.env("python", width=512) as env:     # declare the session width
    with env.sandbox() as sess:                  # one sandbox on it
        group = sess.fork_group(512)             # ... your GRPO fan-out
# exit -> env.release(): capacity relaxes to baseline immediately
```

- `client.env(<catalog ref>, width=N)` bakes the catalog image into *your*
  environment with the width declared up front, so the platform pre-warms it
  (and reserves exec capacity for the full width) before the first rollout.
- `client.env(<your template id>, width=N)` declares the width on an
  environment you already own (`expect_width` under the hood).
- `env.release()` — automatic on context exit, exceptions included — ends the
  session so you stop paying for warmth the moment you're done. It is always
  optional: an idle environment decays back to baseline on its own.

Lower-level: `client.expect_width(template_id, N)` /
`client.release_template(template_id)` map 1:1 to
`POST /v1/templates/{id}/expect|release`.

Key-tier differences are enforced server-side and surface as **typed errors**
(`DemoTierLimit` with the friendly upgrade message, `quota_exceeded`, ...) —
never as silent behavior changes. One thing to size for: the gateway has no
one-shot `/v1/exec`, so `exec_once` composes create → exec → delete
client-side (reap guaranteed even on failure) and each in-flight grade briefly
holds one concurrency slot — keep `grade_batch(concurrency=...)` within your
tier's cap (demo: 8).

## Framework adapters

| framework | adapter | wiring |
|---|---|---|
| **SkyRL** | `collimate_rl.skyrl` `CollimateCodeEnv` | `register("collimate_code")` + `env_config` (add `api_key:` for the gateway) |
| **OpenEnv** (verl / TRL / TorchForge / ...) | `collimate_rl.openenv` `CollimateProvider` | `from_docker_image(image, provider=CollimateProvider(...))` |
| **slime** | `collimate_rl.slime` `reward_func` / `group_reward_func` | `--custom-rm-path collimate_rl.slime.reward_func` (no slime import, zero deps) |

## Install

```bash
pip install collimate-rl              # zero runtime deps (stdlib-only client)
pip install "collimate-rl[skyrl]"     # + skyrl-gym for the SkyRL adapter
pip install "collimate-rl[demo]"      # + rich, for the `collimate` visual CLI
```

## 30-second tour (no GPU, no KVM)

```bash
# an in-process stand-in for the real daemon (⚠️ dev/test only — NOT a sandbox):
python -m collimate_rl.local_backend --port 8080 &
```

```python
from collimate_rl import CollimateClient, PropertyReward, grade, group_seeds

client = CollimateClient("http://127.0.0.1:8080")     # the local backend above —
                                                      # or a real daemon, or
                                                      # connect(api_key="col_...")
spec = PropertyReward(
    reference="    return sum(c in 'aeiouAEIOU' for c in x)",
    generator="''.join(random.choice('abcdEIOUxyz') for _ in range(8))",
    k=16,
)
res = grade(client, "dkr-python311", completion_code, spec, seed="ab"*32)
print(res.reward, res.seed, res.seed_ack)
```

## What's here

```
collimate_rl/
  client.py         CollimateClient — stdlib client for /v1/sessions + /v1/exec
  gateway.py        GatewayClient — the GA gateway (Bearer col_* key, any tier),
                    incl. the daemon-parity session surface
  transport.py      connect() — one entry point, either transport
  template.py       Ready / Template.bake / Session — the gateway bake+session flow
  rewards.py        UnitTestReward / PropertyReward, grade(), grade_batch(),
                    reward_spec_from_ground_truth()
  sandbox.py        Sandbox — the in-worker-pod hot path (fork / exec / attach /
                    byte-exact download + upload_bytes); works over either transport
  seeds.py          common-random-numbers seed helpers (vanilla / crn / deterministic)
  local_backend.py  ⚠️ DEV/TEST in-process stand-in (NOT a sandbox) — runs the
                    exec contract locally so tests + demos work with no VMM
  skyrl/            SkyRL BaseTextEnv adapter (CollimateCodeEnv) + register()
  openenv/          OpenEnv ContainerProvider (CollimateProvider) — Mode A
  slime/            slime custom-RM hooks (reward_func / group_reward_func)
tests/                        230+ tests, all run on a laptop (no KVM)
```

Each adapter's module docstring is its guide; the source tree additionally
carries worked examples (`examples/`) and a throughput benchmark
(`benchmarks/`).

**Validated end-to-end** against a real `serve` daemon forking actual microVMs
on a KVM host: GRPO-group rewards match the local backend exactly, common random
numbers reproduce on the real guest CRNG (`seed_ack: verified`), and every exec
is captured in a bit-for-bit replay manifest. The SkyRL layer is exercised
through the real `skyrl_gym.make()` registry (py3.12).

## File IO (byte-exact, no wire change)

`Sandbox.download(path) -> bytes` and `Sandbox.upload_bytes(path, data)` move
arbitrary bytes in and out of a live clone purely over the existing exec API:

```python
with Sandbox("dkr-python311") as sb:
    sb.upload_bytes("/task/input.tar", tar_bytes)      # verified in-guest (size+sha256)
    r = sb.exec(command="cd /task && ./run.sh")
    patch = sb.download("/task/out.patch")             # bytes, byte-exact
```

- **Byte-exact, both directions.** `upload_bytes` streams base64 chunks through
  `printf | base64 -d` into a temp file, verifies size + sha256 **in-guest**,
  then `mv`s into place — so a failed upload never leaves a truncated target,
  and (unlike `exec(files=[...])`, which writes *text* via heredoc and appends
  a trailing newline to content not ending in one) NULs, newlines and
  missing-EOF-newlines survive exactly. `download` verifies the guest-side
  sha256 against the decoded bytes before returning them.
- **Sentinel framing vs the merged-stderr hazard.** Guest stdout and stderr
  share one pipe (the response's `stderr` is always `""`), so `download` frames
  the base64 payload between per-call random sentinel lines and parses
  *strictly* between them: noise outside the frame is ignored, noise that
  breaks the frame (missing/duplicated sentinel, extra lines) raises, and noise
  spliced into the payload line itself is caught by the sha256 check. A
  download either returns exact bytes or raises — never silently corrupts.
- **Task vs infra errors.** In-guest failures (missing file, unwritable parent,
  disk full, corrupted frame) raise the typed `CollimateTaskError` (a
  `CollimateError` subclass carrying `exit_code` + the guest transcript in
  `.output`); infra failures keep raising plain `CollimateError` exactly as
  `exec` does.
- **Limits.** Downloads are capped at ~11 MiB per file (the guest agent's
  16 MiB exec-output cap ÷ base64's 4/3 inflation; larger files fail with a
  clear in-guest error). Uploads go one ~96 KiB base64 chunk per exec (the
  composed command reaches the guest as a single `sh -c` argv string, bounded
  by the kernel's 128 KiB per-string cap), i.e. ~15 execs — a few ms each on a
  warm clone — per MiB. The guest needs a POSIX shell with a `printf` builtin
  plus `base64` and `sha256sum` (busybox and coreutils both qualify).

## Per-rollout egress (networked templates)

Egress needs a **networked template** — a guest NIC (`no_netns:false` / `NET=1`).
Bake one with `--network` (or `--port` / `--ready-http`, which imply it):

```bash
collimate bake --image my-tool:latest --network --width 64      # a NIC'd template
```

**Networked templates are pro/ent only** — as is baking at all: demo keys run
the ready-made catalog (`client.images()` → `create_sandbox(<catalog id>)`),
while bringing your own image (`Template.bake` / `collimate bake`) is the pro
path. A `demo`-tier request for a NIC (via `--network`, `--port`, or
`--ready-http`) is rejected `403 network_requires_pro`; demo images are
vsock-only. The SDK equivalent is `Ready(network=True)` (or `ports=[...]` /
`http=...`) passed to `Template.bake`.

A networked template gives every fork its own netns, so egress is enforced **per
rollout**, not per host: two forks of the same template can run different
policies at the same time. Pass `egress=` when you open the sandbox:

```python
with Sandbox("py-swebench", egress={"mode": "none"}) as sb: ...             # no egress at all
with Sandbox("py-swebench", egress={"preset": "rl-leakage"}) as sb: ...     # block answer sources
with Sandbox("py-swebench", egress={"allow": ["pypi.org", "10.20.0.0/16"]}) as sb: ...
```

Modes: `open` (the default when unset) | `none` | `allowlist` (`allow`) |
`denylist` (`deny`/`preset`). `allowlist`/`none` are fail-closed; a domain
`denylist` is best-effort (domains resolve at apply time). To default a whole
template, bake the policy in at build with the `EGRESS` knob (it writes the
`egress` object into `collimate.json`); a per-call `egress=` overrides it.
Requires a networked template (`--network`): a vsock-only template rejects a
policy with `400 template is vsock-only`.

## Failure semantics (v0.1.1)

Production RL runs hit load-shed, drains, and long test suites; the SDK makes
each of those explicit instead of folding them into the reward signal:

- **Infra failures never masquerade as task failures.**
  `grade_batch(..., on_infra_error="raise")` (the default) re-raises the first
  backend error so the trainer resamples the group — averaging infra zeros into
  GRPO advantages poisons the gradient. `"zero"` keeps positional alignment and
  flags the miss (`GradeResult.infra_error=True`); `"skip"` drops it.
- **Draining nodes are retry-safe.** A node draining for a rolling upgrade
  refuses new work with 503 *before* admitting it (no VM forked, no exec run),
  so the client retries with backoff and then raises the typed
  `CollimateDrainingError` — catch it and route the request to another node.
- **Retries back off with exponential full jitter,** bounded by a configurable
  total budget (`CollimateClient(..., backoff=0.15, retry_budget=30.0)`), so a
  shed GRPO group doesn't re-arrive as a synchronized thundering herd.
- **Long execs are never abandoned client-side.** The exec socket timeout is
  `max(client.timeout, timeout_seconds + 30)` — a 30-minute test-suite exec
  (`timeout_seconds=1800`; the daemon accepts up to 3600 s) holds the
  connection instead of dying at the flat default.
- **Crashed-worker recovery.** `Sandbox.attach(client, session_id)` adopts an
  already-live session (validated via `GET /v1/sessions/{id}`, nothing forked)
  so a replacement worker or janitor can resume or reap an orphaned clone; the
  context manager reaps on exit as usual.
- **stderr is merged into stdout, everywhere.** The guest agent runs execs with
  both fds on one pipe (`stderr` in the response is always `""`), and the local
  backend now reproduces that — write graders with a sentinel-prefixed reward
  line rather than branching on `stderr` (see `RewardSpec.parse`).

## Tests

```bash
# from a source checkout of the package:
python -m pytest -q                        # ~18s, no network, no VMM
```

The suite runs the real exec/reward contract against the in-process
`LocalCollimateBackend`, including the seed→reproducibility (CRN) property that
the whole common-random-numbers story rests on.

## Design notes

- **Reward stays the env's job.** Collimate only runs the grader faster + isolated;
  it never inspects or alters the reward (the OpenEnv contract, honored here too).
- **The GRPO group fan-out maps to forks.** `grade_batch()` of K candidates = K
  ephemeral VMs forked off one warm template — the after-setup group-fork.
- **CRN is opt-in and explicit.** `seeds.group_seeds(K, "crn")` pins one guest
  seed across a prompt-group so env-noise cancels in the group-relative advantage;
  `"vanilla"` keeps fresh entropy per rollout; `"deterministic"` removes env-noise
  entirely (a control arm). Packaged for reuse.
- **Small public surface.** The integration only ever assumes "fork a warm
  microVM and exec in it"; everything else stays behind the API.
