collimate-rl 0.3.0__tar.gz

collimate_rl-0.3.0/LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

collimate_rl-0.3.0/PKG-INFO

@@ -0,0 +1,256 @@
+Metadata-Version: 2.4
+Name: collimate-rl
+Version: 0.3.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
+
+res = grade(client, "dkr-python311", candidate, spec)   # identical either way
+```
+
+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 (net templates)
+
+A `NET=1` 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 `NET=1` (a vsock-only template rejects a policy with 400).
+
+## 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.