collimate-rl 0.8.0

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.

collimate-rl — Collimate for RL post-training

Run RL rollout execution and reward grading inside forked Collimate 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):

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:

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

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)

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

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:

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 mvs 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):

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:

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.

Connect from outside (host→guest ingress)

Two ways to drive a sandbox. exec-in-guest (sb.exec(...)) ships code into the guest — the RL hot path, one round trip. connect-from-outside reaches a service the guest exposes (a browser's CDP port, an HTTP/WebSocket app, a database, an OpenEnv env-server) from a client on your own machine. Bake the port with ports=[...] (which also gives the fork a NIC), then:

# A CDP browser → a wss:// URL a Playwright/CDP client consumes directly.
with client.session("browser") as sb:
    cdp = sb.cdp_url()                       # == connect_url(9222, protocol="cdp")
    # playwright.chromium.connect_over_cdp(cdp), or the turnkey helper:
    from collimate_rl.browser import connect_browser
    browser = connect_browser(sb)            # a Playwright browser, on your machine

# Any HTTP / WebSocket / TCP service → forward its port to localhost.
with client.session("app") as sb:            # a template that exposes :8000
    with sb.port_forward(8000) as fwd:        # the kubectl-port-forward shape
        requests.get(fwd.url + "/health")     # fwd.url = http://127.0.0.1:<local>

connect_url(port) returns the raw wss://…/connect?port=…&t=<token> URL; cdp_url() mints it with protocol="cdp" for CDP clients; port_forward(port) runs a local listener over the raw (tcp) tunnel so any client hits localhost. All three are Pro/Ent only and need the port declared (ports=[port]); a below-pro key raises, an unexposed port is 409 port_not_exposed. The tunnel is zero-dependency stdlib — importing the SDK adds nothing (Playwright, for connect_browser, is an optional extra).

Failure semantics

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

# 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.