From 535a24f5d8b4bd61312769b69c8aa3f1722cc236 Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Fri, 9 Oct 2026 00:07:34 -0400 Subject: [PATCH 01/18] lab: container jail on Linux; runtime design (RUNTIME.md) srt cannot jail GPU code on Linux: no device passthrough, and each jail gets its own network namespace, so a served engine is unreachable from the clients. On Linux, agent code (tests, profilers, served engines) now runs in one container per command: NVIDIA runtime, --network none, read-only root, non-root user, no capabilities, cpuset leaving cores for the clients. A served port reaches the host through a socat bridge over a unix socket. srt stays on macOS; LAB_JAIL overrides. vm-setup.sh installs docker + NVIDIA Container Toolkit, builds lab-jail, and checks the GPU inside it. lab/RUNTIME.md is the controller/worker design this starts. Co-Authored-By: Claude Opus 5.5 --- CLAUDE.md | 24 +++--- lab/README.md | 7 +- lab/RUNTIME.md | 149 ++++++++++++++++++++++++++++++++++++ lab/safety/container.py | 123 +++++++++++++++++++++++++++++ lab/safety/grader.py | 50 +++++++++--- lab/safety/jail.Dockerfile | 7 ++ lab/safety/jail.py | 6 ++ lab/serve.py | 6 +- lab/vm-setup.sh | 21 ++++- tests/test_lab_container.py | 143 ++++++++++++++++++++++++++++++++++ tests/test_lab_proftools.py | 3 + tests/test_lab_safety.py | 3 +- 12 files changed, 515 insertions(+), 27 deletions(-) create mode 100644 lab/RUNTIME.md create mode 100644 lab/safety/container.py create mode 100644 lab/safety/jail.Dockerfile create mode 100644 tests/test_lab_container.py diff --git a/CLAUDE.md b/CLAUDE.md index 6c85e13..88c9781 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -37,13 +37,15 @@ pytest/ruff config). Each folder with commands has its own `python -m`. with validated args) + `evaltools.py` (bench → `regimes` on the seen split; equiv → `correctness` candidate vs the target's reference; submit → held-out full tier vs the base commit, one aggregate per regime through the Thresholdout guard; all three serve the pristine tree via - `serve.py`, jailed with local binding allowed; every serve also keeps its passive data, `artifacts.py`: device + `serve.py`, jailed (on Linux, port bridged out of a network-less container); every serve also keeps its passive data, `artifacts.py`: device samples via `gpu.DeviceSampler` (DCGM if `dcgmi`, else nvidia-smi), engine telemetry files from `{dir}`, all client rows, the serve log; bench/equiv as ledger blobs under `result.artifacts`, submit only in `/heldout-private/`, never in the ledger; `artifacts.joined` joins client rows to engine rows by trace id), `workspace.py` (exported engine copy + snapshots), `ledger.py` (append-only JSONL, `lab/ledger/` gitignored), `budget.py`, `engine.py` (where the engine repo and its python - are), `safety/` (write surfaces, srt jail, grader: lint/tests run with the engine's python in the jail), + are), `safety/` (write surfaces, grader: lint/tests run with the engine's python in the jail; the jail is `container.py` on + Linux (one container per jailed command, NVIDIA runtime, `--network none`, socat bridge for a served port; tested in + Docker-in-Docker, GPU untested) and srt elsewhere, `LAB_JAIL` overrides), `profile.py` (`--profiler torch|cuda-range|pyspy`) + `bundle.py` + `gpu.py` (profiler harness, staged into the pristine tree under `_harness/` so the jail never opens this repo, which holds `.env` and the ledger; the tool samples its seen-split corpus workload out here, `corpus.sample`, and stages it as `_harness/workload.json`; window-only memory peaks and scheduler counters, warmup sized to the window, op/kernel tables), `canary.py` (honest regressions patched into the engine's process), `vm.py` + `providers/` + `vm-setup.sh` (one GPU VM on Verda/Nebius/Crusoe; pushes both @@ -57,6 +59,8 @@ pytest/ruff config). Each folder with commands has its own `python -m`. measures each regime's run-to-run band from repeated reference runs into `knowledge/noise/.json` (which `submit` reads). Exit 0 only if both hold. Needs a GPU for a real target; tested on fakes. `lab/LEDGER.md` is the record-format spec: the environment's public interface. +- `lab/RUNTIME.md` — design, not built (2026-10-08): controller (laptop: ledger, referee, secrets) and worker (GPU + VM: jails, clients, profilers) behind a job interface; trust zones and rules; why srt cannot jail GPU code on Linux. - `lab/TRACE.md` + `lab/trace_contract.py` — the request-trace contract: per-request engine rows (trace_id = the client's `X-Trace-Id`, `arrival_ts` epoch s, `*_s` monotonic durations, TPOT = (E2E − TTFT)/(n − 1)) in a dir of `*.sqlite` (`requests` + `meta`) or `*.jsonl` + `.meta.json`, with schema version and drop counts; @@ -203,18 +207,20 @@ Unjailed, all four work on the real engine (16 requests × 64 tokens): `profile` `launch_count` honoured; ncu durations are ~1.8× nsys's on a tiny kernel (base clocks), DRAM figures physical. Fixed from the run: the tools now pass the target's engine env (they had loaded the engine's default model); nsys traces CUDA-graph nodes (it had missed every decode kernel: 46,720 vs torch's 860,365 launches); -`kernels.json` drops record_function annotations (they had doubled its total). **Open, blocks every jailed GPU -tool on Linux:** srt always mounts a fresh minimal `/dev` and has no device passthrough, so nothing in the jail -sees the GPU (`No CUDA GPUs are available`); `bench`/`equiv`/`submit` serve jailed too. Needs a design decision -(GPU-aware jail, container with the NVIDIA runtime, or the VM as the boundary). Ubuntu 24.04 also needs -`kernel.apparmor_restrict_unprivileged_userns=0` for bwrap (now in `vm-setup.sh`). +`kernels.json` drops record_function annotations (they had doubled its total). srt always mounts a fresh minimal +`/dev` and has no device passthrough, so nothing in an srt jail sees the GPU (`No CUDA GPUs are available`), and +each srt jail on Linux has its own network namespace, so a jailed server is unreachable. Decided 2026-10-08: on +Linux the jail is a container with the NVIDIA runtime (`lab/safety/container.py`, `lab/RUNTIME.md`); its GPU +check (`vm-setup.sh` prints `jail: gpu ok`) has not run yet. The agent's CLI still runs under srt, which on Ubuntu +24.04 needs `kernel.apparmor_restrict_unprivileged_userns=0` (in `vm-setup.sh`). ## Next GPU (needs spend approval): resolve the BF16 puzzle (plain BF16 on the dev tier at concurrency 32 and 128 on one VM, ~20 min, ~$1.50); first real run of `python -m regimes run all` against the engine and against vLLM. GPU, with approval (and the auto-stop armed): `python -m lab.validate --target targets/inference-server.toml` -(reference, FP8/INT8 good, Int4 bad, noise bands for all 8 regimes; ~3-4 h), then a first `python -m lab.session ---task ...` dry run with a tiny budget. Both engine branches from the 2026-10-06 report (`integ/gpu-session-5`, +(reference, FP8/INT8 good, Int4 bad, noise bands for all 8 regimes; ~3-4 h; serves unjailed, so not blocked), then +a first `python -m lab.session --task ...` dry run with a tiny budget, which waits on the rest of `lab/RUNTIME.md` +(the controller/worker split; the agent's CLI still runs under srt, which has no GPU on Linux). Both engine branches from the 2026-10-06 report (`integ/gpu-session-5`, `engine/gpu-validated`) were judged by the old gate: re-judge with `python -m correctness verdict` on their saved results before deciding anything. lab `session.py --run` re-resolves `base` from HEAD rather than the run's original base (pre-existing; resuming after the engine moved would misaudit); diff --git a/lab/README.md b/lab/README.md index b3686fe..74cfc9c 100644 --- a/lab/README.md +++ b/lab/README.md @@ -159,8 +159,11 @@ A session ends when the agent says `stop`, its per-session cap is spent, or it t run ends on budget, on `stop`, or on a write-surface violation (the one hard rule). Model cost comes from the provider's own accounting; `LAB_MODEL` picks the model, `LAB_AGENT_PROVIDER` the provider (only `claude` today, through the Agent SDK CLI; `lab/agent.py` is the seam for others). -The jail is sandbox-runtime (`npm install -g @anthropic-ai/sandbox-runtime`); `LAB_NO_JAIL=1` -waives it for tests and a box you trust. +Agent code runs jailed. On Linux the jail is a container (`lab/safety/container.py`: NVIDIA runtime, no +network, read-only root, non-root user, no capabilities; a served engine's port reaches the host through a socket +bridge; image from `python -m lab.safety.container build`, which `vm-setup.sh` runs). Elsewhere it is +sandbox-runtime (`npm install -g @anthropic-ai/sandbox-runtime`); `LAB_JAIL` picks one. `LAB_NO_JAIL=1` waives +the jail for tests and a box you trust. The agent's own CLI still runs under srt (`lab/RUNTIME.md`). ## Canaries diff --git a/lab/RUNTIME.md b/lab/RUNTIME.md new file mode 100644 index 0000000..ee799bd --- /dev/null +++ b/lab/RUNTIME.md @@ -0,0 +1,149 @@ +# Where the lab runs + +Design, 2026-10-08. Decided with the user: the lab splits into a controller and a worker; no separate +measurement VM (cost); NVIDIA only for now; the jail on Linux is a container with the NVIDIA runtime (option A +below). Built: the container jail (`lab/safety/container.py`). Not built: the split. Open decisions are marked +**Decide**. + +## Why + +Today every part of the lab runs on one machine and serves the engine on that machine's localhost. On a +Linux GPU VM no jailed GPU tool can work, for two reasons: + +- **No GPU in the jail.** srt mounts a fresh minimal `/dev` and has no device passthrough: jailed code gets + `No CUDA GPUs are available` (H200, Nebius, 2026-10-08; see `fix/profiler-gpu`'s CLAUDE.md). +- **No way in from outside.** On Linux srt runs every jail in its own network namespace (`bwrap + --unshare-net`); `allowLocalBinding` only means something on macOS. A server inside a jail answers only + inside that jail, so `serve.py`'s health check and the measurement clients outside it are refused. + Measured in Docker (Linux, srt from npm, 2026-10-08): same jail serve + request → 200; host → jailed + server, other jail → jailed server, jail → host server: all refused. + +What srt does well on Linux, from the same test: its own PID namespace (a jail sees only its processes and +cannot signal others), and everything in a jail dies when the jail's command exits. + +## Trust zones + +| Zone | Where | Holds | Agent's code can reach it | +|---|---|---|---| +| Controller | laptop (later an always-on box) | session loop, ledger, referee, holdout guard, budget, snapshot store, every secret | no | +| Worker | GPU VM, host side, outside every jail | job runner, measurement clients, profilers, GPU state | only through the engine's HTTP port | +| Jails | GPU VM, inside containers | agent's shell and workspace; each served engine | yes, it is the agent's | + +Rules: + +1. **No secrets on the VM.** No cloud keys, no model API key, no HF token (weights are staged on a read-only + disk). Model calls go through a proxy on the controller that adds the key. +2. **Only the controller writes the ledger.** The worker returns results, raw files and hashes; the controller + records them. +3. **What is measured is what was audited.** The controller pulls each snapshot's changed files, stores them, + audits them, and sends that exact content back for `test`/`bench`/`equiv`/`submit`/`profile`. The worker + rebuilds the tree from base + those files and returns its hash; a mismatch refuses the job. +4. **Nothing of the agent runs during a measurement.** Before every measured job the worker freezes the agent's + jail (cgroup freezer) and kills any of its processes holding a GPU context; it thaws after. The tool result + states what was killed (a fact). +5. **Measurement code is out of reach.** The clients run on the host, outside every container, from a tree the + jails cannot read; their CPU cores are reserved (cpuset) so the engine cannot starve them. +6. **One VM per run, destroyed after** (**Decide**; recommended). Nothing from one run can touch the next one's numbers. Setup time is the + cost; a prebuilt image with the venv plus a reattached read-only weights disk cuts it. +7. **Network.** VM egress: the controller tunnel only. Agent jail: the API proxy only. Engine jail: nothing; its + port is reachable from the host alone. + +Left open by not having a separate measurement VM: the clients and the agent's code share a kernel, so a +kernel or container escape defeats rules 4 and 5. A measurement VM in the same private network closes it later. + +## The jail on Linux + +Needs: the GPU (devices and driver libraries); a port the host can reach and nothing else; write access to the +workspace and a tmp only; weights read-only; an unprivileged user; its own PID namespace; CPU, memory and +process limits; freeze and kill as a unit. + +Options considered; **A was chosen** (2026-10-08): + +- **A. Containers with the NVIDIA runtime.** Docker or Podman plus NVIDIA Container Toolkit: + `--gpus`, an `--internal` network per jail (host reaches the engine's port, no egress), read-only root, + non-root user, all capabilities dropped, `no-new-privileges`, seccomp default, cgroup limits, `pause` as the + freezer. Standard and well-trodden (SWE-bench and OpenHands sandbox this way), and the same unit a scheduler + runs at scale (a pod). Harder later: gVisor (`runsc`, GPU via nvproxy) as a drop-in for a stronger boundary. +- **B. bwrap called directly.** Bind the `/dev/nvidia*` devices and driver libraries in, share the host network, + and restrict egress per uid with nftables. Lighter, but we would own every sharp edge srt hides. +- **C. The VM as the boundary.** One VM per jail. Strongest, and ruled out for cost. + +macOS stays as today (srt; Seatbelt shares the network stack) for development and the fake-engine tests. + +As built: one container per jailed command (`lab-jail` image: Ubuntu 24.04 with a compiler, Python and socat; +the engine's venv, interpreter, weights and `/usr/local/cuda` bind-mounted read-only at their host paths). +`--network none` rather than an internal network: the engine keeps binding 127.0.0.1, a socat inside the +container serves that port on a unix socket in the jail's tmp, and a host socat serves the socket on the host's +127.0.0.1. Nothing in any jail can reach it; the clients pay one forwarding hop per connection, the same for base +and candidate. Run as root (Verda), the jail's user is uid 10001 and the writable paths are handed to it; the +first `CLIENT_CORES` cores are kept out of every jail. Tested in Docker-in-Docker on Ubuntu 24.04 (no network, +no capabilities, non-root, host files invisible, port reaches the host and no other jail, nothing left after +timeout or teardown); `--gpus` and the profilers inside it are untested until a GPU run. + +Profilers under A: nsys runs inside the engine's container; ncu needs GPU performance counters, which +the single-use VM can open to non-admins (`NVreg_RestrictProfilingToAdminUsers=0`) since nothing else runs on it; +py-spy attaches from the host (root) into the container's process. + +## A tool call, end to end (Claude provider) + +1. The agent's CLI runs in the agent container on the VM. Its stdin/stdout reach the Agent SDK on the + controller over SSH, so the lab's tools stay an in-process MCP server on the controller, as today. +2. The agent edits and runs commands in its container (built-in tools; they never leave the VM). +3. It calls a lab tool. The controller asks the worker for the snapshot (changed files), stores and audits + them, and writes the ledger record's snapshot. +4. For a measured tool: the controller sends the job with the snapshot's files and hash. The worker freezes + the agent container, rebuilds and checks the tree, serves it in an engine container, runs the clients from + the host, tears the engine down, thaws, and returns results, raw files and hashes. +5. The controller applies the policy (submit's preconditions, the holdout guard, noise bands), writes the + ledger, and answers the agent. + +The model API path: CLI → `ANTHROPIC_BASE_URL` on the agent network → a host forwarder → SSH tunnel → +controller proxy (adds the key, meters tokens) → api.anthropic.com. Metering at the proxy makes the budget a +measured fact instead of the CLI's report. + +## The worker job interface + +A short fixed list; each job takes a snapshot (files + hash), the target spec and arguments, and returns result +JSON, raw files with hashes, and hardware facts (device, driver, clocks). + +| Job | Does | +|---|---| +| `agent` | start the agent container on the workspace and stream the CLI's stdio (Claude) | +| `exec` | one command in the agent container (providers with their own tool loop) | +| `snapshot` | the workspace's changed files since base | +| `restore` | put the workspace at a snapshot the controller sends | +| `test` | lint and the engine's suite on the rebuilt tree, jailed | +| `measure` | serve the rebuilt tree and run regimes or the correctness gate (`bench`, `equiv`, `submit`) | +| `profile` | `profile`, `trace`, `kernel`, `hostprof` | +| `facts` | GPU, driver, clocks, versions | + +Transport now: SSH from the controller, one worker. Transport later: a queue. The job list does not change. + +## What moves where + +- Controller: `session`, `agent` (SDK side), the tool front ends in `tools`/`evaltools`/`proftools` (policy, + ledger writes), `ledger`, `budget`, `task`, `feedback/`, the holdout guard, `workspace`'s snapshot store, + `grader.audit`, `vm` and `providers/`. +- Worker: a new job runner; `serve`, `regimes/`, `correctness/` (clients), `profile`/`bundle`/`gpu`, + `grader.run_tests`, `canary`, the container jail. + +## Leftover processes, today + +On Linux, srt already isolates them (own PID and network namespace; killed when the jail exits). During a +session they still share the GPU and CPU with the engine being measured, and nothing stops them before a +measurement: `serve.py` tears down only the engine's own process group. Rule 4 closes it. + +## At scale (not built) + +Many single-use workers under a scheduler (SkyPilot, Kubernetes or Ray) instead of `vm.py`; the controller as +an always-on service holding the one holdout guard per corpus version; corpus, tasks and weights in object +storage fetched by hash; the target spec describing multi-GPU and multi-node layouts; ledgers indexed in a +database, still append-only. + +## Order + +1. ~~Decide the jail~~ (A) and build it. Done, minus the GPU check. +2. Without a GPU: the job interface with a local transport, the snapshot round trip and rule 3, the API proxy, + and the container jail's network, filesystem and freeze behaviour on any Linux box (Docker locally). +3. With a GPU (spend approval): `--gpus` and the profilers inside the container, then a first `lab.session` + dry run with a tiny budget. diff --git a/lab/safety/container.py b/lab/safety/container.py new file mode 100644 index 0000000..0c04322 --- /dev/null +++ b/lab/safety/container.py @@ -0,0 +1,123 @@ +"""The Linux jail: one container per jailed command (NVIDIA runtime for the GPU), no network at all, read-only root, +non-root user, no capabilities. A served engine's port reaches the host through a socket bridge. See lab/RUNTIME.md.""" + +from __future__ import annotations + +import os +import shutil +import subprocess +import sys +import uuid +from pathlib import Path + +IMAGE = os.environ.get("LAB_JAIL_IMAGE", "lab-jail") +DOCKERFILE = Path(__file__).parent / "jail.Dockerfile" +JAIL_UID = 10001 # the jail's user when the lab runs as root; otherwise the lab's own uid +CLIENT_CORES = 4 # kept out of every jail's cpuset so the measurement clients cannot be starved +SYSTEM_PREFIXES = {Path("/"), Path("/usr"), Path("/usr/local")} # the image has its own; never mounted over it +BRIDGE = "bridge.sock" + + +def _docker() -> str | None: + return shutil.which("docker") + + +def available() -> bool: + """Docker answers and the jail image is built.""" + d = _docker() + return bool(d) and subprocess.run([d, "image", "inspect", IMAGE], capture_output=True).returncode == 0 + + +def build() -> None: + subprocess.run([_docker(), "build", "-q", "-t", IMAGE, "-f", str(DOCKERFILE), str(DOCKERFILE.parent)], check=True) + + +def gpus() -> bool: + return shutil.which("nvidia-smi") is not None + + +def user() -> tuple[int, int]: + return (JAIL_UID, JAIL_UID) if os.getuid() == 0 else (os.getuid(), os.getgid()) + + +def give(paths: list[Path]) -> None: + """Running as root, hand the writable paths to the jail's user so nothing in the jail runs as root.""" + if os.getuid() != 0: + return + for root in paths: + for p in [root, *root.rglob("*")]: + os.lchown(p, JAIL_UID, JAIL_UID) + + +def cpuset() -> str | None: + n = os.cpu_count() or 1 + return f"{CLIENT_CORES}-{n - 1}" if n > 2 * CLIENT_CORES else None + + +def mounts(writable: list[Path], readonly: list[Path], python: str) -> list[str]: + """Bind mounts at the same absolute paths, so argv and env mean the same inside; interpreter dirs read-only.""" + interpreter = {Path(sys.base_prefix).resolve(), Path(python).resolve().parent.parent} + ro = [p for p in [*readonly, *sorted(interpreter)] if p.exists() and p.resolve() not in SYSTEM_PREFIXES] + if Path("/usr/local/cuda").exists(): + ro.append(Path("/usr/local/cuda").resolve()) # nvcc for JIT-compiling libraries + args = [] + for p in writable: + args += ["--mount", f"type=bind,src={p.resolve()},dst={p.resolve()}"] + for p in dict.fromkeys(q.resolve() for q in ro): + args += ["--mount", f"type=bind,src={p},dst={p},readonly"] + return args + + +def argv(name: str, workdir: Path, command: list[str], env: dict[str, str], writable: list[Path], + readonly: list[Path], python: str, bridge_port: int | None = None) -> list[str]: + """`docker run` for `command`. With `bridge_port`, the engine's 127.0.0.1:port is also served on the unix socket + `/bridge.sock`, which the host side of `Bridge` forwards to the host's 127.0.0.1:port.""" + uid, gid = user() + out = [_docker(), "run", "--rm", "--init", "--name", name, "--network", "none", "--read-only", + "--tmpfs", "/tmp:exec", "--cap-drop", "ALL", "--security-opt", "no-new-privileges", + "--pids-limit", "8192", "--shm-size", "16g", "--user", f"{uid}:{gid}", "--workdir", str(workdir.resolve()), + "--label", "lab.jail=1"] + if gpus(): + out += ["--gpus", "all"] + if (cpus := cpuset()): + out += ["--cpuset-cpus", cpus] + for k, v in env.items(): + out += ["--env", f"{k}={v}"] + out += mounts(writable, readonly, python) + out.append(IMAGE) + if bridge_port is None: + return out + command + sock = writable[-1].resolve() / BRIDGE + return out + ["sh", "-c", f'socat UNIX-LISTEN:{sock},fork,mode=666 TCP:127.0.0.1:{bridge_port} & exec "$@"', + "jail", *command] + + +def new_name() -> str: + return f"lab-jail-{uuid.uuid4().hex[:12]}" + + +def remove(name: str) -> None: + """Kill and remove the container; a no-op if it is already gone.""" + subprocess.run([_docker(), "rm", "-f", name], capture_output=True) + + +class Bridge: + """Host side: 127.0.0.1:port on the host forwarded to the jail's socket. Only host processes can reach it.""" + + def __init__(self, port: int, sock: Path): + self.proc = subprocess.Popen(["socat", f"TCP-LISTEN:{port},bind=127.0.0.1,reuseaddr,fork", + f"UNIX-CONNECT:{sock}"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, + start_new_session=True) + + def stop(self) -> None: + try: + os.killpg(self.proc.pid, 9) + except ProcessLookupError: + pass + self.proc.wait() + + +if __name__ == "__main__": + if sys.argv[1:] != ["build"]: + sys.exit("usage: python -m lab.safety.container build") + build() diff --git a/lab/safety/grader.py b/lab/safety/grader.py index 91e727e..9cd4368 100644 --- a/lab/safety/grader.py +++ b/lab/safety/grader.py @@ -15,7 +15,7 @@ from pathlib import Path from lab import engine, target -from lab.safety import jail +from lab.safety import container, jail from lab.safety.surfaces import ALWAYS_DENY, HIDDEN, may_write IGNORED = ("*/__pycache__/*", "__pycache__/*", "*.pyc", ".pytest_cache/*", "*/.pytest_cache/*", @@ -186,35 +186,59 @@ def run_lint(tree: Path) -> Run: def _jail_argv(tree: Path, argv: list[str], env_extra: dict | None, domains, local_binding: bool = False, - tag: str = "", read: list[Path] = ()) -> tuple[list[str], dict, Path]: + tag: str = "", read: list[Path] = (), port: int | None = None + ) -> tuple[list[str], dict, Path, str | None]: """The wrapped argv, env and private tmp for running `argv` in `tree` as agent code: wiped env, jail around the - tree and the tmp, weights read-only.""" + tree and the tmp, weights read-only. The last item names the container when the container jail is used.""" tmp = Path(tempfile.mkdtemp(prefix="jail-")).resolve() # short: srt's sockets live here env = {"PATH": _engine_path(), "HOME": str(Path.home()), "TMPDIR": str(tmp), "PYTHONPATH": str(tree / "src") + os.pathsep + str(tree), "PYTHONDONTWRITEBYTECODE": "1", "HF_HOME": str(tmp / "hf-home"), "HF_HUB_CACHE": str(HF_HUB), "HF_HUB_OFFLINE": "1", **(env_extra or {})} - config = jail.settings([tree.resolve(), tmp], engine.venv(), list(domains), - readonly=([HF_HUB] if HF_HUB.exists() else []) + list(read), python=engine.python(), - local_binding=local_binding) + readonly = ([HF_HUB] if HF_HUB.exists() else []) + list(read) + if jail.backend() == "container": + if not container.available(): + if jail.required(): + raise jail.JailMissing(f"no container jail: needs docker and the {container.IMAGE} image " + "(python -m lab.safety.container build), or LAB_NO_JAIL=1 on a box you trust") + return argv, env, tmp, None + if domains: + raise ValueError("the container jail has no network") + name = container.new_name() + container.give([tree.resolve(), tmp]) + wrapped = container.argv(name, tree, argv, {**env, "HOME": str(tmp)}, [tree.resolve(), tmp], + [engine.venv(), *readonly], engine.python(), bridge_port=port) + cli_env = {k: os.environ[k] for k in ("PATH", "HOME", "DOCKER_HOST", "DOCKER_CONFIG") if k in os.environ} + return wrapped, cli_env, tmp, name + config = jail.settings([tree.resolve(), tmp], engine.venv(), list(domains), readonly=readonly, + python=engine.python(), local_binding=local_binding) # srt sets TMPDIR to a dir that does not exist inside bwrap (nsys and ncu then refuse to start): restate ours. - return jail.wrap(config, tree.with_suffix(f"{tag}.srt.json"), ["env", f"TMPDIR={tmp}", *argv]), env, tmp + return jail.wrap(config, tree.with_suffix(f"{tag}.srt.json"), ["env", f"TMPDIR={tmp}", *argv]), env, tmp, None def jailed_popen(tree: Path, argv: list[str], *, env_extra: dict | None = None, stdout=None, stderr=None, - local_binding: bool = True) -> tuple[subprocess.Popen, Path]: - """A long-running jailed process (a served engine) in its own session; returns (proc, tmp dir to remove).""" - wrapped, env, tmp = _jail_argv(tree, argv, env_extra, (), local_binding, tag=".serve") + local_binding: bool = True, port: int | None = None): + """A long-running jailed process (a served engine) in its own session, reachable on the host's 127.0.0.1:`port`. + Returns (proc, tmp dir to remove, stop): `stop` removes what lives outside proc's process group.""" + wrapped, env, tmp, name = _jail_argv(tree, argv, env_extra, (), local_binding, tag=".serve", port=port) proc = subprocess.Popen(wrapped, cwd=tree, env=env, stdout=stdout, stderr=stderr, start_new_session=True) - return proc, tmp + if name is None: + return proc, tmp, lambda: None + bridge = container.Bridge(port, tmp / container.BRIDGE) if port else None + + def stop() -> None: + container.remove(name) + if bridge: + bridge.stop() + return proc, tmp, stop def jailed(tree: Path, argv: list[str], *, timeout_s: float, env_extra: dict | None = None, domains: list[str] = (), read: list[Path] = ()) -> subprocess.CompletedProcess: """Run `argv` in `tree` as agent code: wiped env, jail around the tree and a private tmp, weights and `read` read-only.""" - wrapped, env, tmp = _jail_argv(tree, argv, env_extra, domains, read=read) + wrapped, env, tmp, name = _jail_argv(tree, argv, env_extra, domains, read=read) try: return subprocess.run(wrapped, cwd=tree, env=env, capture_output=True, text=True, timeout=timeout_s) except subprocess.TimeoutExpired as e: # a hung suite is a failed run, recorded like any other @@ -222,6 +246,8 @@ def jailed(tree: Path, argv: list[str], *, timeout_s: float, env_extra: dict | N return subprocess.CompletedProcess(argv, -9, out.decode(errors="replace") if isinstance(out, bytes) else out, f"timed out after {timeout_s:.0f}s") finally: + if name: + container.remove(name) shutil.rmtree(tmp, ignore_errors=True) diff --git a/lab/safety/jail.Dockerfile b/lab/safety/jail.Dockerfile new file mode 100644 index 0000000..4397716 --- /dev/null +++ b/lab/safety/jail.Dockerfile @@ -0,0 +1,7 @@ +# The jail image: an empty Ubuntu with what agent code needs to build and run. The engine's venv, the interpreter, +# weights and CUDA are bind-mounted read-only from the host; the NVIDIA runtime adds the driver. +FROM ubuntu:24.04 +RUN apt-get update -qq && DEBIAN_FRONTEND=noninteractive apt-get install -y -qq --no-install-recommends \ + build-essential python3 python3-dev git socat ca-certificates >/dev/null \ + && rm -rf /var/lib/apt/lists/* \ + && chmod 755 /root diff --git a/lab/safety/jail.py b/lab/safety/jail.py index 646d2df..80cb0d3 100644 --- a/lab/safety/jail.py +++ b/lab/safety/jail.py @@ -17,6 +17,12 @@ def available() -> bool: return shutil.which("srt") is not None +def backend() -> str: + """`container` on Linux (srt there has no GPU and gives each jail its own network), `srt` elsewhere; LAB_JAIL + overrides.""" + return os.environ.get("LAB_JAIL") or ("container" if sys.platform.startswith("linux") else "srt") + + def required() -> bool: """The referee runs agent code jailed unless LAB_NO_JAIL=1 (tests, a dev box you trust).""" return os.environ.get("LAB_NO_JAIL") != "1" diff --git a/lab/serve.py b/lab/serve.py index 9d7bf22..4312bae 100644 --- a/lab/serve.py +++ b/lab/serve.py @@ -59,6 +59,7 @@ def __init__(self, tree: Path, server: target.Server, log: Path | None = None, p self.log_offset = 0 self.proc: subprocess.Popen | None = None self.tmp: Path | None = None + self.stop_jail = None self.started_at: float | None = None self.ready_s: float | None = None @@ -74,7 +75,8 @@ def __enter__(self) -> "Served": out = open(self.log, "ab") if self.log else subprocess.DEVNULL self.started_at = time.monotonic() if self.jailed: - self.proc, self.tmp = grader.jailed_popen(self.tree, argv, env_extra=self.env_extra, stdout=out, stderr=out) + self.proc, self.tmp, self.stop_jail = grader.jailed_popen(self.tree, argv, env_extra=self.env_extra, + stdout=out, stderr=out, port=self.port) else: self.proc = subprocess.Popen(argv, cwd=self.tree, env={**os.environ, **self.env_extra}, stdout=out, stderr=out, start_new_session=True) @@ -113,6 +115,8 @@ def __exit__(self, *exc) -> None: os.killpg(self.proc.pid, signal.SIGKILL) except ProcessLookupError: pass + if self.stop_jail: + self.stop_jail() if self.tmp: shutil.rmtree(self.tmp, ignore_errors=True) if self.artifacts: diff --git a/lab/vm-setup.sh b/lab/vm-setup.sh index 0b219ec..9484eb0 100755 --- a/lab/vm-setup.sh +++ b/lab/vm-setup.sh @@ -2,6 +2,7 @@ # One-time bootstrap of a fresh GPU VM for the lab. Idempotent; run via `python -m lab.vm setup`. # Arg: the engine tree on the VM (default ~/inference-server). set -euo pipefail +here=$(cd "$(dirname "$0")" && pwd) # this environment's lab/ cd "${1:-$HOME/inference-server}" # the engine tree: its venv is the one every GPU command runs with SUDO=$([ "$(id -u)" = 0 ] && echo "" || echo sudo) # Verda images log in as root @@ -37,7 +38,7 @@ fi nsys --version 2>/dev/null || true ncu --version 2>/dev/null | head -1 || true -echo "== py-spy (hostprof) and the jail (srt needs Node >= 20, bubblewrap, socat, ripgrep)" +echo "== py-spy (hostprof) and srt, the agent CLI's jail until lab/RUNTIME.md's split (Node >= 20, bubblewrap, socat, ripgrep)" uv pip install -q --python .venv/bin/python py-spy && .venv/bin/py-spy --version if ! command -v srt >/dev/null; then $SUDO apt-get install -y -qq bubblewrap socat ripgrep xz-utils >/dev/null @@ -48,7 +49,23 @@ if ! command -v srt >/dev/null; then fi # Ubuntu 24.04 blocks unprivileged user namespaces, so bwrap cannot start; this VM is single-purpose. $SUDO sysctl -qw kernel.apparmor_restrict_unprivileged_userns=0 2>/dev/null || true -srt -c true >/dev/null 2>&1 && echo "jail: ok" || echo "jail: srt cannot start" +srt -c true >/dev/null 2>&1 && echo "srt: ok" || echo "srt: cannot start" + +echo "== jail: docker + NVIDIA Container Toolkit + the lab-jail image (lab/safety/container.py)" +command -v docker >/dev/null || $SUDO apt-get install -y -qq docker.io socat >/dev/null +command -v socat >/dev/null || $SUDO apt-get install -y -qq socat >/dev/null +if ! command -v nvidia-ctk >/dev/null; then + curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey \ + | $SUDO gpg --dearmor --yes -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg + curl -fsSL https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list \ + | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' \ + | $SUDO tee /etc/apt/sources.list.d/nvidia-container-toolkit.list >/dev/null + $SUDO apt-get update -qq && $SUDO apt-get install -y -qq nvidia-container-toolkit >/dev/null + $SUDO nvidia-ctk runtime configure --runtime=docker >/dev/null && $SUDO systemctl restart docker +fi +[ "$(id -u)" = 0 ] || $SUDO usermod -aG docker "$USER" # takes effect at the next login +$SUDO docker build -q -t lab-jail -f "$here/safety/jail.Dockerfile" "$here/safety" >/dev/null +$SUDO docker run --rm --gpus all --network none lab-jail nvidia-smi -L && echo "jail: gpu ok" || echo "jail: no GPU in the container" echo "== counters: can a non-admin process read GPU performance counters?" # ncu needs NVreg_RestrictProfilingToAdminUsers=0 (or root). Recorded as a fact about this venue. diff --git a/tests/test_lab_container.py b/tests/test_lab_container.py new file mode 100644 index 0000000..420868f --- /dev/null +++ b/tests/test_lab_container.py @@ -0,0 +1,143 @@ +"""The container jail: the docker argv everywhere; isolation and the port bridge where docker and the image exist.""" + +from __future__ import annotations + +import os +import socket +import subprocess +import sys +import textwrap +import time +import urllib.request +from pathlib import Path + +import pytest + +from lab.safety import container, grader, jail +from tests.lab_fixtures import make_repo + + +@pytest.fixture +def docker_named(monkeypatch): + monkeypatch.setattr(container, "_docker", lambda: "docker") + + +def test_argv_has_no_network_no_caps_and_a_non_root_user(docker_named, tmp_path, monkeypatch): + monkeypatch.setattr(container.os, "getuid", lambda: 0) + a = container.argv("n", tmp_path, ["true"], {"K": "v"}, [tmp_path], [], sys.executable) + s = " ".join(a) + for flag in ("--network none", "--read-only", "--cap-drop ALL", "--security-opt no-new-privileges", + f"--user {container.JAIL_UID}:{container.JAIL_UID}", "--env K=v", "--rm"): + assert flag in s + assert a[-1] == "true" and a[-2] == container.IMAGE + + +def test_system_prefixes_are_never_mounted_over_the_image(tmp_path): + m = " ".join(container.mounts([tmp_path], [Path("/usr"), tmp_path / "ro"], "/usr/bin/python3")) + assert "dst=/usr," not in m and "dst=/usr " not in m and not m.endswith("dst=/usr") + + +def test_readonly_mounts_are_readonly(tmp_path): + (tmp_path / "ro").mkdir() + m = container.mounts([tmp_path / "w"], [tmp_path / "ro"], "/usr/bin/python3") + ro = [x for x in m if f"dst={(tmp_path / 'ro').resolve()}" in x] + assert ro and ro[0].endswith(",readonly") + + +def test_bridge_wraps_the_command_and_keeps_it_last(docker_named, tmp_path): + a = container.argv("n", tmp_path, ["python", "-m", "srv"], {}, [tmp_path / "t"], [], sys.executable, + bridge_port=8123) + assert a[-3:] == ["python", "-m", "srv"] and "TCP:127.0.0.1:8123" in a[a.index(container.IMAGE) + 3] + + +def test_linux_defaults_to_the_container_jail(monkeypatch): + monkeypatch.delenv("LAB_JAIL", raising=False) + monkeypatch.setattr(jail.sys, "platform", "linux") + assert jail.backend() == "container" + monkeypatch.setattr(jail.sys, "platform", "darwin") + assert jail.backend() == "srt" + + +live = pytest.mark.skipif(not (jail.backend() == "container" and container.available()), + reason="needs the container jail: Linux, docker and the lab-jail image") + +PROBE = textwrap.dedent(""" + import os, socket, sys + out = {} + out["uid"] = os.getuid() + out["caps"] = [l.split()[1] for l in open("/proc/self/status") if l.startswith("CapEff")][0] + open("written_by_jail", "w").write("x") + def tries(f): + try: + f(); return "yes" + except Exception: + return "no" + out["write_root"] = tries(lambda: open("/etc/x", "w")) + out["read_secret"] = tries(lambda: open(sys.argv[1]).read()) + out["net"] = tries(lambda: socket.create_connection(("1.1.1.1", 53), timeout=2)) + out["host_port"] = tries(lambda: socket.create_connection(("127.0.0.1", int(sys.argv[2])), timeout=2)) + print(out) +""") + + +@live +def test_jail_isolates_files_network_and_privileges(tmp_path, monkeypatch): + monkeypatch.delenv("LAB_NO_JAIL", raising=False) + tree = tmp_path / "t" + grader.export(make_repo(tmp_path), "HEAD", tree) + (tree / "probe.py").write_text(PROBE) + secret = tmp_path / "secret" + secret.write_text("key") + srv = socket.socket() + srv.bind(("127.0.0.1", 0)) + srv.listen() + proc = grader.jailed(tree, ["python3", "probe.py", str(secret), str(srv.getsockname()[1])], timeout_s=60) + srv.close() + assert proc.returncode == 0, proc.stderr + out = eval(proc.stdout.strip().splitlines()[-1]) + assert out["uid"] != 0 and int(out["caps"], 16) == 0 + assert (tree / "written_by_jail").exists() + assert out["write_root"] == "no" and out["read_secret"] == "no" + assert out["net"] == "no" and out["host_port"] == "no" + + +@live +def test_timeout_removes_the_container(tmp_path, monkeypatch): + monkeypatch.delenv("LAB_NO_JAIL", raising=False) + tree = tmp_path / "t" + grader.export(make_repo(tmp_path), "HEAD", tree) + proc = grader.jailed(tree, ["sleep", "60"], timeout_s=3) + assert proc.returncode == -9 + left = subprocess.run(["docker", "ps", "-q", "--filter", "label=lab.jail=1"], capture_output=True, text=True) + assert left.stdout.strip() == "" + + +@live +def test_served_port_reaches_the_host_and_nothing_else(tmp_path, monkeypatch): + from lab import serve + monkeypatch.delenv("LAB_NO_JAIL", raising=False) + tree = tmp_path / "t" + grader.export(make_repo(tmp_path), "HEAD", tree) + (tree / "index.html").write_text("ok") + port = serve.free_port() + proc, tmp, stop = grader.jailed_popen(tree, ["python3", "-m", "http.server", str(port), "--bind", "127.0.0.1"], + port=port) + try: + body, deadline = None, time.monotonic() + 30 + while body is None and time.monotonic() < deadline: + try: + body = urllib.request.urlopen(f"http://127.0.0.1:{port}/index.html", timeout=2).read() + except OSError: + time.sleep(0.3) + assert body == b"ok" + (tree / "probe.py").write_text(PROBE) # another jail cannot reach the served engine + other = grader.jailed(tree, ["python3", "probe.py", "/nonexistent", str(port)], timeout_s=60) + assert eval(other.stdout.strip().splitlines()[-1])["host_port"] == "no" + finally: + os.killpg(proc.pid, 15) + proc.wait(timeout=30) + stop() + left = subprocess.run(["docker", "ps", "-q", "--filter", "label=lab.jail=1"], capture_output=True, text=True) + assert left.stdout.strip() == "" + with pytest.raises(OSError): + socket.create_connection(("127.0.0.1", port), timeout=2) diff --git a/tests/test_lab_proftools.py b/tests/test_lab_proftools.py index 998d3da..667f0bb 100644 --- a/tests/test_lab_proftools.py +++ b/tests/test_lab_proftools.py @@ -13,6 +13,7 @@ from lab import ledger, proftools from lab import tools as labtools from lab.budget import Budget, BudgetExceeded +from lab.safety import container, jail from lab.session import Session from lab.tools import GPU_TOOLS, Toolbox from lab.workspace import Workspace @@ -39,6 +40,8 @@ @pytest.fixture def tb(tmp_path, monkeypatch): monkeypatch.setenv("LAB_NO_JAIL", "1") + monkeypatch.setattr(jail, "available", lambda: False) # the fake instruments log outside any jail + monkeypatch.setattr(container, "available", lambda: False) monkeypatch.setenv("LAB_GPU_USD_PER_HOUR", "3600") # $1 per second: the charge equals the seconds bin_dir = tmp_path / "bin" bin_dir.mkdir() diff --git a/tests/test_lab_safety.py b/tests/test_lab_safety.py index 5a2e111..db6f432 100644 --- a/tests/test_lab_safety.py +++ b/tests/test_lab_safety.py @@ -7,7 +7,7 @@ import pytest -from lab.safety import grader, jail +from lab.safety import container, grader, jail from lab.safety.surfaces import may_write from tests.lab_fixtures import make_repo @@ -119,6 +119,7 @@ def test_pristine_tree_diff_lint_and_tests(tmp_path, monkeypatch): def test_run_tests_refuses_without_the_jail(tmp_path, monkeypatch): monkeypatch.delenv("LAB_NO_JAIL", raising=False) monkeypatch.setattr(jail, "available", lambda: False) + monkeypatch.setattr(container, "available", lambda: False) repo = make_repo(tmp_path) tree = tmp_path / "t" grader.export(repo, "HEAD", tree) From a27dedcf9eb28eb15dc1dcea0758c4819b85f063 Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Fri, 9 Oct 2026 00:23:13 -0400 Subject: [PATCH 02/18] lab: no measurement on a shared GPU (refused / contaminated) A leftover process on the GPU could slow a bench, equiv or submit or stop the engine from starting, and the ledger recorded that as the snapshot's result; a contaminated base measurement was cached and skewed every later submit. serve.Served now refuses to launch while any process holds the GPU (GpuBusy, nothing launched) and polls the GPU's processes every second while serving; anything outside the engine's own process tree (descendants, or the container's processes) makes exclusive() raise Contaminated. bench/equiv/submit/baseline record verdict refused or contaminated with gpu_processes as facts; a contaminated run is never evidence, never cached as the base, and is caught before the holdout guard spends a query. lab.validate's noise bands get the same check. Co-Authored-By: Claude Opus 5.5 --- CLAUDE.md | 4 ++- lab/LEDGER.md | 5 ++- lab/RUNTIME.md | 9 ++++-- lab/evaltools.py | 36 +++++++++++++++------ lab/gpu.py | 34 ++++++++++++++++++++ lab/safety/container.py | 6 ++++ lab/safety/grader.py | 9 +++--- lab/serve.py | 64 +++++++++++++++++++++++++++++++++++-- lab/validate.py | 1 + tests/test_lab_container.py | 4 +-- tests/test_lab_evaltools.py | 44 +++++++++++++++++++++++++ tests/test_lab_gpu.py | 19 +++++++++++ 12 files changed, 212 insertions(+), 23 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 88c9781..4d4789c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -36,7 +36,9 @@ pytest/ruff config). Each folder with commands has its own `python -m`. `proftools.py` (trace = nsys, kernel = ncu, hostprof = py-spy, each wrapping lab.profile's workload in the jail with validated args) + `evaltools.py` (bench → `regimes` on the seen split; equiv → `correctness` candidate vs the target's reference; submit → held-out full tier vs the - base commit, one aggregate per regime through the Thresholdout guard; all three serve the pristine tree via + base commit, one aggregate per regime through the Thresholdout guard; nothing is measured while another process holds + the GPU (`refused`), and a process other than the engine on the GPU mid-run makes it `contaminated` (never evidence, + never the cached base; `serve.Served.exclusive()`); all three serve the pristine tree via `serve.py`, jailed (on Linux, port bridged out of a network-less container); every serve also keeps its passive data, `artifacts.py`: device samples via `gpu.DeviceSampler` (DCGM if `dcgmi`, else nvidia-smi), engine telemetry files from `{dir}`, all client rows, the serve log; bench/equiv as ledger blobs under `result.artifacts`, submit only in diff --git a/lab/LEDGER.md b/lab/LEDGER.md index aecd476..b722f48 100644 --- a/lab/LEDGER.md +++ b/lab/LEDGER.md @@ -21,7 +21,10 @@ on write. Every record: `id`, `at`, `schema` (set by the writer), `kind`, `run`, | `note` | `note`, `restore`, refusals | free text, or `{verdict: refused, reason}` | A tool that could not run writes `{verdict: refused | error, reason}` and the agent sees that text; it never -crashes the session. The integrity verdict (`feedback/lab_verdict.py`) reads only these records: a run is valid +crashes the session. `bench`, `equiv`, `submit` and `baseline` also write `refused` when a process already held +the GPU (no engine launched) and `contaminated` when a process other than the served engine held it during the +measurement; both carry `gpu_processes: [{pid, used_mib, name}]`. A contaminated run is never evidence and never +cached as the base. The integrity verdict (`feedback/lab_verdict.py`) reads only these records: a run is valid iff a completed `submit` exists whose snapshot also has a completed seen-split `bench`, a passing `test`, a passing `equiv` (tier full) and no failing `equiv`, and no record reports a write-surface violation. diff --git a/lab/RUNTIME.md b/lab/RUNTIME.md index ee799bd..8059e13 100644 --- a/lab/RUNTIME.md +++ b/lab/RUNTIME.md @@ -130,8 +130,13 @@ Transport now: SSH from the controller, one worker. Transport later: a queue. Th ## Leftover processes, today On Linux, srt already isolates them (own PID and network namespace; killed when the jail exits). During a -session they still share the GPU and CPU with the engine being measured, and nothing stops them before a -measurement: `serve.py` tears down only the engine's own process group. Rule 4 closes it. +session they could still share the GPU with the engine being measured. Since 2026-10-09 `serve.py` refuses to +launch an engine while any process holds the GPU (`GpuBusy`) and polls the GPU's processes every second while it +serves; anything outside the engine's own processes marks the measurement `contaminated` (`Served.exclusive()`): +never evidence for submit, never cached as the base, and caught before the holdout guard spends a query. The +facts (pid, MiB, name) go to the agent. It refuses rather than kills because the lab cannot yet tell the agent's +processes apart; rule 4's freeze-and-kill comes with the agent container. Not covered: CPU contention, and a GPU +burst shorter than the poll interval. ## At scale (not built) diff --git a/lab/evaltools.py b/lab/evaltools.py index 3f86457..ae19546 100644 --- a/lab/evaltools.py +++ b/lab/evaltools.py @@ -121,6 +121,19 @@ def equiv_rows(out: Path) -> list[dict]: return [{k: v for k, v in json.loads(x).items() if k != "text"} for x in p.read_text().splitlines() if x.strip()] +def _failed(e: BaseException) -> dict: + """The record of a measurement that did not complete; GPU processes that stopped it are kept as facts.""" + if isinstance(e, serve.GpuBusy): + return {"verdict": "refused", "reason": str(e), "gpu_processes": e.apps} + if isinstance(e, serve.Contaminated): + return {"verdict": "contaminated", "reason": str(e), "gpu_processes": e.apps} + return {"verdict": "error", "reason": _why(e)} + + +def _said(tool: str, result: dict) -> str: + return f"{tool} {'failed' if result['verdict'] == 'error' else result['verdict']}: {result['reason']}" + + def _why(e: BaseException) -> str: if isinstance(e, serve.NotReady): return f"engine did not start: {e}" @@ -144,11 +157,11 @@ def bench(self, args: dict) -> str: with serve.Served(tree, t.engine, log=self.s.run_dir / "serve-bench.log", jailed=self.serve_jailed, artifacts=work) as srv: results = run_regimes(srv.url, t, names, "seen", tier, self.seed, rows) + srv.exclusive() except Exception as e: # a tool never crashes the session: the failure is the record - result = {"verdict": "error", "reason": _why(e), "seconds": time.monotonic() - t0, - "artifacts": self._store(work, rows)} + result = {**_failed(e), "seconds": time.monotonic() - t0, "artifacts": self._store(work, rows)} self._record("bench", "bench", args, result, snap, config={"split": "seen", "tier": tier}) - return f"bench failed: {result['reason']}" + return _said("bench", result) result = {"verdict": "ok", "metrics": headline(results), "regimes": results, "tier": tier, "seconds": time.monotonic() - t0, "ready_s": srv.ready_s, "artifacts": self._store(work, rows)} self._record("bench", "bench", args, result, snap, config={"split": "seen", "tier": tier}) @@ -176,11 +189,11 @@ def equiv(self, args: dict) -> str: Thresholds(min_score_ratio=t.min_score_ratio, min_length_ratio=t.min_length_ratio), concurrency=self.equiv_concurrency, api=client.Api(t.engine.api, t.chat_kwargs), config={"tier": tier}) + srv.exclusive() except Exception as e: - self._record("equiv", "equiv", args, {"verdict": "error", "reason": _why(e), - "artifacts": self._store(work, equiv_rows(out))}, - snap, config={"split": "seen", "tier": tier}) - return f"equiv failed: {_why(e)}" + result = {**_failed(e), "artifacts": self._store(work, equiv_rows(out))} + self._record("equiv", "equiv", args, result, snap, config={"split": "seen", "tier": tier}) + return _said("equiv", result) passed = {"pass": True, "fail": False}.get(res["verdict"]) # inconclusive is None: neither passing nor failing record = {"verdict": res["verdict"], "passed": passed, "reasons": res["reasons"], "gates": res["gates"], "metrics": res["metrics"], "thresholds": res["thresholds"], "tier": tier, @@ -219,6 +232,7 @@ def submit(self, args: dict) -> str: with serve.Served(tree, t.engine, log=self.s.run_dir / "serve-submit.log", jailed=self.serve_jailed, artifacts=private) as srv: results = run_regimes(srv.url, t, names, "heldout", "full", self.seed, rows) + srv.exclusive() # before the holdout guard: a contaminated run spends no query finally: artifacts.write_rows(private, rows) new = headline(results) @@ -234,8 +248,9 @@ def submit(self, args: dict) -> str: self._record("submit", "submit", args, {"verdict": "refused", "reason": str(e)}, snap) return f"submit refused: {e}" except Exception as e: - self._record("submit", "submit", args, {"verdict": "error", "reason": _why(e)}, snap) - return f"submit failed: {_why(e)}" + result = _failed(e) + self._record("submit", "submit", args, result, snap) + return _said("submit", result) rec = self._record_heldout("submit", "submit", args, {"tier": "full"}, metrics, snap) out = {"snapshot": snap.id, "split": "heldout", "tier": "full", "metrics": metrics, "record": rec["id"]} task = getattr(self.s, "task", None) @@ -283,6 +298,7 @@ def _base_measure(self, t: target.Target, names: list[str], split: str, tier: st grader.export(self.s.workspace.repo, self.s.workspace.base, tree) with serve.Served(tree, t.engine, log=self.s.run_dir / f"serve-base-{split}.log", jailed=self.serve_jailed) as srv: have.update(headline(run_regimes(srv.url, t, missing, split, tier, self.seed))) + srv.exclusive() # a contaminated base is never cached: every later submit reads it cache.write_text(json.dumps(have, indent=1)) return have @@ -297,7 +313,7 @@ def baseline(self, commit: str) -> dict: seen = self.base_seen(target.load(), names) result = {"verdict": "ok", "metrics": {n: seen[n] for n in names}} except Exception as e: # no GPU or a base that will not serve is itself the fact - result = {"verdict": "error", "reason": _why(e)} + result = _failed(e) return ledger.append({"kind": "baseline", "run": self.s.run_id, "session": self.s.session_id, "config": {"split": "seen", "tier": BENCH_TIER, "commit": commit}, "result": result, "cost": self._cost("baseline")}, self.s.ledger_root) diff --git a/lab/gpu.py b/lab/gpu.py index 577c188..111a118 100644 --- a/lab/gpu.py +++ b/lab/gpu.py @@ -161,3 +161,37 @@ def stop(self) -> None: self._file.close() self._proc = None (self.path.parent / "meta.json").write_text(json.dumps(meta, indent=1)) + + +def compute_apps() -> list[dict] | None: + """Every process holding a GPU context (host pids): pid, used MiB, name. None where there is no nvidia-smi.""" + if not available(): + return None + out = subprocess.run(["nvidia-smi", "--query-compute-apps=pid,used_memory,process_name", + "--format=csv,noheader,nounits"], capture_output=True, text=True) + if out.returncode != 0: + return None + apps = [] + for line in out.stdout.splitlines(): + pid, mib, name = ([p.strip() for p in line.split(",", 2)] + ["", ""])[:3] + if pid.isdigit(): + apps.append({"pid": int(pid), "used_mib": int(mib) if mib.isdigit() else None, "name": name}) + return apps + + +def descendants(pid: int) -> set[int]: + """`pid` and every process below it, from /proc (Linux); `{pid}` where there is no /proc.""" + children: dict[int, list[int]] = {} + for stat in Path("/proc").glob("[0-9]*/stat"): + try: + ppid = int(stat.read_text().rsplit(")", 1)[1].split()[1]) + except (OSError, ValueError, IndexError): + continue + children.setdefault(ppid, []).append(int(stat.parent.name)) + out, todo = set(), [pid] + while todo: + p = todo.pop() + if p not in out: + out.add(p) + todo += children.get(p, []) + return out diff --git a/lab/safety/container.py b/lab/safety/container.py index 0c04322..99fbfe6 100644 --- a/lab/safety/container.py +++ b/lab/safety/container.py @@ -96,6 +96,12 @@ def new_name() -> str: return f"lab-jail-{uuid.uuid4().hex[:12]}" +def pids(name: str) -> set[int]: + """The container's processes as host pids; empty once it is gone.""" + out = subprocess.run([_docker(), "top", name, "-eo", "pid"], capture_output=True, text=True) + return {int(x) for x in out.stdout.split() if x.isdigit()} if out.returncode == 0 else set() + + def remove(name: str) -> None: """Kill and remove the container; a no-op if it is already gone.""" subprocess.run([_docker(), "rm", "-f", name], capture_output=True) diff --git a/lab/safety/grader.py b/lab/safety/grader.py index 9cd4368..04195ad 100644 --- a/lab/safety/grader.py +++ b/lab/safety/grader.py @@ -14,7 +14,7 @@ from fnmatch import fnmatch from pathlib import Path -from lab import engine, target +from lab import engine, gpu, target from lab.safety import container, jail from lab.safety.surfaces import ALWAYS_DENY, HIDDEN, may_write @@ -220,18 +220,19 @@ def _jail_argv(tree: Path, argv: list[str], env_extra: dict | None, domains, loc def jailed_popen(tree: Path, argv: list[str], *, env_extra: dict | None = None, stdout=None, stderr=None, local_binding: bool = True, port: int | None = None): """A long-running jailed process (a served engine) in its own session, reachable on the host's 127.0.0.1:`port`. - Returns (proc, tmp dir to remove, stop): `stop` removes what lives outside proc's process group.""" + Returns (proc, tmp dir to remove, stop, pids): `stop` removes what lives outside proc's process group; `pids()` + is every host pid the jail runs.""" wrapped, env, tmp, name = _jail_argv(tree, argv, env_extra, (), local_binding, tag=".serve", port=port) proc = subprocess.Popen(wrapped, cwd=tree, env=env, stdout=stdout, stderr=stderr, start_new_session=True) if name is None: - return proc, tmp, lambda: None + return proc, tmp, lambda: None, lambda: gpu.descendants(proc.pid) bridge = container.Bridge(port, tmp / container.BRIDGE) if port else None def stop() -> None: container.remove(name) if bridge: bridge.stop() - return proc, tmp, stop + return proc, tmp, stop, lambda: container.pids(name) def jailed(tree: Path, argv: list[str], *, timeout_s: float, env_extra: dict | None = None, diff --git a/lab/serve.py b/lab/serve.py index 4312bae..90037de 100644 --- a/lab/serve.py +++ b/lab/serve.py @@ -2,7 +2,8 @@ free localhost port, with the target's env; ready when its health path answers 200; torn down with its whole process group. `with Served(tree, spec) as s: ... s.url`. With `artifacts`, the run's passive data lands in that dir: `device/` (samples for the whole served window), `telemetry/` (files the engine wrote, when the target asks for them) -and `serve.log` (this run's part of the log).""" +and `serve.log` (this run's part of the log). The GPU must be free when the engine starts (else `GpuBusy`, nothing +launched); while it is served, any other process holding the GPU is recorded and `exclusive()` raises `Contaminated`.""" from __future__ import annotations @@ -12,6 +13,7 @@ import shutil import socket import subprocess +import threading import time import urllib.error import urllib.request @@ -21,6 +23,7 @@ from lab.safety import grader TELEMETRY_SUBDIR = "lab-telemetry" # inside the served tree, the one place a jailed engine may write +POLL_S = 1.0 # how often the GPU's processes are checked while an engine is served def free_port() -> int: @@ -40,6 +43,26 @@ class NotReady(RuntimeError): pass +def _listed(apps: list[dict]) -> str: + return "; ".join(f"pid {a['pid']} ({a['name']}, {a['used_mib']} MiB)" for a in apps) + + +class GpuBusy(RuntimeError): + """Something already held the GPU, so no engine was launched.""" + + def __init__(self, apps: list[dict]): + super().__init__(f"the GPU was in use before the engine started: {_listed(apps)}") + self.apps = apps + + +class Contaminated(RuntimeError): + """A process outside the served engine held the GPU during the measurement.""" + + def __init__(self, apps: list[dict]): + super().__init__(f"another process held the GPU during the measurement: {_listed(apps)}") + self.apps = apps + + class Served: def __init__(self, tree: Path, server: target.Server, log: Path | None = None, port: int | None = None, timeout_s: float | None = None, env_extra: dict | None = None, jailed: bool = True, @@ -60,10 +83,17 @@ def __init__(self, tree: Path, server: target.Server, log: Path | None = None, p self.proc: subprocess.Popen | None = None self.tmp: Path | None = None self.stop_jail = None + self.pids = None # every host pid of the served engine + self.foreign: dict[int, dict] = {} # GPU processes seen while serving that were not the engine's + self._stop = threading.Event() + self._watcher: threading.Thread | None = None self.started_at: float | None = None self.ready_s: float | None = None def __enter__(self) -> "Served": + busy = gpu.compute_apps() + if busy: + raise GpuBusy(busy) argv = argv_for(self.server, self.port) self.log_offset = self.log.stat().st_size if self.log and self.log.exists() else 0 if self.artifacts: @@ -75,11 +105,15 @@ def __enter__(self) -> "Served": out = open(self.log, "ab") if self.log else subprocess.DEVNULL self.started_at = time.monotonic() if self.jailed: - self.proc, self.tmp, self.stop_jail = grader.jailed_popen(self.tree, argv, env_extra=self.env_extra, - stdout=out, stderr=out, port=self.port) + self.proc, self.tmp, self.stop_jail, self.pids = grader.jailed_popen( + self.tree, argv, env_extra=self.env_extra, stdout=out, stderr=out, port=self.port) else: self.proc = subprocess.Popen(argv, cwd=self.tree, env={**os.environ, **self.env_extra}, stdout=out, stderr=out, start_new_session=True) + self.pids = lambda: gpu.descendants(self.proc.pid) + if busy is not None: # None: no nvidia-smi here, nothing to watch + self._watcher = threading.Thread(target=self._watch, daemon=True) + self._watcher.start() try: self._wait_ready() except BaseException: @@ -105,7 +139,31 @@ def _wait_ready(self) -> None: raise NotReady(f"{health} not ready after {self.timeout_s:.0f}s" + (f"; log: {self.log}" if self.log else "")) time.sleep(0.5) + def _watch(self) -> None: + while True: + before = self.pids() + apps = gpu.compute_apps() or [] + mine = before | (self.pids() if apps else set()) # a process the engine starts meanwhile is its own + for a in apps: + if a["pid"] not in mine: + self.foreign.setdefault(a["pid"], a) + if self._stop.wait(POLL_S): + return + + def _stop_watching(self) -> None: + self._stop.set() + if self._watcher is not None: + self._watcher.join() + self._watcher = None + + def exclusive(self) -> None: + """Call once the measurement is done: raises `Contaminated` if anything else held the GPU since launch.""" + self._stop_watching() + if self.foreign: + raise Contaminated(list(self.foreign.values())) + def __exit__(self, *exc) -> None: + self._stop_watching() # before teardown: an engine going away must not read as someone else if self.proc is not None and self.proc.poll() is None: try: os.killpg(self.proc.pid, signal.SIGTERM) diff --git a/lab/validate.py b/lab/validate.py index c1dad84..1093318 100644 --- a/lab/validate.py +++ b/lab/validate.py @@ -90,6 +90,7 @@ def validate(t: target.Target, tier: str = "dev", noise_runs: int = 3, regimes: for n in names: r = suite.REGIMES[n](ctx) values[n].append(r["value"] if r["valid"] else None) + srv.exclusive() # a band from a shared GPU would skew every submit for n in names: report["noise"][n] = write_noise(n, values[n], t, {"tier": "short", "split": "seen"}) report["bands_measured"] = all(d["band_pct"] is not None for d in report["noise"].values()) if names and noise_runs else None diff --git a/tests/test_lab_container.py b/tests/test_lab_container.py index 420868f..72b8457 100644 --- a/tests/test_lab_container.py +++ b/tests/test_lab_container.py @@ -120,7 +120,7 @@ def test_served_port_reaches_the_host_and_nothing_else(tmp_path, monkeypatch): grader.export(make_repo(tmp_path), "HEAD", tree) (tree / "index.html").write_text("ok") port = serve.free_port() - proc, tmp, stop = grader.jailed_popen(tree, ["python3", "-m", "http.server", str(port), "--bind", "127.0.0.1"], + proc, tmp, stop, pids = grader.jailed_popen(tree, ["python3", "-m", "http.server", str(port), "--bind", "127.0.0.1"], port=port) try: body, deadline = None, time.monotonic() + 30 @@ -129,7 +129,7 @@ def test_served_port_reaches_the_host_and_nothing_else(tmp_path, monkeypatch): body = urllib.request.urlopen(f"http://127.0.0.1:{port}/index.html", timeout=2).read() except OSError: time.sleep(0.3) - assert body == b"ok" + assert body == b"ok" and pids() # the jail reports its processes as host pids (tree / "probe.py").write_text(PROBE) # another jail cannot reach the served engine other = grader.jailed(tree, ["python3", "probe.py", "/nonexistent", str(port)], timeout_s=60) assert eval(other.stdout.strip().splitlines()[-1])["host_port"] == "no" diff --git a/tests/test_lab_evaltools.py b/tests/test_lab_evaltools.py index 015e714..3565cd1 100644 --- a/tests/test_lab_evaltools.py +++ b/tests/test_lab_evaltools.py @@ -248,6 +248,50 @@ def test_tools_refuse_cleanly_when_the_engine_does_not_start(lab, monkeypatch): assert rec["result"]["artifacts"]["telemetry"] is None and rec["result"]["artifacts"]["serve_log"] # no section: skipped +AGENT_PROC = {"pid": 4242, "used_mib": 9000, "name": "python"} + + +def test_nothing_is_measured_while_another_process_holds_the_gpu(lab, monkeypatch): + from lab import gpu + tb, s = lab + monkeypatch.setattr(gpu, "compute_apps", lambda: [AGENT_PROC]) + out = tb.bench({"regimes": ["single_stream"]}) + assert out.startswith("bench refused: the GPU was in use") and "pid 4242" in out + rec = list(ledger.records(s.ledger_root, kind="bench"))[-1]["result"] + assert rec["verdict"] == "refused" and rec["gpu_processes"] == [AGENT_PROC] + assert not (s.run_dir / "serve-bench.log").exists() # no engine was launched + + +def test_a_measurement_shared_with_another_gpu_process_is_contaminated_and_never_cached(lab, monkeypatch): + from lab import gpu + tb, s = lab + calls = [] + monkeypatch.setattr(serve, "POLL_S", 0.05) + monkeypatch.setattr(gpu, "compute_apps", lambda: calls.append(1) or ([] if len(calls) == 1 else [AGENT_PROC])) + out = tb.bench({"regimes": ["single_stream"]}) + assert out.startswith("bench contaminated: another process held the GPU") + rec = list(ledger.records(s.ledger_root, kind="bench"))[-1]["result"] + assert rec["verdict"] == "contaminated" and rec["gpu_processes"] == [AGENT_PROC] and "metrics" not in rec + calls.clear() + with pytest.raises(serve.Contaminated): + tb.base_seen(target.load(), ["single_stream"]) + assert not list(s.run_dir.glob("base-*.json")) + + +def test_the_engines_own_gpu_processes_are_not_foreign(lab, monkeypatch): + import time + from lab import gpu + tb, s = lab + srv = serve.Served(tb._pristine(), target.load().engine, jailed=False) + monkeypatch.setattr(serve, "POLL_S", 0.05) + monkeypatch.setattr(gpu, "compute_apps", lambda: [] if srv.proc is None else + [{"pid": srv.proc.pid, "used_mib": 1, "name": "engine"}]) + with srv: + time.sleep(0.3) + srv.exclusive() + assert srv.foreign == {} + + class TwoSessions: """Continues once, then stops; keeps each session's spec.""" name = "scripted" diff --git a/tests/test_lab_gpu.py b/tests/test_lab_gpu.py index 51fdc3c..e88f928 100644 --- a/tests/test_lab_gpu.py +++ b/tests/test_lab_gpu.py @@ -80,3 +80,22 @@ def test_sampler_writes_epoch_seconds(fake_smi, tmp_path): def test_unparseable_timestamp_leaves_epoch_empty(): out = gpu.add_epoch("timestamp, x\nnot a time, 1\n") assert out.splitlines() == ["epoch_s,timestamp,x", ",not a time,1"] + + +def test_compute_apps_parses_pids_memory_and_names(monkeypatch): + import subprocess + monkeypatch.setattr(gpu, "available", lambda: True) + monkeypatch.setattr(gpu.subprocess, "run", lambda *a, **k: subprocess.CompletedProcess( + a[0], 0, "123, 4000, python\n456, [N/A], /usr/bin/vllm, worker\n", "")) + assert gpu.compute_apps() == [{"pid": 123, "used_mib": 4000, "name": "python"}, + {"pid": 456, "used_mib": None, "name": "/usr/bin/vllm, worker"}] + + +@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="reads /proc") +def test_descendants_include_children(): + import subprocess + child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(5)"]) + try: + assert {os.getpid(), child.pid} <= gpu.descendants(os.getpid()) + finally: + child.kill() From c9b43ad97657e559d1d4f0b31a058330734c5f4b Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Fri, 9 Oct 2026 00:42:18 -0400 Subject: [PATCH 03/18] lab: worker job interface (local and SSH transports) The GPU work behind test, profile/trace/kernel/hostprof, bench, equiv, submit and the base measurement now runs as worker jobs (lab/worker.py: test, measure, equiv, profile). A job takes input trees and JSON args, writes an out dir and returns JSON; the tools keep the policy, the ledger and the budget and call self.worker. LAB_WORKER=local (default) runs jobs in process, as before. LAB_WORKER=ssh runs them on the lab.vm VM: each tree goes by content hash at most once (has/put), the worker re-hashes a fresh copy before every job and refuses a mismatch, and the out dir comes back as a tar. GpuBusy, Contaminated, NotReady and ValueError are raised again on the controller. equiv ships the snapshot's earlier answers with the job, so a second equiv never re-rolls the gate (it already did not, in process). Co-Authored-By: Claude Opus 5.5 --- CLAUDE.md | 8 +- lab/README.md | 4 + lab/RUNTIME.md | 18 ++- lab/evaltools.py | 102 ++++++------ lab/proftools.py | 125 +++++++-------- lab/tools.py | 36 +++-- lab/worker.py | 300 ++++++++++++++++++++++++++++++++++++ tests/test_lab_evaltools.py | 11 ++ tests/test_lab_worker.py | 105 +++++++++++++ 9 files changed, 577 insertions(+), 132 deletions(-) create mode 100644 lab/worker.py create mode 100644 tests/test_lab_worker.py diff --git a/CLAUDE.md b/CLAUDE.md index 4d4789c..8e4db71 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -61,8 +61,12 @@ pytest/ruff config). Each folder with commands has its own `python -m`. measures each regime's run-to-run band from repeated reference runs into `knowledge/noise/.json` (which `submit` reads). Exit 0 only if both hold. Needs a GPU for a real target; tested on fakes. `lab/LEDGER.md` is the record-format spec: the environment's public interface. -- `lab/RUNTIME.md` — design, not built (2026-10-08): controller (laptop: ledger, referee, secrets) and worker (GPU - VM: jails, clients, profilers) behind a job interface; trust zones and rules; why srt cannot jail GPU code on Linux. +- `lab/RUNTIME.md` — design (2026-10-08): controller (laptop: ledger, referee, secrets) and worker (GPU VM: jails, + clients, profilers) behind a job interface; trust zones and rules; why srt cannot jail GPU code on Linux. +- `lab/worker.py` — the job interface, built for the measured jobs: `test`, `measure`, `equiv`, `profile` take input + trees + JSON args, write an out dir, return JSON; the tools call `self.worker`. `LAB_WORKER=local` (default, in + process) or `ssh` (the `lab.vm` VM): inputs by content hash, re-hashed on the worker before every job; GpuBusy, + Contaminated, NotReady, ValueError cross intact. Remote is tested through a local shell, not yet over SSH. - `lab/TRACE.md` + `lab/trace_contract.py` — the request-trace contract: per-request engine rows (trace_id = the client's `X-Trace-Id`, `arrival_ts` epoch s, `*_s` monotonic durations, TPOT = (E2E − TTFT)/(n − 1)) in a dir of `*.sqlite` (`requests` + `meta`) or `*.jsonl` + `.meta.json`, with schema version and drop counts; diff --git a/lab/README.md b/lab/README.md index 74cfc9c..5900c70 100644 --- a/lab/README.md +++ b/lab/README.md @@ -165,6 +165,10 @@ bridge; image from `python -m lab.safety.container build`, which `vm-setup.sh` r sandbox-runtime (`npm install -g @anthropic-ai/sandbox-runtime`); `LAB_JAIL` picks one. `LAB_NO_JAIL=1` waives the jail for tests and a box you trust. The agent's own CLI still runs under srt (`lab/RUNTIME.md`). +The GPU work behind `test`, `profile`, `trace`, `kernel`, `hostprof`, `bench`, `equiv` and `submit` runs as worker +jobs (`lab/worker.py`). `LAB_WORKER=local` (default) runs them in process; `LAB_WORKER=ssh` runs them on the VM +`lab.vm` names (start it and run `python -m lab.vm setup` first), shipping each tree by content hash. + ## Canaries python -m lab.canary slow_decode,kv_leak --port 8000 diff --git a/lab/RUNTIME.md b/lab/RUNTIME.md index 8059e13..07b301e 100644 --- a/lab/RUNTIME.md +++ b/lab/RUNTIME.md @@ -2,8 +2,9 @@ Design, 2026-10-08. Decided with the user: the lab splits into a controller and a worker; no separate measurement VM (cost); NVIDIA only for now; the jail on Linux is a container with the NVIDIA runtime (option A -below). Built: the container jail (`lab/safety/container.py`). Not built: the split. Open decisions are marked -**Decide**. +below). Built: the container jail (`lab/safety/container.py`), GPU exclusivity during measurements, and the +measured jobs behind the worker interface (`lab/worker.py`). Not built: the agent on the worker (its container, the +API proxy) and the agent-side jobs. Open decisions are marked **Decide**. ## Why @@ -119,6 +120,15 @@ JSON, raw files with hashes, and hardware facts (device, driver, clocks). Transport now: SSH from the controller, one worker. Transport later: a queue. The job list does not change. +Built (2026-10-09, `lab/worker.py`): `test`, `measure` (bench, submit and the base share it), `equiv` and `profile` +(all four instruments). The tools keep the policy, the ledger and the budget and call `self.worker`; `LAB_WORKER` +picks `local` (default: in process, as before) or `ssh` (the VM `lab.vm` names, started and set up beforehand). +Over SSH each input tree goes by content hash, at most once per worker (`has`, `put`), and the worker re-hashes a +fresh copy right before every job, refusing a mismatch. The job's out dir comes back as a tar; GpuBusy, +Contaminated, NotReady and ValueError are raised again on the controller. equiv ships the snapshot's earlier +answers with the job, so calling it again never re-rolls the gate. Tested through a local shell with the same CLI; +not yet over SSH to a VM. `agent`, `exec`, `snapshot`, `restore` and `facts` come with the agent container. + ## What moves where - Controller: `session`, `agent` (SDK side), the tool front ends in `tools`/`evaltools`/`proftools` (policy, @@ -148,7 +158,7 @@ database, still append-only. ## Order 1. ~~Decide the jail~~ (A) and build it. Done, minus the GPU check. -2. Without a GPU: the job interface with a local transport, the snapshot round trip and rule 3, the API proxy, - and the container jail's network, filesystem and freeze behaviour on any Linux box (Docker locally). +2. Without a GPU: ~~the job interface with a local transport, the snapshot round trip and rule 3~~ (done for the + measured jobs); next the agent container with its freeze, and the API proxy. 3. With a GPU (spend approval): `--gpus` and the profilers inside the container, then a first `lab.session` dry run with a tiny budget. diff --git a/lab/evaltools.py b/lab/evaltools.py index ae19546..6c1277c 100644 --- a/lab/evaltools.py +++ b/lab/evaltools.py @@ -1,5 +1,6 @@ -"""The three tools that can make a change count: bench, equiv and submit. Each serves the agent's pristine tree -from outside the jail (lab/serve.py), measures it with the environment's own instruments, and writes the ledger. +"""The three tools that can make a change count: bench, equiv and submit. Each has the worker (lab/worker.py) serve +the agent's pristine tree outside the jail and measure it with the environment's own instruments; the policy and the +ledger stay here. bench the regimes on the seen split, short tier -> one headline per regime, raw equiv the correctness gate against the target's reference -> pass | fail | inconclusive, with every metric @@ -18,11 +19,12 @@ import json import os +import shutil import tempfile import time from pathlib import Path -from lab import artifacts, ledger, serve, target +from lab import artifacts, ledger, serve, target, worker from validity.holdout import HoldoutGuard, Refused HOLDOUT_BUDGET, HOLDOUT_THRESHOLD_PCT, HOLDOUT_SIGMA_PCT = 100, 2.0, 0.5 @@ -147,23 +149,21 @@ class EvalTools: # -- bench -------------------------------------------------------------------------- def bench(self, args: dict) -> str: snap = self._audited("bench", args) - t = target.load() names = _regimes(args, self.bench_defaults()) tier = args.get("tier") or "short" tree = self._pristine() t0 = time.monotonic() - work, rows = self._work("bench"), [] + work = self._work("bench") try: - with serve.Served(tree, t.engine, log=self.s.run_dir / "serve-bench.log", - jailed=self.serve_jailed, artifacts=work) as srv: - results = run_regimes(srv.url, t, names, "seen", tier, self.seed, rows) - srv.exclusive() + r = self._measure(tree, work, "serve-bench.log", names=names, split="seen", tier=tier, passive=True) except Exception as e: # a tool never crashes the session: the failure is the record - result = {**_failed(e), "seconds": time.monotonic() - t0, "artifacts": self._store(work, rows)} + result = {**_failed(e), "seconds": time.monotonic() - t0, + "artifacts": artifacts.store(work, self.s.ledger_root)} self._record("bench", "bench", args, result, snap, config={"split": "seen", "tier": tier}) return _said("bench", result) - result = {"verdict": "ok", "metrics": headline(results), "regimes": results, "tier": tier, - "seconds": time.monotonic() - t0, "ready_s": srv.ready_s, "artifacts": self._store(work, rows)} + result = {"verdict": "ok", "metrics": headline(r["results"]), "regimes": r["results"], "tier": tier, + "seconds": time.monotonic() - t0, "ready_s": r["ready_s"], + "artifacts": artifacts.store(work, self.s.ledger_root)} self._record("bench", "bench", args, result, snap, config={"split": "seen", "tier": tier}) return json.dumps({"snapshot": snap.id, "split": "seen", "tier": tier, "metrics": result["metrics"]}, indent=1) @@ -177,27 +177,25 @@ def equiv(self, args: dict) -> str: self._record("equiv", "equiv", args, {"verdict": "refused", "reason": f"no {tier}-tier reference outputs; " "run `python -m correctness reference` for this target first"}, snap) return "equiv refused: no reference outputs for this target" - from correctness import client, run as crun - from correctness.gate import Thresholds tree = self._pristine() out = self.s.run_dir / "equiv" / f"{snap.id}-{tier}.json" work = self._work("equiv") + # answers already collected for this snapshot are reused, never re-rolled: the gate is statistical + prior = Path(tempfile.mkdtemp(prefix="equiv-prior-", dir=self.s.run_dir)) + if out.with_suffix(".outputs.jsonl").exists(): + shutil.copyfile(out.with_suffix(".outputs.jsonl"), prior / "equiv.outputs.jsonl") try: - with serve.Served(tree, t.engine, log=self.s.run_dir / "serve-equiv.log", jailed=self.serve_jailed, - artifacts=work) as srv: - res = crun.candidate(ref, srv.url, t.model, out, self.encoder(t), - Thresholds(min_score_ratio=t.min_score_ratio, min_length_ratio=t.min_length_ratio), - concurrency=self.equiv_concurrency, api=client.Api(t.engine.api, t.chat_kwargs), - config={"tier": tier}) - srv.exclusive() + res = self.worker.call("equiv", {"tree": tree, "reference": ref, "prior": prior}, + {"tier": tier, "concurrency": self.equiv_concurrency, "jailed": self.serve_jailed}, + work) except Exception as e: - result = {**_failed(e), "artifacts": self._store(work, equiv_rows(out))} + result = {**_failed(e), "artifacts": self._equiv_store(work, out, prior)} self._record("equiv", "equiv", args, result, snap, config={"split": "seen", "tier": tier}) return _said("equiv", result) passed = {"pass": True, "fail": False}.get(res["verdict"]) # inconclusive is None: neither passing nor failing record = {"verdict": res["verdict"], "passed": passed, "reasons": res["reasons"], "gates": res["gates"], "metrics": res["metrics"], "thresholds": res["thresholds"], "tier": tier, - "artifacts": self._store(work, equiv_rows(out))} + "artifacts": self._equiv_store(work, out, prior)} self._record("equiv", "equiv", args, record, snap, config={"split": "seen", "tier": tier}) return json.dumps({"snapshot": snap.id, "tier": tier, "verdict": res["verdict"], "reasons": res["reasons"], "metrics": res["metrics"]}, indent=1, default=str) @@ -226,16 +224,11 @@ def submit(self, args: dict) -> str: base_seen = self.base_seen(t, names) guard = self.holdout_guard() tree = self._pristine() - # held-out passive data stays operator-only: never a blob, never in the record, never inside the jail - private, rows = self.s.run_dir / "heldout-private" / f"{snap.id}-{time.strftime('%Y%m%dT%H%M%S')}", [] - try: - with serve.Served(tree, t.engine, log=self.s.run_dir / "serve-submit.log", jailed=self.serve_jailed, - artifacts=private) as srv: - results = run_regimes(srv.url, t, names, "heldout", "full", self.seed, rows) - srv.exclusive() # before the holdout guard: a contaminated run spends no query - finally: - artifacts.write_rows(private, rows) - new = headline(results) + # held-out passive data stays operator-only: never a blob, never in the record, never inside the jail. + # A contaminated run raises here, before the holdout guard: it spends no query. + private = self.s.run_dir / "heldout-private" / f"{snap.id}-{time.strftime('%Y%m%dT%H%M%S')}" + r = self._measure(tree, private, "serve-submit.log", names=names, split="heldout", tier="full", passive=True) + new = headline(r["results"]) metrics, better = {}, {} for n in names: held_d, seen_d = delta_pct(base.get(n, {}), new[n]), delta_pct(base_seen.get(n, {}), seen[n]) @@ -270,16 +263,35 @@ def bench_defaults(self) -> tuple[str, ...]: equiv_concurrency = 16 serve_jailed = True - def encoder(self, t: target.Target): - from correctness.run import HFEncoder - return getattr(self.s, "encoder", None) or HFEncoder(t.model, t.chat_kwargs) - def _work(self, tool: str) -> Path: - """A scratch dir outside the jail for one served run's passive data; emptied into blobs by `_store`.""" + """A scratch dir outside the jail for one worker job's output; emptied into blobs or removed after.""" return Path(tempfile.mkdtemp(prefix=f"passive-{tool}-", dir=self.s.run_dir)) - def _store(self, work: Path, rows: list[dict]) -> dict: - artifacts.write_rows(work, rows) + def _measure(self, tree: Path, out: Path, log: str, **args) -> dict: + """The worker's `measure` job into `out`; the engine's log is appended to `/` either way.""" + try: + return self.worker.call("measure", {"tree": tree}, {**args, "seed": self.seed, "jailed": self.serve_jailed}, + out) + finally: + src = out / worker.ENGINE_LOG + if src.exists(): + with open(self.s.run_dir / log, "ab") as f: + f.write(src.read_bytes()) + src.unlink() + + def _equiv_store(self, work: Path, out: Path, prior: Path) -> dict: + """The gate's files to `out` (kept per snapshot for the next equiv), its rows and passive data as blobs.""" + shutil.rmtree(prior, ignore_errors=True) + out.parent.mkdir(parents=True, exist_ok=True) + for name, dest in (("equiv.json", out), ("equiv.outputs.jsonl", out.with_suffix(".outputs.jsonl"))): + if (work / name).exists(): + shutil.move(work / name, dest) + log = work / worker.ENGINE_LOG + if log.exists(): + with open(self.s.run_dir / "serve-equiv.log", "ab") as f: + f.write(log.read_bytes()) + log.unlink() + artifacts.write_rows(work, equiv_rows(out)) return artifacts.store(work, self.s.ledger_root) def holdout_guard(self) -> HoldoutGuard: @@ -296,9 +308,13 @@ def _base_measure(self, t: target.Target, names: list[str], split: str, tier: st tree = self.s.run_dir / f"base-tree-{split}" from lab.safety import grader grader.export(self.s.workspace.repo, self.s.workspace.base, tree) - with serve.Served(tree, t.engine, log=self.s.run_dir / f"serve-base-{split}.log", jailed=self.serve_jailed) as srv: - have.update(headline(run_regimes(srv.url, t, missing, split, tier, self.seed))) - srv.exclusive() # a contaminated base is never cached: every later submit reads it + work = self._work("base") + try: # a contaminated base raises: never cached, every later submit reads it + r = self._measure(tree, work, f"serve-base-{split}.log", names=missing, split=split, tier=tier, + passive=False) + finally: + shutil.rmtree(work, ignore_errors=True) + have.update(headline(r["results"])) cache.write_text(json.dumps(have, indent=1)) return have diff --git a/lab/proftools.py b/lab/proftools.py index 12a739a..e6967d1 100644 --- a/lab/proftools.py +++ b/lab/proftools.py @@ -13,7 +13,6 @@ import time from pathlib import Path -from lab import engine from lab.agent import ToolSpec OUT = "_instr" # in the pristine tree: writable in the jail, rebuilt on every call @@ -56,82 +55,76 @@ def _head(path: Path | None) -> str: return path.read_text(errors="replace")[:SUMMARY_CHARS] if path and path.is_file() and not path.is_symlink() else "" -class ProfTools: - """Mixin for lab.tools.Toolbox: needs self._audited, self._pristine, self._record, self._keep.""" - - def trace(self, args: dict) -> str: - wl = workload(args) - wrap = lambda exe, out: [exe, "profile", "--output", str(out / "trace"), "--force-overwrite", "true", - "--trace", "cuda,nvtx,osrt", "--capture-range", "cudaProfilerApi", - "--capture-range-end", "stop", "--cuda-graph-trace", "node", "--stats", "false"] - post = lambda exe, out: ([exe, "stats", "--report", NSYS_REPORTS, "--format", "csv", "--output", - str(out / "stats"), str(out / "trace.nsys-rep")], None) - return self._instrumented("trace", args, "nsys", wrap, wl + ["--profiler", "cuda-range"], post, - "stats_cuda_gpu_kern_sum.csv") - - def kernel(self, args: dict) -> str: +def instrument(kind: str, args: dict) -> dict: + """One instrument from validated args: its binary, `wrap(exe, out)` around lab.profile, lab.profile `flags`, + `post(exe, out)` -> (argv, file for its stdout) run after, the `summary` glob, and `extra(exe)` flags.""" + wl = workload(args) + if kind == "trace": + return {"binary": "nsys", "flags": wl + ["--profiler", "cuda-range"], "summary": "stats_cuda_gpu_kern_sum.csv", + "wrap": lambda exe, out: [exe, "profile", "--output", str(out / "trace"), "--force-overwrite", "true", + "--trace", "cuda,nvtx,osrt", "--capture-range", "cudaProfilerApi", + "--capture-range-end", "stop", "--cuda-graph-trace", "node", "--stats", "false"], + "post": lambda exe, out: ([exe, "stats", "--report", NSYS_REPORTS, "--format", "csv", "--output", + str(out / "stats"), str(out / "trace.nsys-rep")], None), + "extra": lambda exe: []} + if kind == "kernel": rx = kernel_regex(args.get("kernel_regex")) skip, count = _int(args, "launch_skip", 0, 0, 100000), _int(args, "launch_count", 8, 1, 64) kset = args.get("set", "basic") if kset not in NCU_SETS: raise ValueError(f"set must be one of {NCU_SETS}") - wl = workload(args) - wrap = lambda exe, out: [exe, "--export", str(out / "kernel"), "--force-overwrite", "--kernel-name", - f"regex:{rx}", "--launch-skip", str(skip), "--launch-count", str(count), "--set", kset, - "--profile-from-start", "off", "--target-processes", "all"] - post = lambda exe, out: ([exe, "--import", str(out / "kernel.ncu-rep"), "--csv", "--page", "raw"], - out / "metrics.csv") - return self._instrumented("kernel", args, "ncu", wrap, wl + ["--profiler", "cuda-range"], post, "metrics.csv") + return {"binary": "ncu", "flags": wl + ["--profiler", "cuda-range"], "summary": "metrics.csv", + "wrap": lambda exe, out: [exe, "--export", str(out / "kernel"), "--force-overwrite", "--kernel-name", + f"regex:{rx}", "--launch-skip", str(skip), "--launch-count", str(count), + "--set", kset, "--profile-from-start", "off", "--target-processes", "all"], + "post": lambda exe, out: ([exe, "--import", str(out / "kernel.ncu-rep"), "--csv", "--page", "raw"], + out / "metrics.csv"), + "extra": lambda exe: []} + if kind == "hostprof": + seconds, rate = _int(args, "seconds", 20, 1, 300), _int(args, "rate", 100, 1, 1000) + return {"binary": "py-spy", "flags": wl, "summary": "bundle/*/pyspy-dump.txt", "wrap": lambda exe, out: [], + "post": None, + "extra": lambda exe: ["--profiler", "pyspy", "--seconds", str(seconds), "--rate", str(rate), + "--pyspy-bin", exe]} + raise ValueError(f"unknown instrument {kind!r}") + + +class ProfTools: + """Mixin for lab.tools.Toolbox: needs self._audited, self._pristine, self._record, self._keep, self._work, + self.worker.""" + + def trace(self, args: dict) -> str: + return self._instrumented("trace", args) + + def kernel(self, args: dict) -> str: + return self._instrumented("kernel", args) def hostprof(self, args: dict) -> str: - wl = workload(args) - seconds, rate = _int(args, "seconds", 20, 1, 300), _int(args, "rate", 100, 1, 1000) - extra = lambda exe: ["--profiler", "pyspy", "--seconds", str(seconds), "--rate", str(rate), "--pyspy-bin", exe] - return self._instrumented("hostprof", args, "py-spy", lambda exe, out: [], wl, None, - "bundle/*/pyspy-dump.txt", extra) - - def _instrumented(self, tool: str, args: dict, binary: str, wrap, flags: list[str], post, summary: str, - extra=lambda exe: []) -> str: - """`wrap(exe, out)` + lab.profile `flags`, jailed with the instrument readable; then `post`, jailed too - (the report was written by agent code). Kept as one blob; the head of `summary` is returned.""" - from lab import tools as T + return self._instrumented("hostprof", args) + + def _instrumented(self, tool: str, args: dict) -> str: + """The instrument around lab.profile's workload, then its report step, both jailed on the worker (the report + was written by agent code). Kept as one blob; the head of its summary is returned.""" + instrument(tool, args) # refuse bad arguments before anything is snapshotted or run snap = self._audited(tool, args) - exe = shutil.which(binary) - if not exe: - reason = f"{binary} is not on PATH on this host" - self._record("note", tool, args, {"verdict": "refused", "reason": reason}, snap) - return f"{tool} refused: {reason}" - exe = str(Path(exe).resolve()) tree = self._pristine() - out = tree / OUT - out.mkdir() - harness = T.stage_harness(tree) - read = [Path(exe).parent] - src = T.workload_flags(args, harness, _int(args, "requests", *WORKLOAD["requests"])) - argv = [*wrap(exe, out), engine.python(), "-m", "lab.profile", *flags, *src, *extra(exe), "--out", str(out / "bundle")] + out = self._work(tool) t0 = time.monotonic() - proc = T.default_profile_runner(tree, argv, read=read) - log = proc.stdout + proc.stderr - commands = [argv] - T._drop_links(out) - if post and proc.returncode == 0: - pargv, dest = post(exe, out) - commands.append(pargv) - p = T.default_profile_runner(tree, pargv, read=read) - if dest: - dest.write_text(p.stdout) - log += p.stderr if dest else p.stdout + p.stderr - proc = p if p.returncode else proc - T._drop_links(out) - found = sorted(out.glob(summary)) - head = _head(found[-1] if found else None) - blob, visible = self._keep(out) if any(p.is_file() for p in out.rglob("*")) else ("", "") - result = {"returncode": proc.returncode, "bundle": blob, "workspace_copy": visible, "instrument": exe, - "argv": commands, "seconds": time.monotonic() - t0, "output": log[-3000:]} + try: + r = self.worker.call("profile", {"tree": tree}, {"kind": tool, "args": args}, out) + if r.get("refused"): + self._record("note", tool, args, {"verdict": "refused", "reason": r["refused"]}, snap) + return f"{tool} refused: {r['refused']}" + files = out / "files" + blob, visible = self._keep(files) if files.is_dir() and any(p.is_file() for p in files.rglob("*")) else ("", "") + finally: + shutil.rmtree(out, ignore_errors=True) + result = {"returncode": r["returncode"], "bundle": blob, "workspace_copy": visible, "instrument": r["instrument"], + "argv": r["argv"], "seconds": time.monotonic() - t0, "output": r["output"]} self._record("profile", tool, args, result, snap) - if proc.returncode != 0 or not head: - return f"{tool} failed ({proc.returncode}):\n{result['output']}" - return f"{tool} output at {visible} (in your workspace; left out of your change)\n{head}" + if r["returncode"] != 0 or not r["head"]: + return f"{tool} failed ({r['returncode']}):\n{result['output']}" + return f"{tool} output at {visible} (in your workspace; left out of your change)\n{r['head']}" def prof_specs(self, gpu) -> list[ToolSpec]: """`gpu(time)` renders the cost fact for a description.""" diff --git a/lab/tools.py b/lab/tools.py index 6a51a02..86180a4 100644 --- a/lab/tools.py +++ b/lab/tools.py @@ -33,8 +33,11 @@ class Toolbox(EvalTools, ProfTools): """Bound to one session: knows the workspace, the budget, the ledger root and the run.""" def __init__(self, session: Any): + from lab import worker self.s = session self.violation: str | None = None + self.worker = getattr(session, "worker", None) or worker.load( + encoder=getattr(session, "encoder", None), profile_runner=getattr(session, "profile_runner", None)) # -- plumbing ------------------------------------------------------------------------ def _cost(self, tool: str) -> dict: @@ -97,9 +100,12 @@ def _keep(self, src: Path) -> tuple[str, str]: # -- tools --------------------------------------------------------------------------- def test(self, args: dict) -> str: snap = self._audited("test", args) - tree = self._pristine() - lint = grader.run_lint(tree) - tests = grader.run_tests(tree) if lint.passed else grader.Run(False, -1, "skipped: lint failed") + out = self._work("test") + try: + r = self.worker.call("test", {"tree": self._pristine()}, {}, out) + finally: + shutil.rmtree(out, ignore_errors=True) + lint, tests = grader.Run(**r["lint"]), grader.Run(**r["tests"]) result = {"lint": lint.passed, "tests": tests.passed, "returncode": tests.returncode, "lint_output": lint.output[-2000:], "test_output": tests.output[-4000:], "changed": snap.files + snap.deleted, "scratch_left_out": self.s.workspace.audit().scratch} @@ -110,22 +116,18 @@ def test(self, args: dict) -> str: def profile(self, args: dict) -> str: snap = self._audited("profile", args) - tree = self._pristine() - out = tree / "lab" / "runs" - harness = stage_harness(tree) - n = int(args.get("requests", 8)) - argv = [engine.python(), "-m", "lab.profile", "--requests", str(n), - "--max-tokens", str(args.get("max_tokens", 32)), "--out", str(out), - *workload_flags(args, harness, n)] + out = self._work("profile") t0 = time.monotonic() - proc = (self.s.profile_runner or default_profile_runner)(tree, argv) - bundles = sorted(out.glob("*")) if out.exists() else [] - blob, visible = self._keep(bundles[-1]) if bundles else ("", "") - result = {"returncode": proc.returncode, "bundle": blob, "workspace_copy": visible, - "seconds": time.monotonic() - t0, "output": (proc.stdout + proc.stderr)[-3000:]} + try: + r = self.worker.call("profile", {"tree": self._pristine()}, {"kind": "profile", "args": args}, out) + blob, visible = self._keep(out / "files") if (out / "files").is_dir() else ("", "") + finally: + shutil.rmtree(out, ignore_errors=True) + result = {"returncode": r["returncode"], "bundle": blob, "workspace_copy": visible, + "seconds": time.monotonic() - t0, "output": r["output"]} self._record("profile", "profile", args, result, snap) - if proc.returncode != 0 or not blob: - return f"profile failed ({proc.returncode}):\n{result['output']}" + if r["returncode"] != 0 or not blob: + return f"profile failed ({r['returncode']}):\n{result['output']}" return f"bundle at {visible} (in your workspace; left out of your change)\n" + "\n".join( f" {p.name}" for p in sorted((self.s.workspace.path / visible).iterdir())) diff --git a/lab/worker.py b/lab/worker.py new file mode 100644 index 0000000..1d895ce --- /dev/null +++ b/lab/worker.py @@ -0,0 +1,300 @@ +"""Where the GPU work runs (lab/RUNTIME.md). A job takes input trees and JSON arguments, writes files into an output +dir and returns JSON; the tools keep the policy, the ledger and the budget. `Local` runs a job in this process. +`Remote` ships each input by content hash, runs `python -m lab.worker run JOB` on the worker host over SSH and brings +the output dir back; the worker re-hashes every tree right before using it and refuses one that does not match, so +what is measured is what the controller audited. GpuBusy, Contaminated, NotReady and ValueError cross intact. + + test lint + the engine's suite on `tree` -> {lint, tests} + measure serve `tree`, run regimes; passive data into out -> {results, ready_s} + equiv serve `tree`, the correctness gate against `reference` -> the gate's result; outputs in out + profile lab.profile under torch, nsys, ncu or py-spy on `tree` -> {returncode, output, ...}; files in out + + python -m lab.worker has SHA... | put SHA < tar | run JOB < json > tar (the worker host's side) +""" + +from __future__ import annotations + +import dataclasses +import hashlib +import io +import json +import os +import shutil +import subprocess +import sys +import tarfile +import tempfile +from pathlib import Path + +from lab import engine, serve, target + +RESULT = "_result.json" +ENGINE_LOG = "engine.log" # the served engine's log, in every job's out dir + + +class WorkerError(RuntimeError): + """The worker could not run the job (transport, hash mismatch, or an unexpected failure there).""" + + +def tree_sha(root: Path) -> str: + """Content hash of a directory: every regular file's path, exec bit and bytes; refuses symlinks.""" + h = hashlib.sha256() + for p in sorted(Path(root).rglob("*")): + rel = p.relative_to(root).as_posix() + if p.is_symlink(): + raise WorkerError(f"symlink in a shipped tree: {rel}") + if p.is_file(): + h.update(f"{rel}\0{os.access(p, os.X_OK):d}\0{p.stat().st_size}\0".encode()) + h.update(p.read_bytes()) + return h.hexdigest() + + +# -- jobs (run where the GPU is) --------------------------------------------------------------------------------- + +def job_test(inputs: dict, out: Path, args: dict, hooks: dict) -> dict: + from lab.safety import grader + tree = inputs["tree"] + lint = grader.run_lint(tree) + tests = grader.run_tests(tree) if lint.passed else grader.Run(False, -1, "skipped: lint failed") + return {"lint": dataclasses.asdict(lint), "tests": dataclasses.asdict(tests)} + + +def job_measure(inputs: dict, out: Path, args: dict, hooks: dict) -> dict: + """`args`: names, split, tier, seed, jailed, passive (keep device/telemetry/client rows in `out`).""" + from lab import artifacts + from lab.evaltools import run_regimes + t = target.load() + rows = [] if args["passive"] else None + try: + with serve.Served(inputs["tree"], t.engine, log=out / ENGINE_LOG, jailed=args["jailed"], + artifacts=out if args["passive"] else None) as srv: + results = run_regimes(srv.url, t, args["names"], args["split"], args["tier"], args["seed"], rows) + srv.exclusive() + return {"results": results, "ready_s": srv.ready_s} + finally: + if rows is not None: + artifacts.write_rows(out, rows) + + +def job_equiv(inputs: dict, out: Path, args: dict, hooks: dict) -> dict: + """`args`: tier, concurrency, jailed. The gate's outputs land in `out` as equiv.json and equiv.outputs.jsonl; + answers already in `prior`'s equiv.outputs.jsonl are reused, not asked again.""" + from correctness import client, run as crun + from correctness.gate import Thresholds + t = target.load() + if (inputs["prior"] / "equiv.outputs.jsonl").exists(): + shutil.copyfile(inputs["prior"] / "equiv.outputs.jsonl", out / "equiv.outputs.jsonl") + enc = hooks.get("encoder") or crun.HFEncoder(t.model, t.chat_kwargs) + with serve.Served(inputs["tree"], t.engine, log=out / ENGINE_LOG, jailed=args["jailed"], artifacts=out) as srv: + res = crun.candidate(inputs["reference"], srv.url, t.model, out / "equiv.json", enc, + Thresholds(min_score_ratio=t.min_score_ratio, min_length_ratio=t.min_length_ratio), + concurrency=args["concurrency"], api=client.Api(t.engine.api, t.chat_kwargs), + config={"tier": args["tier"]}) + srv.exclusive() + return res + + +def job_profile(inputs: dict, out: Path, args: dict, hooks: dict) -> dict: + """`args`: kind (profile | trace | kernel | hostprof) and the tool's args. Output files go to `out/files`.""" + from lab import proftools, tools as T + tree, kind, a = inputs["tree"], args["kind"], args["args"] + files = out / "files" + harness = T.stage_harness(tree) + if kind == "profile": + runs = tree / "lab" / "runs" + n = int(a.get("requests", 8)) + argv = [engine.python(), "-m", "lab.profile", "--requests", str(n), "--max-tokens", str(a.get("max_tokens", 32)), + "--out", str(runs), *T.workload_flags(a, harness, n)] + proc = (hooks.get("profile_runner") or T.default_profile_runner)(tree, argv) + bundles = sorted(runs.glob("*")) if runs.exists() else [] + if bundles: + T._drop_links(bundles[-1]) + shutil.copytree(bundles[-1], files) + return {"returncode": proc.returncode, "output": (proc.stdout + proc.stderr)[-3000:]} + ins = proftools.instrument(kind, a) + exe = shutil.which(ins["binary"]) + if not exe: + return {"refused": f"{ins['binary']} is not on PATH on this host"} + exe = str(Path(exe).resolve()) + work = tree / proftools.OUT + work.mkdir() + read = [Path(exe).parent] + src = T.workload_flags(a, harness, proftools._int(a, "requests", *proftools.WORKLOAD["requests"])) + argv = [*ins["wrap"](exe, work), engine.python(), "-m", "lab.profile", *ins["flags"], *src, *ins["extra"](exe), + "--out", str(work / "bundle")] + proc = T.default_profile_runner(tree, argv, read=read) + log, commands = proc.stdout + proc.stderr, [argv] + T._drop_links(work) + if ins["post"] and proc.returncode == 0: + pargv, dest = ins["post"](exe, work) + commands.append(pargv) + p = T.default_profile_runner(tree, pargv, read=read) + if dest: + dest.write_text(p.stdout) + log += p.stderr if dest else p.stdout + p.stderr + proc = p if p.returncode else proc + T._drop_links(work) + found = sorted(work.glob(ins["summary"])) + if any(p.is_file() for p in work.rglob("*")): + shutil.copytree(work, files) + return {"returncode": proc.returncode, "instrument": exe, "argv": commands, "output": log[-3000:], + "head": proftools._head(found[-1] if found else None)} + + +JOBS = {"test": job_test, "measure": job_measure, "equiv": job_equiv, "profile": job_profile} + +# Failures the tools turn into records; anything else is a WorkerError on the controller. +_ERRORS = {"GpuBusy": lambda e: serve.GpuBusy(e["apps"]), "Contaminated": lambda e: serve.Contaminated(e["apps"]), + "NotReady": lambda e: serve.NotReady(e["message"]), "ValueError": lambda e: ValueError(e["message"])} + + +def _error(e: BaseException) -> dict: + return {"type": type(e).__name__, "message": str(e), "apps": getattr(e, "apps", None)} + + +def _raise(err: dict | None) -> None: + if err: + make = _ERRORS.get(err["type"]) + raise make(err) if make else WorkerError(f"{err['type']}: {err['message']}") + + +# -- transports --------------------------------------------------------------------------------------------------- + +class Local: + """Jobs in this process, on the controller's own paths. `hooks`: in-process stand-ins (an encoder, a profile + runner) that tests and fakes provide; a remote worker has none.""" + + def __init__(self, **hooks): + self.hooks = {k: v for k, v in hooks.items() if v is not None} + + def call(self, job: str, inputs: dict[str, Path], args: dict, out: Path) -> dict: + out.mkdir(parents=True, exist_ok=True) + return JOBS[job](inputs, out, args, self.hooks) + + +def _tar(root: Path) -> bytes: + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as t: + t.add(root, arcname=".") + return buf.getvalue() + + +def _untar(data: bytes, dest: Path) -> None: + dest.mkdir(parents=True, exist_ok=True) + with tarfile.open(fileobj=io.BytesIO(data), mode="r:gz") as t: + t.extractall(dest, filter="data") + + +class Remote: + """Jobs on another host. `prefix` is the argv that runs one shell command there (ssh ... user@host); empty runs + it here through sh, which the tests use. `env_dir` is this repo on that host, `python` its interpreter.""" + + def __init__(self, prefix: list[str], env_dir: str, python: str, env: dict[str, str]): + self.prefix, self.env_dir, self.python, self.env = prefix, env_dir, python, env + + def _run(self, command: str, stdin: bytes = b"") -> bytes: + for v in self.env.values(): + if any(c in v for c in '"`\\'): + raise WorkerError(f"unsafe value in the worker's env: {v!r}") + exports = " ".join(f'{k}="{v}"' for k, v in self.env.items()) + line = f"cd {self.env_dir} && {exports} {self.python} -m lab.worker {command}" + argv = [*self.prefix, line] if self.prefix else ["sh", "-c", line] + p = subprocess.run(argv, input=stdin, capture_output=True) + if p.returncode != 0: + raise WorkerError(f"worker `{command}` exited {p.returncode}: {p.stderr.decode(errors='replace')[-2000:]}") + return p.stdout + + def call(self, job: str, inputs: dict[str, Path], args: dict, out: Path) -> dict: + shas = {k: tree_sha(p) for k, p in inputs.items()} + have = set(self._run("has " + " ".join(shas.values())).decode().split()) if shas else set() + for k, p in inputs.items(): + if shas[k] not in have: + self._run(f"put {shas[k]}", _tar(p)) + _untar(self._run(f"run {job}", json.dumps({"inputs": shas, "args": args}).encode()), out) + reply = json.loads((out / RESULT).read_text()) + (out / RESULT).unlink() + _raise(reply.get("error")) + return reply["result"] + + +def load(**hooks): + """`LAB_WORKER`: local (default) or ssh, the VM `lab.vm` names (started and set up beforehand).""" + kind = os.environ.get("LAB_WORKER", "local") + if kind == "local": + return Local(**hooks) + if kind != "ssh": + raise ValueError(f"LAB_WORKER is local or ssh, not {kind!r}") + from lab import vm + v = vm.VM() + record = v.provider.get(v.name) + if record is None or record.state != "running": + raise WorkerError(f"{v.name} is not running: start it with `python -m lab.vm start` and set it up first") + spec = Path(os.environ.get("LAB_TARGET", "")).resolve() + if not spec.is_relative_to(engine.ENV_ROOT): + raise WorkerError(f"LAB_TARGET must be inside this repo to be found on the worker: {spec}") + home = lambda p: p.replace("~", "$HOME", 1) # noqa: E731 - expanded by the worker's shell + return Remote(["ssh", *v.ssh_opts, vm.ssh_target(v, record)], v.env_dir, home(v.remote_dir) + "/.venv/bin/python", + {"LAB_TARGET": spec.relative_to(engine.ENV_ROOT).as_posix(), "LAB_ENGINE_REPO": home(v.remote_dir)}) + + +# -- the worker host's side --------------------------------------------------------------------------------------- + +def cache_dir() -> Path: + return Path(os.environ.get("LAB_WORKER_CACHE") or engine.ENV_ROOT / "lab" / "runs" / "worker-cache") + + +def _put(sha: str, data: bytes) -> None: + dest = cache_dir() / sha + if dest.exists(): + return + part = Path(tempfile.mkdtemp(prefix=f"{sha[:12]}-", dir=cache_dir())) + _untar(data, part) + got = tree_sha(part) + if got != sha: + shutil.rmtree(part) + raise WorkerError(f"received tree hashes to {got}, not {sha}") + part.rename(dest) + + +def _run_job(job: str, request: dict) -> bytes: + """Fresh copies of the inputs (jobs write into them), each re-hashed right before use; out dir as a tar.""" + work = Path(tempfile.mkdtemp(prefix=f"job-{job}-")) + out = work / "out" + out.mkdir() + try: + try: + inputs = {} + for name, sha in request["inputs"].items(): + src = cache_dir() / sha + if not src.is_dir(): + raise WorkerError(f"input {name} ({sha}) was never put") + inputs[name] = work / name + shutil.copytree(src, inputs[name]) + if tree_sha(inputs[name]) != sha: + raise WorkerError(f"input {name} changed after it was put") + reply = {"result": JOBS[job](inputs, out, request["args"], {})} + except Exception as e: # the controller re-raises it; the tar still carries what was written + reply = {"error": _error(e)} + (out / RESULT).write_text(json.dumps(reply, default=str)) + return _tar(out) + finally: + shutil.rmtree(work, ignore_errors=True) + + +def main(argv: list[str] | None = None) -> int: + argv = sys.argv[1:] if argv is None else argv + cache_dir().mkdir(parents=True, exist_ok=True) + if argv[:1] == ["has"]: + print(" ".join(s for s in argv[1:] if (cache_dir() / s).is_dir())) + elif argv[:1] == ["put"] and len(argv) == 2: + _put(argv[1], sys.stdin.buffer.read()) + elif argv[:1] == ["run"] and len(argv) == 2 and argv[1] in JOBS: + sys.stdout.buffer.write(_run_job(argv[1], json.loads(sys.stdin.buffer.read()))) + else: + print(f"usage: python -m lab.worker has SHA... | put SHA | run {{{','.join(JOBS)}}}", file=sys.stderr) + return 2 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_lab_evaltools.py b/tests/test_lab_evaltools.py index 3565cd1..fbd9e04 100644 --- a/tests/test_lab_evaltools.py +++ b/tests/test_lab_evaltools.py @@ -158,6 +158,17 @@ def test_equiv_judges_against_the_reference(lab): assert all(j["engine"] for j in joined) # every gate request has its engine row (X-Trace-Id) +def test_a_second_equiv_on_the_same_snapshot_reuses_its_answers(lab): + tb, s = lab + tb.equiv({}) + answers = next((s.run_dir / "equiv").glob("*-dev.outputs.jsonl")) + marked = "".join(json.dumps({**json.loads(x), "kept": True}) + "\n" for x in answers.read_text().splitlines()) + answers.write_text(marked) + tb.equiv({}) + assert answers.read_text() == marked # nothing asked again: the gate is never re-rolled + assert not list(s.run_dir.glob("equiv-prior-*")) and not list(s.run_dir.glob("passive-*")) + + def test_submit_keeps_heldout_passive_data_out_of_the_ledger(lab): tb, s = lab tb.equiv({"tier": "full"}); tb.bench({"regimes": ["single_stream"]}) diff --git a/tests/test_lab_worker.py b/tests/test_lab_worker.py new file mode 100644 index 0000000..ebf9f27 --- /dev/null +++ b/tests/test_lab_worker.py @@ -0,0 +1,105 @@ +"""lab.worker: tree hashes, and the Remote transport end to end through a local shell (the same CLI ssh would run).""" + +from __future__ import annotations + +import os +import stat +import sys +from pathlib import Path + +import pytest + +from lab import engine, serve, target, worker +from lab.safety import grader +from tests.lab_fixtures import make_repo +from tests.test_lab_evaltools import _target + +MEASURE = {"names": ["single_stream"], "split": "seen", "tier": "short", "seed": 0, "jailed": False, "passive": False} + + +@pytest.fixture +def remote(tmp_path, monkeypatch): + """A Remote whose 'host' is this machine, with its own cache; returns (remote, a pristine tree, calls made).""" + monkeypatch.setenv("LAB_NO_JAIL", "1") + repo = make_repo(tmp_path) + spec = _target(tmp_path, repo, tmp_path / "ref") + monkeypatch.setenv("LAB_TARGET", str(spec)) + monkeypatch.setenv("LAB_ENGINE_REPO", str(repo)) + target._cache.clear() + tree = tmp_path / "tree" + grader.export(repo, "HEAD", tree) + r = worker.Remote([], str(engine.ENV_ROOT), sys.executable, {"LAB_WORKER_CACHE": str(tmp_path / "cache")}) + calls = [] + run = r._run + r._run = lambda command, stdin=b"": calls.append(command.split()[0]) or run(command, stdin) + yield r, tree, calls + target._cache.clear() + + +def test_tree_sha_covers_content_and_exec_bit_and_refuses_links(tmp_path): + (tmp_path / "a").mkdir() + (tmp_path / "a" / "f").write_text("x") + first = worker.tree_sha(tmp_path / "a") + (tmp_path / "a" / "f").chmod(0o755) + second = worker.tree_sha(tmp_path / "a") + (tmp_path / "a" / "f").write_text("y") + assert len({first, second, worker.tree_sha(tmp_path / "a")}) == 3 + (tmp_path / "a" / "l").symlink_to("/etc/passwd") + with pytest.raises(worker.WorkerError): + worker.tree_sha(tmp_path / "a") + + +def test_remote_runs_the_job_on_the_shipped_tree_and_ships_it_once(remote, tmp_path): + r, tree, calls = remote + res = r.call("test", {"tree": tree}, {}, tmp_path / "out1") + assert res["lint"]["passed"] and res["tests"]["passed"], res + assert calls == ["has", "put", "run"] + calls.clear() + r.call("test", {"tree": tree}, {}, tmp_path / "out2") + assert calls == ["has", "run"] # cached by hash: not shipped again + + +def test_remote_carries_gpu_busy_with_its_facts(remote, tmp_path, monkeypatch): + r, tree, _ = remote + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + smi = bin_dir / "nvidia-smi" + smi.write_text(f"#!{sys.executable}\nimport sys\n" + "if any(a.startswith('--query-compute-apps') for a in sys.argv): print('4242, 9000, python')\n") + smi.chmod(smi.stat().st_mode | stat.S_IEXEC) + monkeypatch.setenv("PATH", f"{bin_dir}{os.pathsep}{os.environ['PATH']}") + with pytest.raises(serve.GpuBusy) as e: + r.call("measure", {"tree": tree}, MEASURE, tmp_path / "out") + assert e.value.apps == [{"pid": 4242, "used_mib": 9000, "name": "python"}] + + +def test_remote_carries_an_engine_that_did_not_start(remote, tmp_path, monkeypatch): + r, tree, _ = remote + spec = Path(os.environ["LAB_TARGET"]) + spec.write_text(spec.read_text().replace('launch = "python ', 'launch = "python -c \'import sys; sys.exit(3)\' #', 1)) + with pytest.raises(serve.NotReady): + r.call("measure", {"tree": tree}, MEASURE, tmp_path / "out") + assert (tmp_path / "out" / worker.ENGINE_LOG).exists() # what was written still comes back + + +def test_a_tree_changed_after_it_was_put_is_refused(remote, tmp_path): + r, tree, _ = remote + r.call("test", {"tree": tree}, {}, tmp_path / "out1") + cached = tmp_path / "cache" / worker.tree_sha(tree) + (cached / "src" / "inference_server" / "engine.py").write_text("def speed():\n return 100\n") + with pytest.raises(worker.WorkerError, match="changed after it was put"): + r.call("test", {"tree": tree}, {}, tmp_path / "out2") + + +def test_put_refuses_a_tree_that_does_not_hash_to_its_name(remote): + r, tree, _ = remote + with pytest.raises(worker.WorkerError, match="received tree hashes to"): + r._run(f"put {'0' * 64}", worker._tar(tree)) + + +def test_local_by_default_and_nothing_else_but_ssh(monkeypatch): + monkeypatch.delenv("LAB_WORKER", raising=False) + assert isinstance(worker.load(), worker.Local) + monkeypatch.setenv("LAB_WORKER", "gpu-box") + with pytest.raises(ValueError): + worker.load() From a79167e63cc226428643e030216798107a8c8e98 Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Fri, 9 Oct 2026 01:23:53 -0400 Subject: [PATCH 04/18] RUNTIME.md: the user's access line and the workbench decisions (2026-10-09) Co-Authored-By: Claude Opus 5.5 --- lab/RUNTIME.md | 35 +++++++++++++++++++++++++++++++---- 1 file changed, 31 insertions(+), 4 deletions(-) diff --git a/lab/RUNTIME.md b/lab/RUNTIME.md index 07b301e..6fab5cb 100644 --- a/lab/RUNTIME.md +++ b/lab/RUNTIME.md @@ -6,6 +6,32 @@ below). Built: the container jail (`lab/safety/container.py`), GPU exclusivity d measured jobs behind the worker interface (`lab/worker.py`). Not built: the agent on the worker (its container, the API proxy) and the agent-side jobs. Open decisions are marked **Decide**. +## The line (the user, 2026-10-09) + +Don't restrict the agent's search space; restrict only what would let it cheat or cost money. The agent never +touches (1) secrets and money: API keys, cloud credentials, the HF token; (2) the answer key: held-out workloads, +correctness reference outputs; (3) the scoreboard: ledger, measurement clients, grader, budget; (4) the GPU +during its own measurement. Everything else it can have. Decided with it: + +- **Internet:** the agent's workbench has it. Hence nothing secret on the VM (rule 1). +- **Installing libraries:** the agent can pip- and apt-install. Its room is a looser *workbench* container: GPU, + network, root inside the container (never on the VM host), writable system dirs. The tight container (no + network, no capabilities, read-only root) stays for the measurement rooms only (`test`, `profile`, served + engines). +- **Installed libraries reach measurement only through the change.** A measurement room is rebuilt from the + agent's code change alone, so a dependency change has to be part of it: the engine's dependency files (e.g. + `pyproject.toml`, `uv.lock`) join the target's write surface and are installed from scratch for the + measurement. Patching an installed library in the workbench changes no measured result. + **Decide:** the measurement room has no network, so that install needs its own build step (package indexes + only, no GPU, no secrets, cached by lockfile hash) before the offline measurement. +- **Fresh VM per run:** yes (rule 6). +- **No separate measurement VM** for now: container escape is not a realistic threat at this stage, and it + would double GPU cost. Revisit for a public leaderboard. +- **Refuse rather than kill** (#20) until the workbench can be frozen during a measurement (rule 4). + +Not decided here: held-out workloads are sampled from public datasets, so with internet the agent could +download them; the workload generator's owner decides (rephrased or non-public held-out prompts). + ## Why Today every part of the lab runs on one machine and serves the engine on that machine's localhost. On a @@ -28,7 +54,8 @@ cannot signal others), and everything in a jail dies when the jail's command exi |---|---|---|---| | Controller | laptop (later an always-on box) | session loop, ledger, referee, holdout guard, budget, snapshot store, every secret | no | | Worker | GPU VM, host side, outside every jail | job runner, measurement clients, profilers, GPU state | only through the engine's HTTP port | -| Jails | GPU VM, inside containers | agent's shell and workspace; each served engine | yes, it is the agent's | +| Workbench | GPU VM, the agent's container | agent's shell and workspace; GPU, internet, root inside | yes, it is the agent's | +| Measurement rooms | GPU VM, one tight container per job | a tree under `test`/`profile`, a served engine | its code is the agent's; nothing else reaches in | Rules: @@ -44,10 +71,10 @@ Rules: states what was killed (a fact). 5. **Measurement code is out of reach.** The clients run on the host, outside every container, from a tree the jails cannot read; their CPU cores are reserved (cpuset) so the engine cannot starve them. -6. **One VM per run, destroyed after** (**Decide**; recommended). Nothing from one run can touch the next one's numbers. Setup time is the +6. **One VM per run, created at run start and deleted at the end** (decided 2026-10-09). Nothing from one run can touch the next one's numbers. Setup time is the cost; a prebuilt image with the venv plus a reattached read-only weights disk cuts it. -7. **Network.** VM egress: the controller tunnel only. Agent jail: the API proxy only. Engine jail: nothing; its - port is reachable from the host alone. +7. **Network.** Workbench: the internet (decided 2026-10-09), so the VM holds nothing secret. Measurement rooms: + nothing; a served engine's port is reachable from the host alone. Left open by not having a separate measurement VM: the clients and the agent's code share a kernel, so a kernel or container escape defeats rules 4 and 5. A measurement VM in the same private network closes it later. From a089209e14afda8d40844b595232801fd483832d Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Fri, 9 Oct 2026 23:33:12 -0400 Subject: [PATCH 05/18] RUNTIME.md: open workbench, offline measurement as a fairness rule, the dependency build step (approved 2026-10-09) Co-Authored-By: Claude Opus 5.5 --- lab/RUNTIME.md | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/lab/RUNTIME.md b/lab/RUNTIME.md index 6fab5cb..6ccc544 100644 --- a/lab/RUNTIME.md +++ b/lab/RUNTIME.md @@ -13,7 +13,9 @@ touches (1) secrets and money: API keys, cloud credentials, the HF token; (2) th correctness reference outputs; (3) the scoreboard: ledger, measurement clients, grader, budget; (4) the GPU during its own measurement. Everything else it can have. Decided with it: -- **Internet:** the agent's workbench has it. Hence nothing secret on the VM (rule 1). +- **Internet:** the agent's workbench has it, open: the same capabilities Claude Code has (a shell with network, + web search and fetch, any package install), not gated behind tools. Hence nothing secret on the VM (rule 1). + Agent-facing tools are the user's to design; the lab builds none. - **Installing libraries:** the agent can pip- and apt-install. Its room is a looser *workbench* container: GPU, network, root inside the container (never on the VM host), writable system dirs. The tight container (no network, no capabilities, read-only root) stays for the measurement rooms only (`test`, `profile`, served @@ -22,8 +24,12 @@ during its own measurement. Everything else it can have. Decided with it: agent's code change alone, so a dependency change has to be part of it: the engine's dependency files (e.g. `pyproject.toml`, `uv.lock`) join the target's write surface and are installed from scratch for the measurement. Patching an installed library in the workbench changes no measured result. - **Decide:** the measurement room has no network, so that install needs its own build step (package indexes - only, no GPU, no secrets, cached by lockfile hash) before the offline measurement. + The measurement room stays offline, a fairness rule rather than a safety one: an engine with network during + a measurement could forward requests to a remote GPU or API. So the install is its own build step before the + measurement (decided 2026-10-09): package indexes only, no GPU, no secrets, cached by lockfile hash. Anything the + engine needs at runtime (draft models, prebuilt kernels, files a library downloads on first use) must be in the + change, in the dependency files or saved in the workspace; the tool descriptions state that as a fact once the + build step exists (not before: today measurements still use the engine venv on the VM). - **Fresh VM per run:** yes (rule 6). - **No separate measurement VM** for now: container escape is not a realistic threat at this stage, and it would double GPU cost. Revisit for a public leaderboard. From 125cd97f9b68250e76d8cf5faa2d7ad4582651f4 Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Sat, 10 Oct 2026 00:16:45 -0400 Subject: [PATCH 06/18] lab: the agent's workbench container, its pause during measurements, and the API-key proxy API proxy (lab/apiproxy.py): the agent's CLI holds a placeholder key and talks to a per-session proxy that swaps in the lab's key or OAuth token, forwards only /v1/messages, count_tokens and /v1/models, prices each response's usage (JSON or SSE) at first-party rates, and refuses once the session cap is spent (and any unpriced model or fast mode). The session is charged what the proxy priced. Before, the key sat in the CLI's environment, readable from /proc by the agent's shell. Workbench (Linux): the CLI runs in a per-session container with the GPU, the internet, root inside, writable system dirs; the engine venv, weights and CLI binary read-only so an install never reaches a measurement. vm-setup.sh blocks the cloud metadata service for containers. Pause: measure, equiv and profile jobs kill the workbench's GPU holders, pause it, and thaw it after; what was killed is recorded as workbench_killed and returned as a fact. Co-Authored-By: Claude Opus 5.5 --- CLAUDE.md | 5 +- lab/LEDGER.md | 4 +- lab/README.md | 4 +- lab/RUNTIME.md | 36 ++++++- lab/agent.py | 81 ++++++++++----- lab/apiproxy.py | 190 ++++++++++++++++++++++++++++++++++++ lab/evaltools.py | 25 +++-- lab/proftools.py | 12 ++- lab/safety/container.py | 71 +++++++++++++- lab/session.py | 8 +- lab/tools.py | 18 ++-- lab/vm-setup.sh | 4 + lab/worker.py | 23 +++-- tests/test_lab_agent.py | 36 +++++++ tests/test_lab_apiproxy.py | 129 ++++++++++++++++++++++++ tests/test_lab_container.py | 67 +++++++++++++ tests/test_lab_evaltools.py | 18 ++++ 17 files changed, 666 insertions(+), 65 deletions(-) create mode 100644 lab/apiproxy.py create mode 100644 tests/test_lab_agent.py create mode 100644 tests/test_lab_apiproxy.py diff --git a/CLAUDE.md b/CLAUDE.md index 8e4db71..278458c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -30,7 +30,10 @@ pytest/ruff config). Each folder with commands has its own `python -m`. what `tests/conftest.py` points at. Surfaces, grader commands, `lab/engine.py`, and the `correctness`, `regimes`, `validity` CLIs all read it. - `lab/` — moved from inference-server 2026-10-06 (`lab/README.md`). `session.py` (the loop over a dollar budget), - `agent.py` (the one place a model is called; the agent's jailed shell gets the engine's python on PATH), + `agent.py` (the one place a model is called; the agent's jailed shell gets the engine's python on PATH; on Linux + the CLI runs in a per-session workbench container: GPU, internet, root inside, engine venv read-only; its model + calls go through `apiproxy.py`, which injects the key, prices usage and enforces the session cap; the workbench is + paused by `container.quiet` around measure/equiv/profile jobs, GPU holders killed and recorded), `tools.py` (metered tools: test, profile, ledger, budget, restore, note; every record's `cost` has wall seconds, and GPU tools charge them at `LAB_GPU_USD_PER_HOUR` / the target's `[cost] gpu_usd_per_hour` to the run's budget) + `proftools.py` (trace = nsys, kernel = ncu, hostprof = py-spy, each wrapping lab.profile's workload in the jail diff --git a/lab/LEDGER.md b/lab/LEDGER.md index b722f48..0c4730c 100644 --- a/lab/LEDGER.md +++ b/lab/LEDGER.md @@ -24,7 +24,9 @@ A tool that could not run writes `{verdict: refused | error, reason}` and the ag crashes the session. `bench`, `equiv`, `submit` and `baseline` also write `refused` when a process already held the GPU (no engine launched) and `contaminated` when a process other than the served engine held it during the measurement; both carry `gpu_processes: [{pid, used_mib, name}]`. A contaminated run is never evidence and never -cached as the base. The integrity verdict (`feedback/lab_verdict.py`) reads only these records: a run is valid +cached as the base. A measurement that killed processes in the agent's workbench (they held the GPU) lists them +as `workbench_killed: [{pid, used_mib, name}]` (bench, equiv, profile records; submit's reply). A `session`'s +`cost.usd` is what the API proxy priced, and `cost.tokens` its token counts. The integrity verdict (`feedback/lab_verdict.py`) reads only these records: a run is valid iff a completed `submit` exists whose snapshot also has a completed seen-split `bench`, a passing `test`, a passing `equiv` (tier full) and no failing `equiv`, and no record reports a write-surface violation. diff --git a/lab/README.md b/lab/README.md index 5900c70..291287c 100644 --- a/lab/README.md +++ b/lab/README.md @@ -163,7 +163,9 @@ Agent code runs jailed. On Linux the jail is a container (`lab/safety/container. network, read-only root, non-root user, no capabilities; a served engine's port reaches the host through a socket bridge; image from `python -m lab.safety.container build`, which `vm-setup.sh` runs). Elsewhere it is sandbox-runtime (`npm install -g @anthropic-ai/sandbox-runtime`); `LAB_JAIL` picks one. `LAB_NO_JAIL=1` waives -the jail for tests and a box you trust. The agent's own CLI still runs under srt (`lab/RUNTIME.md`). +the jail for tests and a box you trust. The agent's CLI runs in its own workbench container on Linux (GPU, internet, +root inside; paused while anything is measured) and under srt elsewhere; its model calls go through +`lab/apiproxy.py`, so no credential is ever inside a jail (`lab/RUNTIME.md`). The GPU work behind `test`, `profile`, `trace`, `kernel`, `hostprof`, `bench`, `equiv` and `submit` runs as worker jobs (`lab/worker.py`). `LAB_WORKER=local` (default) runs them in process; `LAB_WORKER=ssh` runs them on the VM diff --git a/lab/RUNTIME.md b/lab/RUNTIME.md index 6ccc544..4bbd0be 100644 --- a/lab/RUNTIME.md +++ b/lab/RUNTIME.md @@ -3,8 +3,9 @@ Design, 2026-10-08. Decided with the user: the lab splits into a controller and a worker; no separate measurement VM (cost); NVIDIA only for now; the jail on Linux is a container with the NVIDIA runtime (option A below). Built: the container jail (`lab/safety/container.py`), GPU exclusivity during measurements, and the -measured jobs behind the worker interface (`lab/worker.py`). Not built: the agent on the worker (its container, the -API proxy) and the agent-side jobs. Open decisions are marked **Decide**. +measured jobs behind the worker interface (`lab/worker.py`), the API proxy (`lab/apiproxy.py`), the workbench +container with its pause during measurements. Not built: the agent on a *remote* worker (its CLI over SSH, the +workspace jobs, the proxy through a tunnel) and the dependency build step. Open decisions are marked **Decide**. ## The line (the user, 2026-10-09) @@ -33,7 +34,8 @@ during its own measurement. Everything else it can have. Decided with it: - **Fresh VM per run:** yes (rule 6). - **No separate measurement VM** for now: container escape is not a realistic threat at this stage, and it would double GPU cost. Revisit for a public leaderboard. -- **Refuse rather than kill** (#20) until the workbench can be frozen during a measurement (rule 4). +- **Refuse rather than kill** (#20) until the workbench can be frozen during a measurement (rule 4). Since the + workbench container (2026-10-10) its GPU holders are killed and it is paused; refusing stays for anything else. Not decided here: held-out workloads are sampled from public datasets, so with internet the agent could download them; the workload generator's owner decides (rephrased or non-public held-out prompts). @@ -162,6 +164,31 @@ Contaminated, NotReady and ValueError are raised again on the controller. equiv answers with the job, so calling it again never re-rolls the gate. Tested through a local shell with the same CLI; not yet over SSH to a VM. `agent`, `exec`, `snapshot`, `restore` and `facts` come with the agent container. +## The workbench and the proxy, as built (2026-10-10) + +- **Workbench** (`container.workbench_argv`, Linux): one container per session, `docker run -i` so the CLI's stdio + stays the Agent SDK's. GPU, Docker's default bridge (the internet), root inside when the lab runs as root (the + lab's uid otherwise, so the workspace stays the lab's), writable system dirs, Docker's default capabilities, + `no-new-privileges`. Mounted: the workspace and the CLI's home and tmp read-write; the engine's venv, the + interpreter, weights, CUDA and the CLI binary read-only, so an install in the workbench never reaches a + measurement. Container root is uid 0 on the host kernel (no user-namespace remap): an escape is host root, + accepted under decision 5; `userns-remap` is the hardening step. `vm-setup.sh` drops container traffic to + 169.254.0.0/16 (the cloud metadata service) and checks it. +- **Pause** (`container.quiet`): around every `measure`, `equiv` and `profile` job, processes in the workbench + holding the GPU are killed, the container is paused, a second sweep catches any that started before the pause + landed, and it waits for their memory to be released; thawed after. The killed processes are in the record and + the tool's result (`workbench_killed`), and every GPU tool's description says so on Linux. +- **API proxy** (`lab/apiproxy.py`): per session, bound to the bridge gateway (Linux) or 127.0.0.1 (srt). The CLI + holds the placeholder key `lab-proxy`; the proxy swaps in the lab's key (`x-api-key`) or OAuth token + (`Authorization: Bearer` + the oauth beta), forwards only `/v1/messages`, `/v1/messages/count_tokens` and + `/v1/models`, prices every response's usage (JSON or SSE) at first-party rates, and refuses once the session's + cap is spent, and refuses any unpriced model or fast mode. The session is charged what the proxy priced, not + what the CLI reports, and its tokens go in the session record. Before this, the key sat in the CLI's + environment, readable from `/proc` by anything in the agent's shell. +- Not tested: the real CLI through the proxy (it may call an endpoint outside the three; it would see a 403), and + the macOS srt path (srt is not installed on the dev box). Tested in Docker-in-Docker: workbench argv, root and + writable system dirs, reaching the gateway, pause and thaw; the proxy against a fake API. + ## What moves where - Controller: `session`, `agent` (SDK side), the tool front ends in `tools`/`evaltools`/`proftools` (policy, @@ -192,6 +219,7 @@ database, still append-only. 1. ~~Decide the jail~~ (A) and build it. Done, minus the GPU check. 2. Without a GPU: ~~the job interface with a local transport, the snapshot round trip and rule 3~~ (done for the - measured jobs); next the agent container with its freeze, and the API proxy. + measured jobs); ~~the agent container with its freeze, and the API proxy~~ (done, lab on one machine); next the + agent on a remote worker, then the dependency build step. 3. With a GPU (spend approval): `--gpus` and the profilers inside the container, then a first `lab.session` dry run with a tiny budget. diff --git a/lab/agent.py b/lab/agent.py index f496d30..80165f6 100644 --- a/lab/agent.py +++ b/lab/agent.py @@ -9,15 +9,14 @@ from pathlib import Path from typing import Any, Callable, Protocol -from lab import engine -from lab.safety import jail +from lab import apiproxy, engine +from lab.safety import container, jail from lab.safety.grader import HF_HUB from lab.safety.hooks import WRITE_TOOLS, write_guard -API_HOST = "api.anthropic.com" BUILTIN_TOOLS = ["Read", "Grep", "Glob", "Edit", "Write", "Bash"] -PASS_ENV = ("PATH", "ANTHROPIC_API_KEY", "CLAUDE_CODE_OAUTH_TOKEN", "CLAUDE_CODE_ENTRYPOINT", - "CLAUDE_AGENT_SDK_VERSION") +PASS_ENV = ("CLAUDE_CODE_ENTRYPOINT", "CLAUDE_AGENT_SDK_VERSION") # what the SDK sets for its CLI +SYSTEM_PATH = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" # What a session says when it ends. `stop` means the agent is done with this run. OUTPUT_SCHEMA = {"type": "object", "additionalProperties": False, "required": ["status", "note"], "properties": {"status": {"type": "string", "enum": ["continue", "stop"]}, @@ -46,6 +45,7 @@ class AgentSpec: max_budget_usd: float = 5.0 timeout_s: float = 3600.0 model: str = field(default_factory=lambda: os.environ.get("LAB_MODEL", "claude-fable-5-1")) + workbench: str | None = None # the container name of the agent's workbench (Linux) @dataclass @@ -55,6 +55,7 @@ class AgentReply: turns: int error: str | None cost_estimated: bool = False # no accounting came back; charged at the session cap + usage: dict | None = None # tokens as the API proxy counted them class Provider(Protocol): @@ -71,33 +72,48 @@ def load(name: str | None = None) -> Provider: class ClaudeAgentSDK: - """Claude through the Agent SDK CLI, which runs inside the srt jail with a wiped env; lab tools run out here.""" + """Claude through the Agent SDK CLI, jailed with a wiped env (on Linux in the workbench container, elsewhere under + srt); its model calls go through the lab's API proxy, so no credential is ever inside; lab tools run out here.""" name = "claude" - def _wrapper(self, spec: AgentSpec, cli: Path) -> Path: + @staticmethod + def containerized() -> bool: + return jail.backend() == "container" and container.available() + + def _wrapper(self, spec: AgentSpec, cli: Path, base_url: str) -> Path: home, tmp = spec.scratch / "home", spec.scratch / "tmp" home.mkdir(parents=True, exist_ok=True) tmp.mkdir(exist_ok=True) ws = spec.workspace.resolve() weights = [HF_HUB] if HF_HUB.exists() else [] # model weights read-only; the token beside them stays hidden - cli_tmp = Path(f"/private/tmp/claude-{os.getuid()}") # the CLI's own per-project scratch; Bash fails without it - cli_tmp.mkdir(parents=True, exist_ok=True) - config = jail.settings([home.resolve(), tmp.resolve(), ws, cli_tmp.resolve()], engine.venv(), - [API_HOST], readonly=weights, python=engine.python()) - argv = jail.wrap(config, spec.scratch / "srt.json", [str(cli)]) - keep = " ".join(f'"{k}=${k}"' for k in PASS_ENV) - path = f'"PATH={Path(engine.python()).parent}:$PATH"' # the agent's `python` is the engine's - env = (f"{keep} {path} HOME={shlex.quote(str(home))} CLAUDE_CONFIG_DIR={shlex.quote(str(home))} " - f"TMPDIR={shlex.quote(str(tmp))} PYTHONPATH={shlex.quote(str(ws / 'src'))} " - f"PYTHONDONTWRITEBYTECODE=1 HF_HUB_CACHE={shlex.quote(str(HF_HUB))} HF_HUB_OFFLINE=1 " - f"HF_HOME={shlex.quote(str(tmp / 'hf-home'))}") + env = {"HOME": str(home), "CLAUDE_CONFIG_DIR": str(home), "TMPDIR": str(tmp), "PYTHONPATH": str(ws / "src"), + "PYTHONDONTWRITEBYTECODE": "1", "HF_HUB_CACHE": str(HF_HUB), "HF_HUB_OFFLINE": "1", + "HF_HOME": str(tmp / "hf-home"), "ANTHROPIC_BASE_URL": base_url, + "ANTHROPIC_API_KEY": apiproxy.PLACEHOLDER, "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1"} + engine_bin = Path(engine.python()).parent # the agent's `python` is the engine's + if self.containerized(): + env["PATH"] = f"{engine_bin}:{SYSTEM_PATH}" + argv = container.workbench_argv(spec.workbench, ws, [str(cli)], env, [home.resolve(), tmp.resolve(), ws], + [engine.venv(), *weights, cli.parent], engine.python()) + at = argv.index(container.IMAGE) + argv[at:at] = [w for k in PASS_ENV for w in ("--env", k)] # the SDK's values, passed through + line = f"exec {shlex.join(argv)} \"$@\"" + else: + cli_tmp = Path(f"/private/tmp/claude-{os.getuid()}") # the CLI's own per-project scratch; Bash fails without it + cli_tmp.mkdir(parents=True, exist_ok=True) + config = jail.settings([home.resolve(), tmp.resolve(), ws, cli_tmp.resolve()], engine.venv(), + [base_url.removeprefix("http://")], readonly=weights, python=engine.python()) + argv = jail.wrap(config, spec.scratch / "srt.json", [str(cli)]) + keep = " ".join(f'"{k}=${k}"' for k in PASS_ENV) + sets = " ".join(f"{k}={shlex.quote(v)}" for k, v in env.items()) + line = f'exec env -i {keep} "PATH={engine_bin}:$PATH" {sets} {shlex.join(argv)} "$@"' script = spec.scratch / "cli.sh" - script.write_text(f"#!/bin/sh\nexec env -i {env} {shlex.join(argv)} \"$@\"\n") + script.write_text(f"#!/bin/sh\n{line}\n") script.chmod(0o700) return script - async def _run(self, spec: AgentSpec) -> AgentReply: + async def _run(self, spec: AgentSpec, base_url: str) -> AgentReply: import claude_agent_sdk from claude_agent_sdk import (ClaudeAgentOptions, HookMatcher, ResultMessage, create_sdk_mcp_server, query, tool) @@ -119,7 +135,7 @@ async def call(args: dict[str, Any]) -> dict[str, Any]: mcp_servers={"lab": server}, permission_mode="dontAsk", setting_sources=[], cwd=str(spec.workspace), max_turns=spec.max_turns, max_budget_usd=spec.max_budget_usd, output_format={"type": "json_schema", "schema": OUTPUT_SCHEMA}, - cli_path=str(self._wrapper(spec, cli)), + cli_path=str(self._wrapper(spec, cli, base_url)), hooks={"PreToolUse": [HookMatcher(matcher=WRITE_TOOLS, hooks=[write_guard(spec.workspace, spec.base_files)])]}, ) @@ -132,11 +148,22 @@ async def call(args: dict[str, Any]) -> dict[str, Any]: return reply def run(self, spec: AgentSpec) -> AgentReply: - """Never raises. Without a result message the spend is unknown and is charged at the session cap.""" + """Never raises. The spend is what the API proxy priced, whatever the CLI reports or however it ended.""" spec.scratch.mkdir(parents=True, exist_ok=True) try: - return asyncio.run(asyncio.wait_for(self._run(spec), spec.timeout_s)) - except TimeoutError: - return AgentReply(None, spec.max_budget_usd, 0, f"timed out after {spec.timeout_s:.0f}s", True) - except Exception as e: - return AgentReply(None, spec.max_budget_usd, 0, f"provider failed: {type(e).__name__}: {e}", True) + host = container.bridge_gateway() if self.containerized() else "127.0.0.1" + proxy = apiproxy.Proxy(spec.max_budget_usd, host) + except Exception as e: # nothing was sent: nothing spent + return AgentReply(None, 0.0, 0, f"provider failed: {type(e).__name__}: {e}") + with proxy: + try: + reply = asyncio.run(asyncio.wait_for(self._run(spec, proxy.url), spec.timeout_s)) + except TimeoutError: + reply = AgentReply(None, 0.0, 0, f"timed out after {spec.timeout_s:.0f}s") + except Exception as e: + reply = AgentReply(None, 0.0, 0, f"provider failed: {type(e).__name__}: {e}") + finally: + if spec.workbench and self.containerized(): + container.remove(spec.workbench) # a CLI killed mid-session leaves its container behind + reply.cost_usd, reply.cost_estimated, reply.usage = proxy.spent_usd, False, dict(proxy.usage) + return reply diff --git a/lab/apiproxy.py b/lab/apiproxy.py new file mode 100644 index 0000000..d0ee5d1 --- /dev/null +++ b/lab/apiproxy.py @@ -0,0 +1,190 @@ +"""The model API key stays out of every jail: the agent's CLI talks to this proxy, which adds the lab's credential, +forwards to the Anthropic API, prices each response's usage, and refuses once the session's dollar cap is spent. + +Only /v1/messages, /v1/messages/count_tokens and /v1/models[/id] are forwarded (not batches, files or anything else); a request naming a model +without a price here, or asking for fast mode, is refused (fail closed: anything unpriced could spend unmetered). +Prices are first-party per-million-token rates (claude-api reference, cached 2026-09-25).""" + +from __future__ import annotations + +import http.client +import json +import os +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from urllib.parse import urlsplit + +UPSTREAM = "https://api.anthropic.com" +PLACEHOLDER = "lab-proxy" # what the jailed CLI holds instead of a key +OAUTH_BETA = "oauth-2025-04-20" +# $/MTok: input, output, cache read, cache write 5 min, cache write 1 h +PRICES = { + "claude-fable-5-1": (10.0, 50.0, 0.25, 12.5, 20.0), + "claude-opus-5-5": (4.0, 20.0, 0.20, 5.0, 8.0), + "claude-sonnet-5-5": (2.0, 10.0, 0.20, 2.5, 4.0), + "claude-haiku-4-5": (1.0, 5.0, 0.10, 1.25, 2.0), +} +ALLOWED = ("/v1/messages", "/v1/messages/count_tokens", "/v1/models") +HOP = {"connection", "keep-alive", "transfer-encoding", "te", "trailer", "upgrade", "proxy-authorization", + "proxy-authenticate", "content-length", "host", "accept-encoding"} + + +def price(model: str) -> tuple | None: + """The model's rates; a `[1m]`-style suffix the CLI adds does not change them.""" + return PRICES.get(model.split("[")[0]) + + +def cost_usd(model: str, usage: dict) -> float: + p_in, p_out, p_read, p_5m, p_1h = price(model) + created = usage.get("cache_creation") or {} + w5, w1 = created.get("ephemeral_5m_input_tokens"), created.get("ephemeral_1h_input_tokens") + if w5 is None and w1 is None: # no breakdown: every write at the 5-minute rate + w5, w1 = usage.get("cache_creation_input_tokens") or 0, 0 + return ((usage.get("input_tokens") or 0) * p_in + (usage.get("output_tokens") or 0) * p_out + + (usage.get("cache_read_input_tokens") or 0) * p_read + (w5 or 0) * p_5m + (w1 or 0) * p_1h) / 1e6 + + +def credential() -> dict[str, str]: + """The headers that authenticate upstream, from the lab's own environment.""" + if os.environ.get("ANTHROPIC_API_KEY"): + return {"x-api-key": os.environ["ANTHROPIC_API_KEY"]} + if os.environ.get("CLAUDE_CODE_OAUTH_TOKEN"): + return {"authorization": f"Bearer {os.environ['CLAUDE_CODE_OAUTH_TOKEN']}"} + raise RuntimeError("no model credential for the proxy: set ANTHROPIC_API_KEY or CLAUDE_CODE_OAUTH_TOKEN") + + +class _Usage: + """Usage read off a response as it streams by: JSON, or SSE message_start + message_delta.""" + + def __init__(self, sse: bool): + self.sse, self.buf, self.model, self.usage = sse, b"", None, {} + + def feed(self, chunk: bytes) -> None: + self.buf += chunk + if not self.sse: + return + *lines, self.buf = self.buf.split(b"\n") + for line in lines: + if line.startswith(b"data:"): + self._event(line[5:].strip()) + + def _event(self, data: bytes) -> None: + try: + e = json.loads(data) + except ValueError: + return + if e.get("type") == "message_start": + self.model = e["message"].get("model") + self.usage.update(e["message"].get("usage") or {}) + elif e.get("type") == "message_delta": + self.usage.update(e.get("usage") or {}) + + def done(self) -> None: + if self.sse: + self.feed(b"\n") + return + try: + body = json.loads(self.buf) + except ValueError: + return + if isinstance(body, dict) and isinstance(body.get("usage"), dict): + self.model, self.usage = body.get("model"), body["usage"] + + +class Proxy: + """`with Proxy(cap_usd, host) as p: ... p.url`; `p.spent_usd` and `p.usage` are what the session cost.""" + + def __init__(self, cap_usd: float, host: str = "127.0.0.1", upstream: str = UPSTREAM): + self.cap_usd, self.spent_usd, self.requests = cap_usd, 0.0, 0 + self.usage = {"input_tokens": 0, "output_tokens": 0, "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0} + self.auth = credential() + self.upstream = urlsplit(upstream) + self.lock = threading.Lock() + proxy = self + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, *a): + pass + + def do_GET(self): + proxy._handle(self, None) + + def do_POST(self): + n = self.headers.get("content-length") + if n is None: + return proxy._refuse(self, 411, "a request body needs a Content-Length") + proxy._handle(self, self.rfile.read(int(n))) + + self.server = ThreadingHTTPServer((host, 0), Handler) + self.server.daemon_threads = True + self.url = f"http://{host}:{self.server.server_address[1]}" + + def __enter__(self) -> "Proxy": + threading.Thread(target=self.server.serve_forever, daemon=True).start() + return self + + def __exit__(self, *exc) -> None: + self.server.shutdown() + self.server.server_close() + + def _refuse(self, h, status: int, why: str) -> None: + body = json.dumps({"type": "error", "error": {"type": "permission_error", "message": f"lab proxy: {why}"}}).encode() + h.send_response(status) + h.send_header("content-type", "application/json") + h.send_header("content-length", str(len(body))) + h.end_headers() + h.wfile.write(body) + + def _handle(self, h, body: bytes | None) -> None: + path = urlsplit(h.path).path + if path not in ALLOWED and not (path.startswith("/v1/models/") and path.count("/") == 3): + return self._refuse(h, 403, f"{path} is not forwarded") + model = None + if body: + try: + req = json.loads(body) + except ValueError: + return self._refuse(h, 400, "the request body is not JSON") + model = req.get("model") + if path == "/v1/messages" and (not isinstance(model, str) or price(model) is None): + return self._refuse(h, 403, f"no price for model {model!r}; priced: {sorted(PRICES)}") + if req.get("speed") == "fast": + return self._refuse(h, 403, "fast mode is not priced here") + with self.lock: + if path == "/v1/messages" and self.spent_usd >= self.cap_usd: + return self._refuse(h, 403, f"this session's ${self.cap_usd:.2f} is spent") + headers = {k: v for k, v in h.headers.items() if k.lower() not in HOP | {"x-api-key", "authorization"}} + headers.update(self.auth, host=self.upstream.netloc, **{"accept-encoding": "identity"}) + if "authorization" in self.auth: + betas = [b for b in headers.pop("anthropic-beta", "").split(",") if b.strip()] + headers["anthropic-beta"] = ",".join(dict.fromkeys([*betas, OAUTH_BETA])) + conn = (http.client.HTTPSConnection if self.upstream.scheme == "https" else http.client.HTTPConnection)( + self.upstream.netloc, timeout=900) + try: + conn.request(h.command, h.path, body=body, headers=headers) + resp = conn.getresponse() + usage = _Usage("text/event-stream" in resp.getheader("content-type", "")) + h.send_response(resp.status) + for k, v in resp.getheaders(): + if k.lower() not in HOP: + h.send_header(k, v) + h.send_header("transfer-encoding", "chunked") + h.end_headers() + while chunk := resp.read1(65536): + usage.feed(chunk) + h.wfile.write(b"%x\r\n%s\r\n" % (len(chunk), chunk)) + h.wfile.flush() + h.wfile.write(b"0\r\n\r\n") + usage.done() + finally: + conn.close() + if path == "/v1/messages" and usage.usage: + spent = cost_usd(usage.model if usage.model and price(usage.model) else model, usage.usage) + with self.lock: + self.spent_usd += spent + self.requests += 1 + for k in self.usage: + self.usage[k] += usage.usage.get(k) or 0 diff --git a/lab/evaltools.py b/lab/evaltools.py index 6c1277c..784ca71 100644 --- a/lab/evaltools.py +++ b/lab/evaltools.py @@ -132,6 +132,11 @@ def _failed(e: BaseException) -> dict: return {"verdict": "error", "reason": _why(e)} +def _killed(r: dict) -> dict: + """Processes of the agent's workbench killed because they held the GPU when the measurement began (a fact).""" + return {"workbench_killed": r["killed"]} if r.get("killed") else {} + + def _said(tool: str, result: dict) -> str: return f"{tool} {'failed' if result['verdict'] == 'error' else result['verdict']}: {result['reason']}" @@ -163,9 +168,10 @@ def bench(self, args: dict) -> str: return _said("bench", result) result = {"verdict": "ok", "metrics": headline(r["results"]), "regimes": r["results"], "tier": tier, "seconds": time.monotonic() - t0, "ready_s": r["ready_s"], - "artifacts": artifacts.store(work, self.s.ledger_root)} + "artifacts": artifacts.store(work, self.s.ledger_root), **_killed(r)} self._record("bench", "bench", args, result, snap, config={"split": "seen", "tier": tier}) - return json.dumps({"snapshot": snap.id, "split": "seen", "tier": tier, "metrics": result["metrics"]}, indent=1) + return json.dumps({"snapshot": snap.id, "split": "seen", "tier": tier, "metrics": result["metrics"], + **_killed(r)}, indent=1) # -- equiv -------------------------------------------------------------------------- def equiv(self, args: dict) -> str: @@ -186,8 +192,8 @@ def equiv(self, args: dict) -> str: shutil.copyfile(out.with_suffix(".outputs.jsonl"), prior / "equiv.outputs.jsonl") try: res = self.worker.call("equiv", {"tree": tree, "reference": ref, "prior": prior}, - {"tier": tier, "concurrency": self.equiv_concurrency, "jailed": self.serve_jailed}, - work) + {"tier": tier, "concurrency": self.equiv_concurrency, "jailed": self.serve_jailed, + "workbench": getattr(self.s, "workbench", None)}, work) except Exception as e: result = {**_failed(e), "artifacts": self._equiv_store(work, out, prior)} self._record("equiv", "equiv", args, result, snap, config={"split": "seen", "tier": tier}) @@ -195,10 +201,10 @@ def equiv(self, args: dict) -> str: passed = {"pass": True, "fail": False}.get(res["verdict"]) # inconclusive is None: neither passing nor failing record = {"verdict": res["verdict"], "passed": passed, "reasons": res["reasons"], "gates": res["gates"], "metrics": res["metrics"], "thresholds": res["thresholds"], "tier": tier, - "artifacts": self._equiv_store(work, out, prior)} + "artifacts": self._equiv_store(work, out, prior), **_killed(res)} self._record("equiv", "equiv", args, record, snap, config={"split": "seen", "tier": tier}) return json.dumps({"snapshot": snap.id, "tier": tier, "verdict": res["verdict"], "reasons": res["reasons"], - "metrics": res["metrics"]}, indent=1, default=str) + "metrics": res["metrics"], **_killed(res)}, indent=1, default=str) # -- submit ------------------------------------------------------------------------- def submit(self, args: dict) -> str: @@ -245,7 +251,8 @@ def submit(self, args: dict) -> str: self._record("submit", "submit", args, result, snap) return _said("submit", result) rec = self._record_heldout("submit", "submit", args, {"tier": "full"}, metrics, snap) - out = {"snapshot": snap.id, "split": "heldout", "tier": "full", "metrics": metrics, "record": rec["id"]} + out = {"snapshot": snap.id, "split": "heldout", "tier": "full", "metrics": metrics, "record": rec["id"], + **_killed(r)} task = getattr(self.s, "task", None) if task is not None: # the task's own win condition, applied to these aggregates from lab.task import score @@ -270,8 +277,8 @@ def _work(self, tool: str) -> Path: def _measure(self, tree: Path, out: Path, log: str, **args) -> dict: """The worker's `measure` job into `out`; the engine's log is appended to `/` either way.""" try: - return self.worker.call("measure", {"tree": tree}, {**args, "seed": self.seed, "jailed": self.serve_jailed}, - out) + return self.worker.call("measure", {"tree": tree}, {**args, "seed": self.seed, "jailed": self.serve_jailed, + "workbench": getattr(self.s, "workbench", None)}, out) finally: src = out / worker.ENGINE_LOG if src.exists(): diff --git a/lab/proftools.py b/lab/proftools.py index e6967d1..8b53c81 100644 --- a/lab/proftools.py +++ b/lab/proftools.py @@ -8,12 +8,14 @@ from __future__ import annotations +import json import re import shutil import time from pathlib import Path from lab.agent import ToolSpec +from lab.evaltools import _killed OUT = "_instr" # in the pristine tree: writable in the jail, rebuilt on every call NSYS_REPORTS = "cuda_gpu_kern_sum,cuda_api_sum,cuda_gpu_mem_time_sum,nvtx_sum" @@ -111,7 +113,8 @@ def _instrumented(self, tool: str, args: dict) -> str: out = self._work(tool) t0 = time.monotonic() try: - r = self.worker.call("profile", {"tree": tree}, {"kind": tool, "args": args}, out) + r = self.worker.call("profile", {"tree": tree}, + {"kind": tool, "args": args, "workbench": getattr(self.s, "workbench", None)}, out) if r.get("refused"): self._record("note", tool, args, {"verdict": "refused", "reason": r["refused"]}, snap) return f"{tool} refused: {r['refused']}" @@ -120,11 +123,12 @@ def _instrumented(self, tool: str, args: dict) -> str: finally: shutil.rmtree(out, ignore_errors=True) result = {"returncode": r["returncode"], "bundle": blob, "workspace_copy": visible, "instrument": r["instrument"], - "argv": r["argv"], "seconds": time.monotonic() - t0, "output": r["output"]} + "argv": r["argv"], "seconds": time.monotonic() - t0, "output": r["output"], **_killed(r)} self._record("profile", tool, args, result, snap) + killed = f"\nkilled in your workbench (held the GPU): {json.dumps(r['killed'])}" if r.get("killed") else "" if r["returncode"] != 0 or not r["head"]: - return f"{tool} failed ({r['returncode']}):\n{result['output']}" - return f"{tool} output at {visible} (in your workspace; left out of your change)\n{r['head']}" + return f"{tool} failed ({r['returncode']}):\n{result['output']}{killed}" + return f"{tool} output at {visible} (in your workspace; left out of your change)\n{r['head']}{killed}" def prof_specs(self, gpu) -> list[ToolSpec]: """`gpu(time)` renders the cost fact for a description.""" diff --git a/lab/safety/container.py b/lab/safety/container.py index 99fbfe6..fcb6382 100644 --- a/lab/safety/container.py +++ b/lab/safety/container.py @@ -1,5 +1,7 @@ -"""The Linux jail: one container per jailed command (NVIDIA runtime for the GPU), no network at all, read-only root, -non-root user, no capabilities. A served engine's port reaches the host through a socket bridge. See lab/RUNTIME.md.""" +"""The Linux jails (lab/RUNTIME.md). A measurement room: one container per jailed command (NVIDIA runtime for the +GPU), no network at all, read-only root, non-root user, no capabilities; a served engine's port reaches the host +through a socket bridge. The agent's workbench: one container per session with the GPU, the internet and root +inside it (writable system dirs, Docker's default capabilities), paused by `quiet` while anything is measured.""" from __future__ import annotations @@ -7,7 +9,10 @@ import shutil import subprocess import sys +import signal +import time import uuid +from contextlib import contextmanager from pathlib import Path IMAGE = os.environ.get("LAB_JAIL_IMAGE", "lab-jail") @@ -92,6 +97,68 @@ def argv(name: str, workdir: Path, command: list[str], env: dict[str, str], writ "jail", *command] +def workbench_argv(name: str, workdir: Path, command: list[str], env: dict[str, str], writable: list[Path], + readonly: list[Path], python: str) -> list[str]: + """`docker run -i` for the agent's CLI: stdio stays the SDK's, the default bridge network gives the internet, + root inside when the lab runs as root (otherwise the lab's uid, so the workspace stays the lab's).""" + uid, gid = (0, 0) if os.getuid() == 0 else (os.getuid(), os.getgid()) + out = [_docker(), "run", "-i", "--rm", "--init", "--name", name, "--security-opt", "no-new-privileges", + "--pids-limit", "16384", "--shm-size", "16g", "--user", f"{uid}:{gid}", "--workdir", str(workdir.resolve()), + "--label", "lab.workbench=1"] + if gpus(): + out += ["--gpus", "all"] + if (cpus := cpuset()): + out += ["--cpuset-cpus", cpus] + for k, v in env.items(): + out += ["--env", f"{k}={v}"] + return out + mounts(writable, readonly, python) + [IMAGE, *command] + + +def bridge_gateway() -> str: + """The host's address on Docker's default bridge: what a workbench reaches the host at, and nothing outside.""" + out = subprocess.run([_docker(), "network", "inspect", "bridge", "-f", "{{(index .IPAM.Config 0).Gateway}}"], + capture_output=True, text=True, check=True) + return out.stdout.strip() + + +def running(name: str) -> bool: + out = subprocess.run([_docker(), "inspect", "-f", "{{.State.Running}}", name], capture_output=True, text=True) + return out.returncode == 0 and out.stdout.strip() == "true" + + +def _kill_gpu_holders(name: str) -> list[dict]: + from lab import gpu + mine = pids(name) + hit = [a for a in (gpu.compute_apps() or []) if a["pid"] in mine] + for a in hit: + try: + os.kill(a["pid"], signal.SIGKILL) + except ProcessLookupError: + pass + return hit + + +@contextmanager +def quiet(name: str | None): + """The workbench paused for a measurement: processes in it holding the GPU are killed (and listed), the rest + frozen; thawed after. A no-op when there is no such running container.""" + if not name or _docker() is None or not running(name): + yield [] + return + from lab import gpu + killed = _kill_gpu_holders(name) + subprocess.run([_docker(), "pause", name], capture_output=True, check=True) + try: + killed += _kill_gpu_holders(name) # anything that took the GPU before the pause landed + deadline = time.monotonic() + 15 + while killed and time.monotonic() < deadline and any( + a["pid"] in {k["pid"] for k in killed} for a in gpu.compute_apps() or []): + time.sleep(0.2) # until the driver has released their memory + yield killed + finally: + subprocess.run([_docker(), "unpause", name], capture_output=True) + + def new_name() -> str: return f"lab-jail-{uuid.uuid4().hex[:12]}" diff --git a/lab/session.py b/lab/session.py index d63eb2f..ac3de12 100644 --- a/lab/session.py +++ b/lab/session.py @@ -33,6 +33,7 @@ class Session: ledger_root: Path profile_runner: object = None # None: lab.tools.default_profile_runner task: Task | None = None + workbench: str | None = None # the agent's container: paused while the tools measure @dataclass @@ -128,13 +129,14 @@ def run(cfg: RunConfig, provider: agent.Provider | None = None) -> dict: if budget.remaining_usd < MIN_SESSION_USD: summary["stopped"] = "budget" break - s = Session(run_id, f"{run_id}-s{n}", run_dir, ws, budget, cfg.ledger_root, task=cfg.task) + s = Session(run_id, f"{run_id}-s{n}", run_dir, ws, budget, cfg.ledger_root, task=cfg.task, + workbench=f"lab-workbench-{run_id}-s{n}") tools = Toolbox(s) snap = ws.snapshot() spec = agent.AgentSpec(system="", prompt=brief(cfg, s, snap.id, n), workspace=ws.path, scratch=run_dir / "scratch" / s.session_id, tools=tools.specs(), base_files=base_files, max_turns=cfg.max_turns, - max_budget_usd=min(cfg.session_usd, budget.remaining_usd)) + max_budget_usd=min(cfg.session_usd, budget.remaining_usd), workbench=s.workbench) spec.system = system_prompt(target.load(), spec) t0 = time.monotonic() try: @@ -154,7 +156,7 @@ def run(cfg: RunConfig, provider: agent.Provider | None = None) -> dict: status = (reply.output or {}).get("status") if reply.output else None ledger.append({"kind": "session", "run": run_id, "session": s.session_id, "provider": provider.name, "model": spec.model, "snapshot": end.id, "snapshot_blob": end.blob, "patch": end.patch, - "cost": {"usd": reply.cost_usd, "estimated": reply.cost_estimated}, + "cost": {"usd": reply.cost_usd, "estimated": reply.cost_estimated, "tokens": reply.usage}, "turns": reply.turns, "seconds": time.monotonic() - t0, "status": status, "error": reply.error, "violation": tools.violation, "claim": {"note": (reply.output or {}).get("note")}}, cfg.ledger_root) diff --git a/lab/tools.py b/lab/tools.py index 86180a4..8943cb7 100644 --- a/lab/tools.py +++ b/lab/tools.py @@ -12,7 +12,7 @@ from lab import corpus, engine, ledger, target from lab.agent import ToolSpec from lab.budget import Budget, BudgetExceeded, gpu_rate -from lab.evaltools import EvalTools +from lab.evaltools import EvalTools, _killed from lab.proftools import ProfTools from lab.safety import grader @@ -119,17 +119,19 @@ def profile(self, args: dict) -> str: out = self._work("profile") t0 = time.monotonic() try: - r = self.worker.call("profile", {"tree": self._pristine()}, {"kind": "profile", "args": args}, out) + r = self.worker.call("profile", {"tree": self._pristine()}, + {"kind": "profile", "args": args, "workbench": getattr(self.s, "workbench", None)}, out) blob, visible = self._keep(out / "files") if (out / "files").is_dir() else ("", "") finally: shutil.rmtree(out, ignore_errors=True) result = {"returncode": r["returncode"], "bundle": blob, "workspace_copy": visible, - "seconds": time.monotonic() - t0, "output": r["output"]} + "seconds": time.monotonic() - t0, "output": r["output"], **_killed(r)} self._record("profile", "profile", args, result, snap) + killed = f"\nkilled in your workbench (held the GPU): {json.dumps(r['killed'])}" if r.get("killed") else "" if r["returncode"] != 0 or not blob: - return f"profile failed ({r['returncode']}):\n{result['output']}" + return f"profile failed ({r['returncode']}):\n{result['output']}{killed}" return f"bundle at {visible} (in your workspace; left out of your change)\n" + "\n".join( - f" {p.name}" for p in sorted((self.s.workspace.path / visible).iterdir())) + f" {p.name}" for p in sorted((self.s.workspace.path / visible).iterdir())) + killed def ledger_tool(self, args: dict) -> str: match = {k: v for k, v in args.items() if k in ("kind", "session", "snapshot", "tool") and v} @@ -163,9 +165,13 @@ def specs(self) -> list[ToolSpec]: rate, known = gpu_rate() free = " Costs no GPU time ($0)." + from lab.agent import ClaudeAgentSDK + paused = (" While it runs, the container your shell runs in is paused, and processes in it holding the GPU are " + "killed and listed in the result." if ClaudeAgentSDK.containerized() else "") + def gpu(t: str) -> str: return (f" Costs GPU time: ~{t} per call, charged at ${rate:g}/h." if known else - f" Costs GPU time: ~{t} per call; no GPU rate is configured on this host, so it is charged $0.") + f" Costs GPU time: ~{t} per call; no GPU rate is configured on this host, so it is charged $0.") + paused short, full = suite.TIERS["short"], suite.TIERS["full"] windows = (f"one engine start plus, per regime, a {short['final_s']:.0f} s measurement window at tier short or " f"{full['final_s']:.0f} s at full, plus {short['probe_s']:.0f} s / {full['probe_s']:.0f} s probes for " diff --git a/lab/vm-setup.sh b/lab/vm-setup.sh index 9484eb0..4c92c33 100755 --- a/lab/vm-setup.sh +++ b/lab/vm-setup.sh @@ -66,6 +66,10 @@ fi [ "$(id -u)" = 0 ] || $SUDO usermod -aG docker "$USER" # takes effect at the next login $SUDO docker build -q -t lab-jail -f "$here/safety/jail.Dockerfile" "$here/safety" >/dev/null $SUDO docker run --rm --gpus all --network none lab-jail nvidia-smi -L && echo "jail: gpu ok" || echo "jail: no GPU in the container" +# The agent's workbench has the internet but not the cloud's metadata service (instance credentials, user-data). +$SUDO iptables -C DOCKER-USER -d 169.254.0.0/16 -j DROP 2>/dev/null || $SUDO iptables -I DOCKER-USER -d 169.254.0.0/16 -j DROP +$SUDO docker run --rm lab-jail python3 -c "import socket; socket.create_connection(('169.254.169.254', 80), timeout=3)" \ + 2>/dev/null && echo "workbench: METADATA REACHABLE" || echo "workbench: metadata blocked" echo "== counters: can a non-admin process read GPU performance counters?" # ncu needs NVreg_RestrictProfilingToAdminUsers=0 (or root). Recorded as a fact about this venue. diff --git a/lab/worker.py b/lab/worker.py index 1d895ce..c9407fd 100644 --- a/lab/worker.py +++ b/lab/worker.py @@ -2,7 +2,8 @@ dir and returns JSON; the tools keep the policy, the ledger and the budget. `Local` runs a job in this process. `Remote` ships each input by content hash, runs `python -m lab.worker run JOB` on the worker host over SSH and brings the output dir back; the worker re-hashes every tree right before using it and refuses one that does not match, so -what is measured is what the controller audited. GpuBusy, Contaminated, NotReady and ValueError cross intact. +what is measured is what the controller audited. GpuBusy, Contaminated, NotReady and ValueError cross intact. measure, equiv and profile take `workbench`, the agent's +container: it is paused for the job, and processes in it holding the GPU are killed and listed under `killed`. test lint + the engine's suite on `tree` -> {lint, tests} measure serve `tree`, run regimes; passive data into out -> {results, ready_s} @@ -27,6 +28,7 @@ from pathlib import Path from lab import engine, serve, target +from lab.safety import container RESULT = "_result.json" ENGINE_LOG = "engine.log" # the served engine's log, in every job's out dir @@ -66,11 +68,12 @@ def job_measure(inputs: dict, out: Path, args: dict, hooks: dict) -> dict: t = target.load() rows = [] if args["passive"] else None try: - with serve.Served(inputs["tree"], t.engine, log=out / ENGINE_LOG, jailed=args["jailed"], - artifacts=out if args["passive"] else None) as srv: + with container.quiet(args.get("workbench")) as killed, serve.Served( + inputs["tree"], t.engine, log=out / ENGINE_LOG, jailed=args["jailed"], + artifacts=out if args["passive"] else None) as srv: results = run_regimes(srv.url, t, args["names"], args["split"], args["tier"], args["seed"], rows) srv.exclusive() - return {"results": results, "ready_s": srv.ready_s} + return {"results": results, "ready_s": srv.ready_s, "killed": killed} finally: if rows is not None: artifacts.write_rows(out, rows) @@ -85,17 +88,23 @@ def job_equiv(inputs: dict, out: Path, args: dict, hooks: dict) -> dict: if (inputs["prior"] / "equiv.outputs.jsonl").exists(): shutil.copyfile(inputs["prior"] / "equiv.outputs.jsonl", out / "equiv.outputs.jsonl") enc = hooks.get("encoder") or crun.HFEncoder(t.model, t.chat_kwargs) - with serve.Served(inputs["tree"], t.engine, log=out / ENGINE_LOG, jailed=args["jailed"], artifacts=out) as srv: + with container.quiet(args.get("workbench")) as killed, serve.Served( + inputs["tree"], t.engine, log=out / ENGINE_LOG, jailed=args["jailed"], artifacts=out) as srv: res = crun.candidate(inputs["reference"], srv.url, t.model, out / "equiv.json", enc, Thresholds(min_score_ratio=t.min_score_ratio, min_length_ratio=t.min_length_ratio), concurrency=args["concurrency"], api=client.Api(t.engine.api, t.chat_kwargs), config={"tier": args["tier"]}) srv.exclusive() - return res + return {**res, "killed": killed} def job_profile(inputs: dict, out: Path, args: dict, hooks: dict) -> dict: - """`args`: kind (profile | trace | kernel | hostprof) and the tool's args. Output files go to `out/files`.""" + """`args`: kind (profile | trace | kernel | hostprof), the tool's args, workbench. Files go to `out/files`.""" + with container.quiet(args.get("workbench")) as killed: + return {**_profile(inputs, out, args, hooks), "killed": killed} + + +def _profile(inputs: dict, out: Path, args: dict, hooks: dict) -> dict: from lab import proftools, tools as T tree, kind, a = inputs["tree"], args["kind"], args["args"] files = out / "files" diff --git a/tests/test_lab_agent.py b/tests/test_lab_agent.py new file mode 100644 index 0000000..1957bda --- /dev/null +++ b/tests/test_lab_agent.py @@ -0,0 +1,36 @@ +"""lab.agent's CLI wrapper: the model credential never enters the jail; the CLI is pointed at the API proxy.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from lab import agent +from lab.safety import container + + +def spec(tmp_path: Path) -> agent.AgentSpec: + (tmp_path / "ws").mkdir() + return agent.AgentSpec(system="", prompt="", workspace=tmp_path / "ws", scratch=tmp_path / "scratch", tools=[], + workbench="lab-workbench-r1-s1") + + +@pytest.mark.parametrize("containerized", [True, False]) +def test_the_wrapper_holds_no_credential_and_points_at_the_proxy(tmp_path, monkeypatch, containerized): + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-secret-key") + monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oat-secret-token") + monkeypatch.setenv("LAB_NO_JAIL", "1") + monkeypatch.setattr(agent.ClaudeAgentSDK, "containerized", staticmethod(lambda: containerized)) + monkeypatch.setattr(container, "_docker", lambda: "docker") + monkeypatch.setattr(agent.jail, "available", lambda: False) + s = spec(tmp_path) + s.scratch.mkdir() + script = agent.ClaudeAgentSDK()._wrapper(s, tmp_path / "bin" / "claude", "http://172.17.0.1:4321").read_text() + assert "sk-secret" not in script and "oat-secret" not in script + assert "ANTHROPIC_BASE_URL=http://172.17.0.1:4321" in script and f"ANTHROPIC_API_KEY={agent.apiproxy.PLACEHOLDER}" in script + if containerized: + assert "docker run -i" in script and "--name lab-workbench-r1-s1" in script + assert "--env CLAUDE_CODE_ENTRYPOINT" in script # the SDK's value, passed through by name + else: + assert script.splitlines()[1].startswith("exec env -i") diff --git a/tests/test_lab_apiproxy.py b/tests/test_lab_apiproxy.py new file mode 100644 index 0000000..bc03393 --- /dev/null +++ b/tests/test_lab_apiproxy.py @@ -0,0 +1,129 @@ +"""lab.apiproxy against a fake Anthropic API: credential injection, usage pricing (JSON and SSE), the cap, refusals.""" + +from __future__ import annotations + +import http.client +import json +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from urllib.parse import urlsplit + +import pytest + +from lab import apiproxy + +USAGE = {"input_tokens": 1000, "output_tokens": 200, "cache_read_input_tokens": 5000, + "cache_creation_input_tokens": 300, + "cache_creation": {"ephemeral_5m_input_tokens": 100, "ephemeral_1h_input_tokens": 200}} +SSE = (b'event: message_start\ndata: {"type":"message_start","message":{"model":"claude-opus-5-5","usage":' + + json.dumps({**USAGE, "output_tokens": 1}).encode() + b'}}\n\n' + b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n' + b'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":200}}\n\n' + b'event: message_stop\ndata: {"type":"message_stop"}\n\n') + + +@pytest.fixture +def upstream(): + seen = [] + + class H(BaseHTTPRequestHandler): + def log_message(self, *a): + pass + + def do_POST(self): + body = json.loads(self.rfile.read(int(self.headers["content-length"]))) + seen.append({"path": self.path, "headers": {k.lower(): v for k, v in self.headers.items()}, "body": body}) + out, kind = ((SSE, "text/event-stream") if body.get("stream") else + (json.dumps({"model": body["model"], "usage": USAGE}).encode(), "application/json")) + self.send_response(200) + self.send_header("content-type", kind) + self.send_header("content-length", str(len(out))) + self.end_headers() + self.wfile.write(out) + + srv = ThreadingHTTPServer(("127.0.0.1", 0), H) + threading.Thread(target=srv.serve_forever, daemon=True).start() + yield f"http://127.0.0.1:{srv.server_address[1]}", seen + srv.shutdown() + + +def post(proxy, path, body, headers=None): + c = http.client.HTTPConnection(urlsplit(proxy.url).netloc, timeout=10) + c.request("POST", path, body=json.dumps(body), headers={"x-api-key": apiproxy.PLACEHOLDER, + "content-type": "application/json", **(headers or {})}) + r = c.getresponse() + return r.status, r.read() + + +def expected(model: str) -> float: + p_in, p_out, p_read, p_5m, p_1h = apiproxy.PRICES[model] + return (1000 * p_in + 200 * p_out + 5000 * p_read + 100 * p_5m + 200 * p_1h) / 1e6 + + +def test_the_lab_key_replaces_the_jails_placeholder(upstream, monkeypatch): + url, seen = upstream + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-real") + with apiproxy.Proxy(5.0, upstream=url) as p: + status, _ = post(p, "/v1/messages?beta=true", {"model": "claude-opus-5-5", "max_tokens": 10}) + assert status == 200 and seen[0]["headers"]["x-api-key"] == "sk-real" + assert seen[0]["path"] == "/v1/messages?beta=true" and "authorization" not in seen[0]["headers"] + + +def test_an_oauth_token_goes_as_bearer_with_its_beta(upstream, monkeypatch): + url, seen = upstream + monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False) + monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oat-real") + with apiproxy.Proxy(5.0, upstream=url) as p: + post(p, "/v1/messages", {"model": "claude-opus-5-5"}, {"anthropic-beta": "fast-x"}) + h = seen[0]["headers"] + assert h["authorization"] == "Bearer oat-real" and "x-api-key" not in h + assert set(h["anthropic-beta"].split(",")) == {"fast-x", apiproxy.OAUTH_BETA} + + +def test_json_and_streamed_usage_are_priced_and_the_stream_passes_through(upstream, monkeypatch): + url, _ = upstream + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-real") + with apiproxy.Proxy(5.0, upstream=url) as p: + post(p, "/v1/messages", {"model": "claude-haiku-4-5"}) + status, body = post(p, "/v1/messages", {"model": "claude-opus-5-5", "stream": True}) + assert status == 200 and body == SSE + assert p.spent_usd == pytest.approx(expected("claude-haiku-4-5") + expected("claude-opus-5-5")) + assert p.requests == 2 and p.usage["output_tokens"] == 400 and p.usage["cache_read_input_tokens"] == 10000 + + +def test_a_spent_cap_refuses_the_next_request(upstream, monkeypatch): + url, seen = upstream + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-real") + with apiproxy.Proxy(expected("claude-opus-5-5") / 2, upstream=url) as p: + assert post(p, "/v1/messages", {"model": "claude-opus-5-5"})[0] == 200 + status, body = post(p, "/v1/messages", {"model": "claude-opus-5-5"}) + assert status == 403 and b"is spent" in body and len(seen) == 1 + + +@pytest.mark.parametrize("path,body,why", [ + ("/v1/messages", {"model": "some-new-model"}, b"no price"), + ("/v1/messages", {"model": "claude-opus-5-5", "speed": "fast"}, b"fast mode"), + ("/v1/messages/batches", {"requests": []}, b"not forwarded"), + ("/v1/files", {}, b"not forwarded"), +]) +def test_anything_unpriced_is_refused_and_never_sent(upstream, monkeypatch, path, body, why): + url, seen = upstream + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-real") + with apiproxy.Proxy(5.0, upstream=url) as p: + status, out = post(p, path, body) + assert status in (403, 400) and why in out and seen == [] + + +def test_count_tokens_is_forwarded_and_costs_nothing(upstream, monkeypatch): + url, seen = upstream + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-real") + with apiproxy.Proxy(5.0, upstream=url) as p: + assert post(p, "/v1/messages/count_tokens", {"model": "claude-opus-5-5"})[0] == 200 + assert len(seen) == 1 and p.spent_usd == 0 + + +def test_no_credential_means_no_proxy(monkeypatch): + monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False) + monkeypatch.delenv("CLAUDE_CODE_OAUTH_TOKEN", raising=False) + with pytest.raises(RuntimeError): + apiproxy.Proxy(5.0) diff --git a/tests/test_lab_container.py b/tests/test_lab_container.py index 72b8457..d65078c 100644 --- a/tests/test_lab_container.py +++ b/tests/test_lab_container.py @@ -141,3 +141,70 @@ def test_served_port_reaches_the_host_and_nothing_else(tmp_path, monkeypatch): assert left.stdout.strip() == "" with pytest.raises(OSError): socket.create_connection(("127.0.0.1", port), timeout=2) + + +def test_the_workbench_has_network_gpu_and_root_inside(docker_named, tmp_path, monkeypatch): + monkeypatch.setattr(container.os, "getuid", lambda: 0) + monkeypatch.setattr(container, "gpus", lambda: True) + a = container.workbench_argv("wb", tmp_path, ["claude"], {"K": "v"}, [tmp_path], [], sys.executable) + s = " ".join(a) + assert "--user 0:0" in s and "--gpus all" in s and "-i" in a and a[-1] == "claude" + for absent in ("--network", "--read-only", "--cap-drop"): + assert absent not in s + + +def test_quiet_kills_the_workbenchs_gpu_holders_then_pauses_and_thaws(monkeypatch): + from lab import gpu + calls, killed = [], [] + monkeypatch.setattr(container, "_docker", lambda: "docker") + monkeypatch.setattr(container, "running", lambda n: True) + monkeypatch.setattr(container, "pids", lambda n: {101, 102}) + apps = [[{"pid": 101, "used_mib": 9000, "name": "python"}, {"pid": 999, "used_mib": 1, "name": "other"}]] + monkeypatch.setattr(gpu, "compute_apps", lambda: apps[0]) + monkeypatch.setattr(container.os, "kill", lambda pid, sig: killed.append(pid) or apps.__setitem__(0, [apps[0][1]])) + monkeypatch.setattr(container.subprocess, "run", lambda argv, **k: calls.append(argv[1:]) or subprocess.CompletedProcess(argv, 0)) + with container.quiet("wb") as got: + assert calls == [["pause", "wb"]] + assert killed == [101] and [a["pid"] for a in got] == [101] # 999 is not the workbench's + assert calls == [["pause", "wb"], ["unpause", "wb"]] + + +def test_quiet_is_a_no_op_without_a_running_workbench(monkeypatch): + monkeypatch.setattr(container, "_docker", lambda: "docker") + monkeypatch.setattr(container, "running", lambda n: False) + with container.quiet("wb") as got: + assert got == [] + with container.quiet(None) as got: + assert got == [] + + +@live +def test_a_paused_workbench_is_frozen_and_thawed_after(tmp_path): + name = container.new_name() + subprocess.run(["docker", "run", "-d", "--rm", "--name", name, container.IMAGE, "sleep", "60"], check=True, + capture_output=True) + state = lambda: subprocess.run(["docker", "inspect", "-f", "{{.State.Paused}}", name], capture_output=True, # noqa: E731 + text=True).stdout.strip() + try: + with container.quiet(name): + assert state() == "true" + assert state() == "false" + finally: + container.remove(name) + + +@live +def test_the_workbench_is_root_with_writable_system_dirs_and_reaches_the_host_gateway(tmp_path): + gw = container.bridge_gateway() + srv = socket.socket() + srv.bind((gw, 0)) + srv.listen() + port = srv.getsockname()[1] + probe = ("import os, socket; open('/usr/local/x', 'w').write('y'); " + f"socket.create_connection(('{gw}', {port}), timeout=5); print(os.getuid())") + argv = container.workbench_argv(container.new_name(), tmp_path, ["python3", "-c", probe], {}, [tmp_path], [], + sys.executable) + out = subprocess.run(argv, capture_output=True, text=True, timeout=60, stdin=subprocess.DEVNULL) + srv.close() + assert out.returncode == 0, out.stderr + assert out.stdout.strip() == ("0" if os.getuid() == 0 else str(os.getuid())) diff --git a/tests/test_lab_evaltools.py b/tests/test_lab_evaltools.py index fbd9e04..84fae0e 100644 --- a/tests/test_lab_evaltools.py +++ b/tests/test_lab_evaltools.py @@ -259,6 +259,24 @@ def test_tools_refuse_cleanly_when_the_engine_does_not_start(lab, monkeypatch): assert rec["result"]["artifacts"]["telemetry"] is None and rec["result"]["artifacts"]["serve_log"] # no section: skipped +def test_the_workbench_is_quieted_for_a_measurement_and_what_was_killed_is_a_fact(lab, monkeypatch): + from contextlib import contextmanager + from lab.safety import container + tb, s = lab + s.workbench = "lab-workbench-r1-s1" + asked = [] + + @contextmanager + def quiet(name): + asked.append(name) + yield [{"pid": 4242, "used_mib": 9000, "name": "python"}] + monkeypatch.setattr(container, "quiet", quiet) + out = json.loads(tb.bench({"regimes": ["single_stream"]})) + assert asked == ["lab-workbench-r1-s1"] and out["workbench_killed"][0]["pid"] == 4242 + rec = list(ledger.records(s.ledger_root, kind="bench"))[-1]["result"] + assert rec["verdict"] == "ok" and rec["workbench_killed"][0]["pid"] == 4242 + + AGENT_PROC = {"pid": 4242, "used_mib": 9000, "name": "python"} From 5c05798608add26aee311fd13841264cb2e77ff8 Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Sat, 10 Oct 2026 00:23:07 -0400 Subject: [PATCH 07/18] agent: the CLI's scratch dir is /tmp off macOS (srt branch) Co-Authored-By: Claude Opus 5.5 --- lab/agent.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/lab/agent.py b/lab/agent.py index 80165f6..35024ae 100644 --- a/lab/agent.py +++ b/lab/agent.py @@ -5,6 +5,7 @@ import asyncio import os import shlex +import sys from dataclasses import dataclass, field from pathlib import Path from typing import Any, Callable, Protocol @@ -100,7 +101,8 @@ def _wrapper(self, spec: AgentSpec, cli: Path, base_url: str) -> Path: argv[at:at] = [w for k in PASS_ENV for w in ("--env", k)] # the SDK's values, passed through line = f"exec {shlex.join(argv)} \"$@\"" else: - cli_tmp = Path(f"/private/tmp/claude-{os.getuid()}") # the CLI's own per-project scratch; Bash fails without it + tmp_root = "/private/tmp" if sys.platform == "darwin" else "/tmp" + cli_tmp = Path(f"{tmp_root}/claude-{os.getuid()}") # the CLI's own per-project scratch; Bash fails without it cli_tmp.mkdir(parents=True, exist_ok=True) config = jail.settings([home.resolve(), tmp.resolve(), ws, cli_tmp.resolve()], engine.venv(), [base_url.removeprefix("http://")], readonly=weights, python=engine.python()) From 28b8c5ed6afb967fbade6c7d71efa97aa03c59b4 Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Sat, 10 Oct 2026 12:59:17 -0400 Subject: [PATCH 08/18] lab: the agent on a remote worker (CLI over SSH, workspace mirror, proxy tunnel) With LAB_WORKER=ssh the agent works on the VM while the loop, ledger, referee, budget and API proxy stay on the controller. The workspace's live copy is on the VM, mirrored by rsync by checksum: pushed before each session and after restore or a bundle copy-in, pulled before every lab tool's audit and after the session (links stay links; the audit refuses them). The VM builds its own workbench `docker run` with the Linux CLI of the controller's Agent SDK version; cli.sh is `ssh -T -R :127.0.0.1: vm ...`, so the CLI's stdio is the SDK's and lab tools stay in process here. The proxy reaches the workbench through a unix socket and a socat on the container's own loopback; it listens on no interface of the VM. Co-Authored-By: Claude Opus 5.5 --- CLAUDE.md | 4 +- lab/RUNTIME.md | 25 +++++++++- lab/agent.py | 92 +++++++++++++++++++++++++++---------- lab/session.py | 19 ++++++-- lab/tools.py | 10 ++++ lab/worker.py | 83 +++++++++++++++++++++++++++++++-- tests/test_lab_agent.py | 33 +++++++++++++ tests/test_lab_evaltools.py | 22 +++++++++ tests/test_lab_worker.py | 32 +++++++++++++ 9 files changed, 284 insertions(+), 36 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 278458c..0d464db 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -69,7 +69,9 @@ pytest/ruff config). Each folder with commands has its own `python -m`. - `lab/worker.py` — the job interface, built for the measured jobs: `test`, `measure`, `equiv`, `profile` take input trees + JSON args, write an out dir, return JSON; the tools call `self.worker`. `LAB_WORKER=local` (default, in process) or `ssh` (the `lab.vm` VM): inputs by content hash, re-hashed on the worker before every job; GpuBusy, - Contaminated, NotReady, ValueError cross intact. Remote is tested through a local shell, not yet over SSH. + Contaminated, NotReady, ValueError cross intact. Remote is tested through a local shell, not yet over SSH. With + `ssh` the agent works on the VM too: its workspace mirrored by rsync (pulled before every audit), its workbench + built there (`workbench`), its CLI's stdio over `ssh -R`, the API proxy tunnelled to a unix socket in the workbench. - `lab/TRACE.md` + `lab/trace_contract.py` — the request-trace contract: per-request engine rows (trace_id = the client's `X-Trace-Id`, `arrival_ts` epoch s, `*_s` monotonic durations, TPOT = (E2E − TTFT)/(n − 1)) in a dir of `*.sqlite` (`requests` + `meta`) or `*.jsonl` + `.meta.json`, with schema version and drop counts; diff --git a/lab/RUNTIME.md b/lab/RUNTIME.md index 4bbd0be..0489eb4 100644 --- a/lab/RUNTIME.md +++ b/lab/RUNTIME.md @@ -4,8 +4,8 @@ Design, 2026-10-08. Decided with the user: the lab splits into a controller and measurement VM (cost); NVIDIA only for now; the jail on Linux is a container with the NVIDIA runtime (option A below). Built: the container jail (`lab/safety/container.py`), GPU exclusivity during measurements, and the measured jobs behind the worker interface (`lab/worker.py`), the API proxy (`lab/apiproxy.py`), the workbench -container with its pause during measurements. Not built: the agent on a *remote* worker (its CLI over SSH, the -workspace jobs, the proxy through a tunnel) and the dependency build step. Open decisions are marked **Decide**. +container with its pause during measurements, and the agent on a remote worker. Not built: the dependency build +step. Open decisions are marked **Decide**. ## The line (the user, 2026-10-09) @@ -189,6 +189,27 @@ not yet over SSH to a VM. `agent`, `exec`, `snapshot`, `restore` and `facts` com the macOS srt path (srt is not installed on the dev box). Tested in Docker-in-Docker: workbench argv, root and writable system dirs, reaching the gateway, pause and thaw; the proxy against a fake API. +## The agent on a remote worker, as built (2026-10-10) + +With `LAB_WORKER=ssh` the session loop, ledger, referee, budget and proxy stay on the controller and the agent works +on the VM: + +- **Workspace**: the live copy is on the VM (`/lab/runs/remote//workspace`); the controller's is its + mirror. rsync by checksum (`--delete`): pushed before each session and after `restore` or a profile bundle copied + in, pulled before every lab tool's snapshot and audit and after the session. Pulled files are untrusted: links stay + links, and the audit refuses them as before. +- **Workbench**: `python -m lab.worker workbench` on the VM builds that host's `docker run` (its GPU, paths, cores) + with the Linux CLI of the controller's Agent SDK version (installed once with uv into the worker cache). The + controller's `cli.sh` is `ssh -T -R :127.0.0.1: vm '' `: the CLI's stdio + is the SDK's over SSH, so the lab tools stay an in-process MCP server on the controller. +- **Proxy**: on the controller, bound to its loopback. SSH forwards it to a unix socket in `/tmp/lab-proxy-/` + on the VM, mounted into the workbench, where a socat serves it on the container's own 127.0.0.1:4000. It listens on + no interface of the VM. +- The write guard checks the VM's workspace path; the workbench is removed over SSH when the session ends. +- Tested without a VM: rsync mirroring (content edits, deletions, links), the VM-side `workbench` command, and the + argument quoting through bash and ssh (an `ssh` stand-in that prints the remote command). Not yet tested against a + real VM. + ## What moves where - Controller: `session`, `agent` (SDK side), the tool front ends in `tools`/`evaltools`/`proftools` (policy, diff --git a/lab/agent.py b/lab/agent.py index 35024ae..3857f52 100644 --- a/lab/agent.py +++ b/lab/agent.py @@ -9,6 +9,7 @@ from dataclasses import dataclass, field from pathlib import Path from typing import Any, Callable, Protocol +from urllib.parse import urlsplit from lab import apiproxy, engine from lab.safety import container, jail @@ -47,6 +48,8 @@ class AgentSpec: timeout_s: float = 3600.0 model: str = field(default_factory=lambda: os.environ.get("LAB_MODEL", "claude-fable-5-1")) workbench: str | None = None # the container name of the agent's workbench (Linux) + remote: Any = None # a lab.worker.Remote: the workbench runs on that host + remote_workspace: str | None = None # the workspace's path there @dataclass @@ -72,6 +75,44 @@ def load(name: str | None = None) -> Provider: raise ValueError(f"unknown agent provider {name!r}; LAB_AGENT_PROVIDER is claude") +PROXY_PORT_IN = 4000 # where a remote workbench serves the tunnelled proxy, inside its own network + + +def _scratch(scratch: Path) -> tuple[Path, Path]: + home, tmp = scratch / "home", scratch / "tmp" + home.mkdir(parents=True, exist_ok=True) + tmp.mkdir(exist_ok=True) + return home, tmp + + +def _weights() -> list[Path]: + return [HF_HUB] if HF_HUB.exists() else [] # model weights read-only; the token beside them stays hidden + + +def cli_env(home: Path, tmp: Path, ws: Path, base_url: str) -> dict[str, str]: + """The CLI's whole environment inside its jail: no credential, only the proxy's URL and its placeholder.""" + return {"HOME": str(home), "CLAUDE_CONFIG_DIR": str(home), "TMPDIR": str(tmp), "PYTHONPATH": str(ws / "src"), + "PYTHONDONTWRITEBYTECODE": "1", "HF_HUB_CACHE": str(HF_HUB), "HF_HUB_OFFLINE": "1", + "HF_HOME": str(tmp / "hf-home"), "ANTHROPIC_BASE_URL": base_url, + "ANTHROPIC_API_KEY": apiproxy.PLACEHOLDER, "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1"} + + +def workbench_command(name: str, ws: Path, scratch: Path, cli: Path, base_url: str | None = None, + sock_dir: Path | None = None) -> list[str]: + """`docker run` for the CLI's workbench on this host; the SDK's arguments go after it. With `sock_dir`, the proxy + is the unix socket `proxy.sock` there (tunnelled from the controller), served inside on 127.0.0.1:PROXY_PORT_IN.""" + home, tmp = _scratch(scratch) + ws = Path(ws).resolve() + env = cli_env(home, tmp, ws, base_url or f"http://127.0.0.1:{PROXY_PORT_IN}") + env["PATH"] = f"{Path(engine.python()).parent}:{SYSTEM_PATH}" # the agent's `python` is the engine's + command = [str(cli)] if sock_dir is None else [ + "sh", "-c", f'socat TCP-LISTEN:{PROXY_PORT_IN},bind=127.0.0.1,fork,reuseaddr UNIX-CONNECT:{sock_dir}/proxy.sock ' + '& exec "$0" "$@"', str(cli)] + writable = [home.resolve(), tmp.resolve(), ws] + ([Path(sock_dir)] if sock_dir else []) + return container.workbench_argv(name, ws, command, env, writable, [engine.venv(), *_weights(), cli.parent], + engine.python()) + + class ClaudeAgentSDK: """Claude through the Agent SDK CLI, jailed with a wiped env (on Linux in the workbench container, elsewhere under srt); its model calls go through the lab's API proxy, so no credential is ever inside; lab tools run out here.""" @@ -83,35 +124,33 @@ def containerized() -> bool: return jail.backend() == "container" and container.available() def _wrapper(self, spec: AgentSpec, cli: Path, base_url: str) -> Path: - home, tmp = spec.scratch / "home", spec.scratch / "tmp" - home.mkdir(parents=True, exist_ok=True) - tmp.mkdir(exist_ok=True) - ws = spec.workspace.resolve() - weights = [HF_HUB] if HF_HUB.exists() else [] # model weights read-only; the token beside them stays hidden - env = {"HOME": str(home), "CLAUDE_CONFIG_DIR": str(home), "TMPDIR": str(tmp), "PYTHONPATH": str(ws / "src"), - "PYTHONDONTWRITEBYTECODE": "1", "HF_HUB_CACHE": str(HF_HUB), "HF_HUB_OFFLINE": "1", - "HF_HOME": str(tmp / "hf-home"), "ANTHROPIC_BASE_URL": base_url, - "ANTHROPIC_API_KEY": apiproxy.PLACEHOLDER, "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1"} - engine_bin = Path(engine.python()).parent # the agent's `python` is the engine's - if self.containerized(): - env["PATH"] = f"{engine_bin}:{SYSTEM_PATH}" - argv = container.workbench_argv(spec.workbench, ws, [str(cli)], env, [home.resolve(), tmp.resolve(), ws], - [engine.venv(), *weights, cli.parent], engine.python()) + script = spec.scratch / "cli.sh" + if spec.remote is not None: # the worker host builds its own workbench; the proxy goes down the tunnel + import claude_agent_sdk + w = spec.remote.workbench(spec.workbench, spec.scratch.name, spec.remote_workspace, + claude_agent_sdk.__version__) + *ssh, target = spec.remote.prefix + line = (f'args=$(printf " %q" "$@")\ncmd={shlex.quote(shlex.join(w["argv"]))}\n' + f'exec {shlex.join(ssh)} -T -o ExitOnForwardFailure=yes -o StreamLocalBindUnlink=yes ' + f'-R {w["sock"]}:127.0.0.1:{urlsplit(base_url).port} {target} "$cmd$args"') + elif self.containerized(): + argv = workbench_command(spec.workbench, spec.workspace, spec.scratch, cli, base_url=base_url) at = argv.index(container.IMAGE) argv[at:at] = [w for k in PASS_ENV for w in ("--env", k)] # the SDK's values, passed through - line = f"exec {shlex.join(argv)} \"$@\"" + line = f'exec {shlex.join(argv)} "$@"' else: + home, tmp = _scratch(spec.scratch) + ws = spec.workspace.resolve() tmp_root = "/private/tmp" if sys.platform == "darwin" else "/tmp" cli_tmp = Path(f"{tmp_root}/claude-{os.getuid()}") # the CLI's own per-project scratch; Bash fails without it cli_tmp.mkdir(parents=True, exist_ok=True) config = jail.settings([home.resolve(), tmp.resolve(), ws, cli_tmp.resolve()], engine.venv(), - [base_url.removeprefix("http://")], readonly=weights, python=engine.python()) + [base_url.removeprefix("http://")], readonly=_weights(), python=engine.python()) argv = jail.wrap(config, spec.scratch / "srt.json", [str(cli)]) keep = " ".join(f'"{k}=${k}"' for k in PASS_ENV) - sets = " ".join(f"{k}={shlex.quote(v)}" for k, v in env.items()) - line = f'exec env -i {keep} "PATH={engine_bin}:$PATH" {sets} {shlex.join(argv)} "$@"' - script = spec.scratch / "cli.sh" - script.write_text(f"#!/bin/sh\n{line}\n") + sets = " ".join(f"{k}={shlex.quote(v)}" for k, v in cli_env(home, tmp, ws, base_url).items()) + line = f'exec env -i {keep} "PATH={Path(engine.python()).parent}:$PATH" {sets} {shlex.join(argv)} "$@"' + script.write_text(f"#!/bin/bash\n{line}\n") script.chmod(0o700) return script @@ -139,7 +178,8 @@ async def call(args: dict[str, Any]) -> dict[str, Any]: output_format={"type": "json_schema", "schema": OUTPUT_SCHEMA}, cli_path=str(self._wrapper(spec, cli, base_url)), hooks={"PreToolUse": [HookMatcher(matcher=WRITE_TOOLS, - hooks=[write_guard(spec.workspace, spec.base_files)])]}, + hooks=[write_guard(Path(spec.remote_workspace or spec.workspace), + spec.base_files)])]}, ) reply = AgentReply(None, 0.0, 0, "no result message") async for msg in query(prompt=spec.prompt, options=options): @@ -153,7 +193,7 @@ def run(self, spec: AgentSpec) -> AgentReply: """Never raises. The spend is what the API proxy priced, whatever the CLI reports or however it ended.""" spec.scratch.mkdir(parents=True, exist_ok=True) try: - host = container.bridge_gateway() if self.containerized() else "127.0.0.1" + host = container.bridge_gateway() if spec.remote is None and self.containerized() else "127.0.0.1" proxy = apiproxy.Proxy(spec.max_budget_usd, host) except Exception as e: # nothing was sent: nothing spent return AgentReply(None, 0.0, 0, f"provider failed: {type(e).__name__}: {e}") @@ -164,8 +204,10 @@ def run(self, spec: AgentSpec) -> AgentReply: reply = AgentReply(None, 0.0, 0, f"timed out after {spec.timeout_s:.0f}s") except Exception as e: reply = AgentReply(None, 0.0, 0, f"provider failed: {type(e).__name__}: {e}") - finally: - if spec.workbench and self.containerized(): - container.remove(spec.workbench) # a CLI killed mid-session leaves its container behind + finally: # a CLI killed mid-session leaves its container behind + if spec.remote is not None: + spec.remote.remove_workbench(spec.workbench) + elif spec.workbench and self.containerized(): + container.remove(spec.workbench) reply.cost_usd, reply.cost_estimated, reply.usage = proxy.spent_usd, False, dict(proxy.usage) return reply diff --git a/lab/session.py b/lab/session.py index ac3de12..4978cd7 100644 --- a/lab/session.py +++ b/lab/session.py @@ -11,7 +11,7 @@ from dataclasses import dataclass, field from pathlib import Path -from lab import agent, engine, ledger, target +from lab import agent, engine, ledger, target, worker from lab.budget import Budget, BudgetExceeded from lab.safety import grader from lab.task import Task @@ -34,6 +34,8 @@ class Session: profile_runner: object = None # None: lab.tools.default_profile_runner task: Task | None = None workbench: str | None = None # the agent's container: paused while the tools measure + worker: object = None # None: lab.worker.load() per toolbox + remote_workspace: str | None = None # the workspace's mirror on a Remote worker, where the agent works @dataclass @@ -122,27 +124,36 @@ def run(cfg: RunConfig, provider: agent.Provider | None = None) -> dict: done = sum(1 for _ in ledger.records(cfg.ledger_root, run=run_id, kind="session")) base_files = grader.base_files(cfg.repo, base) summary = {"run": run_id, "base": base, "sessions": 0, "stopped": None, "spent_usd": budget.spent_usd} + wk = worker.load() + remote = wk if isinstance(wk, worker.Remote) else None # then the agent works on that host, not here + rws = remote.workspace_path(run_id) if remote else None if cfg.baseline: - Toolbox(Session(run_id, f"{run_id}-baseline", run_dir, ws, budget, cfg.ledger_root, task=cfg.task)).baseline(base) + Toolbox(Session(run_id, f"{run_id}-baseline", run_dir, ws, budget, cfg.ledger_root, task=cfg.task, + worker=remote)).baseline(base) for n in range(done + 1, done + cfg.max_sessions + 1): if budget.remaining_usd < MIN_SESSION_USD: summary["stopped"] = "budget" break s = Session(run_id, f"{run_id}-s{n}", run_dir, ws, budget, cfg.ledger_root, task=cfg.task, - workbench=f"lab-workbench-{run_id}-s{n}") + workbench=f"lab-workbench-{run_id}-s{n}", worker=remote, remote_workspace=rws) tools = Toolbox(s) snap = ws.snapshot() spec = agent.AgentSpec(system="", prompt=brief(cfg, s, snap.id, n), workspace=ws.path, scratch=run_dir / "scratch" / s.session_id, tools=tools.specs(), base_files=base_files, max_turns=cfg.max_turns, - max_budget_usd=min(cfg.session_usd, budget.remaining_usd), workbench=s.workbench) + max_budget_usd=min(cfg.session_usd, budget.remaining_usd), workbench=s.workbench, + remote=remote, remote_workspace=rws) spec.system = system_prompt(target.load(), spec) t0 = time.monotonic() try: + if remote: + remote.push(ws.path, rws) reply = provider.run(spec) except Exception as e: # a provider that raises still gets charged and recorded reply = agent.AgentReply(None, spec.max_budget_usd, 0, f"provider raised: {e}", True) + if remote: + remote.pull(rws, ws.path) # what the agent left on the worker is what is snapshotted clean_pristine(run_dir) try: budget.charge(reply.cost_usd, s.session_id) diff --git a/lab/tools.py b/lab/tools.py index 8943cb7..8277829 100644 --- a/lab/tools.py +++ b/lab/tools.py @@ -73,6 +73,7 @@ def _audited(self, tool: str, args: dict): self._t0 = time.monotonic() if tool in GPU_TOOLS and self.s.budget.remaining_usd <= 0 and gpu_rate()[0] > 0: raise BudgetExceeded(f"{tool}: no budget left for GPU time") + self._sync("down") snap = self.s.workspace.snapshot() if snap.violations: self.violation = "; ".join(snap.violations) @@ -80,6 +81,13 @@ def _audited(self, tool: str, args: dict): raise PermissionError(f"workspace violates the write surface: {self.violation}") return snap + def _sync(self, way: str) -> None: + """With the agent on a Remote worker: its workspace there is the live one, this one its mirror.""" + rws = getattr(self.s, "remote_workspace", None) + if rws: + (self.worker.pull(rws, self.s.workspace.path) if way == "down" else + self.worker.push(self.s.workspace.path, rws)) + def _pristine(self) -> Path: a = self.s.workspace.audit() dest = self.s.run_dir / "pristine" @@ -95,6 +103,7 @@ def _keep(self, src: Path) -> tuple[str, str]: dest = self.s.workspace.path / visible shutil.rmtree(dest, ignore_errors=True) shutil.copytree(src, dest) + self._sync("up") return blob, str(visible) # -- tools --------------------------------------------------------------------------- @@ -148,6 +157,7 @@ def restore(self, args: dict) -> str: sid = args["snapshot"] self._audited("restore", args) self.s.workspace.restore(sid) + self._sync("up") snap = self._audited("restore", args) self._record("note", "restore", args, {"restored": sid, "now": snap.id}, snap) return f"workspace restored to {sid}" diff --git a/lab/worker.py b/lab/worker.py index c9407fd..442dde5 100644 --- a/lab/worker.py +++ b/lab/worker.py @@ -11,6 +11,11 @@ profile lab.profile under torch, nsys, ncu or py-spy on `tree` -> {returncode, output, ...}; files in out python -m lab.worker has SHA... | put SHA < tar | run JOB < json > tar (the worker host's side) + python -m lab.worker root | mkdir PATH | workbench < json > json | rm NAME (the agent's side) + +When the agent runs on the worker host too (Remote), its workspace lives there: `push` and `pull` mirror it with +rsync (pulled files are untrusted: links stay links and the audit refuses them), and `workbench` builds that host's +`docker run` for the CLI, with the Linux CLI of the controller's Agent SDK version. """ from __future__ import annotations @@ -20,6 +25,7 @@ import io import json import os +import shlex import shutil import subprocess import sys @@ -198,8 +204,45 @@ class Remote: """Jobs on another host. `prefix` is the argv that runs one shell command there (ssh ... user@host); empty runs it here through sh, which the tests use. `env_dir` is this repo on that host, `python` its interpreter.""" - def __init__(self, prefix: list[str], env_dir: str, python: str, env: dict[str, str]): - self.prefix, self.env_dir, self.python, self.env = prefix, env_dir, python, env + def __init__(self, prefix: list[str], env_dir: str, python: str, env: dict[str, str], rsh: str | None = None): + self.prefix, self.env_dir, self.python, self.env, self.rsh = prefix, env_dir, python, env, rsh + self._root: str | None = None + + def _where(self, path: str) -> str: + return f"{self.prefix[-1]}:{path}" if self.prefix else path + + def _rsync(self, src: str, dst: str) -> None: + """By content, not size and mtime: an edit within the same second must not be missed.""" + p = subprocess.run(["rsync", "-a", "--checksum", "--delete", *(["-e", self.rsh] if self.rsh else []), src, dst], + capture_output=True, text=True) + if p.returncode != 0: + raise WorkerError(f"rsync {src} -> {dst}: {p.stderr.strip()[-1000:]}") + + def root(self) -> str: + """This repo's absolute path on the worker host.""" + if self._root is None: + self._root = self._run("root").decode().strip() + return self._root + + def workspace_path(self, run_id: str) -> str: + return f"{self.root()}/lab/runs/remote/{run_id}/workspace" + + def push(self, local: Path, remote: str) -> None: + self._run(f"mkdir {shlex.quote(remote)}") + self._rsync(f"{local}/", self._where(f"{remote}/")) + + def pull(self, remote: str, local: Path) -> None: + Path(local).mkdir(parents=True, exist_ok=True) + self._rsync(self._where(f"{remote}/"), f"{local}/") + + def workbench(self, name: str, session: str, workspace: str, sdk_version: str) -> dict: + """{argv, sock}: that host's `docker run` for the CLI, and where the proxy's tunnel must land.""" + req = {"name": name, "session": session, "workspace": workspace, "sdk_version": sdk_version} + return json.loads(self._run("workbench", json.dumps(req).encode())) + + def remove_workbench(self, name: str | None) -> None: + if name: + self._run(f"rm {shlex.quote(name)}") def _run(self, command: str, stdin: bytes = b"") -> bytes: for v in self.env.values(): @@ -243,7 +286,8 @@ def load(**hooks): raise WorkerError(f"LAB_TARGET must be inside this repo to be found on the worker: {spec}") home = lambda p: p.replace("~", "$HOME", 1) # noqa: E731 - expanded by the worker's shell return Remote(["ssh", *v.ssh_opts, vm.ssh_target(v, record)], v.env_dir, home(v.remote_dir) + "/.venv/bin/python", - {"LAB_TARGET": spec.relative_to(engine.ENV_ROOT).as_posix(), "LAB_ENGINE_REPO": home(v.remote_dir)}) + {"LAB_TARGET": spec.relative_to(engine.ENV_ROOT).as_posix(), "LAB_ENGINE_REPO": home(v.remote_dir)}, + rsh=f"ssh {shlex.join(v.ssh_opts)}") # -- the worker host's side --------------------------------------------------------------------------------------- @@ -290,6 +334,28 @@ def _run_job(job: str, request: dict) -> bytes: shutil.rmtree(work, ignore_errors=True) +def agent_cli(version: str) -> Path: + """The Agent SDK's bundled CLI for this host at `version`, installed once into the cache.""" + d = cache_dir() / f"cli-{version}" + exe = d / "claude_agent_sdk" / "_bundled" / "claude" + if not exe.exists(): + uv = shutil.which("uv") or str(Path.home() / ".local" / "bin" / "uv") + subprocess.run([uv, "pip", "install", "--quiet", "--no-deps", "--target", str(d), "--python", sys.executable, + f"claude-agent-sdk=={version}"], check=True, stdout=sys.stderr) + return exe + + +def _workbench(req: dict) -> dict: + from lab import agent + sock_dir = Path(f"/tmp/lab-proxy-{req['session']}") + sock_dir.mkdir(mode=0o700, exist_ok=True) + (sock_dir / "proxy.sock").unlink(missing_ok=True) + scratch = engine.ENV_ROOT / "lab" / "runs" / "remote" / "scratch" / req["session"] + argv = agent.workbench_command(req["name"], Path(req["workspace"]), scratch, agent_cli(req["sdk_version"]), + sock_dir=sock_dir) + return {"argv": argv, "sock": str(sock_dir / "proxy.sock")} + + def main(argv: list[str] | None = None) -> int: argv = sys.argv[1:] if argv is None else argv cache_dir().mkdir(parents=True, exist_ok=True) @@ -299,8 +365,17 @@ def main(argv: list[str] | None = None) -> int: _put(argv[1], sys.stdin.buffer.read()) elif argv[:1] == ["run"] and len(argv) == 2 and argv[1] in JOBS: sys.stdout.buffer.write(_run_job(argv[1], json.loads(sys.stdin.buffer.read()))) + elif argv == ["root"]: + print(engine.ENV_ROOT) + elif argv[:1] == ["mkdir"] and len(argv) == 2: + Path(argv[1]).mkdir(parents=True, exist_ok=True) + elif argv == ["workbench"]: + print(json.dumps(_workbench(json.loads(sys.stdin.buffer.read())))) + elif argv[:1] == ["rm"] and len(argv) == 2: + container.remove(argv[1]) else: - print(f"usage: python -m lab.worker has SHA... | put SHA | run {{{','.join(JOBS)}}}", file=sys.stderr) + print(f"usage: python -m lab.worker has SHA... | put SHA | run {{{','.join(JOBS)}}} | root | mkdir PATH | " + "workbench | rm NAME", file=sys.stderr) return 2 return 0 diff --git a/tests/test_lab_agent.py b/tests/test_lab_agent.py index 1957bda..4b1fa92 100644 --- a/tests/test_lab_agent.py +++ b/tests/test_lab_agent.py @@ -34,3 +34,36 @@ def test_the_wrapper_holds_no_credential_and_points_at_the_proxy(tmp_path, monke assert "--env CLAUDE_CODE_ENTRYPOINT" in script # the SDK's value, passed through by name else: assert script.splitlines()[1].startswith("exec env -i") + + +def test_a_remote_workbench_is_launched_over_ssh_with_the_proxy_tunnelled_and_args_intact(tmp_path, monkeypatch): + import shlex + import subprocess + + class FakeRemote: + prefix = ["ssh", "-i", "key", "root@203.0.113.5"] + + def workbench(self, name, session, workspace, version): + return {"argv": ["docker", "run", "-i", "--name", name, "lab-jail", "sh", "-c", 'exec "$0" "$@"', "/c/claude"], + "sock": f"/tmp/lab-proxy-{session}/proxy.sock"} + + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-secret-key") + import sys + import types + monkeypatch.setitem(sys.modules, "claude_agent_sdk", types.SimpleNamespace(__version__="9.9.9")) + s = spec(tmp_path) + s.scratch.mkdir() + s.remote, s.remote_workspace = FakeRemote(), "/root/BlameGraph/lab/runs/remote/r1/workspace" + script = agent.ClaudeAgentSDK()._wrapper(s, Path("/unused"), "http://127.0.0.1:5555") + text = script.read_text() + assert "sk-secret" not in text + assert "-R /tmp/lab-proxy-scratch/proxy.sock:127.0.0.1:5555 root@203.0.113.5" in text + fake = tmp_path / "bin" / "ssh" # an ssh that prints the remote command it was given + fake.parent.mkdir() + fake.write_text('#!/bin/sh\nfor a; do last="$a"; done\nprintf "%s" "$last"\n') + fake.chmod(0o755) + out = subprocess.run(["bash", str(script), "--output-format", "stream-json", "it's \"quoted\" $HOME"], + capture_output=True, text=True, env={"PATH": f"{fake.parent}:/usr/bin:/bin"}) + remote_argv = shlex.split(out.stdout) + assert remote_argv == ["docker", "run", "-i", "--name", "lab-workbench-r1-s1", "lab-jail", "sh", "-c", + 'exec "$0" "$@"', "/c/claude", "--output-format", "stream-json", "it's \"quoted\" $HOME"] diff --git a/tests/test_lab_evaltools.py b/tests/test_lab_evaltools.py index 84fae0e..4dc8afc 100644 --- a/tests/test_lab_evaltools.py +++ b/tests/test_lab_evaltools.py @@ -387,3 +387,25 @@ def test_bench_defaults_follow_the_task(): tb.s.task = Task("g", ("long_prompt_short_output",), constraints={"single_stream": {"max_regression_pct": 5}, "long_prompt_short_output": {"max_regression_pct": 1}}) assert tb.bench_defaults() == ("long_prompt_short_output", "single_stream") + + +def test_with_the_agent_on_a_remote_worker_its_workspace_is_pulled_before_every_audit_and_pushed_after_a_restore(lab): + tb, s = lab + moves = [] + local = tb.worker + + class Mirror: + def call(self, *a, **k): + return local.call(*a, **k) + + def pull(self, remote, path): + moves.append(("pull", remote)) + + def push(self, path, remote): + moves.append(("push", remote)) + tb.worker, s.remote_workspace = Mirror(), "/far/ws" + tb.test({}) + assert moves == [("pull", "/far/ws")] + moves.clear() + tb.restore({"snapshot": "base"}) + assert ("push", "/far/ws") in moves and moves[0] == ("pull", "/far/ws") diff --git a/tests/test_lab_worker.py b/tests/test_lab_worker.py index ebf9f27..89fa359 100644 --- a/tests/test_lab_worker.py +++ b/tests/test_lab_worker.py @@ -103,3 +103,35 @@ def test_local_by_default_and_nothing_else_but_ssh(monkeypatch): monkeypatch.setenv("LAB_WORKER", "gpu-box") with pytest.raises(ValueError): worker.load() + + +def test_push_and_pull_mirror_the_workspace_and_keep_links_as_links(remote, tmp_path): + r, _, _ = remote + local = tmp_path / "local" + (local / "src").mkdir(parents=True) + (local / "src" / "a.py").write_text("x = 1\n") + far = str(tmp_path / "far" / "ws") + r.push(local, far) + assert (Path(far) / "src" / "a.py").read_text() == "x = 1\n" + (Path(far) / "src" / "a.py").write_text("x = 2\n") + (Path(far) / "src" / "b.py").write_text("y = 1\n") + (Path(far) / "leak").symlink_to("/etc/passwd") + r.pull(far, local) + assert (local / "src" / "a.py").read_text() == "x = 2\n" and (local / "src" / "b.py").exists() + assert (local / "leak").is_symlink() # never followed; the audit refuses it + (Path(far) / "src" / "b.py").unlink() + r.pull(far, local) + assert not (local / "src" / "b.py").exists() + + +def test_the_worker_builds_its_own_workbench_with_the_tunnelled_proxy(remote, tmp_path, monkeypatch): + r, _, _ = remote + cli = tmp_path / "cache" / "cli-9.9.9" / "claude_agent_sdk" / "_bundled" / "claude" + cli.parent.mkdir(parents=True) + cli.write_text("") # already installed: nothing is downloaded + (tmp_path / "ws").mkdir() + w = r.workbench("lab-workbench-r1-s1", "r1-s1", str(tmp_path / "ws"), "9.9.9") + argv = w["argv"] + assert argv[1:3] == ["run", "-i"] and "lab-workbench-r1-s1" in argv and str(cli) == argv[-1] + assert w["sock"] == "/tmp/lab-proxy-r1-s1/proxy.sock" and "UNIX-CONNECT:/tmp/lab-proxy-r1-s1/proxy.sock" in argv[-2] + assert "ANTHROPIC_BASE_URL=http://127.0.0.1:4000" in argv and not any("sk-" in a for a in argv) From 258cf8026f4055ffc5c447bcf09a0e88f03bce0a Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Sat, 10 Oct 2026 13:23:59 -0400 Subject: [PATCH 09/18] lab: dependency build step (deps installed from scratch, index-only network) The target names its dependency files (deps) and the command that installs them (build). The agent may change them; a pyproject.toml only in its dependency tables (checked by the audit; test and lint config stay protected). A tree whose change touched them gets a venv built for it in a no-GPU, --network none container whose only way out is a host-side CONNECT proxy to the package-index hosts, cached by the files' hash; an unchanged tree keeps the engine's venv, so base and candidate run identically. Jobs run under engine.using(build.python_for(tree)); a failed build is a recorded failure. Tool descriptions now state the fact: only the change reaches the engine, and it runs offline. Co-Authored-By: Claude Opus 5.5 --- CLAUDE.md | 4 + lab/RUNTIME.md | 22 ++++- lab/build.py | 179 ++++++++++++++++++++++++++++++++++ lab/engine.py | 16 +++ lab/safety/container.py | 4 +- lab/safety/grader.py | 11 ++- lab/safety/surfaces.py | 32 +++++- lab/target.py | 4 +- lab/tools.py | 12 ++- lab/worker.py | 29 ++++-- targets/inference-server.toml | 4 + tests/test_lab_build.py | 148 ++++++++++++++++++++++++++++ tests/test_lab_safety.py | 38 +++++++- 13 files changed, 485 insertions(+), 18 deletions(-) create mode 100644 lab/build.py create mode 100644 tests/test_lab_build.py diff --git a/CLAUDE.md b/CLAUDE.md index 0d464db..0a5a35d 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -66,6 +66,10 @@ pytest/ruff config). Each folder with commands has its own `python -m`. record-format spec: the environment's public interface. - `lab/RUNTIME.md` — design (2026-10-08): controller (laptop: ledger, referee, secrets) and worker (GPU VM: jails, clients, profilers) behind a job interface; trust zones and rules; why srt cannot jail GPU code on Linux. +- `lab/build.py` — the dependency build step: a tree whose change touched the target's `deps` (pyproject: dependency + tables only, checked by the audit) gets a venv built by the target's `build` command in a no-GPU, `--network none` + container whose only way out is a CONNECT proxy to the package-index hosts; cached by the files' hash; jobs run + under `engine.using(build.python_for(tree))`. - `lab/worker.py` — the job interface, built for the measured jobs: `test`, `measure`, `equiv`, `profile` take input trees + JSON args, write an out dir, return JSON; the tools call `self.worker`. `LAB_WORKER=local` (default, in process) or `ssh` (the `lab.vm` VM): inputs by content hash, re-hashed on the worker before every job; GpuBusy, diff --git a/lab/RUNTIME.md b/lab/RUNTIME.md index 0489eb4..086dff5 100644 --- a/lab/RUNTIME.md +++ b/lab/RUNTIME.md @@ -4,8 +4,8 @@ Design, 2026-10-08. Decided with the user: the lab splits into a controller and measurement VM (cost); NVIDIA only for now; the jail on Linux is a container with the NVIDIA runtime (option A below). Built: the container jail (`lab/safety/container.py`), GPU exclusivity during measurements, and the measured jobs behind the worker interface (`lab/worker.py`), the API proxy (`lab/apiproxy.py`), the workbench -container with its pause during measurements, and the agent on a remote worker. Not built: the dependency build -step. Open decisions are marked **Decide**. +container with its pause during measurements, the agent on a remote worker, and the dependency build step +(`lab/build.py`). Open decisions are marked **Decide**. ## The line (the user, 2026-10-09) @@ -210,6 +210,24 @@ on the VM: argument quoting through bash and ssh (an `ssh` stand-in that prints the remote command). Not yet tested against a real VM. +## The dependency build step, as built (2026-10-10) + +- The target names its dependency files (`deps`) and the command that installs them (`build`, run with + `$UV_PROJECT_ENVIRONMENT` set). For inference-server that is vm-setup.sh's recipe: `uv sync --locked`, then the + CUDA torch wheel at the locked version. +- The agent may change those files; a `pyproject.toml` only in its dependency tables (`[project] dependencies`, + `optional-dependencies`, `[dependency-groups]`, `[tool.uv] sources` and `index`): any other change, its test and + lint config included, is a violation, and deleting a dependency file is one too. +- Each job asks `build.python_for(tree)`: a tree whose change left the dependency files as in its base (HEAD~1 of + the two-commit pristine tree) runs on the engine's venv; a changed one gets a venv built for it, cached by the + files' hash. The build runs in a measurement-room container without the GPU and with `--network none`; its only way + out is a host-side CONNECT proxy, over a unix socket, to pypi.org, files.pythonhosted.org and the https hosts its + files and command name. A failed build is a failed `test`/`profile` or an `error` record carrying the build's + output. The tool descriptions state it as a fact: only the change reaches the engine, installed from scratch, and + the engine then runs offline. +- Tested in Docker-in-Docker: the build room gets 200 from PyPI through the proxy, 403 for another host, and no + direct route out. The real inference-server build (torch, CUDA) is untested until the GPU run. + ## What moves where - Controller: `session`, `agent` (SDK side), the tool front ends in `tools`/`evaltools`/`proftools` (policy, diff --git a/lab/build.py b/lab/build.py new file mode 100644 index 0000000..8681745 --- /dev/null +++ b/lab/build.py @@ -0,0 +1,179 @@ +"""The dependency build step (lab/RUNTIME.md). A measurement room is rebuilt from the agent's change alone and runs +offline, so a change to the engine's dependency files (the target's `deps`) is installed from scratch for it: a +fresh venv built by the target's `build` command in a container with no GPU, no secrets and no network but a +CONNECT proxy to the package-index hosts, cached by the dependency files' hash. A tree whose change leaves the dependency +files as they were in its base keeps the engine's venv (built from that base by vm-setup.sh), so the base and an +unchanged candidate run identically. A measured tree is a two-commit repo (base, change; grader.pristine_tree); one +without that history, such as the exported base itself, has no change.""" + +from __future__ import annotations + +import hashlib +import os +import re +import select +import shutil +import socket +import socketserver +import subprocess +import tempfile +import threading +from pathlib import Path + +from lab import engine, target +from lab.safety import container + +INDEX_HOSTS = ("pypi.org", "files.pythonhosted.org") # plus every host the dependency files and build name +PROXY_PORT_IN = 3128 +LOG_TAIL = 4000 +TIMEOUT_S = 3600 + + +class BuildFailed(RuntimeError): + """The tree's dependencies could not be installed; the message carries the build's output.""" + + +def deps_sha(tree: Path) -> str: + h = hashlib.sha256() + for rel in target.load().deps: + p = Path(tree) / rel + h.update(rel.encode() + b"\0" + (p.read_bytes() if p.is_file() else b"") + b"\0") + return h.hexdigest() + + +def hosts(tree: Path, build: str) -> set[str]: + """The index hosts a build may reach: the defaults and every https host its dependency files and command name.""" + text = build + "".join((Path(tree) / r).read_text(errors="replace") for r in target.load().deps + if (Path(tree) / r).is_file()) + return set(INDEX_HOSTS) | {m.lower() for m in re.findall(r"https://([A-Za-z0-9.-]+)", text)} + + +def changed(tree: Path) -> bool: + """Whether the tree's change touched a dependency file: its working files against its base commit (HEAD~1).""" + tree = Path(tree) + if not target.load().deps or not (tree / ".git").exists(): + return False + for rel in target.load().deps: + old = subprocess.run(["git", "-C", str(tree), "show", f"HEAD~1:{rel}"], capture_output=True) + if old.returncode != 0 and b"exists on disk, but not in" not in old.stderr and b"does not exist" not in old.stderr: + return False # no base commit: nothing to compare against + now = (tree / rel).read_bytes() if (tree / rel).is_file() else None + if (old.stdout if old.returncode == 0 else None) != now: + return True + return False + + +def python_for(tree: Path) -> str: + """The interpreter that runs `tree`: the engine's own, or a venv built from the tree's changed dependency files.""" + if not changed(tree): + return engine.python() + venv = cache_dir() / deps_sha(tree) + if not (venv / ".lab-built").exists(): + build(tree, venv) + return str(venv / "env" / "bin" / "python") + + +def cache_dir() -> Path: + from lab import worker + return worker.cache_dir() / "venvs" + + +class _Connect(socketserver.BaseRequestHandler): + """One CONNECT host:443 to an allowed host, then bytes both ways; anything else is refused.""" + allowed: set[str] = set() + + def handle(self): + head = b"" + while b"\r\n\r\n" not in head and len(head) < 8192: + chunk = self.request.recv(4096) + if not chunk: + return + head += chunk + line = head.split(b"\r\n", 1)[0].decode(errors="replace").split() + host, _, port = (line[1] if len(line) > 1 else "").rpartition(":") + if len(line) < 2 or line[0] != "CONNECT" or port != "443" or host.lower() not in self.allowed: + self.request.sendall(b"HTTP/1.1 403 Forbidden\r\ncontent-length: 0\r\n\r\n") + return + try: + up = socket.create_connection((host, 443), timeout=30) + except OSError: + self.request.sendall(b"HTTP/1.1 502 Bad Gateway\r\ncontent-length: 0\r\n\r\n") + return + self.request.sendall(b"HTTP/1.1 200 Connection established\r\n\r\n") + with up: + pair = (self.request, up) + while True: + ready, _, _ = select.select(pair, [], [], 300) + if not ready: + return + for s in ready: + data = s.recv(65536) + if not data: + return + (up if s is self.request else self.request).sendall(data) + + +class IndexProxy: + """`with IndexProxy(hosts, sock):` a CONNECT-only proxy on unix socket `sock` that reaches `hosts` and nothing else.""" + + def __init__(self, allowed: set[str], sock: Path): + handler = type("Handler", (_Connect,), {"allowed": {h.lower() for h in allowed}}) + self.server = socketserver.ThreadingUnixStreamServer(str(sock), handler) + self.server.daemon_threads = True + os.chmod(sock, 0o666) # the build runs as the jail's user + + def __enter__(self): + threading.Thread(target=self.server.serve_forever, daemon=True).start() + return self + + def __exit__(self, *exc): + self.server.shutdown() + self.server.server_close() + + +def build(tree: Path, venv: Path) -> None: + """The target's `build` in a container: the tree's dependency files only, the engine's interpreter and uv read-only, + uv's cache and `venv` writable (uv makes the environment at `venv/env`), the index proxy as the only way out. + A failed build leaves no venv behind.""" + t = target.load() + if not t.build: + raise BuildFailed("the dependency files changed, and the target has no `build` command to install them") + if not container.available(): + raise BuildFailed("a dependency build needs the container jail (docker and the lab-jail image)") + uv = shutil.which("uv") or str(Path.home() / ".local" / "bin" / "uv") + interpreter = Path(engine.python()).resolve() + venv.parent.mkdir(parents=True, exist_ok=True) + uv_cache = cache_dir().parent / "uv-cache" + uv_cache.mkdir(parents=True, exist_ok=True) + work = Path(tempfile.mkdtemp(prefix="build-", dir=cache_dir())) + src, sock_dir = work / "src", Path(tempfile.mkdtemp(prefix="lab-index-")) # short: a unix socket path + sock_dir.chmod(0o711) # the jail's user reaches the socket, lists nothing + src.mkdir() + for rel in t.deps: + if (Path(tree) / rel).is_file(): + (src / rel).parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(Path(tree) / rel, src / rel) + shutil.rmtree(venv, ignore_errors=True) # an interrupted earlier build + venv.mkdir() + container.give([src, venv, uv_cache]) + name = container.new_name() + env = {"UV_PROJECT_ENVIRONMENT": str(venv / "env"), "UV_CACHE_DIR": str(uv_cache), "UV_PYTHON": str(interpreter), + "UV_PYTHON_DOWNLOADS": "never", "HOME": str(src), "PATH": f"{Path(uv).parent}:/usr/bin:/bin", + "HTTPS_PROXY": f"http://127.0.0.1:{PROXY_PORT_IN}", "HTTP_PROXY": f"http://127.0.0.1:{PROXY_PORT_IN}"} + argv = container.argv(name, src, ["sh", "-c", f"socat TCP-LISTEN:{PROXY_PORT_IN},bind=127.0.0.1,fork,reuseaddr " + f"UNIX-CONNECT:{sock_dir}/index.sock & {t.build}"], + env, [src, venv, uv_cache, sock_dir], [Path(uv).parent, interpreter.parent.parent], str(interpreter), + gpu=False) + try: + with IndexProxy(hosts(tree, t.build), sock_dir / "index.sock"): + p = subprocess.run(argv, capture_output=True, text=True, timeout=TIMEOUT_S) + except subprocess.TimeoutExpired: + p = subprocess.CompletedProcess(argv, -9, "", f"timed out after {TIMEOUT_S}s") + finally: + container.remove(name) + shutil.rmtree(work, ignore_errors=True) + shutil.rmtree(sock_dir, ignore_errors=True) + if p.returncode != 0 or not (venv / "env" / "bin" / "python").exists(): + shutil.rmtree(venv, ignore_errors=True) + raise BuildFailed(f"dependency build exited {p.returncode}: {(p.stdout + p.stderr)[-LOG_TAIL:]}") + (venv / ".lab-built").write_text(deps_sha(tree)) diff --git a/lab/engine.py b/lab/engine.py index 9210cff..1445286 100644 --- a/lab/engine.py +++ b/lab/engine.py @@ -4,6 +4,7 @@ from __future__ import annotations import sys +from contextlib import contextmanager from pathlib import Path ENV_ROOT = Path(__file__).resolve().parents[1] @@ -15,9 +16,24 @@ def repo() -> Path: return target.load().engine_repo +_using: list[str] = [] + + +@contextmanager +def using(py: str): + """`python()` is `py` inside the block: the venv a tree's own dependency files were built into (lab/build.py).""" + _using.append(py) + try: + yield + finally: + _using.pop() + + def python() -> str: """The target's engine interpreter (LAB_ENGINE_PYTHON overrides); this interpreter if that file does not exist. The engine's venv needs what its test and lint commands import (e.g. `uv sync --extra dev`).""" + if _using: + return _using[-1] from lab import target py = target.load().engine_python return py if Path(py).exists() else sys.executable diff --git a/lab/safety/container.py b/lab/safety/container.py index fcb6382..1fc180f 100644 --- a/lab/safety/container.py +++ b/lab/safety/container.py @@ -74,7 +74,7 @@ def mounts(writable: list[Path], readonly: list[Path], python: str) -> list[str] def argv(name: str, workdir: Path, command: list[str], env: dict[str, str], writable: list[Path], - readonly: list[Path], python: str, bridge_port: int | None = None) -> list[str]: + readonly: list[Path], python: str, bridge_port: int | None = None, gpu: bool | None = None) -> list[str]: """`docker run` for `command`. With `bridge_port`, the engine's 127.0.0.1:port is also served on the unix socket `/bridge.sock`, which the host side of `Bridge` forwards to the host's 127.0.0.1:port.""" uid, gid = user() @@ -82,7 +82,7 @@ def argv(name: str, workdir: Path, command: list[str], env: dict[str, str], writ "--tmpfs", "/tmp:exec", "--cap-drop", "ALL", "--security-opt", "no-new-privileges", "--pids-limit", "8192", "--shm-size", "16g", "--user", f"{uid}:{gid}", "--workdir", str(workdir.resolve()), "--label", "lab.jail=1"] - if gpus(): + if gpus() if gpu is None else gpu: out += ["--gpus", "all"] if (cpus := cpuset()): out += ["--cpuset-cpus", cpus] diff --git a/lab/safety/grader.py b/lab/safety/grader.py index 04195ad..947d456 100644 --- a/lab/safety/grader.py +++ b/lab/safety/grader.py @@ -16,7 +16,7 @@ from lab import engine, gpu, target from lab.safety import container, jail -from lab.safety.surfaces import ALWAYS_DENY, HIDDEN, may_write +from lab.safety.surfaces import ALWAYS_DENY, HIDDEN, deps_only, may_write IGNORED = ("*/__pycache__/*", "__pycache__/*", "*.pyc", ".pytest_cache/*", "*/.pytest_cache/*", ".ruff_cache/*", "*/.ruff_cache/*", "*.egg-info/*", ".DS_Store", "*/.DS_Store", @@ -121,6 +121,15 @@ def audit(repo: Path, base: str, workspace: Path) -> Audit: for rel in a.modified + a.deleted: if not may_write(rel, new_file=False): a.violations.append(f"may not change: {rel}") + deps = target.load().deps + for rel in a.deleted: + if rel in deps: + a.violations.append(f"may not delete a dependency file: {rel}") + for rel in a.modified: + if rel in deps and rel.endswith(".toml"): + old = subprocess.run(["git", "-C", str(repo), "show", f"{base}:{rel}"], capture_output=True, text=True).stdout + if not deps_only(old, (workspace / rel).read_text(errors="replace")): + a.violations.append(f"may change only the dependency tables of {rel}") return a diff --git a/lab/safety/surfaces.py b/lab/safety/surfaces.py index caf520c..2b10286 100644 --- a/lab/safety/surfaces.py +++ b/lab/safety/surfaces.py @@ -22,11 +22,41 @@ def write_surface() -> tuple[tuple[str, ...], tuple[str, ...]]: HIDDEN = ("corpus/*/heldout.jsonl",) +# The parts of a pyproject.toml that are dependencies; everything else in it (test and lint config) is protected. +DEP_TABLES = (("project", "dependencies"), ("project", "optional-dependencies"), ("dependency-groups",), + ("tool", "uv", "sources"), ("tool", "uv", "index")) + + def may_write(path: str, *, new_file: bool) -> bool: - """Whether the agent may touch repo-relative `path`.""" + """Whether the agent may touch repo-relative `path`. The target's dependency files may be changed (a + pyproject.toml only in its dependency tables: `deps_only`), though some are on the deny list for the rest.""" + if not new_file and path in target.load().deps: + return True if any(fnmatch(path.lower(), p.lower()) for p in ALWAYS_DENY + HIDDEN): # macOS ignores case return False write, add_only = write_surface() if any(fnmatch(path, p) for p in write): return True return new_file and any(fnmatch(path, p) for p in add_only) + + +def deps_only(old: str, new: str) -> bool: + """Whether a pyproject.toml changed nothing but its dependency tables (and still parses).""" + import tomllib + + def rest(text: str) -> dict: + d = tomllib.loads(text) + for keys in DEP_TABLES: + path = [d] + for k in keys[:-1]: + path.append(path[-1].get(k) if isinstance(path[-1], dict) else None) + if isinstance(path[-1], dict): + path[-1].pop(keys[-1], None) + for parent, k, child in reversed(list(zip(path, keys, path[1:]))): # tables emptied by it go too + if child == {}: + parent.pop(k) + return d + try: + return rest(old) == rest(new) + except tomllib.TOMLDecodeError: + return False diff --git a/lab/target.py b/lab/target.py index affe662..4f1d850 100644 --- a/lab/target.py +++ b/lab/target.py @@ -51,6 +51,8 @@ class Target: add_only: tuple[str, ...] test: str lint: str + deps: tuple[str, ...] # dependency files the agent may change (pyproject: its dependency tables only) + build: str | None # installs `deps` into $UV_PROJECT_ENVIRONMENT (lab/build.py) reference: Server | None reference_dir: Path | None interactive: Limits @@ -96,7 +98,7 @@ def parse(path: Path) -> Target: path=path, model=m["name"], chat_kwargs=dict(m.get("chat_kwargs") or {}), max_model_len=m.get("max_model_len"), engine_repo=repo, engine_python=py, engine=_server(e), write=tuple(e.get("write") or ()), add_only=tuple(e.get("add_only") or ()), - test=e["test"], lint=e["lint"], + test=e["test"], lint=e["lint"], deps=tuple(e.get("deps") or ()), build=e.get("build"), reference=_server(r), reference_dir=(ENV_ROOT / r["dir"]) if r and r.get("dir") else None, interactive=Limits(**lim.get("interactive", {"ttft_s": 0.5, "tpot_s": 0.030})), conversational=Limits(**lim.get("conversational", {"ttft_s": 2.0, "tpot_s": 0.100})), diff --git a/lab/tools.py b/lab/tools.py index 8277829..cc84306 100644 --- a/lab/tools.py +++ b/lab/tools.py @@ -179,16 +179,24 @@ def specs(self) -> list[ToolSpec]: paused = (" While it runs, the container your shell runs in is paused, and processes in it holding the GPU are " "killed and listed in the result." if ClaudeAgentSDK.containerized() else "") + deps = target.load().deps + offline = (f" Only your change reaches the engine here: if it changes {' or '.join(deps)}, those are installed " + "from scratch for it (package indexes only), and the engine then runs with no network, so anything it " + "needs at runtime (draft models, prebuilt kernels, files a library downloads when it is used) has to be " + "in your change." if deps else " The engine runs with no network.") + def gpu(t: str) -> str: return (f" Costs GPU time: ~{t} per call, charged at ${rate:g}/h." if known else - f" Costs GPU time: ~{t} per call; no GPU rate is configured on this host, so it is charged $0.") + paused + f" Costs GPU time: ~{t} per call; no GPU rate is configured on this host, so it is charged $0." + ) + offline + paused short, full = suite.TIERS["short"], suite.TIERS["full"] windows = (f"one engine start plus, per regime, a {short['final_s']:.0f} s measurement window at tier short or " f"{full['final_s']:.0f} s at full, plus {short['probe_s']:.0f} s / {full['probe_s']:.0f} s probes for " f"regimes that search for a load") return [ ToolSpec("test", "Lint and run the test suite on a pristine copy of your current workspace, inside the " - "referee's jail. Returns PASS or FAIL, lint and test status, and the tail of each output." + free, + "referee's jail. Returns PASS or FAIL, lint and test status, and the tail of each output." + offline + + free, obj, self.test), ToolSpec("profile", "Run the engine in process under the profiler, jailed, on `requests` sampled from a seen " "corpus class (`corpus_class`, default steady_interactive; `seed`) or on synthetic token ids " diff --git a/lab/worker.py b/lab/worker.py index 442dde5..42b8a7c 100644 --- a/lab/worker.py +++ b/lab/worker.py @@ -4,6 +4,8 @@ the output dir back; the worker re-hashes every tree right before using it and refuses one that does not match, so what is measured is what the controller audited. GpuBusy, Contaminated, NotReady and ValueError cross intact. measure, equiv and profile take `workbench`, the agent's container: it is paused for the job, and processes in it holding the GPU are killed and listed under `killed`. +Every job runs the tree with the interpreter `build.python_for` gives it: the engine's, or a venv built from the +tree's own dependency files. test lint + the engine's suite on `tree` -> {lint, tests} measure serve `tree`, run regimes; passive data into out -> {results, ready_s} @@ -60,21 +62,28 @@ def tree_sha(root: Path) -> str: # -- jobs (run where the GPU is) --------------------------------------------------------------------------------- def job_test(inputs: dict, out: Path, args: dict, hooks: dict) -> dict: + from lab import build from lab.safety import grader tree = inputs["tree"] - lint = grader.run_lint(tree) - tests = grader.run_tests(tree) if lint.passed else grader.Run(False, -1, "skipped: lint failed") + try: + py = build.python_for(tree) + except build.BuildFailed as e: # a failed test run, recorded like any other + skipped = grader.Run(False, -1, str(e)) + return {"lint": dataclasses.asdict(skipped), "tests": dataclasses.asdict(skipped)} + with engine.using(py): + lint = grader.run_lint(tree) + tests = grader.run_tests(tree) if lint.passed else grader.Run(False, -1, "skipped: lint failed") return {"lint": dataclasses.asdict(lint), "tests": dataclasses.asdict(tests)} def job_measure(inputs: dict, out: Path, args: dict, hooks: dict) -> dict: """`args`: names, split, tier, seed, jailed, passive (keep device/telemetry/client rows in `out`).""" - from lab import artifacts + from lab import artifacts, build from lab.evaltools import run_regimes t = target.load() rows = [] if args["passive"] else None try: - with container.quiet(args.get("workbench")) as killed, serve.Served( + with engine.using(build.python_for(inputs["tree"])), container.quiet(args.get("workbench")) as killed, serve.Served( inputs["tree"], t.engine, log=out / ENGINE_LOG, jailed=args["jailed"], artifacts=out if args["passive"] else None) as srv: results = run_regimes(srv.url, t, args["names"], args["split"], args["tier"], args["seed"], rows) @@ -90,12 +99,13 @@ def job_equiv(inputs: dict, out: Path, args: dict, hooks: dict) -> dict: answers already in `prior`'s equiv.outputs.jsonl are reused, not asked again.""" from correctness import client, run as crun from correctness.gate import Thresholds + from lab import build t = target.load() if (inputs["prior"] / "equiv.outputs.jsonl").exists(): shutil.copyfile(inputs["prior"] / "equiv.outputs.jsonl", out / "equiv.outputs.jsonl") enc = hooks.get("encoder") or crun.HFEncoder(t.model, t.chat_kwargs) - with container.quiet(args.get("workbench")) as killed, serve.Served( - inputs["tree"], t.engine, log=out / ENGINE_LOG, jailed=args["jailed"], artifacts=out) as srv: + with engine.using(build.python_for(inputs["tree"])), container.quiet(args.get("workbench")) as killed, \ + serve.Served(inputs["tree"], t.engine, log=out / ENGINE_LOG, jailed=args["jailed"], artifacts=out) as srv: res = crun.candidate(inputs["reference"], srv.url, t.model, out / "equiv.json", enc, Thresholds(min_score_ratio=t.min_score_ratio, min_length_ratio=t.min_length_ratio), concurrency=args["concurrency"], api=client.Api(t.engine.api, t.chat_kwargs), @@ -106,7 +116,12 @@ def job_equiv(inputs: dict, out: Path, args: dict, hooks: dict) -> dict: def job_profile(inputs: dict, out: Path, args: dict, hooks: dict) -> dict: """`args`: kind (profile | trace | kernel | hostprof), the tool's args, workbench. Files go to `out/files`.""" - with container.quiet(args.get("workbench")) as killed: + from lab import build + try: + py = build.python_for(inputs["tree"]) + except build.BuildFailed as e: # a failed profile run, recorded like any other + return {"returncode": -1, "output": str(e)[-3000:], "instrument": None, "argv": [], "head": "", "killed": []} + with engine.using(py), container.quiet(args.get("workbench")) as killed: return {**_profile(inputs, out, args, hooks), "killed": killed} diff --git a/targets/inference-server.toml b/targets/inference-server.toml index de30fc0..2c9d8cd 100644 --- a/targets/inference-server.toml +++ b/targets/inference-server.toml @@ -18,6 +18,10 @@ write = ["src/inference_server/*"] # what the agent may add, modify or add_only = ["tests/test_*.py"] # what it may add but never change test = "python -m pytest -q -p no:cacheprovider -m 'not heavy and not needs_host'" lint = "python -m ruff check ." +deps = ["pyproject.toml", "uv.lock"] # the agent may change these (pyproject: its dependency tables only) +# Installs a changed `deps` into $UV_PROJECT_ENVIRONMENT (lab/build.py): vm-setup.sh's recipe, the lock with the +# CUDA torch wheel at the locked version (the lock pins the CPU index on Linux). +build = '''uv sync --locked --extra dev --no-install-project && uv pip install --python "$UV_PROJECT_ENVIRONMENT/bin/python" --reinstall "torch==$(awk '/^name = "torch"$/{getline; sub(/version = "/,""); sub(/"/,""); print; exit}' uv.lock)" --index-url https://download.pytorch.org/whl/cu130''' [engine.telemetry] # optional: engine-native files serve collects after teardown env = { TELEMETRY_DIR = "{dir}", TIMELINE_DIR = "{dir}/timeline" } # request rows (lab/TRACE.md); events.jsonl apart diff --git a/tests/test_lab_build.py b/tests/test_lab_build.py new file mode 100644 index 0000000..8c3e070 --- /dev/null +++ b/tests/test_lab_build.py @@ -0,0 +1,148 @@ +"""lab.build: when a tree gets its own venv, the cache, the index-only proxy, and (Linux + docker + internet) a build +room that reaches the package index through the proxy and nothing else.""" + +from __future__ import annotations + +import dataclasses +import socket +import threading +from pathlib import Path + +import pytest + +from lab import build, engine, target +from lab.safety import container, jail + + +@pytest.fixture +def deps_target(tmp_path, monkeypatch): + """The committed target with its engine repo here, holding one dependency file.""" + repo = tmp_path / "engine" + repo.mkdir() + (repo / "pyproject.toml").write_text('[project]\nname = "e"\ndependencies = ["a"]\n') + t = dataclasses.replace(target.load(), engine_repo=repo, deps=("pyproject.toml",)) + monkeypatch.setattr(target, "load", lambda path=None: t) + monkeypatch.setenv("LAB_WORKER_CACHE", str(tmp_path / "cache")) + return t, repo + + +def pristine(tmp_path, repo: Path, change: dict[str, str]) -> Path: + """A two-commit tree like grader.pristine_tree's: base, then `change`.""" + from lab.safety import grader + from tests.lab_fixtures import make_repo + tree = tmp_path / "tree" + grader.export(make_repo(tmp_path), "HEAD", tree) + (tree / "pyproject.toml").write_text((repo / "pyproject.toml").read_text()) + grader._commit(tree, "base", init=True) + for rel, text in change.items(): + (tree / rel).write_text(text) + grader._commit(tree, "change") + return tree + + +def test_unchanged_dependencies_keep_the_engines_venv_and_changed_ones_build_once(deps_target, tmp_path, monkeypatch): + t, repo = deps_target + tree = pristine(tmp_path, repo, {"src/inference_server/engine.py": "def speed():\n return 2\n"}) + assert not build.changed(tree) and build.python_for(tree) == engine.python() + flat = tmp_path / "flat" # an exported base: no history, no change + flat.mkdir() + (flat / "pyproject.toml").write_text("anything") + assert not build.changed(flat) + built = [] + + def fake(tree, venv): + built.append(venv) + (venv / "env" / "bin").mkdir(parents=True) + (venv / ".lab-built").write_text("") + monkeypatch.setattr(build, "build", fake) + (tree / "pyproject.toml").write_text('[project]\nname = "e"\ndependencies = ["a", "b"]\n') + assert build.changed(tree) + py = build.python_for(tree) + assert py == str(built[0] / "env" / "bin" / "python") and build.python_for(tree) == py and len(built) == 1 + + +def test_the_allowed_hosts_are_the_index_defaults_plus_what_the_files_and_build_name(deps_target, tmp_path): + tree = tmp_path / "tree" + tree.mkdir() + (tree / "pyproject.toml").write_text('[[tool.uv.index]]\nurl = "https://download.pytorch.org/whl/cu130"\n') + got = build.hosts(tree, "uv sync --index-url https://mirror.example.org/simple") + assert got == {"pypi.org", "files.pythonhosted.org", "download.pytorch.org", "mirror.example.org"} + + +def proxy_request(sock_path: Path, line: bytes) -> bytes: + c = socket.socket(socket.AF_UNIX) + c.connect(str(sock_path)) + c.sendall(line) + out = c.recv(4096) + return c, out + + +def test_the_index_proxy_tunnels_to_allowed_hosts_only(tmp_path, monkeypatch): + echo = socket.socket() + echo.bind(("127.0.0.1", 0)) + echo.listen() + + def serve(): + conn, _ = echo.accept() + conn.sendall(conn.recv(100).upper()) + conn.close() + threading.Thread(target=serve, daemon=True).start() + real = socket.create_connection + monkeypatch.setattr(build.socket, "create_connection", + lambda addr, timeout=None: real(echo.getsockname(), timeout=timeout)) # "pypi.org:443" here + sock = Path("/tmp") / f"lab-test-{id(tmp_path)}.sock" + with build.IndexProxy({"pypi.org"}, sock): + c, head = proxy_request(sock, b"CONNECT pypi.org:443 HTTP/1.1\r\nHost: pypi.org\r\n\r\n") + assert head.startswith(b"HTTP/1.1 200") + c.sendall(b"hello") + assert c.recv(100) == b"HELLO" + for line in (b"CONNECT example.com:443 HTTP/1.1\r\n\r\n", b"CONNECT pypi.org:80 HTTP/1.1\r\n\r\n", + b"GET http://pypi.org/ HTTP/1.1\r\n\r\n"): + assert proxy_request(sock, line)[1].startswith(b"HTTP/1.1 403") + sock.unlink(missing_ok=True) + + +def _online() -> bool: + try: + socket.create_connection(("pypi.org", 443), timeout=5).close() + return True + except OSError: + return False + + +PROBE = r'''import os, socket, urllib.request, urllib.error +def get(url): + try: + return urllib.request.urlopen(url, timeout=20).status + except urllib.error.HTTPError as e: + return e.code + except Exception as e: # a refused tunnel is a URLError carrying the proxy's 403 + return "403" if "403" in str(e) else type(e).__name__ +direct = "no" +try: + socket.create_connection(("1.1.1.1", 443), timeout=5); direct = "yes" +except OSError: + pass +print("index", get("https://pypi.org/simple/pip/"), "other", get("https://example.com/"), "direct", direct) +''' + + +@pytest.mark.skipif(not (jail.backend() == "container" and container.available()), reason="needs the container jail") +@pytest.mark.skipif(not _online(), reason="needs the internet") +def test_a_build_room_reaches_the_index_through_the_proxy_and_nothing_else(deps_target, tmp_path, monkeypatch): + t, repo = deps_target + tree = tmp_path / "tree" + tree.mkdir() + (tree / "pyproject.toml").write_text('[project]\nname = "e"\ndependencies = ["b"]\n') + (tree / "probe.py").write_text(PROBE) + out = tmp_path / "probe.out" + out.touch() + cmd = (f'python3 -c "$(cat {tree}/probe.py)" > "$(dirname "$UV_PROJECT_ENVIRONMENT")/probe.out" 2>&1; ' + 'mkdir -p "$UV_PROJECT_ENVIRONMENT/bin" && ln -s /usr/bin/python3 "$UV_PROJECT_ENVIRONMENT/bin/python"') + monkeypatch.setattr(target, "load", lambda path=None: dataclasses.replace(t, build=cmd)) + real = container.argv + monkeypatch.setattr(container, "argv", lambda *a, **k: real(*a[:5], [*a[5], tree], *a[6:], **k)) # probe readable + venv = build.cache_dir() / "probe" + build.build(tree, venv) + said = (venv / "probe.out").read_text() + assert said.split() == ["index", "200", "other", "403", "direct", "no"], said diff --git a/tests/test_lab_safety.py b/tests/test_lab_safety.py index db6f432..06227cd 100644 --- a/tests/test_lab_safety.py +++ b/tests/test_lab_safety.py @@ -19,7 +19,9 @@ ("tests/test_engine.py", False, False), # existing tests are evidence ("tests/conftest.py", False, False), ("src/inference_server/sitecustomize.py", True, False), - ("pyproject.toml", False, False), + ("pyproject.toml", False, True), # a dependency file: its dependency tables only (the audit) + ("uv.lock", False, True), + ("sub/pyproject.toml", False, False), ("lab/ledger.py", False, False), ("corpus/a/heldout.jsonl", True, False), ("CORPUS/A/HELDOUT.JSONL", True, False), @@ -82,7 +84,7 @@ def test_write_hook_lets_the_agent_edit_its_own_new_test(tmp_path): (ws / "tests/test_new.py").write_text("x") assert allowed(ws, base, "tests/test_new.py") # still "new": not in the base tree assert not allowed(ws, base, "tests/test_engine.py") - assert not allowed(ws, base, "../outside.py") and not allowed(ws, base, "pyproject.toml") + assert not allowed(ws, base, "../outside.py") and not allowed(ws, base, "tests/conftest.py") def test_jailed_timeout_is_a_failed_run(tmp_path, monkeypatch): @@ -135,3 +137,35 @@ def test_pristine_refuses_a_failed_audit(tmp_path): a = grader.audit(repo, "HEAD", ws) with pytest.raises(ValueError): grader.pristine_tree(repo, "HEAD", ws, a, tmp_path / "p") + + +PYPROJECT = """[project] +name = "x" +dependencies = ["a"] + +[tool.pytest.ini_options] +addopts = "-q" +""" + + +def test_a_pyproject_may_change_its_dependencies_and_nothing_else(tmp_path): + from lab.safety.surfaces import deps_only + assert deps_only(PYPROJECT, PYPROJECT.replace('["a"]', '["a", "b>=2"]')) + assert deps_only(PYPROJECT, PYPROJECT + '\n[tool.uv.sources]\nb = { index = "x" }\n') + assert not deps_only(PYPROJECT, PYPROJECT.replace('"-q"', '"-q -k nothing"')) # test config is the evaluator's + assert not deps_only(PYPROJECT, PYPROJECT + "[broken") + + +def test_the_audit_holds_a_pyproject_to_its_dependency_tables(tmp_path): + repo = make_repo(tmp_path) + ws = tmp_path / "ws" + grader.export(repo, "HEAD", ws) + text = (ws / "pyproject.toml").read_text() + (ws / "pyproject.toml").write_text(text.replace("[tool.ruff]", '[project]\nname = "t"\ndependencies = ["b"]\n\n[tool.ruff]')) + assert grader.audit(repo, "HEAD", ws).ok is False # [project].name is not a dependency + (ws / "pyproject.toml").write_text(text.replace("line-length = 100", "line-length = 300")) + assert "may change only the dependency tables of pyproject.toml" in grader.audit(repo, "HEAD", ws).violations + (ws / "pyproject.toml").write_text(text + '\n[dependency-groups]\nextra = ["c"]\n') + assert grader.audit(repo, "HEAD", ws).ok + (ws / "pyproject.toml").unlink() + assert "may not delete a dependency file: pyproject.toml" in grader.audit(repo, "HEAD", ws).violations From 22482f9aa1ec6f304f94ae4c57db6e59676167eb Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Sat, 10 Oct 2026 14:00:57 -0400 Subject: [PATCH 10/18] lab: fixes from the first H200 dry run and the stack review From the dry run (Nebius H200, LAB_WORKER=ssh): - cli.sh removes the tunnel's stale socket before forwarding (the SDK launches the CLI more than once; that sshd does not unlink an existing socket), with `ssh -n` so that call does not swallow the SDK's first message. - A non-root login (Nebius, Crusoe) runs the worker side through sudo, keeping HOME: the workbench is root inside, measurement rooms drop to uid 10001, rsync via sudo, and the proxy's socket dir is handed to the login user for sshd. - Profiler jails mount the instrument's install root (nsys needs host-linux-x64's importer). - vm-setup.sh caches the target's model and the tokenizers its tests read (new target key `hub`), world-readable: every jail is offline. - vm.push excludes lab/runs explicitly: --delete otherwise reached the worker cache and the agent's live workspaces. From the review: - The API proxy charges in a finally; an incomplete response is charged its input and all of max_tokens; each request reserves its worst-case output first, so concurrent requests cannot pass the cap. - When the workbench shares the lab's host, it is frozen for every lab tool call (container.paused), closing the swap-a-file-for-a-link race on snapshot, restore and copies; quiet() nests in it. - Scratch (outside the write surface, not a base file) is left out whatever it is: a venv's links and large files no longer end the run. Co-Authored-By: Claude Opus 5.5 --- lab/agent.py | 12 ++++++++- lab/apiproxy.py | 46 ++++++++++++++++++++----------- lab/safety/container.py | 22 ++++++++++++--- lab/safety/grader.py | 10 +++++++ lab/vm-setup.sh | 12 ++++++++- lab/vm.py | 4 ++- lab/worker.py | 19 ++++++++----- targets/inference-server.toml | 1 + tests/test_lab_agent.py | 6 +++-- tests/test_lab_apiproxy.py | 51 +++++++++++++++++++++++++++++++++++ tests/test_lab_container.py | 20 +++++++++++--- tests/test_lab_safety.py | 14 ++++++++++ 12 files changed, 184 insertions(+), 33 deletions(-) diff --git a/lab/agent.py b/lab/agent.py index 3857f52..5afbef0 100644 --- a/lab/agent.py +++ b/lab/agent.py @@ -130,7 +130,10 @@ def _wrapper(self, spec: AgentSpec, cli: Path, base_url: str) -> Path: w = spec.remote.workbench(spec.workbench, spec.scratch.name, spec.remote_workspace, claude_agent_sdk.__version__) *ssh, target = spec.remote.prefix + # The SDK may launch this more than once, and that sshd need not unlink a socket a launch left behind; + # -n: that call must not read the SDK's stdin. line = (f'args=$(printf " %q" "$@")\ncmd={shlex.quote(shlex.join(w["argv"]))}\n' + f'{shlex.join(ssh)} -n -T {target} {shlex.quote("rm -f " + shlex.quote(w["sock"]))} || exit 1\n' f'exec {shlex.join(ssh)} -T -o ExitOnForwardFailure=yes -o StreamLocalBindUnlink=yes ' f'-R {w["sock"]}:127.0.0.1:{urlsplit(base_url).port} {target} "$cmd$args"') elif self.containerized(): @@ -159,10 +162,17 @@ async def _run(self, spec: AgentSpec, base_url: str) -> AgentReply: from claude_agent_sdk import (ClaudeAgentOptions, HookMatcher, ResultMessage, create_sdk_mcp_server, query, tool) + # Sharing this host, the workbench is frozen for every lab tool call: nothing in it can swap a file for a link + # while the tool snapshots, restores or copies into the workspace (with a Remote worker only rsync writes here). + shared = spec.remote is None and self.containerized() + def adapt(t: ToolSpec): async def call(args: dict[str, Any]) -> dict[str, Any]: try: - text = await asyncio.to_thread(t.fn, args) # a long test run must not block the transport + def run(): + with container.paused(spec.workbench if shared else None): + return t.fn(args) + text = await asyncio.to_thread(run) # a long test run must not block the transport except Exception as e: # the agent sees the refusal, the harness keeps going text = f"{t.name} refused: {e}" return {"content": [{"type": "text", "text": text}]} diff --git a/lab/apiproxy.py b/lab/apiproxy.py index d0ee5d1..0beff8c 100644 --- a/lab/apiproxy.py +++ b/lab/apiproxy.py @@ -54,10 +54,11 @@ def credential() -> dict[str, str]: class _Usage: - """Usage read off a response as it streams by: JSON, or SSE message_start + message_delta.""" + """Usage read off a response as it streams by: JSON, or SSE message_start + message_delta. `complete` once the + final output count is known.""" def __init__(self, sse: bool): - self.sse, self.buf, self.model, self.usage = sse, b"", None, {} + self.sse, self.buf, self.model, self.usage, self.complete = sse, b"", None, {}, False def feed(self, chunk: bytes) -> None: self.buf += chunk @@ -78,6 +79,7 @@ def _event(self, data: bytes) -> None: self.usage.update(e["message"].get("usage") or {}) elif e.get("type") == "message_delta": self.usage.update(e.get("usage") or {}) + self.complete = True def done(self) -> None: if self.sse: @@ -88,14 +90,14 @@ def done(self) -> None: except ValueError: return if isinstance(body, dict) and isinstance(body.get("usage"), dict): - self.model, self.usage = body.get("model"), body["usage"] + self.model, self.usage, self.complete = body.get("model"), body["usage"], True class Proxy: """`with Proxy(cap_usd, host) as p: ... p.url`; `p.spent_usd` and `p.usage` are what the session cost.""" def __init__(self, cap_usd: float, host: str = "127.0.0.1", upstream: str = UPSTREAM): - self.cap_usd, self.spent_usd, self.requests = cap_usd, 0.0, 0 + self.cap_usd, self.spent_usd, self.requests, self.reserved_usd = cap_usd, 0.0, 0, 0.0 self.usage = {"input_tokens": 0, "output_tokens": 0, "cache_read_input_tokens": 0, "cache_creation_input_tokens": 0} self.auth = credential() @@ -142,20 +144,24 @@ def _handle(self, h, body: bytes | None) -> None: path = urlsplit(h.path).path if path not in ALLOWED and not (path.startswith("/v1/models/") and path.count("/") == 3): return self._refuse(h, 403, f"{path} is not forwarded") - model = None + model, max_tokens = None, 0 if body: try: req = json.loads(body) except ValueError: return self._refuse(h, 400, "the request body is not JSON") - model = req.get("model") + model, max_tokens = req.get("model"), int(req.get("max_tokens") or 0) if path == "/v1/messages" and (not isinstance(model, str) or price(model) is None): return self._refuse(h, 403, f"no price for model {model!r}; priced: {sorted(PRICES)}") if req.get("speed") == "fast": return self._refuse(h, 403, "fast mode is not priced here") + metered = path == "/v1/messages" + # the request's worst-case output is reserved before it goes, so concurrent requests cannot overshoot the cap + hold = max_tokens * price(model)[1] / 1e6 if metered else 0.0 with self.lock: - if path == "/v1/messages" and self.spent_usd >= self.cap_usd: - return self._refuse(h, 403, f"this session's ${self.cap_usd:.2f} is spent") + if metered and self.spent_usd + self.reserved_usd + hold > self.cap_usd: + return self._refuse(h, 403, f"this session's ${self.cap_usd:.2f} is spent (or would be by this request)") + self.reserved_usd += hold headers = {k: v for k, v in h.headers.items() if k.lower() not in HOP | {"x-api-key", "authorization"}} headers.update(self.auth, host=self.upstream.netloc, **{"accept-encoding": "identity"}) if "authorization" in self.auth: @@ -163,6 +169,7 @@ def _handle(self, h, body: bytes | None) -> None: headers["anthropic-beta"] = ",".join(dict.fromkeys([*betas, OAUTH_BETA])) conn = (http.client.HTTPSConnection if self.upstream.scheme == "https" else http.client.HTTPConnection)( self.upstream.netloc, timeout=900) + usage = None try: conn.request(h.command, h.path, body=body, headers=headers) resp = conn.getresponse() @@ -178,13 +185,22 @@ def _handle(self, h, body: bytes | None) -> None: h.wfile.write(b"%x\r\n%s\r\n" % (len(chunk), chunk)) h.wfile.flush() h.wfile.write(b"0\r\n\r\n") - usage.done() - finally: + finally: # charged however it ended: a dropped client still spent conn.close() - if path == "/v1/messages" and usage.usage: - spent = cost_usd(usage.model if usage.model and price(usage.model) else model, usage.usage) - with self.lock: - self.spent_usd += spent + if usage is not None: + usage.done() + self._charge(metered and usage is not None, usage, model, max_tokens, hold) + + def _charge(self, metered: bool, usage: "_Usage | None", model: str | None, max_tokens: int, hold: float) -> None: + """An incomplete response is charged its input as far as seen and its whole `max_tokens` of output.""" + got = dict(usage.usage) if usage is not None else {} + if metered and got and not usage.complete: + got["output_tokens"] = max(got.get("output_tokens") or 0, max_tokens) + spent = cost_usd(usage.model if usage.model and price(usage.model) else model, got) if metered and got else 0.0 + with self.lock: + self.reserved_usd -= hold + self.spent_usd += spent + if metered and got: self.requests += 1 for k in self.usage: - self.usage[k] += usage.usage.get(k) or 0 + self.usage[k] += got.get(k) or 0 diff --git a/lab/safety/container.py b/lab/safety/container.py index 1fc180f..d6bc1ac 100644 --- a/lab/safety/container.py +++ b/lab/safety/container.py @@ -138,6 +138,23 @@ def _kill_gpu_holders(name: str) -> list[dict]: return hit +@contextmanager +def paused(name: str | None): + """The container frozen for the block (a no-op if it is not running); left paused if it already was.""" + if not name or _docker() is None or not running(name): + yield + return + already = subprocess.run([_docker(), "inspect", "-f", "{{.State.Paused}}", name], capture_output=True, + text=True).stdout.strip() == "true" + if not already: + subprocess.run([_docker(), "pause", name], capture_output=True, check=True) + try: + yield + finally: + if not already: + subprocess.run([_docker(), "unpause", name], capture_output=True) + + @contextmanager def quiet(name: str | None): """The workbench paused for a measurement: processes in it holding the GPU are killed (and listed), the rest @@ -147,16 +164,13 @@ def quiet(name: str | None): return from lab import gpu killed = _kill_gpu_holders(name) - subprocess.run([_docker(), "pause", name], capture_output=True, check=True) - try: + with paused(name): killed += _kill_gpu_holders(name) # anything that took the GPU before the pause landed deadline = time.monotonic() + 15 while killed and time.monotonic() < deadline and any( a["pid"] in {k["pid"] for k in killed} for a in gpu.compute_apps() or []): time.sleep(0.2) # until the driver has released their memory yield killed - finally: - subprocess.run([_docker(), "unpause", name], capture_output=True) def new_name() -> str: diff --git a/lab/safety/grader.py b/lab/safety/grader.py index 947d456..87d320b 100644 --- a/lab/safety/grader.py +++ b/lab/safety/grader.py @@ -94,6 +94,10 @@ def audit(repo: Path, base: str, workspace: Path) -> Audit: for name in dirs + files: p = Path(root) / name rel = p.relative_to(workspace).as_posix() + if _scratch(rel, blobs) and (p.is_symlink() or not (p.is_dir() or p.is_file()) + or (p.is_file() and p.stat().st_size > MAX_FILE_BYTES)): + a.scratch.append(rel) # never reaches the pristine tree: a venv's links, big wheels + continue if p.is_symlink(): a.violations.append(f"symlink: {rel}") elif p.is_dir() or any(fnmatch(rel, g) for g in IGNORED): @@ -133,6 +137,12 @@ def audit(repo: Path, base: str, workspace: Path) -> Audit: return a +def _scratch(rel: str, blobs: dict) -> bool: + """Outside what the agent may write and not a base file: left out of every tree, whatever it is.""" + return (rel not in blobs and not may_write(rel, new_file=True) + and not any(fnmatch(rel.lower(), p.lower()) for p in ALWAYS_DENY + HIDDEN)) + + def pristine_tree(repo: Path, base: str, workspace: Path, result: Audit, dest: Path) -> None: """A fresh export of `base` plus only the audited changes, as a two-commit repo (base, change).""" if not result.ok: diff --git a/lab/vm-setup.sh b/lab/vm-setup.sh index 4c92c33..824d17c 100755 --- a/lab/vm-setup.sh +++ b/lab/vm-setup.sh @@ -87,8 +87,18 @@ else echo "clock lock: refused" fi -echo "== model cache" +echo "== model cache: the target's model and the Hub repos its tests read (every jail is offline)" mkdir -p "$HOME/.cache/huggingface" +.venv/bin/python - "$here/../${LAB_TARGET:-targets/inference-server.toml}" <<'PY' \ + || echo "model cache: download failed (finding, not failure)" +import sys, tomllib +from huggingface_hub import snapshot_download +t = tomllib.load(open(sys.argv[1], "rb")) +print("cached", snapshot_download(t["model"]["name"])) +for repo in t["engine"].get("hub", []): # what tests read is tokenizers and configs, not weights + print("cached", snapshot_download(repo, allow_patterns=["*.json", "*.txt", "*.model", "*.tiktoken"])) +PY +chmod -R a+rX "$HOME/.cache/huggingface" # the jails' own user reads it (hub metadata is written 0600) du -sh "$HOME/.cache/huggingface" 2>/dev/null || true echo "== smoke" diff --git a/lab/vm.py b/lab/vm.py index d9170f4..7d0aee4 100644 --- a/lab/vm.py +++ b/lab/vm.py @@ -190,7 +190,9 @@ def push(vm: VM, record: Record) -> None: """Both working trees, minus .git and everything .gitignore excludes (secrets, weights, venvs, ledger, runs).""" # .venv is named explicitly: a gitignore-only exclusion did not protect it from --delete on the VM. for local, remote in ((engine.repo(), vm.remote_dir), (engine.ENV_ROOT, vm.env_dir)): - _rsync(vm, "--delete", "--exclude=.git", "--exclude=.venv", "--exclude=__pycache__", "--filter=:- .gitignore", + # lab/runs is named explicitly too: it holds the worker's cache and the agent's live workspaces there. + _rsync(vm, "--delete", "--exclude=.git", "--exclude=.venv", "--exclude=__pycache__", "--exclude=/lab/runs", + "--filter=:- .gitignore", f"{local}/", f"{ssh_target(vm, record)}:{remote}/") diff --git a/lab/worker.py b/lab/worker.py index 42b8a7c..d5fe124 100644 --- a/lab/worker.py +++ b/lab/worker.py @@ -148,7 +148,7 @@ def _profile(inputs: dict, out: Path, args: dict, hooks: dict) -> dict: exe = str(Path(exe).resolve()) work = tree / proftools.OUT work.mkdir() - read = [Path(exe).parent] + read = [Path(exe).parent, Path(exe).parent.parent] # the install root: nsys's importer lives beside its binary dir src = T.workload_flags(a, harness, proftools._int(a, "requests", *proftools.WORKLOAD["requests"])) argv = [*ins["wrap"](exe, work), engine.python(), "-m", "lab.profile", *ins["flags"], *src, *ins["extra"](exe), "--out", str(work / "bundle")] @@ -219,8 +219,11 @@ class Remote: """Jobs on another host. `prefix` is the argv that runs one shell command there (ssh ... user@host); empty runs it here through sh, which the tests use. `env_dir` is this repo on that host, `python` its interpreter.""" - def __init__(self, prefix: list[str], env_dir: str, python: str, env: dict[str, str], rsh: str | None = None): - self.prefix, self.env_dir, self.python, self.env, self.rsh = prefix, env_dir, python, env, rsh + def __init__(self, prefix: list[str], env_dir: str, python: str, env: dict[str, str], rsh: str | None = None, + sudo: bool = False): + """`sudo`: the login user is not root, so the worker's side runs as root (the workbench is root inside, the + measurement rooms drop to their own uid), keeping the login user's HOME (weights, venv, uv).""" + self.prefix, self.env_dir, self.python, self.env, self.rsh, self.sudo = prefix, env_dir, python, env, rsh, sudo self._root: str | None = None def _where(self, path: str) -> str: @@ -228,7 +231,8 @@ def _where(self, path: str) -> str: def _rsync(self, src: str, dst: str) -> None: """By content, not size and mtime: an edit within the same second must not be missed.""" - p = subprocess.run(["rsync", "-a", "--checksum", "--delete", *(["-e", self.rsh] if self.rsh else []), src, dst], + p = subprocess.run(["rsync", "-a", "--checksum", "--delete", *(["-e", self.rsh] if self.rsh else []), + *(["--rsync-path", "sudo -n rsync"] if self.sudo else []), src, dst], capture_output=True, text=True) if p.returncode != 0: raise WorkerError(f"rsync {src} -> {dst}: {p.stderr.strip()[-1000:]}") @@ -264,7 +268,8 @@ def _run(self, command: str, stdin: bytes = b"") -> bytes: if any(c in v for c in '"`\\'): raise WorkerError(f"unsafe value in the worker's env: {v!r}") exports = " ".join(f'{k}="{v}"' for k, v in self.env.items()) - line = f"cd {self.env_dir} && {exports} {self.python} -m lab.worker {command}" + root = 'sudo -n env HOME="$HOME" PATH="$PATH" ' if self.sudo else "" + line = f"cd {self.env_dir} && {root}{exports} {self.python} -m lab.worker {command}" argv = [*self.prefix, line] if self.prefix else ["sh", "-c", line] p = subprocess.run(argv, input=stdin, capture_output=True) if p.returncode != 0: @@ -302,7 +307,7 @@ def load(**hooks): home = lambda p: p.replace("~", "$HOME", 1) # noqa: E731 - expanded by the worker's shell return Remote(["ssh", *v.ssh_opts, vm.ssh_target(v, record)], v.env_dir, home(v.remote_dir) + "/.venv/bin/python", {"LAB_TARGET": spec.relative_to(engine.ENV_ROOT).as_posix(), "LAB_ENGINE_REPO": home(v.remote_dir)}, - rsh=f"ssh {shlex.join(v.ssh_opts)}") + rsh=f"ssh {shlex.join(v.ssh_opts)}", sudo=v.login != "root") # -- the worker host's side --------------------------------------------------------------------------------------- @@ -364,6 +369,8 @@ def _workbench(req: dict) -> dict: from lab import agent sock_dir = Path(f"/tmp/lab-proxy-{req['session']}") sock_dir.mkdir(mode=0o700, exist_ok=True) + if os.environ.get("SUDO_UID"): # sshd binds the tunnel's socket here as the login user + os.chown(sock_dir, int(os.environ["SUDO_UID"]), int(os.environ["SUDO_GID"])) (sock_dir / "proxy.sock").unlink(missing_ok=True) scratch = engine.ENV_ROOT / "lab" / "runs" / "remote" / "scratch" / req["session"] argv = agent.workbench_command(req["name"], Path(req["workspace"]), scratch, agent_cli(req["sdk_version"]), diff --git a/targets/inference-server.toml b/targets/inference-server.toml index 2c9d8cd..b668fb3 100644 --- a/targets/inference-server.toml +++ b/targets/inference-server.toml @@ -18,6 +18,7 @@ write = ["src/inference_server/*"] # what the agent may add, modify or add_only = ["tests/test_*.py"] # what it may add but never change test = "python -m pytest -q -p no:cacheprovider -m 'not heavy and not needs_host'" lint = "python -m ruff check ." +hub = ["gpt2"] # Hub repos whose tokenizer the engine's tests read; cached at setup (jails are offline) deps = ["pyproject.toml", "uv.lock"] # the agent may change these (pyproject: its dependency tables only) # Installs a changed `deps` into $UV_PROJECT_ENVIRONMENT (lab/build.py): vm-setup.sh's recipe, the lock with the # CUDA torch wheel at the locked version (the lock pins the CPU index on Linux). diff --git a/tests/test_lab_agent.py b/tests/test_lab_agent.py index 4b1fa92..b96a8aa 100644 --- a/tests/test_lab_agent.py +++ b/tests/test_lab_agent.py @@ -60,10 +60,12 @@ def workbench(self, name, session, workspace, version): assert "-R /tmp/lab-proxy-scratch/proxy.sock:127.0.0.1:5555 root@203.0.113.5" in text fake = tmp_path / "bin" / "ssh" # an ssh that prints the remote command it was given fake.parent.mkdir() - fake.write_text('#!/bin/sh\nfor a; do last="$a"; done\nprintf "%s" "$last"\n') + fake.write_text('#!/bin/sh\nfor a; do last="$a"; done\nprintf "%s\\n" "$last"\n') fake.chmod(0o755) out = subprocess.run(["bash", str(script), "--output-format", "stream-json", "it's \"quoted\" $HOME"], capture_output=True, text=True, env={"PATH": f"{fake.parent}:/usr/bin:/bin"}) - remote_argv = shlex.split(out.stdout) + first, launch = out.stdout.splitlines() + assert first == "rm -f /tmp/lab-proxy-scratch/proxy.sock" # a stale socket from an earlier launch goes + remote_argv = shlex.split(launch) assert remote_argv == ["docker", "run", "-i", "--name", "lab-workbench-r1-s1", "lab-jail", "sh", "-c", 'exec "$0" "$@"', "/c/claude", "--output-format", "stream-json", "it's \"quoted\" $HOME"] diff --git a/tests/test_lab_apiproxy.py b/tests/test_lab_apiproxy.py index bc03393..e3b2c32 100644 --- a/tests/test_lab_apiproxy.py +++ b/tests/test_lab_apiproxy.py @@ -127,3 +127,54 @@ def test_no_credential_means_no_proxy(monkeypatch): monkeypatch.delenv("CLAUDE_CODE_OAUTH_TOKEN", raising=False) with pytest.raises(RuntimeError): apiproxy.Proxy(5.0) + + +def test_a_client_that_drops_mid_stream_is_charged_its_input_and_all_of_max_tokens(monkeypatch): + import time + seen = [] + + class Slow(BaseHTTPRequestHandler): + def log_message(self, *a): + pass + + def do_POST(self): + self.rfile.read(int(self.headers["content-length"])) + self.send_response(200) + self.send_header("content-type", "text/event-stream") + self.end_headers() + start = b'data: {"type":"message_start","message":{"model":"claude-opus-5-5","usage":{"input_tokens":1000,"output_tokens":1}}}\n\n' + self.wfile.write(start) + self.wfile.flush() + try: + for _ in range(50): # the rest never comes before the client hangs up + time.sleep(0.05) + self.wfile.write(b'data: {"type":"ping"}\n\n') + self.wfile.flush() + except OSError: + pass + seen.append(True) + + srv = ThreadingHTTPServer(("127.0.0.1", 0), Slow) + threading.Thread(target=srv.serve_forever, daemon=True).start() + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-real") + with apiproxy.Proxy(50.0, upstream=f"http://127.0.0.1:{srv.server_address[1]}") as p: + c = http.client.HTTPConnection(urlsplit(p.url).netloc, timeout=10) + c.request("POST", "/v1/messages", body=json.dumps({"model": "claude-opus-5-5", "max_tokens": 4000, "stream": True}), + headers={"content-type": "application/json"}) + r = c.getresponse() + r.read1(200) + c.close() # hang up + deadline = time.time() + 10 + while not p.requests and time.time() < deadline: + time.sleep(0.1) + srv.shutdown() + p_in, p_out = apiproxy.PRICES["claude-opus-5-5"][:2] + assert p.spent_usd == pytest.approx((1000 * p_in + 4000 * p_out) / 1e6) and p.reserved_usd == pytest.approx(0) + + +def test_a_request_whose_worst_case_would_pass_the_cap_is_refused_before_it_goes(upstream, monkeypatch): + url, seen = upstream + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-real") + with apiproxy.Proxy(0.10, upstream=url) as p: # 10k output tokens of opus is $0.20 + status, body = post(p, "/v1/messages", {"model": "claude-opus-5-5", "max_tokens": 10000}) + assert status == 403 and b"would be" in body and seen == [] diff --git a/tests/test_lab_container.py b/tests/test_lab_container.py index d65078c..fee4b4e 100644 --- a/tests/test_lab_container.py +++ b/tests/test_lab_container.py @@ -162,11 +162,25 @@ def test_quiet_kills_the_workbenchs_gpu_holders_then_pauses_and_thaws(monkeypatc apps = [[{"pid": 101, "used_mib": 9000, "name": "python"}, {"pid": 999, "used_mib": 1, "name": "other"}]] monkeypatch.setattr(gpu, "compute_apps", lambda: apps[0]) monkeypatch.setattr(container.os, "kill", lambda pid, sig: killed.append(pid) or apps.__setitem__(0, [apps[0][1]])) - monkeypatch.setattr(container.subprocess, "run", lambda argv, **k: calls.append(argv[1:]) or subprocess.CompletedProcess(argv, 0)) + monkeypatch.setattr(container.subprocess, "run", lambda argv, **k: calls.append(argv[1:]) or + subprocess.CompletedProcess(argv, 0, "false", "")) + states = lambda: [c for c in calls if c[0] in ("pause", "unpause")] # noqa: E731 with container.quiet("wb") as got: - assert calls == [["pause", "wb"]] + assert states() == [["pause", "wb"]] assert killed == [101] and [a["pid"] for a in got] == [101] # 999 is not the workbench's - assert calls == [["pause", "wb"], ["unpause", "wb"]] + assert states() == [["pause", "wb"], ["unpause", "wb"]] + + +def test_an_already_paused_workbench_is_left_paused(monkeypatch): + calls = [] + monkeypatch.setattr(container, "_docker", lambda: "docker") + monkeypatch.setattr(container, "running", lambda n: True) + monkeypatch.setattr(container.subprocess, "run", lambda argv, **k: calls.append(argv[1]) or + subprocess.CompletedProcess(argv, 0, "true", "")) + with container.paused("wb"): + with container.paused("wb"): + pass + assert "pause" not in calls and "unpause" not in calls def test_quiet_is_a_no_op_without_a_running_workbench(monkeypatch): diff --git a/tests/test_lab_safety.py b/tests/test_lab_safety.py index 06227cd..ad6c1c2 100644 --- a/tests/test_lab_safety.py +++ b/tests/test_lab_safety.py @@ -169,3 +169,17 @@ def test_the_audit_holds_a_pyproject_to_its_dependency_tables(tmp_path): assert grader.audit(repo, "HEAD", ws).ok (ws / "pyproject.toml").unlink() assert "may not delete a dependency file: pyproject.toml" in grader.audit(repo, "HEAD", ws).violations + + +def test_a_venv_in_scratch_is_left_out_but_a_link_in_the_surface_is_a_violation(tmp_path): + repo = make_repo(tmp_path) + ws = tmp_path / "ws" + grader.export(repo, "HEAD", ws) + venv = ws / ".venv" / "bin" + venv.mkdir(parents=True) + (venv / "python").symlink_to("/usr/bin/python3") + (ws / ".venv" / "big.whl").write_bytes(b"x" * (grader.MAX_FILE_BYTES + 1)) + a = grader.audit(repo, "HEAD", ws) + assert a.ok and {".venv/bin/python", ".venv/big.whl"} <= set(a.scratch) + (ws / "src/inference_server/sneaky.py").symlink_to("/etc/passwd") + assert "symlink: src/inference_server/sneaky.py" in grader.audit(repo, "HEAD", ws).violations From 4dbe4b4480e31626d90e3f8c6bde0b3e48d604c4 Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Sat, 10 Oct 2026 15:07:05 -0400 Subject: [PATCH 11/18] grader: ignored paths stay ignored before the scratch check Co-Authored-By: Claude Opus 5.5 --- lab/safety/grader.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/lab/safety/grader.py b/lab/safety/grader.py index 87d320b..ac03976 100644 --- a/lab/safety/grader.py +++ b/lab/safety/grader.py @@ -94,6 +94,8 @@ def audit(repo: Path, base: str, workspace: Path) -> Audit: for name in dirs + files: p = Path(root) / name rel = p.relative_to(workspace).as_posix() + if not p.is_symlink() and any(fnmatch(rel, g) for g in IGNORED): + continue if _scratch(rel, blobs) and (p.is_symlink() or not (p.is_dir() or p.is_file()) or (p.is_file() and p.stat().st_size > MAX_FILE_BYTES)): a.scratch.append(rel) # never reaches the pristine tree: a venv's links, big wheels From 0ba019f3151817b6fc431244e06d5b9328506c87 Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Sat, 10 Oct 2026 15:13:03 -0400 Subject: [PATCH 12/18] CLAUDE.md: the first agent dry run on an H200 Co-Authored-By: Claude Opus 5.5 --- CLAUDE.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/CLAUDE.md b/CLAUDE.md index 0131674..b2bf3da 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -231,6 +231,16 @@ Linux the jail is a container with the NVIDIA runtime (`lab/safety/container.py` check (`vm-setup.sh` prints `jail: gpu ok`) has not run yet. The agent's CLI still runs under srt, which on Ubuntu 24.04 needs `kernel.apparmor_restrict_unprivileged_userns=0` (in `vm-setup.sh`). +## First agent dry run (Nebius H200, 2026-10-10) +`LAB_WORKER=ssh`, laptop controller, `claude-fable-5-1`, one session ($2.53 as the proxy priced it, 24 turns, 465 s; +~1 h 20 min of GPU for the whole run). Validated on the real GPU: the GPU in the network-less jail, metadata blocked, +base and candidate served jailed through the port bridge, `profile`, `trace` (after mounting nsys's install root), +`test` (after caching the tests' gpt2 tokenizer), the remote workbench over `ssh -R`, the proxy's pricing. The agent +fused RMSNorm/RoPE in Triton and added a batch-1 decode graph: seen-split short-tier `single_stream` tpot p99 21.4 -> +13.1 ms; not submitted, correctness not run (equiv needs `data/equiv/ref`, absent in a worktree). A base dev-tier equiv +on this engine took > 46 min for 774 of 1,597 items at concurrency 16 (vLLM: ~10 min): the tool descriptions' "10 min +at tier dev" is wrong for this engine. Workbench processes cannot change clocks or power limits on that platform. + ## Next GPU (needs spend approval): resolve the BF16 puzzle (plain BF16 on the dev tier at concurrency 32 and 128 on one VM, ~20 min, ~$1.50); first real run of `python -m regimes run all` against the engine and against vLLM. From b44a153478387e0540e9c16e263d8e7268e3a621 Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Tue, 13 Oct 2026 13:45:42 -0400 Subject: [PATCH 13/18] lab.gpucheck: the GPU-host checks left open after the dry run, in one command Co-Authored-By: Claude Opus 5.5 --- CLAUDE.md | 3 + lab/gpucheck.py | 188 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 191 insertions(+) create mode 100644 lab/gpucheck.py diff --git a/CLAUDE.md b/CLAUDE.md index b2bf3da..4b4955e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -72,6 +72,9 @@ pytest/ruff config). Each folder with commands has its own `python -m`. tables only, checked by the audit) gets a venv built by the target's `build` command in a no-GPU, `--network none` container whose only way out is a CONNECT proxy to the package-index hosts; cached by the files' hash; jobs run under `engine.using(build.python_for(tree))`. +- `lab/gpucheck.py` — `python -m lab.gpucheck` on a GPU VM (as root, like the worker): what the jails need from a real + GPU host, PASS/FAIL each: jail GPU, workbench (GPU, root, apt/pip, metadata blocked), quiet, exclusivity, ncu, + hostprof, the dependency build. - `lab/worker.py` — the job interface, built for the measured jobs: `test`, `measure`, `equiv`, `profile` take input trees + JSON args, write an out dir, return JSON; the tools call `self.worker`. `LAB_WORKER=local` (default, in process) or `ssh` (the `lab.vm` VM): inputs by content hash, re-hashed on the worker before every job; GpuBusy, diff --git a/lab/gpucheck.py b/lab/gpucheck.py new file mode 100644 index 0000000..917548f --- /dev/null +++ b/lab/gpucheck.py @@ -0,0 +1,188 @@ +"""What lab/RUNTIME.md's jails need from a real GPU host, checked on one: run on the GPU VM, as root the way the remote +worker runs (`sudo`), with the engine's venv first on PATH. Each check prints PASS or FAIL with its evidence; the exit +status is the number of failures. `--only NAME,...` runs a subset; `--skip-build` leaves out the dependency build. + + sudo -n env HOME="$HOME" PATH="$HOME/inference-server/.venv/bin:$PATH" LAB_ENGINE_REPO="$HOME/inference-server" \\ + LAB_TARGET=targets/inference-server.toml python -m lab.gpucheck +""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +import subprocess +import sys +import tempfile +import time +import traceback +from pathlib import Path + +from lab import engine, gpu, serve, worker +from lab.safety import container, grader + +HOLD = "import torch, time; x = torch.empty(2**30, dtype=torch.uint8, device='cuda'); print('holding', flush=True); time.sleep(600)" +MEASURE = {"names": ["single_stream"], "split": "seen", "tier": "short", "seed": 0, "jailed": True, "passive": False} +results: list[tuple[str, bool, str]] = [] + + +def check(name: str): + def wrap(fn): + fn.check = name + return fn + return wrap + + +def base_tree(root: Path) -> Path: + """The worker's engine checkout as a two-commit tree (base, empty change): the shape every job gets.""" + tree = root / "base" + shutil.copytree(engine.repo(), tree, ignore=shutil.ignore_patterns(".venv", ".git", "__pycache__", "*.pyc")) + grader._commit(tree, "base", init=True) + grader._commit(tree, "change") + return tree + + +def holder(name: str, wb: bool, tmp: Path) -> subprocess.Popen: + """A process holding 1 GiB of GPU memory: in a workbench container (`wb`) or on the host.""" + py = engine.python() + if wb: + argv = container.workbench_argv(name, tmp, [py, "-c", HOLD], {"HOME": str(tmp)}, [tmp], [engine.venv()], py) + else: + argv = [py, "-c", HOLD] + p = subprocess.Popen(argv, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) + deadline = time.monotonic() + 120 + while time.monotonic() < deadline and "holding" not in (p.stdout.readline() or ""): + if p.poll() is not None: + raise RuntimeError(f"the GPU holder exited {p.returncode}") + return p + + +@check("jail_gpu") +def jail_gpu(tmp: Path) -> str: + """A measurement room (no network, read-only root, uid 10001) sees the GPU.""" + argv = container.argv(container.new_name(), tmp, ["nvidia-smi", "-L"], {}, [tmp], [], engine.python(), gpu=True) + out = subprocess.run(argv, capture_output=True, text=True, timeout=120) + assert out.returncode == 0 and "GPU 0" in out.stdout, out.stdout + out.stderr + return out.stdout.strip() + + +@check("workbench") +def workbench(tmp: Path) -> str: + """The workbench: the GPU, root, apt and pip over the internet, and no cloud metadata service.""" + script = ("set -e; nvidia-smi -L; id -u; apt-get update -qq; apt-get install -y -qq --no-install-recommends " + "python3-pip >/dev/null; pip3 install --quiet --break-system-packages six; python3 -c 'import six'; " + "echo pip-ok; python3 -c \"import socket; socket.create_connection(('169.254.169.254', 80), timeout=3)\" " + "&& echo METADATA-REACHABLE || echo metadata-blocked") + argv = container.workbench_argv(container.new_name(), tmp, ["sh", "-c", script], {"HOME": str(tmp)}, [tmp], [], + engine.python()) + out = subprocess.run(argv, capture_output=True, text=True, timeout=600, stdin=subprocess.DEVNULL) + text = out.stdout + out.stderr + assert out.returncode == 0 and "GPU 0" in text and "pip-ok" in text and "metadata-blocked" in text, text[-2000:] + assert "\n0\n" in f"\n{out.stdout}", "not root inside the workbench (run this check as root, like the worker)" + return "GPU, root, apt, pip; metadata blocked" + + +@check("quiet") +def quiet(tmp: Path) -> str: + """`container.quiet` kills the workbench's GPU holder, freezes the rest, and thaws it after.""" + name = container.new_name() + p = holder(name, True, tmp) + try: + mine = container.pids(name) + assert any(a["pid"] in mine for a in gpu.compute_apps() or []), "the holder is not on the GPU" + with container.quiet(name) as killed: + state = subprocess.run(["docker", "inspect", "-f", "{{.State.Paused}}", name], capture_output=True, + text=True).stdout.strip() + left = [a for a in gpu.compute_apps() or [] if a["pid"] in mine] + after = subprocess.run(["docker", "inspect", "-f", "{{.State.Paused}}", name], capture_output=True, + text=True).stdout.strip() + assert killed and not left and state == "true" and after == "false", (killed, left, state, after) + return f"killed {killed}; paused during, thawed after" + finally: + p.kill() + container.remove(name) + + +@check("exclusive") +def exclusive(tmp: Path) -> str: + """A clean measurement is not contaminated (only the engine's processes on the GPU); a foreign one is refused.""" + tree = base_tree(tmp) + clean = worker.Local().call("measure", {"tree": tree}, MEASURE, tmp / "clean") + p = holder(None, False, tmp) + try: + worker.Local().call("measure", {"tree": tree}, MEASURE, tmp / "busy") + raise AssertionError("measured with a foreign process on the GPU") + except serve.GpuBusy as e: + refused = e.apps + finally: + p.kill() + v = clean["results"][0] + return f"clean: {v['objective']} {v['value']:.4f} valid={v['valid']}; refused with {refused}" + + +@check("ncu") +def ncu(tmp: Path) -> str: + """Nsight Compute inside a measurement room (needs the GPU's performance counters for the jail's user).""" + r = worker.Local().call("profile", {"tree": base_tree(tmp)}, {"kind": "kernel", "args": {"launch_count": 2}}, + tmp / "out") + assert not r.get("refused") and r["returncode"] == 0 and r["head"], r.get("refused") or r["output"][-1500:] + return r["head"][:300] + + +@check("hostprof") +def hostprof(tmp: Path) -> str: + """py-spy inside a measurement room.""" + r = worker.Local().call("profile", {"tree": base_tree(tmp)}, {"kind": "hostprof", "args": {"seconds": 5}}, + tmp / "out") + assert not r.get("refused") and r["returncode"] == 0 and r["head"], r.get("refused") or r["output"][-1500:] + return r["head"][:300] + + +@check("build") +def build_step(tmp: Path) -> str: + """A dependency change is built from scratch (index-only network) and the built venv runs torch on the GPU.""" + from lab import build + tree = tmp / "base" + shutil.copytree(engine.repo(), tree, ignore=shutil.ignore_patterns(".venv", ".git", "__pycache__", "*.pyc")) + grader._commit(tree, "base", init=True) + uv = shutil.which("uv") or str(Path.home() / ".local" / "bin" / "uv") + subprocess.run([uv, "add", "--no-sync", "six"], cwd=tree, check=True, capture_output=True) + grader._commit(tree, "change") + assert build.changed(tree), "the change to pyproject.toml/uv.lock was not seen" + t0 = time.monotonic() + py = build.python_for(tree) + out = subprocess.run([py, "-c", "import six, torch; print(torch.__version__, torch.cuda.is_available())"], + capture_output=True, text=True) + assert out.returncode == 0 and out.stdout.strip().endswith("True"), out.stdout + out.stderr + return f"built in {time.monotonic() - t0:.0f}s: {py}: torch {out.stdout.strip()}" + + +CHECKS = [jail_gpu, workbench, quiet, exclusive, ncu, hostprof, build_step] + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(prog="python -m lab.gpucheck", description=__doc__.split("\n\n")[0]) + ap.add_argument("--only", help="comma-separated: " + ",".join(c.check for c in CHECKS)) + ap.add_argument("--skip-build", action="store_true") + a = ap.parse_args(argv) + want = set(a.only.split(",")) if a.only else {c.check for c in CHECKS} - ({"build"} if a.skip_build else set()) + print(f"as uid {os.getuid()}; GPU: {gpu.query()}", flush=True) + for c in CHECKS: + if c.check not in want: + continue + tmp = Path(tempfile.mkdtemp(prefix=f"gpucheck-{c.check}-")) + t0 = time.monotonic() + try: + detail, ok = c(tmp), True + except Exception as e: + detail, ok = f"{type(e).__name__}: {e}\n{traceback.format_exc()[-1500:]}", False + results.append((c.check, ok, detail)) + print(f"{'PASS' if ok else 'FAIL'} {c.check} ({time.monotonic() - t0:.0f}s): {detail}", flush=True) + shutil.rmtree(tmp, ignore_errors=True) + print(json.dumps({n: ok for n, ok, _ in results})) + return sum(not ok for _, ok, _ in results) + + +if __name__ == "__main__": + sys.exit(main()) From 5b30cc2235805d8082695c3348cb7e399ce27d7c Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Tue, 13 Oct 2026 15:15:09 -0400 Subject: [PATCH 14/18] lab: fixes from the gpucheck run (workbench DNS, ncu on PATH, NVIDIA's wheel index) - vm-setup.sh drops only 169.254.169.254 for containers, not 169.254.0.0/16 (a cloud's DNS can sit there; the workbench had no names), and checks the workbench's internet. - vm-setup.sh links /usr/local/cuda/bin/ncu into /usr/local/bin when ncu is on no PATH. - The build's index hosts include download.pytorch.org, download-r2.pytorch.org and pypi.nvidia.com: CUDA torch's nvidia-* wheels come from NVIDIA's index. - gpucheck's workbench check prints resolv.conf. Co-Authored-By: Claude Opus 5.5 --- lab/RUNTIME.md | 3 ++- lab/build.py | 5 ++++- lab/gpucheck.py | 2 +- lab/vm-setup.sh | 7 ++++++- tests/test_lab_build.py | 2 +- 5 files changed, 14 insertions(+), 5 deletions(-) diff --git a/lab/RUNTIME.md b/lab/RUNTIME.md index 086dff5..7a05aaa 100644 --- a/lab/RUNTIME.md +++ b/lab/RUNTIME.md @@ -173,7 +173,8 @@ not yet over SSH to a VM. `agent`, `exec`, `snapshot`, `restore` and `facts` com interpreter, weights, CUDA and the CLI binary read-only, so an install in the workbench never reaches a measurement. Container root is uid 0 on the host kernel (no user-namespace remap): an escape is host root, accepted under decision 5; `userns-remap` is the hardening step. `vm-setup.sh` drops container traffic to - 169.254.0.0/16 (the cloud metadata service) and checks it. + 169.254.169.254 (the cloud metadata service; only that address, a cloud's DNS can sit elsewhere in 169.254/16) and + checks both that and the workbench's internet. - **Pause** (`container.quiet`): around every `measure`, `equiv` and `profile` job, processes in the workbench holding the GPU are killed, the container is paused, a second sweep catches any that started before the pause landed, and it waits for their memory to be released; thawed after. The killed processes are in the record and diff --git a/lab/build.py b/lab/build.py index 8681745..b3e3ef3 100644 --- a/lab/build.py +++ b/lab/build.py @@ -23,7 +23,10 @@ from lab import engine, target from lab.safety import container -INDEX_HOSTS = ("pypi.org", "files.pythonhosted.org") # plus every host the dependency files and build name +# PyPI, plus the wheel hosts that CUDA builds of torch pull from (its nvidia-* dependencies live on NVIDIA's index); +# every https host the dependency files and the build command name is added per build. +INDEX_HOSTS = ("pypi.org", "files.pythonhosted.org", "download.pytorch.org", "download-r2.pytorch.org", + "pypi.nvidia.com") PROXY_PORT_IN = 3128 LOG_TAIL = 4000 TIMEOUT_S = 3600 diff --git a/lab/gpucheck.py b/lab/gpucheck.py index 917548f..ebb7bed 100644 --- a/lab/gpucheck.py +++ b/lab/gpucheck.py @@ -70,7 +70,7 @@ def jail_gpu(tmp: Path) -> str: @check("workbench") def workbench(tmp: Path) -> str: """The workbench: the GPU, root, apt and pip over the internet, and no cloud metadata service.""" - script = ("set -e; nvidia-smi -L; id -u; apt-get update -qq; apt-get install -y -qq --no-install-recommends " + script = ("cat /etc/resolv.conf; set -e; nvidia-smi -L; id -u; apt-get update -qq; apt-get install -y -qq --no-install-recommends " "python3-pip >/dev/null; pip3 install --quiet --break-system-packages six; python3 -c 'import six'; " "echo pip-ok; python3 -c \"import socket; socket.create_connection(('169.254.169.254', 80), timeout=3)\" " "&& echo METADATA-REACHABLE || echo metadata-blocked") diff --git a/lab/vm-setup.sh b/lab/vm-setup.sh index 824d17c..f0ebbd5 100755 --- a/lab/vm-setup.sh +++ b/lab/vm-setup.sh @@ -30,6 +30,8 @@ echo "== CUDA toolkit" nvcc --version 2>/dev/null | tail -1 || ls /usr/local/cuda/bin/nvcc 2>/dev/null || echo "nvcc: not on PATH" echo "== Nsight" +# CUDA images keep ncu in /usr/local/cuda/bin, which is on no login PATH; the tools look it up by name. +command -v ncu >/dev/null || { [ -x /usr/local/cuda/bin/ncu ] && $SUDO ln -sf /usr/local/cuda/bin/ncu /usr/local/bin/ncu; } if ! command -v nsys >/dev/null || ! command -v ncu >/dev/null; then $SUDO apt-get update -qq $SUDO apt-get install -y -qq nsight-systems nsight-compute 2>/dev/null \ @@ -67,9 +69,12 @@ fi $SUDO docker build -q -t lab-jail -f "$here/safety/jail.Dockerfile" "$here/safety" >/dev/null $SUDO docker run --rm --gpus all --network none lab-jail nvidia-smi -L && echo "jail: gpu ok" || echo "jail: no GPU in the container" # The agent's workbench has the internet but not the cloud's metadata service (instance credentials, user-data). -$SUDO iptables -C DOCKER-USER -d 169.254.0.0/16 -j DROP 2>/dev/null || $SUDO iptables -I DOCKER-USER -d 169.254.0.0/16 -j DROP +# Only that address: a cloud's DNS resolver can sit in 169.254.0.0/16 too, and the workbench needs names. +$SUDO iptables -C DOCKER-USER -d 169.254.169.254 -j DROP 2>/dev/null || $SUDO iptables -I DOCKER-USER -d 169.254.169.254 -j DROP $SUDO docker run --rm lab-jail python3 -c "import socket; socket.create_connection(('169.254.169.254', 80), timeout=3)" \ 2>/dev/null && echo "workbench: METADATA REACHABLE" || echo "workbench: metadata blocked" +$SUDO docker run --rm lab-jail python3 -c "import socket; socket.create_connection(('pypi.org', 443), timeout=10)" \ + 2>/dev/null && echo "workbench: internet ok (DNS and egress)" || echo "workbench: NO INTERNET (DNS or egress)" echo "== counters: can a non-admin process read GPU performance counters?" # ncu needs NVreg_RestrictProfilingToAdminUsers=0 (or root). Recorded as a fact about this venue. diff --git a/tests/test_lab_build.py b/tests/test_lab_build.py index 8c3e070..7893977 100644 --- a/tests/test_lab_build.py +++ b/tests/test_lab_build.py @@ -66,7 +66,7 @@ def test_the_allowed_hosts_are_the_index_defaults_plus_what_the_files_and_build_ tree.mkdir() (tree / "pyproject.toml").write_text('[[tool.uv.index]]\nurl = "https://download.pytorch.org/whl/cu130"\n') got = build.hosts(tree, "uv sync --index-url https://mirror.example.org/simple") - assert got == {"pypi.org", "files.pythonhosted.org", "download.pytorch.org", "mirror.example.org"} + assert got == {*build.INDEX_HOSTS, "mirror.example.org"} and "pypi.nvidia.com" in got def proxy_request(sock_path: Path, line: bytes) -> bytes: From 59fcf861f225182a00ac5cf86b11f1a1a51f3996 Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Tue, 13 Oct 2026 15:28:52 -0400 Subject: [PATCH 15/18] ncu: link the real install, and mount /opt/nvidia into profiler jails CUDA's bin/ncu is a launcher for /opt/nvidia/nsight-compute/; the jail mounted neither (gpucheck on an H200: 'nsight-compute directory is not found'). Co-Authored-By: Claude Opus 5.5 --- lab/vm-setup.sh | 7 +++++-- lab/worker.py | 2 ++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/lab/vm-setup.sh b/lab/vm-setup.sh index f0ebbd5..3503c3e 100755 --- a/lab/vm-setup.sh +++ b/lab/vm-setup.sh @@ -30,8 +30,11 @@ echo "== CUDA toolkit" nvcc --version 2>/dev/null | tail -1 || ls /usr/local/cuda/bin/nvcc 2>/dev/null || echo "nvcc: not on PATH" echo "== Nsight" -# CUDA images keep ncu in /usr/local/cuda/bin, which is on no login PATH; the tools look it up by name. -command -v ncu >/dev/null || { [ -x /usr/local/cuda/bin/ncu ] && $SUDO ln -sf /usr/local/cuda/bin/ncu /usr/local/bin/ncu; } +# CUDA images keep ncu in /usr/local/cuda/bin, on no login PATH, and that one is only a launcher for the install under +# /opt/nvidia/nsight-compute/: link the real one, so the tool's jail mounts its whole install root. +ncu_real=$(ls -d /opt/nvidia/nsight-compute/*/ncu 2>/dev/null | sort -V | tail -1) +if [ -n "$ncu_real" ]; then $SUDO ln -sf "$ncu_real" /usr/local/bin/ncu +elif ! command -v ncu >/dev/null && [ -x /usr/local/cuda/bin/ncu ]; then $SUDO ln -sf /usr/local/cuda/bin/ncu /usr/local/bin/ncu; fi if ! command -v nsys >/dev/null || ! command -v ncu >/dev/null; then $SUDO apt-get update -qq $SUDO apt-get install -y -qq nsight-systems nsight-compute 2>/dev/null \ diff --git a/lab/worker.py b/lab/worker.py index d5fe124..cc12ebe 100644 --- a/lab/worker.py +++ b/lab/worker.py @@ -149,6 +149,8 @@ def _profile(inputs: dict, out: Path, args: dict, hooks: dict) -> dict: work = tree / proftools.OUT work.mkdir() read = [Path(exe).parent, Path(exe).parent.parent] # the install root: nsys's importer lives beside its binary dir + if Path("/opt/nvidia").is_dir(): + read.append(Path("/opt/nvidia")) # where launchers like CUDA's bin/ncu find the real install src = T.workload_flags(a, harness, proftools._int(a, "requests", *proftools.WORKLOAD["requests"])) argv = [*ins["wrap"](exe, work), engine.python(), "-m", "lab.profile", *ins["flags"], *src, *ins["extra"](exe), "--out", str(work / "bundle")] From 099ebf422130c731dfe2fddd99e6731333c4ac66 Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Tue, 13 Oct 2026 15:44:03 -0400 Subject: [PATCH 16/18] CLAUDE.md: lab.gpucheck passes on an H200 Co-Authored-By: Claude Opus 5.5 --- CLAUDE.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CLAUDE.md b/CLAUDE.md index 4b4955e..588b022 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -243,6 +243,10 @@ fused RMSNorm/RoPE in Triton and added a batch-1 decode graph: seen-split short- 13.1 ms; not submitted, correctness not run (equiv needs `data/equiv/ref`, absent in a worktree). A base dev-tier equiv on this engine took > 46 min for 774 of 1,597 items at concurrency 16 (vLLM: ~10 min): the tool descriptions' "10 min at tier dev" is wrong for this engine. Workbench processes cannot change clocks or power limits on that platform. +`python -m lab.gpucheck` then passed every item on Nebius H200s (2026-10-13, three runs, fixes between them): +jail GPU, quiet, exclusivity, hostprof; workbench (after blocking only 169.254.169.254: the /16 took DNS with it); +the dependency build (after allowing pypi.nvidia.com); ncu (after linking the real /opt/nvidia install). The jail's +user can read GPU performance counters on that image. ## Next GPU (needs spend approval): resolve the BF16 puzzle (plain BF16 on the dev tier at concurrency 32 and 128 on one VM, From 46547e0fcda4f5f4f0f8dafa67512f69ba73c4b8 Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Tue, 13 Oct 2026 20:08:36 -0400 Subject: [PATCH 17/18] session: a provider that works locally is refused with a Remote worker With LAB_WORKER=ssh the workspace is mirrored from the VM before every audit; a provider whose agent edits this machine's copy (openai, #23) would have its edits overwritten. Providers say whether they can work on the worker's host (ClaudeAgentSDK.remote). Co-Authored-By: Claude Opus 5.5 --- lab/agent.py | 1 + lab/session.py | 2 ++ tests/test_lab_runtime.py | 14 ++++++++++++++ 3 files changed, 17 insertions(+) diff --git a/lab/agent.py b/lab/agent.py index ac3fcbb..b48cf10 100644 --- a/lab/agent.py +++ b/lab/agent.py @@ -121,6 +121,7 @@ class ClaudeAgentSDK: srt); its model calls go through the lab's API proxy, so no credential is ever inside; lab tools run out here.""" name = "claude" + remote = True # can work on a Remote worker's host (its CLI over ssh) @staticmethod def containerized() -> bool: diff --git a/lab/session.py b/lab/session.py index aa6f96b..39341a2 100644 --- a/lab/session.py +++ b/lab/session.py @@ -128,6 +128,8 @@ def run(cfg: RunConfig, provider: agent.Provider | None = None) -> dict: summary = {"run": run_id, "base": base, "sessions": 0, "stopped": None, "spent_usd": budget.spent_usd} wk = worker.load() remote = wk if isinstance(wk, worker.Remote) else None # then the agent works on that host, not here + if remote and not getattr(provider, "remote", False): # its edits here would be overwritten by the mirror + raise ValueError(f"provider {provider.name!r} works on this machine; it cannot run with LAB_WORKER=ssh") rws = remote.workspace_path(run_id) if remote else None if cfg.baseline: Toolbox(Session(run_id, f"{run_id}-baseline", run_dir, ws, budget, cfg.ledger_root, task=cfg.task, diff --git a/tests/test_lab_runtime.py b/tests/test_lab_runtime.py index 97ec719..ec54c11 100644 --- a/tests/test_lab_runtime.py +++ b/tests/test_lab_runtime.py @@ -255,3 +255,17 @@ def script(n, ws: Path): assert t.model in system and t.engine_repo.name in system and "`stop`" in system and "`note`" in system assert "{" not in system.replace(t.engine.launch, "").replace(json.dumps(t.engine.env), "") # all filled assert not ADVICE.search(system), ADVICE.search(system) + + +def test_a_provider_that_works_locally_is_refused_with_a_remote_worker(cfg, monkeypatch): + from lab import worker + + class Far(worker.Remote): + def __init__(self): + pass + + def workspace_path(self, run_id): + return "/far" + monkeypatch.setattr(worker, "load", lambda **k: Far()) + with pytest.raises(ValueError, match="cannot run with LAB_WORKER=ssh"): + session.run(cfg, ScriptedProvider(lambda n, ws: iter(()))) From e0b1d5a8c5f856a15ebd02500c8162205e684109 Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Tue, 13 Oct 2026 20:38:25 -0400 Subject: [PATCH 18/18] lab: freeze the remote workbench for every tool call; srt reaches the proxy; proxy guards From the stack review (BlameGraph session): - A (high): in LAB_WORKER=ssh mode the workbench ran while the worker's root rsync pulled and pushed its workspace; swapping a directory for a link mid-transfer could carry the held-out split out, or let root write anywhere on the VM. The workbench is now frozen on the VM (`lab.worker pause`) for the whole of every lab tool call, as in local mode; Remote.paused thaws only what it froze. - B: srt on macOS reaches the proxy on 127.0.0.1 only with allowLocalBinding (verified with srt 0.0.78); set on that path, which opens every local port: dev only. - The proxy reserves the input too (at most a token per byte of body), and refuses a /v1/messages request without a JSON body naming a priced model instead of crashing. Co-Authored-By: Claude Opus 5.5 --- lab/RUNTIME.md | 3 +++ lab/agent.py | 16 +++++++++++----- lab/apiproxy.py | 7 +++++-- lab/safety/container.py | 27 ++++++++++++++++++--------- lab/worker.py | 17 ++++++++++++++++- tests/test_lab_apiproxy.py | 20 ++++++++++++++++++++ tests/test_lab_worker.py | 10 ++++++++++ 7 files changed, 83 insertions(+), 17 deletions(-) diff --git a/lab/RUNTIME.md b/lab/RUNTIME.md index 7a05aaa..fd3a748 100644 --- a/lab/RUNTIME.md +++ b/lab/RUNTIME.md @@ -207,6 +207,9 @@ on the VM: on the VM, mounted into the workbench, where a socat serves it on the container's own 127.0.0.1:4000. It listens on no interface of the VM. - The write guard checks the VM's workspace path; the workbench is removed over SSH when the session ends. +- The workbench is frozen on the VM (`python -m lab.worker pause`) for the whole of every lab tool call, so the root + rsync of a pull or push never races a process in it swapping a file or a directory for a link (2026-10-13 review: + a swapped directory could have carried the held-out split out, or let root write anywhere on the VM). - Tested without a VM: rsync mirroring (content edits, deletions, links), the VM-side `workbench` command, and the argument quoting through bash and ssh (an `ssh` stand-in that prints the remote command). Not yet tested against a real VM. diff --git a/lab/agent.py b/lab/agent.py index b48cf10..8debf11 100644 --- a/lab/agent.py +++ b/lab/agent.py @@ -152,7 +152,10 @@ def _wrapper(self, spec: AgentSpec, cli: Path, base_url: str) -> Path: cli_tmp = Path(f"{tmp_root}/claude-{os.getuid()}") # the CLI's own per-project scratch; Bash fails without it cli_tmp.mkdir(parents=True, exist_ok=True) config = jail.settings([home.resolve(), tmp.resolve(), ws, cli_tmp.resolve()], engine.venv(), - [base_url.removeprefix("http://")], readonly=_weights(), python=engine.python()) + [base_url.removeprefix("http://")], readonly=_weights(), python=engine.python(), + local_binding=True) # srt on macOS reaches 127.0.0.1 only with it (and then + # every local port: dev only; measurements that count run + # on Linux, where each jail has its own network) argv = jail.wrap(config, spec.scratch / "srt.json", [str(cli)]) keep = " ".join(f'"{k}=${k}"' for k in PASS_ENV) sets = " ".join(f"{k}={shlex.quote(v)}" for k, v in cli_env(home, tmp, ws, base_url).items()) @@ -166,15 +169,18 @@ async def _run(self, spec: AgentSpec, base_url: str) -> AgentReply: from claude_agent_sdk import (ClaudeAgentOptions, HookMatcher, ResultMessage, create_sdk_mcp_server, query, tool) - # Sharing this host, the workbench is frozen for every lab tool call: nothing in it can swap a file for a link - # while the tool snapshots, restores or copies into the workspace (with a Remote worker only rsync writes here). - shared = spec.remote is None and self.containerized() + # The workbench is frozen for every lab tool call, wherever it runs: nothing in it can swap a file or a dir for a + # link while the tool snapshots, restores, copies into the workspace, or (Remote) rsyncs it as root on the VM. + if spec.remote is not None: + frozen = lambda: spec.remote.paused(spec.workbench) # noqa: E731 + else: + frozen = lambda: container.paused(spec.workbench if self.containerized() else None) # noqa: E731 def adapt(t: ToolSpec): async def call(args: dict[str, Any]) -> dict[str, Any]: try: def run(): - with container.paused(spec.workbench if shared else None): + with frozen(): return t.fn(args) text = await asyncio.to_thread(run) # a long test run must not block the transport except Exception as e: # the agent sees the refusal, the harness keeps going diff --git a/lab/apiproxy.py b/lab/apiproxy.py index 0beff8c..aed95a6 100644 --- a/lab/apiproxy.py +++ b/lab/apiproxy.py @@ -156,8 +156,11 @@ def _handle(self, h, body: bytes | None) -> None: if req.get("speed") == "fast": return self._refuse(h, 403, "fast mode is not priced here") metered = path == "/v1/messages" - # the request's worst-case output is reserved before it goes, so concurrent requests cannot overshoot the cap - hold = max_tokens * price(model)[1] / 1e6 if metered else 0.0 + if metered and (not body or not isinstance(model, str) or price(model) is None): + return self._refuse(h, 403, "a /v1/messages request needs a JSON body naming a priced model") + # Its worst case is reserved before it goes (input: at most a token per byte of body; output: all of max_tokens), + # so concurrent requests cannot overshoot the cap. + hold = (len(body) * price(model)[0] + max_tokens * price(model)[1]) / 1e6 if metered else 0.0 with self.lock: if metered and self.spent_usd + self.reserved_usd + hold > self.cap_usd: return self._refuse(h, 403, f"this session's ${self.cap_usd:.2f} is spent (or would be by this request)") diff --git a/lab/safety/container.py b/lab/safety/container.py index d6bc1ac..9fb5ec7 100644 --- a/lab/safety/container.py +++ b/lab/safety/container.py @@ -138,21 +138,30 @@ def _kill_gpu_holders(name: str) -> list[dict]: return hit +def pause(name: str | None) -> bool: + """Freeze a running container; True only if this call froze it (False: absent, stopped, or already paused).""" + if not name or _docker() is None or not running(name): + return False + if subprocess.run([_docker(), "inspect", "-f", "{{.State.Paused}}", name], capture_output=True, + text=True).stdout.strip() == "true": + return False + subprocess.run([_docker(), "pause", name], capture_output=True, check=True) + return True + + +def unpause(name: str) -> None: + subprocess.run([_docker(), "unpause", name], capture_output=True) + + @contextmanager def paused(name: str | None): """The container frozen for the block (a no-op if it is not running); left paused if it already was.""" - if not name or _docker() is None or not running(name): - yield - return - already = subprocess.run([_docker(), "inspect", "-f", "{{.State.Paused}}", name], capture_output=True, - text=True).stdout.strip() == "true" - if not already: - subprocess.run([_docker(), "pause", name], capture_output=True, check=True) + did = pause(name) try: yield finally: - if not already: - subprocess.run([_docker(), "unpause", name], capture_output=True) + if did: + unpause(name) @contextmanager diff --git a/lab/worker.py b/lab/worker.py index add4d8c..c07d794 100644 --- a/lab/worker.py +++ b/lab/worker.py @@ -23,6 +23,7 @@ from __future__ import annotations import dataclasses +from contextlib import contextmanager import hashlib import io import json @@ -264,6 +265,16 @@ def workbench(self, name: str, session: str, workspace: str, sdk_version: str) - req = {"name": name, "session": session, "workspace": workspace, "sdk_version": sdk_version} return json.loads(self._run("workbench", json.dumps(req).encode())) + @contextmanager + def paused(self, name: str | None): + """The workbench on the worker host frozen for the block (as `container.paused` there).""" + did = bool(name) and self._run(f"pause {shlex.quote(name)}").decode().strip() == "paused" + try: + yield + finally: + if did: + self._run(f"unpause {shlex.quote(name)}") + def remove_workbench(self, name: str | None) -> None: if name: self._run(f"rm {shlex.quote(name)}") @@ -400,9 +411,13 @@ def main(argv: list[str] | None = None) -> int: print(json.dumps(_workbench(json.loads(sys.stdin.buffer.read())))) elif argv[:1] == ["rm"] and len(argv) == 2: container.remove(argv[1]) + elif argv[:1] == ["pause"] and len(argv) == 2: # prints "paused" only if this call froze it + print("paused" if container.pause(argv[1]) else "") + elif argv[:1] == ["unpause"] and len(argv) == 2: + container.unpause(argv[1]) else: print(f"usage: python -m lab.worker has SHA... | put SHA | run {{{','.join(JOBS)}}} | root | mkdir PATH | " - "workbench | rm NAME", file=sys.stderr) + "workbench | rm NAME | pause NAME | unpause NAME", file=sys.stderr) return 2 return 0 diff --git a/tests/test_lab_apiproxy.py b/tests/test_lab_apiproxy.py index e3b2c32..5b0c036 100644 --- a/tests/test_lab_apiproxy.py +++ b/tests/test_lab_apiproxy.py @@ -178,3 +178,23 @@ def test_a_request_whose_worst_case_would_pass_the_cap_is_refused_before_it_goes with apiproxy.Proxy(0.10, upstream=url) as p: # 10k output tokens of opus is $0.20 status, body = post(p, "/v1/messages", {"model": "claude-opus-5-5", "max_tokens": 10000}) assert status == 403 and b"would be" in body and seen == [] + + +def test_a_bodyless_or_unpriced_messages_request_is_refused_not_crashed(upstream, monkeypatch): + url, seen = upstream + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-real") + with apiproxy.Proxy(5.0, upstream=url) as p: + c = http.client.HTTPConnection(urlsplit(p.url).netloc, timeout=10) + c.request("GET", "/v1/messages") + r = c.getresponse() + assert r.status == 403 and b"priced model" in r.read() + assert post(p, "/v1/messages", {"max_tokens": 10})[0] == 403 + assert seen == [] + + +def test_the_input_is_reserved_too(upstream, monkeypatch): + url, seen = upstream + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-real") + with apiproxy.Proxy(0.05, upstream=url) as p: # 20k bytes of opus input is up to $0.08 + status, body = post(p, "/v1/messages", {"model": "claude-opus-5-5", "max_tokens": 1, "pad": "x" * 20000}) + assert status == 403 and seen == [] diff --git a/tests/test_lab_worker.py b/tests/test_lab_worker.py index 89fa359..b007197 100644 --- a/tests/test_lab_worker.py +++ b/tests/test_lab_worker.py @@ -135,3 +135,13 @@ def test_the_worker_builds_its_own_workbench_with_the_tunnelled_proxy(remote, tm assert argv[1:3] == ["run", "-i"] and "lab-workbench-r1-s1" in argv and str(cli) == argv[-1] assert w["sock"] == "/tmp/lab-proxy-r1-s1/proxy.sock" and "UNIX-CONNECT:/tmp/lab-proxy-r1-s1/proxy.sock" in argv[-2] assert "ANTHROPIC_BASE_URL=http://127.0.0.1:4000" in argv and not any("sk-" in a for a in argv) + + +@pytest.mark.parametrize("said,unpaused", [(b"paused\n", True), (b"\n", False)]) +def test_remote_pause_thaws_only_what_it_froze(said, unpaused): + r = worker.Remote(["ssh", "host"], "~/BlameGraph", "python", {}) + calls = [] + r._run = lambda command, stdin=b"": calls.append(command) or (said if command.startswith("pause") else b"") + with r.paused("lab-workbench-r1-s1"): + assert calls == ["pause lab-workbench-r1-s1"] + assert (calls[-1] == "unpause lab-workbench-r1-s1") is unpaused