collimate-rl 0.3.0__py3-none-any.whl

collimate_rl/__init__.py

@@ -0,0 +1,107 @@
+# Copyright 2026 Omer Gero (Gero Consulting)
+# SPDX-License-Identifier: Apache-2.0
+"""Collimate RL integration toolkit.
+
+Thin, dependency-free Python building blocks for using a Collimate control
+plane (the ``serve.rs`` daemon: ``/v1/sessions`` + ``/v1/exec``) as the
+sandboxed *reward grader* in an RL post-training loop.
+
+This is the "exec-grader" integration mode: the environment logic lives in the
+trainer, and Collimate runs each candidate rollout in a forked microVM to
+produce an isolated, reproducible reward. It needs nothing from the guest
+beyond the existing exec API — no host->guest ingress — so it runs against the
+control plane that exists today.
+
+Public surface:
+
+    from collimate_rl import (
+        connect,                # ONE entry point, either transport (see below)
+        CollimateClient,        # the daemon /v1 HTTP client (stdlib urllib only)
+        GatewayClient,          # the GA gateway client (Bearer col_* key, any tier)
+        CollimateDrainingError, # node draining (retry-safe; route elsewhere)
+        PropertyReward,         # property-based / fuzz grader (CRN via seed)
+        UnitTestReward,         # explicit unit-test grader (deterministic)
+        grade, grade_batch,     # run a candidate (or a GRPO group) -> reward
+        GradeResult,
+        extract_code,           # pull code out of an LLM completion
+        fresh_seed, group_seeds, DET_SEED,  # common-random-numbers helpers
+    )
+
+``connect()`` returns the right client for the target you name (a daemon URL /
+UDS, or a ``col_*`` API key of ANY tier → the gateway) — and both clients
+expose the same session surface, so ``grade``/``grade_batch``, ``Sandbox`` and
+every framework adapter run over either transport unchanged.
+
+The framework adapters (``collimate_rl.skyrl``, ``collimate_rl.openenv``,
+``collimate_rl.slime``) are thin shims over this.
+"""
+
+from .client import CollimateClient, CollimateDrainingError, CollimateError
+from .fleet import Fleet, RolloutHandle
+from .gateway import DemoTierLimit, GatewayClient, GatewayError
+from .transport import SessionClient, connect
+# template.py installs GatewayClient.session and re-exports the high-level flow;
+# importing it here wires that method and keeps the public surface discoverable.
+from .template import (
+    Ready,
+    Session,
+    Template,
+    TemplateHandle,
+    TemplateStatus,
+    open_session,
+)
+from .rewards import (
+    GradeResult,
+    PropertyReward,
+    RewardSpec,
+    UnitTestReward,
+    extract_code,
+    grade,
+    grade_batch,
+    reward_spec_from_ground_truth,
+)
+from .sandbox import CollimateTaskError, ExecResult, Sandbox, rollout_once
+from .seeds import DET_SEED, fresh_seed, group_seeds, is_valid_seed, seed_from_key
+
+__all__ = [
+    # transport selection (one interface, any key tier)
+    "connect",
+    "SessionClient",
+    "CollimateClient",
+    "CollimateError",
+    "CollimateDrainingError",
+    # GA gateway surface (Bearer col_* keys) + high-level bake/session flow
+    "GatewayClient",
+    "GatewayError",
+    "DemoTierLimit",
+    "Ready",
+    "Template",
+    "TemplateHandle",
+    "TemplateStatus",
+    "Session",
+    "open_session",
+    # fleet / control-plane surface
+    "Fleet",
+    "RolloutHandle",
+    # in-pod sandbox hot path
+    "Sandbox",
+    "ExecResult",
+    "CollimateTaskError",
+    "rollout_once",
+    # reward grading
+    "RewardSpec",
+    "PropertyReward",
+    "UnitTestReward",
+    "GradeResult",
+    "grade",
+    "grade_batch",
+    "reward_spec_from_ground_truth",
+    "extract_code",
+    "fresh_seed",
+    "group_seeds",
+    "seed_from_key",
+    "is_valid_seed",
+    "DET_SEED",
+]
+
+__version__ = "0.3.0"

collimate_rl/burst.py

