Skip to content
Collimate

Pro quickstart: GRPO rollouts

Pro is the bring-your-own-image tier: you bake your container image into a ready-state artifact once, then fork it at rollout width for as long as you train. This page is the GRPO-shaped recipe — one shared prefix, N isolated branches, group-relative rewards.

Two things you will not find in this guide, because they are not your job:

  • Capacity management. There are no pools to size and no fleet knobs — capacity follows your workload.
  • Backpressure handling. No 503-retry loops. A burst that momentarily outruns ready capacity is absorbed transparently while capacity materializes; your loop just runs. The only limit is your spend setting in Billing.

1. Bake your image

Section titled “1. Bake your image”

A bake turns your OCI image into the thing sandboxes fork from: a booted, warmed, checkpointed artifact. You do this once per image version.

  1. Get a Pro key and install the SDK:

    Terminal window
    pip install collimate-rl
    export COLLIMATE_API_KEY="col_pro_a1b2c3d4e5f60718_..."

  2. Bake, declaring the ready state — the point your rollouts should start from, with the heavy imports already paid:

    from collimate_rl import Ready, Template, connect
    

    client = connect()   # reads COLLIMATE_API_KEY
    

    handle = Template.bake(
        client,
        name="swe-rollout",
        image="ghcr.io/you/swe-rollout@sha256:…",     # pin by digest
        ready=Ready(cmd="python -c 'import torch'"),  # warm state into the artifact
        width=64,        # how wide you expect to fan out — a hint, not a knob
    )

  3. Wait for the fleet to report it ready (bad images fail fast, with the build error attached):

    status = handle.wait(timeout=600)
    print(status.phase)   # Live

Every rollout now starts from that ready state — import torch costs zero at episode time, on every branch.

2. The rollout loop

Section titled “2. The rollout loop”

One GRPO group = one parent session, forked into N live children. Each child inherits the parent's process memory — the shared prefix is computed once, never replayed per branch.

WIDTH = 64


def rollout_group(task_cmd: str, candidates: list[str]) -> list[float]:
    """Grade one GRPO group of len(candidates) == WIDTH completions."""
    with client.session(handle) as sess:
        # Shared prefix: paid ONCE for the whole group.
        sess.exec(command=task_cmd, timeout_seconds=120)


        # Fork the mid-episode state into WIDTH isolated live children.
        with sess.fork_group(WIDTH) as grp:
            def run(i, child):
                r = child.exec(code=candidates[i], timeout_seconds=120)
                return 1.0 if r.ok else 0.0


            return grp.map(run)          # parallel, positional
    # exiting the blocks tears down every child, then the parent — even on error




for step in range(1000):
    prompts, groups = sample_batch()      # your trainer
    for task_cmd, candidates in zip(prompts, groups):
        rewards = rollout_group(task_cmd, candidates)
        update_policy(candidates, rewards)

What the SDK is doing for you in that loop:

  • Common random numbers. Each child carries its own per-rollout CRN seed (child.default_seed), so environment noise cancels in the group-relative advantage. See Deterministic replay.
  • Strict group width. A partial group is never handed to your trainer — if a child drops, the group is torn down and the error surfaces so you resample, instead of averaging an infra failure into the gradient.
  • Guaranteed teardown. Context-manager exit deletes every child even when an exec raises; nothing leaks between steps.

3. Scope a session with client.env

Section titled “3. Scope a session with client.env”

A training run is a session: a bounded period where you hammer one environment at a known width. Declare it up front with client.env(ref, width=N) and the platform has the full width ready before your first burst — and hands the capacity back the moment you're done:

with client.env("swe-rollout", width=WIDTH) as env:   # width ready up front
    with env.sandbox() as sess:
        sess.exec(command=task_cmd, timeout_seconds=120)
        with sess.fork_group(WIDTH) as grp:
            rewards = grp.map(run)
# exit → release: capacity relaxes immediately (idle decay is the fallback)

ref is a template id you baked or a catalog image name — pointing env at a catalog image bakes it into your own environment at that width in one step. Under the hood this is POST /v1/templates/{id}/expect on entry and .../release on exit (the SDK's expect_width / release_template); both are optional — an undeclared session just warms up from observed demand, and an idle environment always relaxes on its own. Releasing is never destructive: the environment, its artifact, and any live sandboxes are untouched.

4. Run it as long as you like

Section titled “4. Run it as long as you like”

There is nothing else to operate. No warm-up scripts before a training run, no draining after, no capacity to hand back. Set a spend limit (and optional auto top-up) in Billing; the platform provisions to your workload underneath it and the meter only runs while sandboxes are live — suspending stops it.

print(client.usage())   # GET /v1/usage — your metered totals, any window

Next steps

Section titled “Next steps”