@@ -0,0 +1,243 @@
+"""Lifecycle load burst: start -> exec -> stop, N times at concurrency C.
+
+The classic load-tool view (histogram + latency distribution + summary), applied
+to the full sandbox lifecycle instead of a bare HTTP endpoint: every sample is a
+CREATE (a live fork off the warm pool), one EXEC inside it, and a DELETE. Each
+phase is timed separately so the report can attribute the milliseconds honestly.
+
+Zero dependencies beyond the optional ``rich`` extra the harness already uses.
+"""
+from __future__ import annotations
+
+import argparse
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+from .gateway import GatewayClient, GatewayError
+
+#: Default exec payload — deliberately the same *class* of work as the classic
+#: "hello sandbox" launch demos (a trivial interpreter one-liner), so lifecycle
+#: numbers compare apples-to-apples. The flex is what the sandbox IS (a warm
+#: fork of the full template, deps resident), not what the one-liner does.
+DEFAULT_EXEC = 'python3 -c "import uuid, time; print(str(uuid.uuid4()) + \' \' + str(time.time()))"'
+
+PERCENTILES = (10.0, 25.0, 50.0, 75.0, 90.0, 95.0, 99.0, 99.9, 100.0)
+HISTOGRAM_BUCKETS = 11
+BAR_WIDTH = 40
+
+
+# --------------------------------------------------------------------------- #
+# pure helpers (unit-tested)                                                   #
+# --------------------------------------------------------------------------- #
+
+def percentile(sorted_vals: Sequence[float], pct: float) -> float:
+    """Nearest-rank percentile of an ascending-sorted sequence."""
+    if not sorted_vals:
+        return 0.0
+    if pct >= 100.0:
+        return sorted_vals[-1]
+    rank = max(1, int(round(pct / 100.0 * len(sorted_vals) + 0.5)))
+    return sorted_vals[min(rank, len(sorted_vals)) - 1]
+
+
+def histogram_buckets(
+    sorted_vals: Sequence[float], n_buckets: int = HISTOGRAM_BUCKETS
+) -> List[Tuple[float, int]]:
+    """Linear buckets over [min, p99.9] (upper-edge, count) — outliers clamp to
+    the last bucket so one straggler cannot flatten the whole picture."""
+    if not sorted_vals:
+        return []
+    lo = sorted_vals[0]
+    hi = max(percentile(sorted_vals, 99.9), lo + 1e-9)
+    step = (hi - lo) / n_buckets
+    buckets = [[lo + step * (i + 1), 0] for i in range(n_buckets)]
+    for v in sorted_vals:
+        idx = min(int((v - lo) / step) if step > 0 else 0, n_buckets - 1)
+        buckets[idx][1] += 1
+    return [(edge, count) for edge, count in buckets]
+
+
+# --------------------------------------------------------------------------- #
+# the burst                                                                    #
+# --------------------------------------------------------------------------- #
+
+class _Sample(Dict[str, Any]):
+    """One lifecycle sample: phase timings (s) or an error code."""
+
+
+def _one_lifecycle(client: GatewayClient, template: str, exec_cmd: str,
+                   timeout_s: int) -> _Sample:
+    s = _Sample()
+    t0 = time.perf_counter()
+    try:
+        created = client.create_sandbox(template)
+        s["create"] = time.perf_counter() - t0
+        sid = created["id"]
+    except GatewayError as e:
+        s["error"] = getattr(e, "code", "") or "create_failed"
+        s["total"] = time.perf_counter() - t0
+        return s
+    try:
+        t1 = time.perf_counter()
+        r = client.exec(sid, command=exec_cmd, timeout_seconds=timeout_s)
+        s["exec"] = time.perf_counter() - t1
+        if r.get("exit_code", 0) != 0:
+            s["error"] = "exec_nonzero"
+    except GatewayError as e:
+        s["error"] = getattr(e, "code", "") or "exec_failed"
+    finally:
+        t2 = time.perf_counter()
+        try:
+            client.delete_sandbox(sid)
+            s["delete"] = time.perf_counter() - t2
+        except GatewayError:
+            s.setdefault("error", "delete_failed")
+    s["total"] = time.perf_counter() - t0
+    return s
+
+
+def run_burst(client: GatewayClient, template: str, *, n: int, concurrency: int,
+              exec_cmd: str = DEFAULT_EXEC, timeout_s: int = 30,
+              on_progress=None) -> Dict[str, Any]:
+    """Run the burst; returns {samples, wall_s}. Thread-per-inflight-lifecycle —
+    each worker owns its create/exec/delete chain end to end."""
+    samples: List[_Sample] = []
+    lock = threading.Lock()
+    done = 0
+    t0 = time.perf_counter()
+    with ThreadPoolExecutor(max_workers=concurrency) as pool:
+        futs = [pool.submit(_one_lifecycle, client, template, exec_cmd, timeout_s)
+                for _ in range(n)]
+        for fut in as_completed(futs):
+            sample = fut.result()
+            with lock:
+                samples.append(sample)
+                done += 1
+                if on_progress is not None:
+                    on_progress(done, n, time.perf_counter() - t0)
+    return {"samples": samples, "wall_s": time.perf_counter() - t0}
+
+
+# --------------------------------------------------------------------------- #
+# rendering (rich — same optional extra as the rest of the harness)            #
+# --------------------------------------------------------------------------- #
+
+def render_report(console, result: Dict[str, Any], *, template: str, n: int,
+                  concurrency: int, note: str = "") -> None:
+    from rich.text import Text
+
+    samples: List[_Sample] = result["samples"]
+    ok = [s for s in samples if "error" not in s]
+    totals = sorted(s["total"] for s in ok)
+    wall = result["wall_s"]
+
+    console.print()
+    console.print(Text.assemble(
+        ("  collimate burst", "bold cyan"),
+        (f"  ·  start → exec → stop × {n:,}  ·  concurrency {concurrency}", "bold"),
+    ))
+    console.print(Text(f"  template {template} — every sample forks a live sandbox "
+                       f"off the warm pool, runs, and is destroyed", "dim"))
+    console.print()
+
+    if not totals:
+        console.print(Text("  no successful lifecycles — see errors below", "bold red"))
+    else:
+        console.print(Text("Response time histogram (full lifecycle):", "bold"))
+        buckets = histogram_buckets(totals)
+        top = max(count for _, count in buckets) or 1
+        for edge, count in buckets:
+            bar = "■" * max(0, round(count / top * BAR_WIDTH))
+            console.print(Text.assemble(
+                (f"  {edge:7.3f}s ", ""), (f"[{count:5d}] ", "dim"), (bar, "cyan"),
+            ))
+        console.print()
+        console.print(Text("Latency distribution:", "bold"))
+        for pct in PERCENTILES:
+            console.print(f"  {pct:6.2f}% in {percentile(totals, pct):8.4f}s")
+        console.print()
+        console.print(Text("Per-phase (p50 / p90):", "bold"))
+        for phase in ("create", "exec", "delete"):
+            vals = sorted(s[phase] for s in ok if phase in s)
+            if vals:
+                console.print(
+                    f"  {phase:<7} {percentile(vals, 50)*1000:7.1f} ms / "
+                    f"{percentile(vals, 90)*1000:7.1f} ms"
+                )
+        console.print()
+
+    rate = (len(ok) / wall) if wall > 0 else 0.0
+    console.print(Text("Summary:", "bold"))
+    console.print(Text.assemble(
+        (f"  total {wall:.2f}s", "bold"),
+        (f"  ·  {rate:,.1f} lifecycles/sec  ·  success ", ""),
+        (f"{len(ok):,}/{n:,}", "bold green" if len(ok) == n else "bold red"),
+    ))
+    errors: Dict[str, int] = {}
+    for s in samples:
+        if "error" in s:
+            errors[s["error"]] = errors.get(s["error"], 0) + 1
+    for code, count in sorted(errors.items()):
+        console.print(Text(f"  error {code}: {count}", "red"))
+    if note:
+        console.print(Text(f"  {note}", "dim"))
+    console.print()
+
+
+def cmd_burst(args: argparse.Namespace) -> int:
+    from rich.console import Console
+    from rich.progress import (BarColumn, Progress, TextColumn,
+                               TimeElapsedColumn)
+
+    console = Console()
+    # retries=6: a lifecycle burst that outruns the advertised warm pool sees
+    # placer-level 503 at_capacity (retryable, no side effect); six full-jitter
+    # attempts (~10 s worst-case) ride one node-agent report/refill window
+    # instead of failing the sample.
+    client = GatewayClient(api_key=args.api_key, base_url=args.base_url,
+                           retries=6)
+    progress = Progress(
+        TextColumn("[bold cyan]burst[/bold cyan]"),
+        BarColumn(bar_width=40, complete_style="cyan"),
+        TextColumn("{task.completed:,}/{task.total:,} lifecycles"),
+        TextColumn("[dim]{task.fields[rate]:,.0f}/s[/dim]"),
+        TimeElapsedColumn(),
+        console=console,
+    )
+    with progress:
+        task = progress.add_task("burst", total=args.n, rate=0.0)
+
+        def on_progress(done: int, total: int, elapsed: float) -> None:
+            progress.update(task, completed=done,
+                            rate=(done / elapsed) if elapsed > 0 else 0.0)
+
+        result = run_burst(client, args.template, n=args.n,
+                           concurrency=args.concurrency, exec_cmd=args.exec_cmd,
+                           timeout_s=args.timeout, on_progress=on_progress)
+    render_report(console, result, template=args.template, n=args.n,
+                  concurrency=args.concurrency, note=args.note)
+    ok = sum(1 for s in result["samples"] if "error" not in s)
+    return 0 if ok == args.n else 1
+
+
+def add_burst_parser(sub, add_common_args) -> None:
+    p = sub.add_parser(
+        "burst",
+        help="start → exec → stop N sandboxes at concurrency C; "
+             "histogram + latency distribution",
+    )
+    add_common_args(p)
+    p.add_argument("template", help="template id to fork from")
+    p.add_argument("-n", type=int, default=1000, dest="n",
+                   help="number of full lifecycles (default 1000)")
+    p.add_argument("-c", type=int, default=100, dest="concurrency",
+                   help="concurrent in-flight lifecycles (default 100)")
+    p.add_argument("--exec-cmd", default=DEFAULT_EXEC,
+                   help="command each sandbox runs (default: trivial python one-liner)")
+    p.add_argument("--timeout", type=int, default=30,
+                   help="per-exec timeout seconds")
+    p.add_argument("--note", default="",
+                   help="footer annotation (e.g. client region / env description)")
+    p.set_defaults(func=cmd_burst)