diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 399ecbfb..f3457920 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -1,9 +1,7 @@ # Dependabot: weekly, grouped, one PR per ecosystem. The Substrate module is # kept out of the Go group on purpose: its bumps change the API the manifests # and the client are generated from and go through the bump procedure, not -# through an automated PR. integrations/nemo-gym is not listed: its -# ate-env-client dependency is a path source Dependabot cannot resolve, and -# nothing in it is pinned. +# through an automated PR. version: 2 updates: - package-ecosystem: gomod diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 1fa209ce..adf9163a 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -49,7 +49,7 @@ jobs: - run: hack/verify/boilerplate.sh python: - name: Python client and integration tests + name: Python client tests runs-on: ubuntu-24.04 # pinned: ubuntu-latest moves to 26.04 from 2026-10-19; bump deliberately steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 @@ -58,8 +58,6 @@ jobs: - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 with: python-version: "3.12" - - run: python -m pip install -e 'clients/python[dev]' -e 'integrations/nemo-gym[test]' + - run: python -m pip install -e 'clients/python[dev]' - run: python -m pytest clients/python/tests -q - - run: python -m pytest internal/benchmarking/swebench/trajectories -q - run: ruff check clients/python - - run: python -m pytest integrations/nemo-gym/tests -q diff --git a/.gitignore b/.gitignore index b08a4d05..3640f8f8 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,3 @@ deployment.yaml # uv writes a lockfile next to each pyproject it runs; the Python projects here are libraries and ship none. uv.lock -cmd/benchmarking/swebench-loader/kodata/ diff --git a/Makefile b/Makefile index fe85bd25..4ceee7d3 100644 --- a/Makefile +++ b/Makefile @@ -60,12 +60,6 @@ python-protos: python-test: cd clients/python && python3 -m pytest -# Run the NeMo Gym provider tests (requires an environment with the client and -# the provider's test extra installed: pip install clients/python -e -# 'integrations/nemo-gym[test]'). -nemo-gym-test: - cd integrations/nemo-gym && python3 -m pytest - images: @echo "Building and pushing container images to $(ATE_ENV_IMAGE_REPO)..." @guest_img=$$(KO_DOCKER_REPO=$(ATE_ENV_IMAGE_REPO)/ate-env-guest ko build --bare ./cmd/ate-env-guest | tail -n 1); \ diff --git a/README.md b/README.md index cf18493e..45e85add 100644 --- a/README.md +++ b/README.md @@ -29,7 +29,6 @@ while this project adds the environment-shaped API on top. - **`cmd/ate-env-guest`** — The daemon server running inside each actor serving command executions, file read/write, and built-in MCP tools. - **`clients/go`** — The Go client library to manage environments, run commands, and perform file operations. - **`clients/python`** — Python client for environment lifecycle management and batch sandbox orchestration ([README](clients/python/README.md)). -- **`integrations/nemo-gym`** — A [NeMo Gym](https://github.com/NVIDIA-NeMo/Gym) sandbox provider that runs rollout sandboxes as environments, built on the Python client ([README](integrations/nemo-gym/README.md)). ## Installation @@ -77,12 +76,11 @@ Substrate install. An older Substrate rejects the string sandbox class (from [`c58f374a`](https://github.com/agent-substrate/substrate/commit/c58f374a)); use an ate-env from before this bump with it. -When you upgrade an existing cluster across that boundary, register your base templates again. -Substrate cannot update an ActorTemplate in place, and a base registered earlier has no CPU and memory -limits, thus each template `ate-env-api` derives from it for a rootfs image fails with -`resources: Required`. Delete the base and its derived templates (named `-<12 hex digits>`) with -`kubectl-ate delete actor-template`, then register the base again with `ate-env manifest template`. -Each environment created after that derives its template again, with the limits. +When you upgrade an existing cluster across that boundary, register your templates again. +Substrate cannot update an ActorTemplate in place, and a template registered earlier has no CPU and +memory limits. Delete it with `kubectl-ate delete actor-template`, then register it again with +`ate-env manifest template`. Templates that older `ate-env-api` versions derived per rootfs image +(named `-<12 hex digits>`) are no longer used and can be deleted the same way. ### 1. Deploy the Kubernetes Resources @@ -121,8 +119,7 @@ ate-env manifest template \ ``` Each actor gets 2 CPUs and 4 GiB unless you pass `--cpu` and `--memory` (Kubernetes quantities such as -`500m` or `8Gi`; CPU at most `999.999`, as Substrate requires). Substrate requires both limits and books them against worker capacity; templates -derived per rootfs image inherit them. +`500m` or `8Gi`; CPU at most `999.999`, as Substrate requires). Substrate requires both limits and books them against worker capacity. That template's container is the guest image itself. To run another image unmodified, name it with `--rootfs-image`: the guest is then mounted into it as a read-only OCI image volume at @@ -135,30 +132,6 @@ ate-env manifest template --template py312 \ --snapshots-bucket | kubectl-ate create actor-template -f - ``` -A second runtime rides along the same way: `--layer` mounts its image read-only next to the -guest, `--sidecar` has the guest start and supervise a process from it, and `--sidecar-readyz` -holds the actor's readiness until that process answers. Here a static busybox stands in for a -runtime daemon and serves HTTP on port 44772 inside the environment: - -```bash -ate-env manifest template --template layered \ - --rootfs-image docker.io/library/python@sha256: \ - --guest-image /ate-env-guest@sha256: \ - --layer web=docker.io/library/busybox@sha256:=/opt/web \ - --sidecar '/opt/web/bin/busybox httpd -f -p 44772 -h /' \ - --sidecar-readyz http://127.0.0.1:44772/etc/hostname \ - --snapshots-bucket | kubectl-ate create actor-template -f - -``` - -Layers must be self-contained: the `1.37-musl` busybox is static, while the glibc build fails -on images with an older libc. The runnable version is -[`examples/runtime-layers/demo.sh`](examples/runtime-layers/demo.sh). - -See [docs/rootfs-images/README.md](docs/rootfs-images/README.md) for the full guide, including -creating environments from an image on demand, [docs/rootfs-images/RUNTIMES.md](docs/rootfs-images/RUNTIMES.md) -for injecting a second runtime as a layer, and [docs/rootfs-images/DESIGN.md](docs/rootfs-images/DESIGN.md) -for the design. - Then create and use an environment: ```bash @@ -168,17 +141,6 @@ kubectl port-forward -n ate-env svc/ate-env-api 7777:7777 & # Create an environment. ate-env create dev1 -# Or run any digest-pinned image unmodified. ate-env-api derives a template -# from default-template on first use (guest mounted in as an image volume) -# and reuses it for later environments on the same image. -ate-env create py1 --rootfs-image docker.io/library/python@sha256: - -# Or from the layered template above: the second runtime answers inside the -# environment, and any other image on that base inherits it. -ate-env create web1 --template layered -ate-env web1 shell '/opt/web/bin/busybox wget -qO- http://127.0.0.1:44772/etc/os-release' -ate-env create node1 --template layered --rootfs-image docker.io/library/node@sha256: - # Execute a shell command inside the environment. ate-env dev1 shell 'echo hello > /note.txt' @@ -260,7 +222,7 @@ Manages the lifecycle of isolated execution environments (defined in [`proto/ate | RPC | Description | | --- | ----------- | -| `CreateEnvironment` | Creates and starts a new environment actor from an ActorTemplate, or from a digest-pinned `rootfs_image` on top of one: the template becomes the base, the image the container, and the guest is mounted in as a read-only image volume | +| `CreateEnvironment` | Creates and starts a new environment actor from an ActorTemplate | | `GetEnvironment` | Retrieves environment details and status | | `SuspendEnvironment` | Suspends and checkpoints the environment to snapshot storage | | `DeleteEnvironment` | Deletes the environment permanently | @@ -289,56 +251,6 @@ Provides chunked streaming file reading and writing within the environment conta | `WriteFile` | Streams raw binary or text chunks directly to a target file | -## Reaching a service inside an environment - -A daemon that listens on a port inside an environment, such as one a runtime layer starts next to -the guest ([docs/rootfs-images/RUNTIMES.md](docs/rootfs-images/RUNTIMES.md)), is reachable through -the API's listener. A plain HTTP request that carries `X-Env-Id` and `X-Env-Port` (and -`X-Env-Atespace` outside `ate-env`) is forwarded through the Substrate router to that port, whatever -its path; the API strips the three headers. The router resumes a suspended environment first, -and the proxy waits for the wake up to `--wake-budget` (the same wait as a guest call) before it -answers 503. - -```bash -curl -H 'X-Env-Id: dev1' -H 'X-Env-Port: 6000' localhost:7777/health -``` - -Clients that cannot set headers, such as WebSocket clients and browsers, put the target in the path -instead: `/v1alpha/envs/{id}/ports/{port}/...`, or `/v1alpha/atespaces/{atespace}/envs/{id}/ports/{port}/...` -outside `ate-env`. The API strips that prefix and forwards the rest, so a client given -`http://localhost:7777/v1alpha/envs/dev1/ports/6000` as its base URL works unchanged. WebSocket upgrades pass -through either form; gRPC to the daemon is not supported yet. - -```bash -curl localhost:7777/v1alpha/envs/dev1/ports/6000/health -``` - -The path form reaches the daemon exactly as sent: a path without the trailing slash is the daemon's -root, and `//` or `..` after the prefix are not cleaned or redirected. - -**Who can use it.** The forwarding does not authenticate callers. Whoever can reach the API's -listener can reach any TCP port of any actor the Substrate router serves, in any atespace, -including actors that ate-env did not create. The guest API already reaches each actor that runs -the guest; the forwarding extends that reach to every port. Keep the listener inside the cluster -(the manifest's Service is `ClusterIP`), restrict it with a NetworkPolicy where other workloads -share the cluster, and make each daemon check its own credential. - -**Browsers.** The path form refuses a state-changing request or a WebSocket handshake that a browser -sends from a page of another origin, so a web page cannot drive a daemon through a `kubectl -port-forward`. A page served by the forwarded daemon itself is same-origin and passes. All daemons -behind one listener share that origin, thus a page from one daemon can call another. -`--port-proxy-trusted-origins` adds origins, such as a UI on its own host. The header form needs no -check: a browser sends custom headers to another origin only after a CORS preflight, which the API -never approves. DNS rebinding, where a page on an attacker's name resolves to the forwarded -listener, is not blocked yet; it needs a check of the `Host` header across the whole listener, the -MCP endpoint included. - -| Flag | Default | Effect | -|---|---|---| -| `--atenet-connect` | `atenet-router.ate-system.svc.cluster.local:8081` | The router's CONNECT listener. An empty value turns the forwarding off. | -| `--port-proxy-trusted-origins` | none | Comma-separated browser origins allowed to use the path form. | -| `--port-proxy-response-header-timeout` | `0`, no limit | Bound on the wait for a daemon's response headers. Without it, the caller's deadline ends the wait. | - ## Built-in MCP Server The API exposes a streamable MCP endpoint at `POST /v1alpha/envs/{id}/mcp` serving built-in tools. @@ -412,8 +324,6 @@ For complete runnable Go programs: - **MCP**: See [mcp](examples/mcp/main.go) to connect to an environment's MCP endpoint, discover tools, and execute tool calls. - **Guest Daemon**: See [guest-daemon](examples/guest-daemon/main.go) to run a standalone in-actor gRPC service for asynchronous process execution and chunked file transfer. -For using environments from a framework, see [integrations/nemo-gym](integrations/nemo-gym/README.md): a NeMo Gym sandbox provider registered through the `nemo_gym.sandbox_providers` entry point. - ## Cleanup ```bash diff --git a/benchmarking/swebench/README.md b/benchmarking/swebench/README.md deleted file mode 100644 index 37ba9451..00000000 --- a/benchmarking/swebench/README.md +++ /dev/null @@ -1,150 +0,0 @@ -# Fast Substrate RL benchmark (SWE-bench Verified) - -A loader that drives real SWE-bench task images through the ate-env API the -way a simulated RL job does, and a report over what every phase cost. It is -the "Fast" rung of the RL benchmark ladder: regular runs that catch -regressions in the sandbox SDK, Substrate or the controller, on a workload -that is the real thing rather than a Python image printing READY. - -Code: `internal/benchmarking/swebench` (library, tested against an in-memory -backend) and `cmd/benchmarking/swebench-loader` (binary). The layout follows -substrate's `internal/benchmarking/boomer` so a user class can be lifted into -that harness later; Go does not allow importing another module's `internal` -packages, so nothing is imported from it. - -## Workloads - -Both are parameters of one binary. Images, fan-out, width, steps and the -trajectory are all flags, so the smoke test and the 500-wide run differ only -in numbers. - -| Shape | Flags | What happens | -|---|---|---| -| eval (SWE-bench Verified evaluation) | `-mode eval -images 500 -tasks-per-image 6 -width 500` | 500 in flight. Each task: create from the image, replay the trajectory, record the reward, delete. With the offline exclusion list the 500 become 445 images and 2,670 tasks; the loader prints the reduced count. | -| steps (training loop) | `-mode steps -images 500 -rollouts 2 -steps 8 -park suspend` | 1,000 environments created once. Each step replays the trajectory on all of them, waits for the last (the barrier), then parks the fleet. `-park` is `suspend`, `delete` or `none`. | - -The trajectory is YAML (`internal/benchmarking/swebench/trajectories/`, -or `-trajectory FILE`). The built-in `swebench-gold` does what a SWE agent -does on a task, with the gold patch as the model's edit: explore, read, write -the patch, `git apply`, write the harness eval script, run it. The reward is -the SWE-bench verdict: `swebench-grade.py` (written into the sandbox by the -trajectory) parses the log with the repo's harness parser and checks the -instance's FAIL_TO_PASS and PASS_TO_PASS tests, the same rules the harness -applies. Neither the script's nor the test command's exit code is the -verdict: unrelated tests in the same files fail routinely. A healthy run -rewards every task and a drop in reward rate is a correctness regression. -`python3 -m pytest benchmarking/swebench/tests` covers the grader. Steps carry think times, the -simulated policy latency between commands. - -## What is measured - -One JSONL event per operation, wall time as the loader sees it (network and -API server included, which is what a trainer sees): - -| Event | Meaning | -|---|---| -| `create` | CreateEnvironment call | -| `ttfi` | create start to the first command's return (time to first instruction) | -| `wake` | first command after a suspend park, including the wake it triggers | -| `step` (by name) | each trajectory step, with exit code; the reward step carries the reward and an output tail | -| `task` | one replay end to end, ok and reward | -| `suspend`, `delete` | the calls | -| `step_barrier` | steps mode: first task start to last task end in a training step | -| `park` | steps mode: parking the whole fleet | -| `run` | the whole run | - -`swebench-loader report RUN.jsonl -json summary.json` prints count, failed, -p50, p95 and max per event, a per-step table (tasks, barrier, park, wake -p50/p95) and the most frequent errors. The JSON is the artifact to diff -between runs. - -## Prepare once per project - -1. Put the task images in a registry the nodes can pull from. If the nodes - can read your mirror directly, skip this step. Otherwise copy them by - digest with crane; `tags.txt` has one tag per line, for example every - `sweb.eval.x86_64.*` tag of the source repository without the `-zran` - variants: - - ```bash - SRC=SOURCE_REGISTRY/swebench-verified - DST=us-central1-docker.pkg.dev/PROJECT/swebench/swebench-verified - crane ls $SRC | grep -v -- '-zran$' > tags.txt - benchmarking/swebench/mirror-images.sh $SRC $DST tags.txt 8 - ``` - -2. Build the manifest: gold patch, test patch, harness eval script, log - parser and the digest-pinned image per instance. `--registry` is the - repository the images live in, the same `$DST` as above (or the source - repository when the nodes can read it). - - ```bash - pip install datasets 'swebench>=5' - benchmarking/swebench/make-dataset.py --registry $DST \ - --revision --out swebench-verified.jsonl - ``` - - `--image-ref` sets how an instance maps to an image reference (default: - one tag per instance, `sweb.eval.x86_64.`, in one - repository); the script lists instances whose image it cannot resolve and - writes the dataset revision and swebench version next to the manifest. - -3. Install ate-env with a base template (`ate-env manifest template`). Per - image, the API derives one template from it on first use and bakes the - golden snapshot then; the first create of each image pays that. - -## Run - -From a machine with a port-forward (fine for the smoke test and the step -loop): - -```bash -kubectl -n ate-env port-forward svc/ate-env-api 7777:7777 & -go run ./cmd/benchmarking/swebench-loader run -dataset swebench-verified.jsonl \ - -images 5 -tasks-per-image 1 -width 5 -run-id smoke -go run ./cmd/benchmarking/swebench-loader report smoke.jsonl -``` - -`-dry-run` loads and renders everything and prints the plan without touching -the API. `-keep` leaves the fleet in place for inspection. - -For the 500-wide shapes run inside the cluster so the loader sits next to the -API and spreads over its replicas (`-connections`, default 4): - -```bash -KO_DOCKER_REPO=gcr.io/PROJECT/repo benchmarking/swebench/run-in-cluster.sh swe-eval500x6 swebench-verified.jsonl \ - -mode eval -images 500 -tasks-per-image 6 -width 500 -connections 8 -``` - -The script bakes the dataset and the exclusion list into the loader image, -runs a Job next to ate-env-api, follows the full loader log (step failures -included), copies the events file back to the current directory, prints the -report and deletes the Job. - -## Network and the offline exclusion list - -Actors have no network until an egress policy is created for them, which the -env SDK cannot do yet, and the SWE-bench harness assumes network (matplotlib's -eval scripts pip-install, requests and django have tests that open sockets). -`offline-exclusions.txt` lists the Verified instances that do not resolve -without egress, with the grader's verdict as the reason; pass it with -`-exclude-file` and every remaining instance should resolve. Regenerate it -from a full pass when the dataset, the images or the sandbox network change, -and review the diff: a new entry is either a regression or a new -network-dependent test, never silently accepted. Egress at scale is a -benchmark of its own. - -## Two things the first runs showed - -- Commands run as start, poll, read rather than through the SDK's streaming - shell: the router's route timeout (five minutes by default) cuts a single - streamed request, and SWE-bench test runs regularly exceed it. -- On a large pool a new actor rarely lands on a node that already has its - image, so time to first instruction is the cold pull plus the golden bake - (about 20 s) for nearly every task. Report cold and warm passes separately. - -## Record with every run - -Substrate and env commits, cluster shape (machine type, node count, sandbox -class), bucket region, the dataset manifest's digest, and the loader flags. -The report alone is not reproducible without them. diff --git a/benchmarking/swebench/make-dataset.py b/benchmarking/swebench/make-dataset.py deleted file mode 100755 index c98f1b95..00000000 --- a/benchmarking/swebench/make-dataset.py +++ /dev/null @@ -1,184 +0,0 @@ -#!/usr/bin/env python3 -# -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Build the JSONL manifest swebench-loader reads from a SWE-bench dataset. - -For every instance it records the gold patch, the test patch, the harness -eval script (the exact script the SWE-bench harness runs: apply the test -patch, run the repo's test command, restore the test files, print the exit -code marker the grader reads) and the task image pinned by digest in the -registry the cluster can pull from. - -The eval script comes from the installed swebench's make_test_spec, which -wraps the dataset row's own eval_script with the harness header and the -"Test Exit Code" marker; the two are checked to agree. The dataset revision -and the swebench version are recorded in OUT.meta.json so two operators can -build the same manifest; pass --revision to pin the dataset. - - pip install datasets 'swebench>=5' - ./make-dataset.py --registry us-central1-docker.pkg.dev/PROJECT/swebench \ - --out swebench-verified.jsonl - -The image reference per instance comes from --image-ref, a format string -over {registry}, {instance_id} and {instance_id_1776}. The default matches a -mirror that keeps one tag per instance in a single repository; the public -SWE-bench images on Docker Hub are one repository per instance: - - --image-ref '{registry}/sweb.eval.x86_64.{instance_id_1776}:latest' - -Instances whose image cannot be resolved are reported and left out, so run -mirror-images.sh first and re-run this until none are missing. -""" - -import argparse -import datetime -import json -import subprocess -import sys - - -def image_ref(registry: str, instance_id: str, fmt: str) -> str: - return fmt.format( - registry=registry, - instance_id=instance_id, - instance_id_1776=instance_id.replace("__", "_1776_"), - ) - - -def resolve_digest(ref: str) -> str | None: - """Return repo@sha256:... for ref, or None when crane cannot see it.""" - proc = subprocess.run(["crane", "digest", ref], capture_output=True, text=True) - if proc.returncode != 0: - return None - repo = ref.split("@", 1)[0] - if "@" not in ref: - head, _, last = repo.rpartition("/") - if ":" in last: - repo = f"{head}/{last.split(':', 1)[0]}" if head else last.split(":", 1)[0] - return f"{repo}@{proc.stdout.strip()}" - - -def main() -> int: - ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - ap.add_argument("--dataset", default="SWE-bench/SWE-bench_Verified", help="Hugging Face dataset name or a local JSON/JSONL path. The rows must carry eval_script (the SWE-bench org copies do; the older princeton-nlp copies do not).") - ap.add_argument("--split", default="test") - ap.add_argument("--revision", default=None, help="Dataset revision (commit or tag) to pin; recorded in the meta file.") - ap.add_argument("--registry", required=True, help="Registry path that holds the task images, e.g. us-docker.pkg.dev/PROJECT/swebench-mirror/swebench-verified.") - ap.add_argument("--image-ref", default="{registry}:sweb.eval.x86_64.{instance_id}", help="Image reference per instance; fields: {registry}, {instance_id}, {instance_id_1776}.") - ap.add_argument("--limit", type=int, default=0, help="Stop after N instances (0 = all).") - ap.add_argument("--instances", default="", help="Comma-separated instance ids to include; default all.") - ap.add_argument("--out", required=True, help="JSONL manifest to write.") - args = ap.parse_args() - - try: - from datasets import load_dataset - except ImportError as e: - print(f"missing dependency: {e}; pip install datasets", file=sys.stderr) - return 2 - try: - make_test_spec = _load_make_test_spec() - except ImportError as e: - print(f"cannot import make_test_spec from swebench: {e}; pip install -U 'swebench>=5'", file=sys.stderr) - return 2 - - if args.dataset.endswith((".json", ".jsonl")): - rows = load_dataset("json", data_files=args.dataset, split="train") - else: - rows = load_dataset(args.dataset, split=args.split, revision=args.revision) - import swebench - swebench_version = getattr(swebench, "__version__", "unknown") - want = {s for s in args.instances.split(",") if s} - - written = missing = 0 - with open(args.out, "w") as out: - for row in rows: - iid = row["instance_id"] - if want and iid not in want: - continue - if args.limit and written >= args.limit: - break - ref = image_ref(args.registry, iid, args.image_ref) - pinned = resolve_digest(ref) - if pinned is None: - print(f"missing image for {iid}: {ref}", file=sys.stderr) - missing += 1 - continue - try: - spec = make_test_spec(row) - except KeyError as e: - print(f"{iid}: dataset row lacks field {e} that this swebench release needs; " - "use a dataset revision that carries eval_script (swebench >= 5) or an older swebench", file=sys.stderr) - return 2 - if row.get("eval_script") and row["eval_script"].strip() not in spec.eval_script: - print(f"{iid}: the installed swebench produced an eval script that does not contain the dataset row's; " - "check the swebench version against the dataset revision", file=sys.stderr) - return 2 - rec = { - "instance_id": iid, - "repo": row["repo"], - "image": pinned, - "patch": row["patch"], - "test_patch": row["test_patch"], - "eval_script": spec.eval_script, - "fail_to_pass": _as_list(row["FAIL_TO_PASS"]), - "pass_to_pass": _as_list(row["PASS_TO_PASS"]), - "log_parser": row["log_parser"], - "eval_type": row["eval_type"], - } - out.write(json.dumps(rec) + "\n") - written += 1 - meta = { - "dataset": args.dataset, - "split": args.split, - "revision": args.revision, - "swebench_version": swebench_version, - "registry": args.registry, - "image_ref": args.image_ref, - "instances": written, - "missing_images": missing, - "generated_at": datetime.datetime.now(datetime.timezone.utc).isoformat(timespec="seconds"), - } - with open(args.out + ".meta.json", "w") as f: - json.dump(meta, f, indent=2) - f.write("\n") - print(f"wrote {written} instances to {args.out}; {missing} missing images; provenance in {args.out}.meta.json", file=sys.stderr) - return 1 if missing else 0 - - -def _load_make_test_spec(): - """make_test_spec moved between swebench releases; try the known homes.""" - import importlib - - import swebench - - errors = [] - for mod in ("swebench.harness.utils", "swebench.harness.test_spec.test_spec", "swebench.harness.test_spec"): - try: - return getattr(importlib.import_module(mod), "make_test_spec") - except (ImportError, AttributeError) as e: - errors.append(f"{mod}: {e}") - version = getattr(swebench, "__version__", "unknown") - raise ImportError(f"swebench {version} at {swebench.__file__}: " + "; ".join(errors)) - - -def _as_list(v): - if isinstance(v, str): - return json.loads(v) - return list(v) - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/benchmarking/swebench/mirror-images.sh b/benchmarking/swebench/mirror-images.sh deleted file mode 100755 index b362f33d..00000000 --- a/benchmarking/swebench/mirror-images.sh +++ /dev/null @@ -1,64 +0,0 @@ -#!/usr/bin/env bash -# -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# Copy SWE-bench task images into a registry the cluster's nodes can pull -# from, with crane. Each tag is resolved to its digest in SRC_REPO and copied -# by that digest to the same tag in DST_REPO, so a tag that moves upstream -# cannot change what was copied; a tag whose digest already matches in -# DST_REPO is skipped. Nodes pull with their own identity, and shared mirrors -# may deny it, so every image the benchmark uses must live in a registry the -# nodes can read. -# -# mirror-images.sh SRC_REPO DST_REPO TAGS_FILE [PARALLEL] -# -# TAGS_FILE has one tag per line, for example -# crane ls SRC_REPO | grep -v -- '-zran$' > tags.txt -# Needs crane on PATH and credentials for both registries -# (gcloud auth configure-docker). - -set -o errexit -o nounset -o pipefail - -if [[ $# -lt 3 ]]; then - echo "usage: $0 SRC_REPO DST_REPO TAGS_FILE [PARALLEL]" >&2 - exit 2 -fi - -SRC="$1" -DST="$2" -TAGS="$3" -PARALLEL="${4:-8}" - -command -v crane >/dev/null || { echo "crane not found on PATH" >&2; exit 1; } - -copy_one() { - local tag="$1" src_digest dst_digest - src_digest="$(crane digest "${SRC}:${tag}" 2>/dev/null)" || { echo "FAILED ${tag}: not in ${SRC}" >&2; return 1; } - dst_digest="$(crane digest "${DST}:${tag}" 2>/dev/null || true)" - if [[ "${dst_digest}" == "${src_digest}" ]]; then - echo "skip ${tag} (same digest)" - return 0 - fi - if crane copy "${SRC}@${src_digest}" "${DST}:${tag}" >/dev/null 2>&1; then - echo "copied ${tag} ${src_digest}" - else - echo "FAILED ${tag}" >&2 - return 1 - fi -} -export -f copy_one -export SRC DST - -grep -v '^\s*$' "${TAGS}" | xargs -P "${PARALLEL}" -I{} bash -c 'copy_one "$@"' _ {} diff --git a/benchmarking/swebench/offline-exclusions.txt b/benchmarking/swebench/offline-exclusions.txt deleted file mode 100644 index 356cfe1b..00000000 --- a/benchmarking/swebench/offline-exclusions.txt +++ /dev/null @@ -1,57 +0,0 @@ -# SWE-bench Verified instances whose evaluation does not resolve without network egress (m2-density-a, substrate ac5c5ee0, 2026-10-15). -# Generated from run swe-500b: 55 of 500 tasks unresolved. -astropy__astropy-8707 f2p 0/1 p2p 4/11 -django__django-10097 f2p 431/438 p2p 1427/1427 -django__django-13343 f2p 1/1 p2p 129/130 -django__django-13658 f2p 1/1 p2p 165/181 -django__django-13809 f2p 1/1 p2p 229/245 -django__django-14011 f2p 0/17 p2p 1/2 -django__django-16145 f2p 1/1 p2p 155/174 -django__django-16493 f2p 1/1 p2p 90/91 -django__django-16502 f2p 1/1 p2p 3/4 -matplotlib__matplotlib-13989 f2p 0/1 p2p 0/411 -matplotlib__matplotlib-14623 f2p 0/1 p2p 0/400 -matplotlib__matplotlib-20488 f2p 0/1 p2p 0/111 -matplotlib__matplotlib-20676 f2p 0/2 p2p 0/32 -matplotlib__matplotlib-20826 f2p 0/1 p2p 0/673 -matplotlib__matplotlib-20859 f2p 0/1 p2p 0/88 -matplotlib__matplotlib-21568 f2p 0/4 p2p 0/73 -matplotlib__matplotlib-22719 f2p 0/1 p2p 0/68 -matplotlib__matplotlib-22865 f2p 0/3 p2p 0/57 -matplotlib__matplotlib-22871 f2p 0/1 p2p 0/88 -matplotlib__matplotlib-23299 f2p 0/1 p2p 0/192 -matplotlib__matplotlib-23314 f2p 0/1 p2p 0/862 -matplotlib__matplotlib-23412 f2p 0/1 p2p 0/46 -matplotlib__matplotlib-23476 f2p 0/1 p2p 0/104 -matplotlib__matplotlib-24026 f2p 0/2 p2p 0/764 -matplotlib__matplotlib-24149 f2p 0/1 p2p 0/767 -matplotlib__matplotlib-24177 f2p 0/1 p2p 0/768 -matplotlib__matplotlib-24570 f2p 0/2 p2p 0/267 -matplotlib__matplotlib-24627 f2p 0/1 p2p 0/776 -matplotlib__matplotlib-24637 f2p 0/1 p2p 0/29 -matplotlib__matplotlib-24870 f2p 0/1 p2p 0/65 -matplotlib__matplotlib-24970 f2p 0/1 p2p 0/253 -matplotlib__matplotlib-25122 f2p 0/12 p2p 0/2476 -matplotlib__matplotlib-25287 f2p 0/2 p2p 0/786 -matplotlib__matplotlib-25311 f2p 0/1 p2p 0/181 -matplotlib__matplotlib-25332 f2p 0/1 p2p 0/181 -matplotlib__matplotlib-25479 f2p 0/2 p2p 0/263 -matplotlib__matplotlib-25775 f2p 0/7 p2p 0/89 -matplotlib__matplotlib-25960 f2p 0/1 p2p 0/137 -matplotlib__matplotlib-26113 f2p 0/1 p2p 0/812 -matplotlib__matplotlib-26208 f2p 0/2 p2p 0/813 -matplotlib__matplotlib-26291 f2p 0/1 p2p 0/49 -matplotlib__matplotlib-26342 f2p 0/1 p2p 0/81 -matplotlib__matplotlib-26466 f2p 0/1 p2p 0/100 -psf__requests-1724 f2p 0/6 p2p 44/79 -psf__requests-1766 f2p 0/6 p2p 45/79 -psf__requests-1921 f2p 1/6 p2p 69/107 -psf__requests-2317 f2p 0/8 p2p 81/133 -psf__requests-2931 f2p 1/1 p2p 82/84 -psf__requests-5414 f2p 1/1 p2p 126/130 -pylint-dev__pylint-4661 invalid run -sphinx-doc__sphinx-10323 f2p 1/1 p2p 38/40 -sphinx-doc__sphinx-10435 f2p 0/1 p2p 73/74 -sphinx-doc__sphinx-7985 f2p 0/2 p2p 2/3 -sphinx-doc__sphinx-8269 f2p 0/1 p2p 2/5 -sphinx-doc__sphinx-8475 f2p 0/1 p2p 0/17 diff --git a/benchmarking/swebench/run-in-cluster.sh b/benchmarking/swebench/run-in-cluster.sh deleted file mode 100755 index a4f669b6..00000000 --- a/benchmarking/swebench/run-in-cluster.sh +++ /dev/null @@ -1,165 +0,0 @@ -#!/usr/bin/env bash -# -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# Run swebench-loader inside the cluster, next to ate-env-api, and bring the -# events file back. The dataset and offline-exclusions.txt are baked into the -# loader image (ko's kodata directory), so the pod needs no bucket access; a -# sidecar keeps the pod alive after the loader exits so the results can be -# copied out with kubectl cp. The Job is deleted when the script exits, on -# success or failure. -# -# KO_DOCKER_REPO=gcr.io/PROJECT/repo \ -# benchmarking/swebench/run-in-cluster.sh RUN_ID DATASET.jsonl [swebench-loader run flags...] -# -# RUN_ID must be a DNS label (lowercase letters, digits, '-', at most 20 -# characters). Results land in ./RUN_ID.jsonl and ./RUN_ID.json in the -# directory the script is run from, which can be outside the repository. -# NAMESPACE defaults to ate-env; START_TIMEOUT (seconds, default 600) bounds -# the wait for the pod to start. - -set -o errexit -o nounset -o pipefail - -if [[ $# -lt 2 ]]; then - echo "usage: $0 RUN_ID DATASET.jsonl [swebench-loader run flags...]" >&2 - exit 2 -fi -RUN_ID="$1" -DATASET="$(cd "$(dirname "$2")" && pwd)/$(basename "$2")" -shift 2 -NS="${NAMESPACE:-ate-env}" -START_TIMEOUT="${START_TIMEOUT:-600}" -: "${KO_DOCKER_REPO:?set KO_DOCKER_REPO to the registry path the cluster can pull from}" -[[ "${RUN_ID}" =~ ^[a-z0-9]([-a-z0-9]*[a-z0-9])?$ && ${#RUN_ID} -le 20 ]] || { echo "RUN_ID must be a DNS label of at most 20 characters" >&2; exit 2; } -[[ -f "${DATASET}" ]] || { echo "dataset ${DATASET} not found" >&2; exit 2; } -command -v python3 >/dev/null || { echo "python3 is needed to encode the loader arguments" >&2; exit 2; } -OUT_DIR="$(pwd)" - -# The repository is found from the script's own location, so the script can -# run from any directory and the results land where it was invoked. -ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" -cd "${ROOT}" -KODATA=cmd/benchmarking/swebench-loader/kodata -rm -rf "${KODATA}" && mkdir -p "${KODATA}" -cp "${DATASET}" "${KODATA}/dataset.jsonl" -cp benchmarking/swebench/offline-exclusions.txt "${KODATA}/offline-exclusions.txt" -IMAGE="$(KO_DOCKER_REPO="${KO_DOCKER_REPO}/swebench-loader" ko build --bare ./cmd/benchmarking/swebench-loader)" -rm -rf "${KODATA}" -echo "loader image ${IMAGE}" - -JOB="swebench-loader-${RUN_ID}" -# Every argument is encoded as a JSON string, so quotes and other YAML -# metacharacters in a flag value cannot break the manifest. -ARGS_JSON="$(python3 -c 'import json, sys; print(json.dumps(sys.argv[1:]))' \ - run "-endpoint=ate-env-api.${NS}.svc:7777" "-dataset=/var/run/ko/dataset.jsonl" \ - "-exclude-file=/var/run/ko/offline-exclusions.txt" "-run-id=${RUN_ID}" "-out=/work/${RUN_ID}.jsonl" "$@")" - -cleanup() { - kubectl -n "${NS}" delete job "${JOB}" --ignore-not-found --wait=false >/dev/null 2>&1 || true -} -trap cleanup EXIT - -kubectl -n "${NS}" delete job "${JOB}" --ignore-not-found --wait=true >/dev/null -# busybox 1.36, pinned by digest; Docker Hub rate-limits unauthenticated pulls. -cat </dev/null || true)" - if [[ -n "${POD}" ]] && kubectl -n "${NS}" get pod "${POD}" -o jsonpath='{.status.containerStatuses[?(@.name=="loader")].state}' 2>/dev/null | grep -q 'running\|terminated'; then - break - fi - sleep 5 -done -if [[ -z "${POD}" ]] || ! kubectl -n "${NS}" get pod "${POD}" -o jsonpath='{.status.containerStatuses[?(@.name=="loader")].state}' 2>/dev/null | grep -q 'running\|terminated'; then - echo "loader did not start within ${START_TIMEOUT}s" >&2 - [[ -n "${POD}" ]] && kubectl -n "${NS}" describe pod "${POD}" | tail -15 >&2 - exit 1 -fi -echo "pod ${POD}; following the loader log" -kubectl -n "${NS}" logs -f "${POD}" -c loader || true - -# The log stream can end before the container does; wait for the exit code, -# with the pod's own state as the deadline check. -until EXIT="$(kubectl -n "${NS}" get pod "${POD}" -o jsonpath='{.status.containerStatuses[?(@.name=="loader")].state.terminated.exitCode}')" && [[ -n "${EXIT}" ]]; do - PHASE="$(kubectl -n "${NS}" get pod "${POD}" -o jsonpath='{.status.phase}' 2>/dev/null || echo gone)" - if [[ "${PHASE}" == "gone" || "${PHASE}" == "Failed" ]]; then - echo "pod ${POD} is ${PHASE} before the loader reported an exit code" >&2 - exit 1 - fi - sleep 10 -done -kubectl -n "${NS}" cp -c keep "${POD}:/work/${RUN_ID}.jsonl" "${OUT_DIR}/${RUN_ID}.jsonl" -echo "loader exit code ${EXIT}; events in ${OUT_DIR}/${RUN_ID}.jsonl" -go run ./cmd/benchmarking/swebench-loader report "${OUT_DIR}/${RUN_ID}.jsonl" -json "${OUT_DIR}/${RUN_ID}.json" -exit "${EXIT}" diff --git a/clients/python/README.md b/clients/python/README.md index e47be4ec..0ad008d8 100644 --- a/clients/python/README.md +++ b/clients/python/README.md @@ -145,17 +145,6 @@ env = await client.create("dev1", template_name="my-template", template_atespace="my-atespace") ``` -To run an arbitrary OCI image, pass it pinned by digest. The server derives -an ActorTemplate from the template above, which acts as the base: the image -becomes the container and the `ate-env-guest` is mounted into it as a -read-only image volume, so the image runs unmodified. The derived template -is created on first use and shared by every environment on that image: - -```python -env = await client.create("py1", rootfs_image="docker.io/library/python@sha256:…") -print((await env.info()).template.name) # default-template-<12 hex of the digest> -``` - To get a handle to an environment that already exists (no RPC is made): ```python @@ -391,7 +380,4 @@ ATE_ENV_API_TARGET=127.0.0.1:17777 .venv/bin/pytest tests/e2e/test_full_stack.py `ATE_ENV_TEMPLATE` optionally overrides the ActorTemplate used for the test environment; `ATE_ENV_READY_TIMEOUT` (default 180s) bounds the wait -for the environment to start serving. `ATE_ENV_ROOTFS_IMAGE`, a digest-pinned -image that has `sh`, enables the create-from-image test, which runs that -image unmodified with the guest injected and checks the derived template is -reused by a second environment. +for the environment to start serving. diff --git a/clients/python/src/ate_env/_gen/ateenv/v1alpha/env_pb2.py b/clients/python/src/ate_env/_gen/ateenv/v1alpha/env_pb2.py index 96bd9280..c77a3e57 100644 --- a/clients/python/src/ate_env/_gen/ateenv/v1alpha/env_pb2.py +++ b/clients/python/src/ate_env/_gen/ateenv/v1alpha/env_pb2.py @@ -38,7 +38,7 @@ -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x18\x61teenv/v1alpha/env.proto\x12\x0e\x61teenv.v1alpha\"*\n\x08Template\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x10\n\x08\x61tespace\x18\x02 \x01(\t\"\x8a\x01\n\x0b\x45nvironment\x12\n\n\x02id\x18\x01 \x01(\t\x12\x10\n\x08\x61tespace\x18\x02 \x01(\t\x12*\n\x08template\x18\x03 \x01(\x0b\x32\x18.ateenv.v1alpha.Template\x12\x31\n\x06status\x18\x04 \x01(\x0e\x32!.ateenv.v1alpha.EnvironmentStatus\"z\n\x18\x43reateEnvironmentRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x10\n\x08\x61tespace\x18\x02 \x01(\t\x12*\n\x08template\x18\x03 \x01(\x0b\x32\x18.ateenv.v1alpha.Template\x12\x14\n\x0crootfs_image\x18\x04 \x01(\t\"M\n\x19\x43reateEnvironmentResponse\x12\x30\n\x0b\x65nvironment\x18\x01 \x01(\x0b\x32\x1b.ateenv.v1alpha.Environment\"5\n\x15GetEnvironmentRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x10\n\x08\x61tespace\x18\x02 \x01(\t\"J\n\x16GetEnvironmentResponse\x12\x30\n\x0b\x65nvironment\x18\x01 \x01(\x0b\x32\x1b.ateenv.v1alpha.Environment\"9\n\x19SuspendEnvironmentRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x10\n\x08\x61tespace\x18\x02 \x01(\t\"\x1c\n\x1aSuspendEnvironmentResponse\"8\n\x18\x44\x65leteEnvironmentRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x10\n\x08\x61tespace\x18\x02 \x01(\t\"\x1b\n\x19\x44\x65leteEnvironmentResponse*\xbd\x02\n\x11\x45nvironmentStatus\x12\"\n\x1e\x45NVIRONMENT_STATUS_UNSPECIFIED\x10\x00\x12\x1f\n\x1b\x45NVIRONMENT_STATUS_RESUMING\x10\x01\x12\x1e\n\x1a\x45NVIRONMENT_STATUS_RUNNING\x10\x02\x12!\n\x1d\x45NVIRONMENT_STATUS_SUSPENDING\x10\x03\x12 \n\x1c\x45NVIRONMENT_STATUS_SUSPENDED\x10\x04\x12\x1e\n\x1a\x45NVIRONMENT_STATUS_PAUSING\x10\x05\x12\x1d\n\x19\x45NVIRONMENT_STATUS_PAUSED\x10\x06\x12\x1e\n\x1a\x45NVIRONMENT_STATUS_CRASHED\x10\x07\x12\x1f\n\x1b\x45NVIRONMENT_STATUS_DELETING\x10\x08\x32\xb6\x03\n\x12\x45nvironmentService\x12h\n\x11\x43reateEnvironment\x12(.ateenv.v1alpha.CreateEnvironmentRequest\x1a).ateenv.v1alpha.CreateEnvironmentResponse\x12_\n\x0eGetEnvironment\x12%.ateenv.v1alpha.GetEnvironmentRequest\x1a&.ateenv.v1alpha.GetEnvironmentResponse\x12k\n\x12SuspendEnvironment\x12).ateenv.v1alpha.SuspendEnvironmentRequest\x1a*.ateenv.v1alpha.SuspendEnvironmentResponse\x12h\n\x11\x44\x65leteEnvironment\x12(.ateenv.v1alpha.DeleteEnvironmentRequest\x1a).ateenv.v1alpha.DeleteEnvironmentResponseBCZAgithub.com/agent-substrate/env/proto/ateenv/v1alpha;ateenvv1alphab\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x18\x61teenv/v1alpha/env.proto\x12\x0e\x61teenv.v1alpha\"*\n\x08Template\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x10\n\x08\x61tespace\x18\x02 \x01(\t\"\x8a\x01\n\x0b\x45nvironment\x12\n\n\x02id\x18\x01 \x01(\t\x12\x10\n\x08\x61tespace\x18\x02 \x01(\t\x12*\n\x08template\x18\x03 \x01(\x0b\x32\x18.ateenv.v1alpha.Template\x12\x31\n\x06status\x18\x04 \x01(\x0e\x32!.ateenv.v1alpha.EnvironmentStatus\"x\n\x18\x43reateEnvironmentRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x10\n\x08\x61tespace\x18\x02 \x01(\t\x12*\n\x08template\x18\x03 \x01(\x0b\x32\x18.ateenv.v1alpha.TemplateJ\x04\x08\x04\x10\x05R\x0crootfs_image\"M\n\x19\x43reateEnvironmentResponse\x12\x30\n\x0b\x65nvironment\x18\x01 \x01(\x0b\x32\x1b.ateenv.v1alpha.Environment\"5\n\x15GetEnvironmentRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x10\n\x08\x61tespace\x18\x02 \x01(\t\"J\n\x16GetEnvironmentResponse\x12\x30\n\x0b\x65nvironment\x18\x01 \x01(\x0b\x32\x1b.ateenv.v1alpha.Environment\"9\n\x19SuspendEnvironmentRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x10\n\x08\x61tespace\x18\x02 \x01(\t\"\x1c\n\x1aSuspendEnvironmentResponse\"8\n\x18\x44\x65leteEnvironmentRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x10\n\x08\x61tespace\x18\x02 \x01(\t\"\x1b\n\x19\x44\x65leteEnvironmentResponse*\xbd\x02\n\x11\x45nvironmentStatus\x12\"\n\x1e\x45NVIRONMENT_STATUS_UNSPECIFIED\x10\x00\x12\x1f\n\x1b\x45NVIRONMENT_STATUS_RESUMING\x10\x01\x12\x1e\n\x1a\x45NVIRONMENT_STATUS_RUNNING\x10\x02\x12!\n\x1d\x45NVIRONMENT_STATUS_SUSPENDING\x10\x03\x12 \n\x1c\x45NVIRONMENT_STATUS_SUSPENDED\x10\x04\x12\x1e\n\x1a\x45NVIRONMENT_STATUS_PAUSING\x10\x05\x12\x1d\n\x19\x45NVIRONMENT_STATUS_PAUSED\x10\x06\x12\x1e\n\x1a\x45NVIRONMENT_STATUS_CRASHED\x10\x07\x12\x1f\n\x1b\x45NVIRONMENT_STATUS_DELETING\x10\x08\x32\xb6\x03\n\x12\x45nvironmentService\x12h\n\x11\x43reateEnvironment\x12(.ateenv.v1alpha.CreateEnvironmentRequest\x1a).ateenv.v1alpha.CreateEnvironmentResponse\x12_\n\x0eGetEnvironment\x12%.ateenv.v1alpha.GetEnvironmentRequest\x1a&.ateenv.v1alpha.GetEnvironmentResponse\x12k\n\x12SuspendEnvironment\x12).ateenv.v1alpha.SuspendEnvironmentRequest\x1a*.ateenv.v1alpha.SuspendEnvironmentResponse\x12h\n\x11\x44\x65leteEnvironment\x12(.ateenv.v1alpha.DeleteEnvironmentRequest\x1a).ateenv.v1alpha.DeleteEnvironmentResponseBCZAgithub.com/agent-substrate/env/proto/ateenv/v1alpha;ateenvv1alphab\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) @@ -46,28 +46,28 @@ if not _descriptor._USE_C_DESCRIPTORS: _globals['DESCRIPTOR']._loaded_options = None _globals['DESCRIPTOR']._serialized_options = b'ZAgithub.com/agent-substrate/env/proto/ateenv/v1alpha;ateenvv1alpha' - _globals['_ENVIRONMENTSTATUS']._serialized_start=740 - _globals['_ENVIRONMENTSTATUS']._serialized_end=1057 + _globals['_ENVIRONMENTSTATUS']._serialized_start=738 + _globals['_ENVIRONMENTSTATUS']._serialized_end=1055 _globals['_TEMPLATE']._serialized_start=44 _globals['_TEMPLATE']._serialized_end=86 _globals['_ENVIRONMENT']._serialized_start=89 _globals['_ENVIRONMENT']._serialized_end=227 _globals['_CREATEENVIRONMENTREQUEST']._serialized_start=229 - _globals['_CREATEENVIRONMENTREQUEST']._serialized_end=351 - _globals['_CREATEENVIRONMENTRESPONSE']._serialized_start=353 - _globals['_CREATEENVIRONMENTRESPONSE']._serialized_end=430 - _globals['_GETENVIRONMENTREQUEST']._serialized_start=432 - _globals['_GETENVIRONMENTREQUEST']._serialized_end=485 - _globals['_GETENVIRONMENTRESPONSE']._serialized_start=487 - _globals['_GETENVIRONMENTRESPONSE']._serialized_end=561 - _globals['_SUSPENDENVIRONMENTREQUEST']._serialized_start=563 - _globals['_SUSPENDENVIRONMENTREQUEST']._serialized_end=620 - _globals['_SUSPENDENVIRONMENTRESPONSE']._serialized_start=622 - _globals['_SUSPENDENVIRONMENTRESPONSE']._serialized_end=650 - _globals['_DELETEENVIRONMENTREQUEST']._serialized_start=652 - _globals['_DELETEENVIRONMENTREQUEST']._serialized_end=708 - _globals['_DELETEENVIRONMENTRESPONSE']._serialized_start=710 - _globals['_DELETEENVIRONMENTRESPONSE']._serialized_end=737 - _globals['_ENVIRONMENTSERVICE']._serialized_start=1060 - _globals['_ENVIRONMENTSERVICE']._serialized_end=1498 + _globals['_CREATEENVIRONMENTREQUEST']._serialized_end=349 + _globals['_CREATEENVIRONMENTRESPONSE']._serialized_start=351 + _globals['_CREATEENVIRONMENTRESPONSE']._serialized_end=428 + _globals['_GETENVIRONMENTREQUEST']._serialized_start=430 + _globals['_GETENVIRONMENTREQUEST']._serialized_end=483 + _globals['_GETENVIRONMENTRESPONSE']._serialized_start=485 + _globals['_GETENVIRONMENTRESPONSE']._serialized_end=559 + _globals['_SUSPENDENVIRONMENTREQUEST']._serialized_start=561 + _globals['_SUSPENDENVIRONMENTREQUEST']._serialized_end=618 + _globals['_SUSPENDENVIRONMENTRESPONSE']._serialized_start=620 + _globals['_SUSPENDENVIRONMENTRESPONSE']._serialized_end=648 + _globals['_DELETEENVIRONMENTREQUEST']._serialized_start=650 + _globals['_DELETEENVIRONMENTREQUEST']._serialized_end=706 + _globals['_DELETEENVIRONMENTRESPONSE']._serialized_start=708 + _globals['_DELETEENVIRONMENTRESPONSE']._serialized_end=735 + _globals['_ENVIRONMENTSERVICE']._serialized_start=1058 + _globals['_ENVIRONMENTSERVICE']._serialized_end=1496 # @@protoc_insertion_point(module_scope) diff --git a/clients/python/src/ate_env/_gen/ateenv/v1alpha/env_pb2.pyi b/clients/python/src/ate_env/_gen/ateenv/v1alpha/env_pb2.pyi index 6dea132d..08ab40eb 100644 --- a/clients/python/src/ate_env/_gen/ateenv/v1alpha/env_pb2.pyi +++ b/clients/python/src/ate_env/_gen/ateenv/v1alpha/env_pb2.pyi @@ -62,16 +62,14 @@ class Environment(_message.Message): def __init__(self, id: _Optional[str] = ..., atespace: _Optional[str] = ..., template: _Optional[_Union[Template, _Mapping]] = ..., status: _Optional[_Union[EnvironmentStatus, str]] = ...) -> None: ... class CreateEnvironmentRequest(_message.Message): - __slots__ = ("id", "atespace", "template", "rootfs_image") + __slots__ = ("id", "atespace", "template") ID_FIELD_NUMBER: _ClassVar[int] ATESPACE_FIELD_NUMBER: _ClassVar[int] TEMPLATE_FIELD_NUMBER: _ClassVar[int] - ROOTFS_IMAGE_FIELD_NUMBER: _ClassVar[int] id: str atespace: str template: Template - rootfs_image: str - def __init__(self, id: _Optional[str] = ..., atespace: _Optional[str] = ..., template: _Optional[_Union[Template, _Mapping]] = ..., rootfs_image: _Optional[str] = ...) -> None: ... + def __init__(self, id: _Optional[str] = ..., atespace: _Optional[str] = ..., template: _Optional[_Union[Template, _Mapping]] = ...) -> None: ... class CreateEnvironmentResponse(_message.Message): __slots__ = ("environment",) diff --git a/clients/python/src/ate_env/client.py b/clients/python/src/ate_env/client.py index dab33f1f..1da8e734 100644 --- a/clients/python/src/ate_env/client.py +++ b/clients/python/src/ate_env/client.py @@ -138,21 +138,13 @@ async def create( atespace: str = DEFAULT_ATESPACE, template_name: str | None = None, template_atespace: str | None = None, - rootfs_image: str | None = None, ) -> Env: """Register and start a new environment; returns a handle to it. The server fills defaults for the template (name "default-template" in atespace "ate-env") when none is given. - - With rootfs_image (a digest-pinned OCI reference, repo@sha256:...), the - environment runs that image unmodified: the server derives an - ActorTemplate from the template, which then acts as the base, with the - image as the container and the ate-env-guest mounted into it as an - image volume. The derived template is created on first use and reused - for later environments on the same image; info() reports its name. """ - req = env_pb2.CreateEnvironmentRequest(id=id, atespace=atespace, rootfs_image=rootfs_image or "") + req = env_pb2.CreateEnvironmentRequest(id=id, atespace=atespace) if template_name or template_atespace: req.template.name = template_name or "" req.template.atespace = template_atespace or "" diff --git a/clients/python/tests/e2e/test_full_stack.py b/clients/python/tests/e2e/test_full_stack.py index 9bdbba65..7705bcb8 100644 --- a/clients/python/tests/e2e/test_full_stack.py +++ b/clients/python/tests/e2e/test_full_stack.py @@ -25,11 +25,6 @@ Unlike test_guest_daemon.py, this exercises the whole path: environment lifecycle against the Substrate control plane, and guest operations proxied by ate-env-api through the atenet router into the environment's actor. - -Set ATE_ENV_ROOTFS_IMAGE to a digest-pinned image that has `sh` (for example -docker.io/library/python@sha256:...) to also run the create-from-image test, -which checks that an unmodified image runs with the guest injected as an -image volume and that a second environment reuses the derived template. """ from __future__ import annotations @@ -45,8 +40,6 @@ Client, EnvError, EnvironmentStatus, - FailedPreconditionError, - InvalidArgumentError, NotFoundError, ProcessState, Signal, @@ -56,8 +49,6 @@ READY_TIMEOUT = float(os.environ.get("ATE_ENV_READY_TIMEOUT", "180")) # Optional ActorTemplate override; the server default is "default-template". TEMPLATE = os.environ.get("ATE_ENV_TEMPLATE") -# Optional digest-pinned rootfs image for the create-from-image test. -ROOTFS_IMAGE = os.environ.get("ATE_ENV_ROOTFS_IMAGE") pytestmark = pytest.mark.skipif( not TARGET, @@ -180,78 +171,3 @@ async def test_full_lifecycle(client): return await asyncio.sleep(2) pytest.fail(f"environment {env_id} still exists after delete") - - -async def test_create_from_image_rejects_unpinned_image(client): - # Validated before anything touches the control plane. - with pytest.raises(InvalidArgumentError): - await client.create(f"pye2e-img-{uuid.uuid4().hex[:8]}", rootfs_image="python:3.12-slim") - - -async def test_create_from_image_needs_a_base_template(client): - missing = f"pye2e-nobase-{uuid.uuid4().hex[:8]}" - with pytest.raises(FailedPreconditionError) as excinfo: - await client.create( - f"pye2e-img-{uuid.uuid4().hex[:8]}", - template_name=missing, - rootfs_image="docker.io/library/busybox@sha256:" + "0" * 64, - ) - assert missing in str(excinfo.value) - - -@pytest.mark.skipif( - not ROOTFS_IMAGE, - reason="set ATE_ENV_ROOTFS_IMAGE=repo@sha256:... (an image with sh) for the create-from-image test", -) -async def test_create_from_image(client): - base = TEMPLATE or "default-template" - want_template = f"{base}-{ROOTFS_IMAGE.split('@sha256:', 1)[1][:12]}" - envs = [] - try: - env = await client.create( - f"pye2e-img-{uuid.uuid4().hex[:8]}", template_name=TEMPLATE, rootfs_image=ROOTFS_IMAGE - ) - envs.append(env) - - # The derived template is reported at once, before the actor serves. - info = await env.info() - assert info.template is not None - assert info.template.name == want_template - assert info.template.atespace == env.atespace - - await _wait_until_serving(env) - - # The process runs in the rootfs image's rootfs, with the guest coming - # from the image volume rather than from the image itself. - result = await env.shell( - "test -x /ate/ko-app/ate-env-guest && echo guest:volume; " - "test -e /ko-app/ate-env-guest && echo guest:baked; " - "test -r /etc/os-release && echo rootfs:ok" - ) - assert result.exit_code == 0, result - lines = result.stdout.splitlines() - assert "guest:volume" in lines, result - assert "guest:baked" not in lines, "the rootfs image should not carry the guest" - assert "rootfs:ok" in lines, result - - # The workspace is writable and the guest file path works unchanged. - content = os.urandom(64 * 1024 + 1) - path = f"/tmp/pye2e-img-{uuid.uuid4().hex[:8]}.bin" - assert await env.write_file(path, content) == len(content) - assert await env.read_file_bytes(path) == content - - # A second environment on the same image reuses the derived template - # instead of minting another one. - env2 = await client.create( - f"pye2e-img-{uuid.uuid4().hex[:8]}", template_name=TEMPLATE, rootfs_image=ROOTFS_IMAGE - ) - envs.append(env2) - assert (await env2.info()).template.name == want_template - await _wait_until_serving(env2) - assert (await env2.shell("echo second")).stdout == "second\n" - finally: - for e in envs: - try: - await e.delete() - except EnvError as exc: - pytest.fail(f"cleanup delete of {e.id} failed: {exc}") diff --git a/clients/python/tests/fakes.py b/clients/python/tests/fakes.py index 5e7f255f..5bf5f96f 100644 --- a/clients/python/tests/fakes.py +++ b/clients/python/tests/fakes.py @@ -62,13 +62,6 @@ async def CreateEnvironment(self, request, context): template.name = request.template.name if request.template.atespace: template.atespace = request.template.atespace - if request.rootfs_image: - # Mirror ate-env-api: the template becomes the base of one derived - # per image, named after the digest. - if "@sha256:" not in request.rootfs_image: - await context.abort(grpc.StatusCode.INVALID_ARGUMENT, "image is not pinned by digest") - digest = request.rootfs_image.split("@sha256:", 1)[1] - template.name = f"{template.name}-{digest[:12]}" environment = env_pb2.Environment( id=request.id, atespace=atespace, diff --git a/clients/python/tests/test_client.py b/clients/python/tests/test_client.py index 780eac67..2b154d78 100644 --- a/clients/python/tests/test_client.py +++ b/clients/python/tests/test_client.py @@ -55,30 +55,6 @@ async def test_create_omits_template_when_not_given(fake_stack): assert not fakes.environments.last_create_request.HasField("template") -async def test_create_with_image_reports_derived_template(fake_stack): - client, fakes = fake_stack - image = "docker.io/library/python@sha256:" + "0123456789abcdef" * 4 - env = await client.create("dev3", rootfs_image=image) - assert fakes.environments.last_create_request.rootfs_image == image - info = await env.info() - assert info.template is not None - assert info.template.name == "default-template-0123456789ab" - - -async def test_create_with_image_and_base_template(fake_stack): - client, fakes = fake_stack - image = "docker.io/library/python@sha256:" + "0123456789abcdef" * 4 - env = await client.create("dev4", template_name="py-base", rootfs_image=image) - info = await env.info() - assert info.template.name == "py-base-0123456789ab" - - -async def test_create_omits_image_when_not_given(fake_stack): - client, fakes = fake_stack - await client.create("dev5") - assert fakes.environments.last_create_request.rootfs_image == "" - - async def test_create_duplicate_maps_to_rpc_error_with_code(fake_stack): client, _ = fake_stack await client.create("dev1") diff --git a/clients/python/tests/test_config.py b/clients/python/tests/test_config.py index 7036f66f..e67910c1 100644 --- a/clients/python/tests/test_config.py +++ b/clients/python/tests/test_config.py @@ -18,7 +18,7 @@ def test_unknown_fields_are_rejected(): - # Previously silently dropped by pydantic (e.g. the NeMo/TML doc snippets). + # Previously silently dropped by pydantic (e.g. stale doc snippets). with pytest.raises(ValidationError, match="warm_pool_size"): FleetConfig(warm_pool_size=64) with pytest.raises(ValidationError, match="warmpool_replicas"): diff --git a/cmd/ate-env-api/main.go b/cmd/ate-env-api/main.go index 666543ff..ad1a5f0b 100644 --- a/cmd/ate-env-api/main.go +++ b/cmd/ate-env-api/main.go @@ -23,14 +23,12 @@ import ( "net/http" "os" "os/signal" - "strings" "syscall" "time" "github.com/agent-substrate/env/internal/apiservice" "github.com/agent-substrate/env/internal/ate" "github.com/agent-substrate/env/internal/mcp" - "github.com/agent-substrate/env/internal/portproxy" ateenvv1alpha "github.com/agent-substrate/env/proto/ateenv/v1alpha" "google.golang.org/grpc" ) @@ -40,13 +38,10 @@ func main() { listen = flag.String("listen", "0.0.0.0:7777", "address to serve the API on") ateapi = flag.String("ateapi", ate.DefaultControlAddr, "address of the ateapi gRPC control plane") atenet = flag.String("atenet", ate.DefaultRouterAddr, "address of the atenet HTTP router") - atenetConn = flag.String("atenet-connect", ate.DefaultRouterConnectAddr, "address of the atenet router's plaintext CONNECT listener; HTTP requests carrying X-Env-Id and X-Env-Port, or to /v1alpha/envs/{id}/ports/{port}/..., are forwarded through it to that port in the environment. The forwarding does not authenticate callers and reaches any port of any actor the router serves. Empty disables it") - ppOrigins = flag.String("port-proxy-trusted-origins", "", "comma-separated browser origins (scheme://host[:port]) whose pages may use the port proxy's path form, next to the listener's own origin") - ppHeaderTimeout = flag.Duration("port-proxy-response-header-timeout", 0, "how long the port proxy waits for a daemon's response headers; 0 means no limit, and the caller's deadline applies") hostSuffix = flag.String("host-suffix", ate.DefaultHostSuffix, "atenet router host suffix for actor routing") skipVerify = flag.Bool("skip-verify", true, "skip TLS certificate verification on the control plane connection") tokenFile = flag.String("ateapi-token-file", "", "file with a bearer token for the control plane, e.g. a projected ServiceAccount token with the ateapi audience") - wakeBudget = flag.Duration("wake-budget", apiservice.DefaultWakeBudget, "how long a guest call or a port-proxy request waits for a suspended environment to wake before failing; cold image pulls and cross-region snapshot reads exceed the router's own parking budget") + wakeBudget = flag.Duration("wake-budget", apiservice.DefaultWakeBudget, "how long a guest call waits for a suspended environment to wake before failing; cold image pulls and cross-region snapshot reads exceed the router's own parking budget") wakeRunningBudget = flag.Duration("wake-running-budget", apiservice.DefaultWakeRunningBudget, "how long a guest call waits when the environment is already RUNNING but the router still answers Unavailable (guest not yet listening vs. path down)") ) flag.Parse() @@ -91,31 +86,6 @@ func main() { mux.Handle("/v1alpha/envs/{id}/mcp", mcp.NewHandler(client)) mux.Handle("/", grpcServer) - // Requests that name an environment port go to the port proxy before - // anything else, whatever their path: the path belongs to the daemon. - var handler http.Handler = mux - if *atenetConn != "" { - var origins []string - for o := range strings.SplitSeq(*ppOrigins, ",") { - if o = strings.TrimSpace(o); o != "" { - origins = append(origins, o) - } - } - pp, err := portproxy.New(portproxy.Options{ - ConnectAddr: *atenetConn, - HostSuffix: *hostSuffix, - DefaultAtespace: ate.DefaultAtespace, - ResponseHeaderTimeout: *ppHeaderTimeout, - TrustedOrigins: origins, - WakeBudget: *wakeBudget, - }) - if err != nil { - log.Fatalf("--port-proxy-trusted-origins: %v", err) - } - // Header form and path form, /v1alpha/envs/{id}/ports/{port}/... - handler = portproxy.Dispatch(pp, mux) - } - // Enable both HTTP/1.1 (for /healthz and /v1alpha/ HTTP endpoints) and // unencrypted HTTP/2 (h2c, for gRPC services) on the same plaintext TCP listener. var protocols http.Protocols @@ -124,7 +94,7 @@ func main() { srv := &http.Server{ Addr: *listen, - Handler: handler, + Handler: mux, Protocols: &protocols, } @@ -137,7 +107,7 @@ func main() { _ = srv.Shutdown(context.Background()) }() - log.Printf("ate-env-api listening on %s (ateapi %s, atenet %s, atenet CONNECT %q)", *listen, *ateapi, *atenet, *atenetConn) + log.Printf("ate-env-api listening on %s (ateapi %s, atenet %s)", *listen, *ateapi, *atenet) if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) { log.Fatalf("server error: %v", err) } diff --git a/cmd/ate-env-guest/main.go b/cmd/ate-env-guest/main.go index 0af0a467..d6f952f9 100644 --- a/cmd/ate-env-guest/main.go +++ b/cmd/ate-env-guest/main.go @@ -15,8 +15,6 @@ // Command ate-env-guest is the daemon that runs inside a Substrate actor // and exposes command execution and filesystem access over gRPC // (ProcessService and FileSystemService) with HTTP readiness probe support. -// It can also start and supervise extra runtimes next to itself (-sidecar), -// folding their readiness into /readyz (-sidecar-readyz). package main import ( @@ -29,10 +27,7 @@ import ( "net/http" "os" "os/signal" - "strings" - "sync" "syscall" - "time" "github.com/agent-substrate/env/guest" ) @@ -41,23 +36,8 @@ func main() { listen := flag.String("listen", ":80", "address to serve the guest API on") logDir := flag.String("log-dir", "", "directory for process logs (defaults to /var/log/ate-jobs or temporary dir)") workspace := flag.String("workspace", "/", "workspace root directory") - var sidecars, sidecarReadyz multiFlag - flag.Var(&sidecars, "sidecar", "extra runtime to start next to the guest and keep running: an absolute binary path followed by its arguments, whitespace-separated (repeatable)") - flag.Var(&sidecarReadyz, "sidecar-readyz", "URL that must answer 2xx before /readyz reports ready, e.g. http://127.0.0.1:44772/ready (repeatable)") flag.Parse() - // Reject bad sidecar values before listening, so a misconfigured template - // fails at actor start instead of looping on a restarting sidecar. - var sidecarArgv [][]string - for _, s := range sidecars { - argv, err := parseSidecar(s) - if err != nil { - log.Fatalf("invalid -sidecar: %v", err) - } - sidecarArgv = append(sidecarArgv, argv) - } - gate := newReadyGate(sidecarReadyz) - addr := *listen if addr == "" { addr = ":80" @@ -79,10 +59,6 @@ func main() { mux := http.NewServeMux() mux.HandleFunc("GET /readyz", func(w http.ResponseWriter, r *http.Request) { - if ok, why := gate.Ready(r.Context()); !ok { - http.Error(w, why, http.StatusServiceUnavailable) - return - } io.WriteString(w, "ok\n") }) mux.Handle("/", grpcServer) @@ -114,37 +90,9 @@ func main() { _ = srv.Shutdown(context.Background()) }() - var supervisors sync.WaitGroup - for _, argv := range sidecarArgv { - log.Printf("starting sidecar %s", strings.Join(argv, " ")) - supervisors.Add(1) - go func() { - defer supervisors.Done() - runSidecar(ctx, argv, log.Printf) - }() - } - log.Printf("ate-env-guest listening on %s (gRPC services=%s, workspace=%s, logdir=%s)", lis.Addr(), guest.FormatEnabledServices(cfg), cfg.Workspace, cfg.LogDir) if err := srv.Serve(lis); err != nil && !errors.Is(err, http.ErrServerClosed) { log.Fatalf("ate-env-guest server error: %v", err) } - // The supervisors stop their process groups with SIGTERM, a grace period - // and SIGKILL; that sequence needs this process to stay alive until it - // has run. Inside an actor the guest is PID 1 and the container would end - // with it anyway; outside one, this is what stops the sidecars. - waitBounded(&supervisors, 2*sidecarStopGrace+time.Second) -} - -// waitBounded waits for wg, but at most d. -func waitBounded(wg *sync.WaitGroup, d time.Duration) { - done := make(chan struct{}) - go func() { - wg.Wait() - close(done) - }() - select { - case <-done: - case <-time.After(d): - } } diff --git a/cmd/ate-env-guest/sidecar.go b/cmd/ate-env-guest/sidecar.go deleted file mode 100644 index 60fd12be..00000000 --- a/cmd/ate-env-guest/sidecar.go +++ /dev/null @@ -1,259 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package main - -import ( - "bytes" - "context" - "errors" - "fmt" - "net/http" - "os" - "os/exec" - "path/filepath" - "strings" - "sync" - "syscall" - "time" -) - -// Sidecars are extra runtimes started next to the guest, from the same -// container: typically a second data-plane daemon that was mounted into the -// actor as an image volume (see docs/rootfs-images/RUNTIMES.md). The guest is -// the container's only process by contract, so it is the one that starts -// them, restarts them if they exit, and folds their readiness into /readyz. - -// multiFlag collects the values of a repeatable flag. -type multiFlag []string - -func (m *multiFlag) String() string { return strings.Join(*m, "; ") } - -func (m *multiFlag) Set(v string) error { - if strings.TrimSpace(v) == "" { - return errors.New("value must not be empty") - } - *m = append(*m, v) - return nil -} - -// parseSidecar turns a -sidecar value into argv. Values are split on -// whitespace and no quoting is interpreted: sidecars run without a shell, so -// a runtime that needs one ships a wrapper script in its layer. The binary -// must be an absolute path, which is what a mounted layer provides and what -// keeps the guest independent of the rootfs image's PATH. -func parseSidecar(v string) ([]string, error) { - argv := strings.Fields(v) - if len(argv) == 0 { - return nil, errors.New("sidecar command is empty") - } - if !filepath.IsAbs(argv[0]) { - return nil, fmt.Errorf("sidecar %q: the binary must be an absolute path (sidecars run without a shell or PATH lookup)", argv[0]) - } - return argv, nil -} - -const ( - sidecarMinBackoff = time.Second - sidecarMaxBackoff = 30 * time.Second - // sidecarStopGrace is how long a stopped sidecar's process group gets to exit - // and release the output pipes after SIGTERM before it is SIGKILLed. - sidecarStopGrace = 5 * time.Second - // sidecarStableAfter is how long a sidecar must have run for its next - // restart to start from the minimum backoff again. - sidecarStableAfter = time.Minute -) - -// runSidecar runs argv until ctx is done, restarting it whenever it exits. -// The backoff doubles from one second to thirty while the sidecar keeps -// failing quickly and resets once it has run for a while. Output goes to -// logf, one line at a time, prefixed with the binary's name. -func runSidecar(ctx context.Context, argv []string, logf func(string, ...any)) { - name := filepath.Base(argv[0]) - backoff := sidecarMinBackoff - for { - cmd := exec.CommandContext(ctx, argv[0], argv[1:]...) - cmd.Env = os.Environ() - out := &lineLogger{prefix: name, logf: logf} - cmd.Stdout = out - cmd.Stderr = out - // The sidecar runs in its own process group so that stopping it also - // stops what it spawned. On shutdown the group gets SIGTERM, then - // sidecarStopGrace to exit and release the output pipes, then SIGKILL. - // Killing only the direct child is not enough: a wrapper script (the - // documented way to pass quoted arguments) would die alone while its - // child kept the pipes open, and Wait would block until that child - // exited. On Linux `sh -c` forks its command; bash on macOS execs it, - // which is why the test passed on one and hung on the other. - cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} - cmd.Cancel = func() error { - logf("sidecar %s: context done, SIGTERM to process group %d", name, cmd.Process.Pid) - return signalGroup(cmd.Process.Pid, syscall.SIGTERM) - } - cmd.WaitDelay = sidecarStopGrace - start := time.Now() - err := cmd.Run() - out.flush() - // One stop sequence for every path. On shutdown it finishes what - // Cancel started; before a restart it clears what a crashed sidecar - // left behind, so the new one does not compete with its own orphans - // for ports, pipes or lock files. - if cmd.Process != nil { - stopGroup(name, cmd.Process.Pid, sidecarStopGrace, logf) - } - if ctx.Err() != nil { - logf("sidecar %s stopped", name) - return - } - ran := time.Since(start) - logf("sidecar %s exited after %s: %v; restarting in %s", name, ran.Round(time.Millisecond), err, backoff) - select { - case <-ctx.Done(): - return - case <-time.After(backoff): - } - if ran >= sidecarStableAfter { - backoff = sidecarMinBackoff - } else { - backoff = min(backoff*2, sidecarMaxBackoff) - } - } -} - -// stopGroup is the one way a sidecar's process group is taken down: SIGTERM, -// up to grace for the group to disappear, then SIGKILL for what is left. Each -// step is logged so a shutdown or a restart can be followed in the guest log. -// A group that is already gone (ESRCH) is success. -func stopGroup(name string, pid int, grace time.Duration, logf func(string, ...any)) { - if !groupAlive(pid) { - logf("sidecar %s: process group %d already exited", name, pid) - return - } - logf("sidecar %s: SIGTERM to process group %d, waiting up to %s", name, pid, grace) - _ = signalGroup(pid, syscall.SIGTERM) - deadline := time.Now().Add(grace) - for groupAlive(pid) { - if time.Now().After(deadline) { - logf("sidecar %s: process group %d still alive after %s, SIGKILL", name, pid, grace) - _ = signalGroup(pid, syscall.SIGKILL) - return - } - time.Sleep(20 * time.Millisecond) - } - logf("sidecar %s: process group %d exited after SIGTERM", name, pid) -} - -// groupAlive reports whether any process remains in the group led by pid. -// Signal 0 performs the permission and existence checks without delivering -// anything; ESRCH means the group is gone. -func groupAlive(pid int) bool { - return !errors.Is(syscall.Kill(-pid, 0), syscall.ESRCH) -} - -// signalGroup sends sig to the process group led by pid. A group that no -// longer exists is not an error: everything in it has already exited. The -// calls after Wait name a reaped pid; a reused pid leading another group is -// possible in theory, but the window is microseconds in a sandbox whose pid -// space is nearly empty. -func signalGroup(pid int, sig syscall.Signal) error { - if err := syscall.Kill(-pid, sig); err != nil && !errors.Is(err, syscall.ESRCH) { - return err - } - return nil -} - -// lineLogger writes complete lines to logf with a prefix; a partial trailing -// line is kept until the next write or flush. -type lineLogger struct { - prefix string - logf func(string, ...any) - mu sync.Mutex - buf bytes.Buffer -} - -func (l *lineLogger) Write(p []byte) (int, error) { - l.mu.Lock() - defer l.mu.Unlock() - l.buf.Write(p) - for { - line, err := l.buf.ReadString('\n') - if err != nil { - // No newline yet: put the partial line back. - l.buf.Reset() - l.buf.WriteString(line) - break - } - l.logf("[%s] %s", l.prefix, strings.TrimRight(line, "\r\n")) - } - return len(p), nil -} - -func (l *lineLogger) flush() { - l.mu.Lock() - defer l.mu.Unlock() - if l.buf.Len() > 0 { - l.logf("[%s] %s", l.prefix, l.buf.String()) - l.buf.Reset() - } -} - -// readyGate makes /readyz wait for the sidecars' own readiness endpoints. -// Each URL is polled until it answers 2xx once; after that it is not asked -// again. Readiness is a start-up gate, so Substrate's wakeup probe admits -// traffic only when every runtime in the actor is serving, not a liveness -// check on the sidecars. -type readyGate struct { - urls []string - client *http.Client - - mu sync.Mutex - seen map[string]bool -} - -func newReadyGate(urls []string) *readyGate { - return &readyGate{ - urls: urls, - client: &http.Client{Timeout: 2 * time.Second}, - seen: make(map[string]bool, len(urls)), - } -} - -// Ready reports whether every URL has answered 2xx, and if not, which one is -// still pending and why. -func (g *readyGate) Ready(ctx context.Context) (bool, string) { - for _, u := range g.urls { - g.mu.Lock() - ok := g.seen[u] - g.mu.Unlock() - if ok { - continue - } - req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil) - if err != nil { - return false, fmt.Sprintf("sidecar readiness %s: %v", u, err) - } - resp, err := g.client.Do(req) - if err != nil { - return false, fmt.Sprintf("sidecar readiness %s: %v", u, err) - } - resp.Body.Close() - if resp.StatusCode/100 != 2 { - return false, fmt.Sprintf("sidecar readiness %s: HTTP %d", u, resp.StatusCode) - } - g.mu.Lock() - g.seen[u] = true - g.mu.Unlock() - } - return true, "" -} diff --git a/cmd/ate-env-guest/sidecar_test.go b/cmd/ate-env-guest/sidecar_test.go deleted file mode 100644 index d2cc6734..00000000 --- a/cmd/ate-env-guest/sidecar_test.go +++ /dev/null @@ -1,159 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package main - -import ( - "context" - "fmt" - "net/http" - "net/http/httptest" - "os/exec" - "strings" - "sync" - "sync/atomic" - "testing" - "time" -) - -func TestParseSidecar(t *testing.T) { - argv, err := parseSidecar(" /opt/osb/execd --port 44772 ") - if err != nil || strings.Join(argv, " ") != "/opt/osb/execd --port 44772" { - t.Errorf("parseSidecar = %v, %v", argv, err) - } - for _, bad := range []string{"", " ", "execd --port 1", "sh -c /opt/osb/execd"} { - if _, err := parseSidecar(bad); err == nil { - t.Errorf("parseSidecar(%q) accepted a non-absolute or empty command", bad) - } - } - var m multiFlag - if err := m.Set(""); err == nil { - t.Error("empty flag value accepted") - } - _ = m.Set("/a") - _ = m.Set("/b x") - if len(m) != 2 || m.String() != "/a; /b x" { - t.Errorf("multiFlag = %v", m) - } -} - -func TestReadyGateWaitsOnceForEachURL(t *testing.T) { - var state atomic.Int32 // 0: 503, 1: 200 - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if state.Load() == 0 { - http.Error(w, "booting", http.StatusServiceUnavailable) - return - } - w.WriteHeader(http.StatusOK) - })) - defer srv.Close() - closed := httptest.NewServer(http.NotFoundHandler()) - closedURL := closed.URL - closed.Close() - - g := newReadyGate([]string{srv.URL + "/ready"}) - ctx := context.Background() - if ok, why := g.Ready(ctx); ok || !strings.Contains(why, "HTTP 503") { - t.Errorf("not ready expected, got ok=%v why=%q", ok, why) - } - state.Store(1) - if ok, why := g.Ready(ctx); !ok { - t.Errorf("ready expected, got %q", why) - } - // Once seen, the URL is not re-checked: readiness is a start-up gate. - state.Store(0) - if ok, _ := g.Ready(ctx); !ok { - t.Error("readiness flipped back after the sidecar had answered once") - } - - unreachable := newReadyGate([]string{closedURL + "/ready"}) - if ok, why := unreachable.Ready(ctx); ok || why == "" { - t.Errorf("an unreachable sidecar must report not ready with a reason, got ok=%v why=%q", ok, why) - } - if ok, _ := newReadyGate(nil).Ready(ctx); !ok { - t.Error("no sidecars means ready") - } -} - -func TestRunSidecarRestartsAndStopsWithContext(t *testing.T) { - sh, err := exec.LookPath("sh") - if err != nil { - t.Skip("no sh on this host") - } - var mu sync.Mutex - var lines []string - logf := func(format string, args ...any) { - mu.Lock() - defer mu.Unlock() - lines = append(lines, fmt.Sprintf(format, args...)) - } - - // A sidecar that prints and exits non-zero is restarted with backoff. - ctx, cancel := context.WithTimeout(context.Background(), 2500*time.Millisecond) - defer cancel() - done := make(chan struct{}) - go func() { - runSidecar(ctx, []string{sh, "-c", "echo hello from sidecar; exit 3"}, logf) - close(done) - }() - select { - case <-done: - case <-time.After(5 * time.Second): - t.Fatal("runSidecar did not return after its context ended") - } - mu.Lock() - joined := strings.Join(lines, "\n") - mu.Unlock() - if !strings.Contains(joined, "[sh] hello from sidecar") { - t.Errorf("sidecar output was not logged with its prefix:\n%s", joined) - } - if n := strings.Count(joined, "restarting in"); n < 2 { - t.Errorf("expected at least two restarts within the window (1s then 2s backoff), got %d:\n%s", n, joined) - } - if !strings.Contains(joined, "restarting in 1s") || !strings.Contains(joined, "restarting in 2s") { - t.Errorf("backoff did not double:\n%s", joined) - } - - // A long-running sidecar is killed promptly when the guest shuts down, - // together with what it forked. `sleep 60 & wait` forks on every - // platform; a bare `sh -c "sleep 60"` execs on macOS and would not - // exercise the process-group kill there. The forked sleep holds the - // output pipe, and WaitDelay (sidecarStopGrace) would not release it - // before 5 s, so a return within the bound below proves the child died. - bound := sidecarStopGrace - time.Second - ctx2, cancel2 := context.WithCancel(context.Background()) - done2 := make(chan struct{}) - go func() { - runSidecar(ctx2, []string{sh, "-c", "sleep 60 & wait"}, logf) - close(done2) - }() - time.Sleep(200 * time.Millisecond) - cancel2() - select { - case <-done2: - case <-time.After(bound): - t.Fatal("sidecar (or its child) was not stopped when the context was cancelled") - } - mu.Lock() - joined = strings.Join(lines, "\n") - mu.Unlock() - for _, want := range []string{"SIGTERM to process group", "sidecar sh stopped"} { - if !strings.Contains(joined, want) { - t.Errorf("termination sequence not logged (%q missing):\n%s", want, joined) - } - } - if !strings.Contains(joined, "exited after SIGTERM") && !strings.Contains(joined, "already exited") { - t.Errorf("the group's end was not logged:\n%s", joined) - } -} diff --git a/cmd/ate-env/main.go b/cmd/ate-env/main.go index 2a67c734..2f0cb31a 100644 --- a/cmd/ate-env/main.go +++ b/cmd/ate-env/main.go @@ -137,7 +137,6 @@ func newCreateCommand() *cobra.Command { atespace string createTemplate string createTemplateAtespace string - rootfsImage string ) cmd := &cobra.Command{ Use: "create ", @@ -153,9 +152,8 @@ func newCreateCommand() *cobra.Command { defer client.Close() req := &ateenvv1alpha.CreateEnvironmentRequest{ - Id: args[0], - Atespace: atespace, - RootfsImage: rootfsImage, + Id: args[0], + Atespace: atespace, } if createTemplate != "" || createTemplateAtespace != "" { tmplAtespace := createTemplateAtespace @@ -175,7 +173,6 @@ func newCreateCommand() *cobra.Command { cmd.Flags().StringVar(&atespace, "atespace", apiservice.DefaultAtespace, "Substrate atespace") cmd.Flags().StringVar(&createTemplate, "template", "", "ActorTemplate name (defaults to server default)") cmd.Flags().StringVar(&createTemplateAtespace, "template-atespace", "", "Substrate atespace of the ActorTemplate (defaults to environment atespace)") - cmd.Flags().StringVar(&rootfsImage, "rootfs-image", "", "digest-pinned rootfs image to run unmodified (repo@sha256:...); the template becomes the base the guest is taken from") return cmd } diff --git a/cmd/ate-env/main_test.go b/cmd/ate-env/main_test.go index e2e19795..80d0ec2d 100644 --- a/cmd/ate-env/main_test.go +++ b/cmd/ate-env/main_test.go @@ -203,21 +203,14 @@ func TestHelpOutput(t *testing.T) { }) } -func TestCreateImageFlag(t *testing.T) { - args := []string{"create", "dev1", "--rootfs-image", "example.com/task@sha256:abc"} - root := newRootCommand(args) - cmd, _, err := root.Find(args[:2]) +func TestCreateHasNoRootfsImageFlag(t *testing.T) { + root := newRootCommand([]string{"create", "dev1"}) + cmd, _, err := root.Find([]string{"create", "dev1"}) if err != nil { t.Fatalf("root.Find failed: %v", err) } - if cmd.Flags().Lookup("rootfs-image") == nil { - t.Fatal("create has no --rootfs-image flag") - } - if err := cmd.ParseFlags(args[2:]); err != nil { - t.Fatalf("parsing --rootfs-image: %v", err) - } - if got, _ := cmd.Flags().GetString("rootfs-image"); got != "example.com/task@sha256:abc" { - t.Errorf("--rootfs-image = %q", got) + if cmd.Flags().Lookup("rootfs-image") != nil { + t.Error("create still has a --rootfs-image flag; on-demand image templates were removed") } } diff --git a/cmd/ate-env/manifest.go b/cmd/ate-env/manifest.go index 9552aec2..56f68482 100644 --- a/cmd/ate-env/manifest.go +++ b/cmd/ate-env/manifest.go @@ -19,7 +19,6 @@ import ( "errors" "fmt" "io" - "strings" "github.com/agent-substrate/env/internal/apiservice" "github.com/spf13/cobra" @@ -95,64 +94,12 @@ type templateConfig struct { rootfsImage string // workspace, when set, is passed to the guest as -workspace. workspace string - // layers are extra runtimes mounted as read-only image volumes, each - // "name=image@sha256:...=/mount/path". - layers []string - // sidecars are command lines the guest starts and supervises next to - // itself (absolute binary path first), one per -sidecar flag. - sidecars []string - // sidecarReadyz are URLs the guest's /readyz waits for. - sidecarReadyz []string // cpu and memory are the actor's compute limits as Kubernetes quantities. // Empty means defaultActorCPU and defaultActorMemory. cpu string memory string } -// runtimeLayer is a parsed --layer value. -type runtimeLayer struct { - name, image, mountPath string -} - -// parseLayer parses "name=image@sha256:...=/mount/path". The image reference -// carries no '=' and the mount path is absolute, so splitting on the first two -// '=' is unambiguous. -func parseLayer(v string) (runtimeLayer, error) { - parts := strings.SplitN(v, "=", 3) - if len(parts) != 3 || parts[0] == "" || parts[1] == "" || parts[2] == "" { - return runtimeLayer{}, fmt.Errorf("--layer %q: want name=image@sha256:...=/mount/path", v) - } - l := runtimeLayer{name: parts[0], image: parts[1], mountPath: parts[2]} - if l.name == apiservice.GuestVolumeName { - return runtimeLayer{}, fmt.Errorf("--layer %q: the name %q is reserved for the guest", v, apiservice.GuestVolumeName) - } - if _, err := apiservice.ImageDigest(l.image); err != nil { - return runtimeLayer{}, fmt.Errorf("--layer %q: %w", l.name, err) - } - if !strings.HasPrefix(l.mountPath, "/") || l.mountPath == apiservice.GuestMountPath { - return runtimeLayer{}, fmt.Errorf("--layer %q: mount path must be absolute and not %s", l.name, apiservice.GuestMountPath) - } - return l, nil -} - -// parseLayers parses and cross-checks every --layer value. -func parseLayers(values []string) ([]runtimeLayer, error) { - var layers []runtimeLayer - names, mounts := map[string]bool{}, map[string]bool{} - for _, v := range values { - l, err := parseLayer(v) - if err != nil { - return nil, err - } - if names[l.name] || mounts[l.mountPath] { - return nil, fmt.Errorf("--layer %q: duplicate name or mount path", v) - } - names[l.name], mounts[l.mountPath] = true, true - layers = append(layers, l) - } - return layers, nil -} - // Bounds Substrate applies to an ActorTemplate's limits // (go.ate.dev/substrate/internal/resources.ValidateLimit and the // Limits.quantity maxLength), checked here so the CLI never prints a template @@ -208,20 +155,6 @@ binaries default it to the ate-env-guest image published with the release return fmt.Errorf("--guest-image must be digest-pinned to be mounted as an image volume: %w", err) } } - if _, err := parseLayers(c.layers); err != nil { - return err - } - for _, s := range c.sidecars { - argv := strings.Fields(s) - if len(argv) == 0 || !strings.HasPrefix(argv[0], "/") { - return fmt.Errorf("--sidecar %q: want an absolute binary path followed by its arguments", s) - } - } - for _, u := range c.sidecarReadyz { - if !strings.HasPrefix(u, "http://") && !strings.HasPrefix(u, "https://") { - return fmt.Errorf("--sidecar-readyz %q: want an http(s) URL", u) - } - } return nil } @@ -308,9 +241,6 @@ It prints YAML to stdout without touching the cluster.`, cmd.Flags().StringVar(&tCfg.snapshotsBucket, "snapshots-bucket", "", "object-storage bucket (with optional prefix) for actor snapshots, e.g. gs://bucket/prefix/") cmd.Flags().StringVar(&tCfg.rootfsImage, "rootfs-image", "", "digest-pinned image to run unmodified; the guest is mounted into it as a read-only image volume at "+apiservice.GuestMountPath) cmd.Flags().StringVar(&tCfg.workspace, "workspace", "", "workspace root inside the actor, passed to the guest as -workspace (default: the guest's default)") - cmd.Flags().StringArrayVar(&tCfg.layers, "layer", nil, "extra runtime mounted as a read-only image volume, as name=image@sha256:...=/mount/path (repeatable)") - cmd.Flags().StringArrayVar(&tCfg.sidecars, "sidecar", nil, "command the guest starts and keeps running next to itself, an absolute binary path plus arguments, e.g. '/opt/opensandbox/execd --port 44772' (repeatable)") - cmd.Flags().StringArrayVar(&tCfg.sidecarReadyz, "sidecar-readyz", nil, "URL the actor's readiness waits for, e.g. http://127.0.0.1:44772/ready (repeatable)") cmd.Flags().StringVar(&tCfg.cpu, "cpu", defaultActorCPU, "CPU limit of each actor, a Kubernetes quantity such as 2 or 500m; Substrate requires one") cmd.Flags().StringVar(&tCfg.memory, "memory", defaultActorMemory, "memory limit of each actor, a Kubernetes quantity such as 4Gi; Substrate requires one") @@ -517,27 +447,6 @@ func buildTemplate(cfg templateConfig) (*ateapipb.ActorTemplate, error) { if cfg.workspace != "" { command = append(command, "-workspace", cfg.workspace) } - // Extra runtimes: each layer is an image volume, each sidecar a flag the - // guest acts on. They are added before any rootfs-image derivation so the - // derived template inherits them, and so do environments created from - // it with an image of their own. - layers, err := parseLayers(cfg.layers) - if err != nil { - return nil, err - } - for _, l := range layers { - tmpl.Volumes = append(tmpl.Volumes, &ateapipb.Volume{ - Name: l.name, - Image: &ateapipb.ImageVolumeSource{Reference: l.image}, - }) - c.VolumeMounts = append(c.VolumeMounts, &ateapipb.VolumeMount{Name: l.name, MountPath: l.mountPath}) - } - for _, s := range cfg.sidecars { - command = append(command, "-sidecar", s) - } - for _, u := range cfg.sidecarReadyz { - command = append(command, "-sidecar-readyz", u) - } c.Command = command if cfg.rootfsImage == "" { return tmpl, nil diff --git a/cmd/ate-env/manifest_test.go b/cmd/ate-env/manifest_test.go index 95f26594..5353d139 100644 --- a/cmd/ate-env/manifest_test.go +++ b/cmd/ate-env/manifest_test.go @@ -438,64 +438,6 @@ func TestTemplateConfigResolveImagesRootfsImage(t *testing.T) { } } -func TestBuildTemplateWithRuntimeLayer(t *testing.T) { - cfg := testTemplateConfig() - cfg.guestImage = "example.com/guest@sha256:" + testHex64 - cfg.rootfsImage = "docker.io/library/python@sha256:" + testHex64 - cfg.layers = []string{"execd=example.com/execd@sha256:" + testHex64 + "=/opt/opensandbox"} - cfg.sidecars = []string{"/opt/opensandbox/execd --port 44772"} - cfg.sidecarReadyz = []string{"http://127.0.0.1:44772/ready"} - - tmpl, err := buildTemplate(cfg) - if err != nil { - t.Fatalf("buildTemplate: %v", err) - } - vols := tmpl.GetVolumes() - if len(vols) != 2 || vols[0].GetName() != "execd" || vols[0].GetImage().GetReference() != "example.com/execd@sha256:"+testHex64 || vols[1].GetName() != apiservice.GuestVolumeName { - t.Errorf("volumes = %v, want the execd layer then the guest volume", vols) - } - c := tmpl.GetContainers()[0] - if len(c.GetVolumeMounts()) != 2 || c.GetVolumeMounts()[0].GetMountPath() != "/opt/opensandbox" || c.GetVolumeMounts()[1].GetMountPath() != apiservice.GuestMountPath { - t.Errorf("mounts = %v", c.GetVolumeMounts()) - } - want := "/ate/ko-app/ate-env-guest -sidecar /opt/opensandbox/execd --port 44772 -sidecar-readyz http://127.0.0.1:44772/ready" - if got := strings.Join(c.GetCommand(), " "); got != want { - t.Errorf("command = %q, want %q", got, want) - } - if c.GetImage() != cfg.rootfsImage { - t.Errorf("image = %q", c.GetImage()) - } -} - -func TestParseLayerRejects(t *testing.T) { - cases := map[string]string{ - "missing parts": "execd=example.com/execd@sha256:" + testHex64, - "unpinned image": "execd=example.com/execd:latest=/opt/execd", - "relative mount": "execd=example.com/execd@sha256:" + testHex64 + "=opt/execd", - "guest mount path": "execd=example.com/execd@sha256:" + testHex64 + "=/ate", - "reserved name": "guest=example.com/execd@sha256:" + testHex64 + "=/opt/execd", - } - for name, v := range cases { - if _, err := parseLayer(v); err == nil { - t.Errorf("%s: %q accepted", name, v) - } - } - dup := "a=example.com/x@sha256:" + testHex64 + "=/opt/a" - if _, err := parseLayers([]string{dup, dup}); err == nil { - t.Error("duplicate layer accepted") - } - cfg := testTemplateConfig() - cfg.sidecars = []string{"execd --port 1"} - if err := cfg.resolveImages(); err == nil { - t.Error("relative sidecar binary accepted") - } - cfg = testTemplateConfig() - cfg.sidecarReadyz = []string{"127.0.0.1:44772/ready"} - if err := cfg.resolveImages(); err == nil { - t.Error("sidecar-readyz without a scheme accepted") - } -} - const relHex64 = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef" func TestImageFlagsDefaultToReleaseImages(t *testing.T) { diff --git a/cmd/benchmarking/swebench-loader/main.go b/cmd/benchmarking/swebench-loader/main.go deleted file mode 100644 index 1be90cbd..00000000 --- a/cmd/benchmarking/swebench-loader/main.go +++ /dev/null @@ -1,246 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// swebench-loader drives SWE-bench task images through the ate-env API as a -// simulated RL job and writes one JSONL event per operation. -// -// swebench-loader run -dataset manifest.jsonl -mode eval -images 500 -tasks-per-image 6 -width 500 -out run.jsonl -// swebench-loader run -dataset manifest.jsonl -mode steps -images 500 -rollouts 2 -steps 8 -park suspend -out run.jsonl -// swebench-loader report run.jsonl [-json summary.json] -package main - -import ( - "context" - "encoding/json" - "flag" - "fmt" - "log/slog" - "os" - "os/signal" - "strings" - "syscall" - "time" - - "github.com/agent-substrate/env/internal/benchmarking/swebench" -) - -func main() { - if len(os.Args) < 2 { - usage() - os.Exit(2) - } - var err error - switch os.Args[1] { - case "run": - err = runCmd(os.Args[2:]) - case "report": - err = reportCmd(os.Args[2:]) - case "-h", "--help", "help": - usage() - return - default: - usage() - os.Exit(2) - } - if err != nil { - fmt.Fprintln(os.Stderr, "error:", err) - os.Exit(1) - } -} - -func usage() { - fmt.Fprintln(os.Stderr, "usage: swebench-loader run [flags] | swebench-loader report FILE [-json OUT]") -} - -func runCmd(args []string) (retErr error) { - fs := flag.NewFlagSet("run", flag.ContinueOnError) - var ( - endpoint = fs.String("endpoint", "localhost:7777", "ate-env API address (host:port).") - connections = fs.Int("connections", 4, "gRPC connections to the API, handed out round-robin per environment; a Service balances per connection, so use at least the number of API replicas.") - atespace = fs.String("atespace", "ate-env", "Atespace the environments are created in.") - baseTemplate = fs.String("base-template", "default-template", "Base ActorTemplate the per-image templates derive from (ate-env manifest template).") - dataset = fs.String("dataset", "", "JSONL manifest of SWE-bench instances; a path or an https URL. Required.") - images = fs.Int("images", 0, "Use the first N instances of the dataset; 0 means all.") - instances = fs.String("instances", "", "Comma-separated instance ids to use instead of the first N.") - excludeFile = fs.String("exclude-file", "", "File of instance ids to leave out, one per line; text after the id is a comment.") - trajectory = fs.String("trajectory", swebench.DefaultTrajectory, fmt.Sprintf("Trajectory file or built-in name (%s).", strings.Join(swebench.Names(), ", "))) - mode = fs.String("mode", string(swebench.ModeEval), "eval (images x tasks-per-image, each create/replay/delete) or steps (images x rollouts kept across steps).") - width = fs.Int("width", 50, "Most tasks in flight at once.") - tasksPerImage = fs.Int("tasks-per-image", 1, "eval: replays per image.") - rollouts = fs.Int("rollouts", 2, "steps: environments per image.") - steps = fs.Int("steps", 8, "steps: training steps.") - park = fs.String("park", string(swebench.ParkSuspend), "steps: what happens between steps: suspend, delete or none.") - createTimeout = fs.Duration("create-timeout", 5*time.Minute, "Deadline for one create call.") - stepTimeout = fs.Duration("step-timeout", 10*time.Minute, "Default deadline for a run step without its own timeout.") - wakeBudget = fs.Duration("wake-budget", 3*time.Minute, "How long a task's first command is retried while the environment is still waking or its image is still being baked.") - wakeBackoff = fs.Duration("wake-backoff", swebench.DefaultWakeBackoff, "Initial pause between wake retries; doubles up to five times this.") - runID = fs.String("run-id", "", "Run id, prefix of every environment name: lowercase letters, digits and '-', at most 20 characters; default swe-.") - out = fs.String("out", "", "JSONL events file; default .jsonl in the working directory.") - keep = fs.Bool("keep", false, "Leave the environments in place when the run ends.") - dryRun = fs.Bool("dry-run", false, "Load and render everything, print the plan and exit without touching the API.") - ) - if err := fs.Parse(args); err != nil { - return err - } - if *dataset == "" { - return fmt.Errorf("-dataset is required") - } - if *runID == "" { - *runID = fmt.Sprintf("swe-%d", time.Now().Unix()) - } - if *out == "" { - *out = *runID + ".jsonl" - } - log := slog.New(slog.NewJSONHandler(os.Stderr, nil)) - - all, err := swebench.LoadDataset(*dataset) - if err != nil { - return err - } - var ids []string - if *instances != "" { - ids = strings.Split(*instances, ",") - } - var exclude []string - if *excludeFile != "" { - if exclude, err = swebench.ReadIDList(*excludeFile); err != nil { - return err - } - } - selected, err := swebench.Select(all, ids, exclude, *images) - if err != nil { - return err - } - if *images > 0 && len(selected) < *images && len(ids) == 0 { - fmt.Fprintf(os.Stderr, "note: %d instances requested, %d available after the exclude list\n", *images, len(selected)) - } - traj, err := swebench.LoadTrajectory(*trajectory) - if err != nil { - return err - } - for _, in := range selected { - if _, err := traj.Render(in); err != nil { - return err - } - } - - fanout := *tasksPerImage - if swebench.Mode(*mode) == swebench.ModeSteps { - fanout = *rollouts - } - log.Info("plan", slog.String("run", *runID), slog.String("mode", *mode), slog.Int("images", len(selected)), - slog.Int("environments", len(selected)*fanout), slog.Int("width", *width), slog.Int("steps", *steps), - slog.String("park", *park), slog.String("trajectory", traj.Name), slog.String("out", *out)) - if *dryRun { - return nil - } - - backend, err := swebench.NewEnvBackend(*endpoint, *atespace, *baseTemplate, *connections) - if err != nil { - return err - } - defer backend.Close() - - f, err := os.Create(*out) - if err != nil { - return err - } - rec := swebench.NewRecorder(f) - // Events lost to a write error, or a file that does not flush, must fail - // the run rather than produce a report with holes. - defer func() { - if cerr := f.Close(); cerr != nil && retErr == nil { - retErr = fmt.Errorf("closing %s: %w", *out, cerr) - } - if rerr := rec.Err(); rerr != nil && retErr == nil { - retErr = fmt.Errorf("writing events to %s: %w", *out, rerr) - } - }() - - ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) - defer stop() - - sum, err := swebench.Run(ctx, swebench.Config{ - RunID: *runID, - Backend: backend, - Instances: selected, - Trajectory: traj, - Recorder: rec, - Logger: log, - Mode: swebench.Mode(*mode), - Width: *width, - TasksPerImage: *tasksPerImage, - Rollouts: *rollouts, - Steps: *steps, - Park: swebench.Park(*park), - CreateTimeout: *createTimeout, - StepTimeout: *stepTimeout, - WakeBudget: *wakeBudget, - WakeBackoff: *wakeBackoff, - Keep: *keep, - }) - if sum != nil { - log.Info("run finished", slog.Int("tasks", sum.Tasks), slog.Int("failed", sum.Failed), - slog.Int("rewarded", sum.Rewarded), slog.Duration("duration", sum.Duration)) - } - if err != nil { - return err - } - if sum.Failed > 0 { - return fmt.Errorf("%d of %d tasks failed", sum.Failed, sum.Tasks) - } - return nil -} - -func reportCmd(args []string) error { - if len(args) < 1 { - return fmt.Errorf("report needs the JSONL file") - } - path := args[0] - fs := flag.NewFlagSet("report", flag.ContinueOnError) - jsonOut := fs.String("json", "", "Also write the summary as JSON to this path.") - unresolvedOut := fs.String("unresolved", "", "Write the unresolved instances to this path in -exclude-file format.") - unresolvedNote := fs.String("unresolved-note", "Instances graded unresolved", "First comment line of the -unresolved file.") - if err := fs.Parse(args[1:]); err != nil { - return err - } - f, err := os.Open(path) - if err != nil { - return err - } - defer f.Close() - rep, err := swebench.Summarize(f) - if err != nil { - return err - } - rep.WriteText(os.Stdout) - if *unresolvedOut != "" { - uf, err := os.Create(*unresolvedOut) - if err != nil { - return err - } - rep.WriteExclusions(uf, *unresolvedNote) - if err := uf.Close(); err != nil { - return err - } - } - if *jsonOut != "" { - b, err := json.MarshalIndent(rep, "", " ") - if err != nil { - return err - } - return os.WriteFile(*jsonOut, b, 0o644) - } - return nil -} diff --git a/docs/rootfs-images/DESIGN.md b/docs/rootfs-images/DESIGN.md deleted file mode 100644 index fd9dfcf6..00000000 --- a/docs/rootfs-images/DESIGN.md +++ /dev/null @@ -1,224 +0,0 @@ -# Design: injecting the guest into unmodified rootfs images - -Status: implemented. Companion to [README.md](README.md), which is the how-to. - -## Problem - -Every ate-env operation goes through `ate-env-guest`, a daemon inside the -actor. Until now the only way to run an image as an environment was to make the -guest the container image, which for a rootfs image meant rebuilding it with the -guest layered on top and set as the entrypoint. For one image that is a `ko` -invocation. For an evaluation or RL fleet it is a build and a push per image, -duplicated guest bytes in the registry, a rebuild of every image for every guest -change, and a registry pipeline the harness has to own. - -The two things a container sources, its rootfs and its process, can be supplied -independently. Substrate's ActorTemplate already allows a second OCI image to be -mounted read-only into the actor. This design uses that to inject the guest. - -## Goals - -- Run any digest-pinned image as an environment with no rebuild. -- Keep the guest behavior, API and client code identical in both modes. -- Make the on-demand path idempotent and safe under concurrency, so harnesses - can call it per task without coordinating. -- Keep one place that knows the mount layout, so a second injected runtime can - reuse it. - -Non-goals: validating that an image exists or starts (the actor reports that), -per-create resource overrides, template garbage collection, and multi-container -templates. These are listed under future work. - -## Design - -### Template derivation - -`internal/apiservice/imagetemplate.go` holds `DeriveImageTemplate(base, name, -image)`. Given a base template it returns a new template that: - -- keeps the base's worker selector, snapshot configuration, sandbox - configuration, resources, container env and readiness probe; -- replaces the first container's image with the rootfs image; -- if the base runs the guest as its image, adds a read-only image volume named - `guest` whose reference is the base's image, mounts it at `/ate`, and moves - the command's first element under `/ate` (`/ko-app/ate-env-guest` becomes - `/ate/ko-app/ate-env-guest`); -- if the base already has an image volume, leaves volumes, mounts and command - alone; -- drops server-assigned metadata and status. - -It refuses an unpinned task or guest image, a base without containers, and a -base command that is not an absolute path (a shell wrapper cannot be re-rooted -and would produce a template whose guest is silently missing). - -Both entry points use it: `ate-env manifest template --rootfs-image` derives from -the guest-image template it would otherwise print, and `ate-env-api` derives -from a template fetched from the control plane. - -### Naming and idempotency - -A derived template is named `-`, in the base's atespace. The name is a pure function of the inputs a -caller controls, so: - -- the same image always maps to the same template, and a lookup by name is the - reuse check; -- a human can find the template for an image from its digest; -- 12 hex digits (48 bits) make accidental collisions a non-concern at any - realistic image count, while leaving room under the 63-character resource - name limit for a base name of up to 50 characters. Longer bases are rejected - with a message naming the limit. - -The guest is deliberately not part of the name. Templates are immutable in -Substrate, so a guest change is a new base template with a new name, which -yields new derived names. Encoding the guest digest would only lengthen names -without adding information. - -### Server flow - -`CreateEnvironment` with `rootfs_image` set: - -1. Validate the digest and compute the derived name; invalid input is - `InvalidArgument` and touches nothing. -2. `GetActorTemplate(derived)`. If it exists and its container image is the - requested image, use it. If it exists with a different image, return - `FailedPrecondition` naming both images: someone else owns that name. -3. Otherwise `GetActorTemplate(base)`; a missing base is `FailedPrecondition` - with the atespace and the command that registers one. -4. Derive and `CreateActorTemplate`. `AlreadyExists` is success: a concurrent - create won the race and built the same template. -5. Create the actor from the derived template and report it in the response. - -The check in step 2 compares the container image only. Volumes and command are -not compared because the only way for them to differ is a template created by -other means under the derived name, which the image comparison already flags -in practice; a stricter comparison would cost a base fetch on every reuse. - -### API shape - -`rootfs_image` is a field on `CreateEnvironmentRequest` rather than a new RPC, and the -existing `template` field doubles as the base. This keeps one create path, lets -clients that already pass a template start passing an image with one more -argument, and makes "which base" an explicit, per-request choice rather than -server configuration. The derived template is surfaced through the existing -`Environment.template`, so nothing new is needed to observe it. - -### Placement - -Derived templates live in the environment's atespace, next to the base. This is -also where `ate-env-api` already creates actors, so no new permissions are -involved beyond template creation, which the API server's identity needs for -this feature. - -### Runtime layers - -A second runtime is expressed with the same primitives rather than a new -template concept: an extra image volume for its binary, and two guest flags. -`-sidecar` makes the guest start and supervise a process (restart with -backoff, output into the guest log), and `-sidecar-readyz` makes the guest's -`/readyz`, which Substrate uses as the wakeup probe, wait for the runtime's own -health endpoint once. The guest is the right supervisor because it is already -the container's only process and the readiness authority. - -Derivation treats layers as part of what carries over. The guest volume is -recognized by its name, not by being an image volume, so a base with a layer -still gets the guest added or kept correctly, and `create --rootfs-image` on a layered -base yields environments with the second runtime present. Layers are -documented in [RUNTIMES.md](RUNTIMES.md); the verified example uses busybox as -a stand-in runtime. - -What this does not solve is reaching a second runtime from outside the actor. -The router forwards to port 80 unless a CONNECT authority names another port, -and `ate-env-api` proxies only the guest. Exposing a runtime's port is the next -piece, and it belongs with the OpenSandbox backend that needs it. - -## Failure semantics - -| Situation | Result | -|---|---| -| unpinned rootfs image | `InvalidArgument`, nothing created | -| base missing | `FailedPrecondition`, nothing created | -| base command not re-rootable | `InvalidArgument`, nothing created | -| derived name held by another image | `FailedPrecondition`, nothing created | -| two creates race on a new image | both succeed, one template | -| image cannot be pulled or guest cannot start | create succeeds; actor never becomes ready; golden bake fails and the template reports it | -| control plane unavailable mid-flow | the underlying gRPC status is returned; a half-created template is reused next time because the name is deterministic | - -## Security - -- The rootfs image is run with the same sandbox class and configuration as the - base, so injection changes what runs, not how it is isolated. -- The guest volume is read-only. The rootfs image cannot alter the guest. -- Digest pinning is enforced for both images at every entry point. A tag could - be re-pointed between the golden bake and a later cold boot, which would - invalidate snapshots and make two environments on "the same image" differ. -- Callers of `CreateEnvironment` can now cause template creation. An - `ate-env-api` that is reachable by untrusted callers should already be behind - authentication; this adds template records and golden snapshots to what such - a caller can accumulate. - -## Performance - -The first environment on an image pays the rootfs image pull on the worker it -lands on and a cold boot; the derived template's golden bake runs concurrently -on another worker. Later environments restore from the golden. In the -verification run below the first environment answered a shell command 37 s -after the create call on a small two-worker pool, most of it image pull. - -Costs that scale with the number of distinct images: one template record and -one golden snapshot per image. Neither is reclaimed automatically. - -## Alternatives considered - -- **Bake the guest into each image.** Works today and remains supported; it is - the per-image build cost this design removes. -- **An init container that copies the guest into a shared volume.** Needs a - writable volume and an extra container per actor, and the copy happens on - every cold boot. An image volume is shared by the node's cache and needs no - copy. -- **A server-side image-to-template map.** Equivalent to the derived name, but - state that can drift from the control plane. The name derivation is stateless - and recoverable from the templates themselves. -- **Including the guest digest in the derived name.** Rejected because template - immutability already forces a new base name for a new guest; see Naming. -- **A separate `CreateEnvironmentFromImage` RPC.** More surface for the same - behavior; a field keeps one path. - -## Future work - -- **Derived-template garbage collection** is the first follow-up. A derived - template is shared by every environment on its image and carries the - golden snapshot, so deleting it when one environment is deleted would send - the next create on that image back through the cold path; the count grows - with distinct images, not with environments. The plan is a last-used - annotation updated on each create, collection by age or by last use, and - the golden snapshot removed alongside. -- **gRPC and raw TCP to a second runtime.** `ate-env-api` forwards HTTP to a - runtime's port through the router's CONNECT listener (see - [RUNTIMES.md](RUNTIMES.md#who-can-reach-the-second-runtime)); a runtime that - speaks gRPC or a non-HTTP protocol needs an h2c upstream or a raw tunnel. -- **Per-create overrides** for resources and workspace, so one base can serve - images with different needs. -- **Image validation at create time**, so a bad reference fails the call - instead of the actor. -- **Template atespace.** `CreateEnvironment` takes `template.atespace` but the - server creates the actor and resolves the template in the environment's - atespace; this predates the change and is unchanged by it. - -## Verification - -- Unit tests cover derivation from both base kinds, naming and its limits, - every rejection, and that runtime layers and sidecar flags survive - derivation. The guest's sidecar supervisor and readiness gate have their own - tests (restart with backoff, stop on shutdown, gate once per URL). -- Server tests against the in-process fake control plane cover first create, - reuse, a named base in another atespace, and the failure table. -- `clients/python/tests/e2e/test_full_stack.py` has a create-from-image test - gated on `ATE_ENV_ROOTFS_IMAGE`, which checks that the guest comes from the - volume and not the image, that the rootfs is the rootfs image's, that files - round-trip, and that a second environment reuses the derived template. -- Run once end to end on a Kubernetes cluster running Substrate with a fresh - two-worker ate-env: `ate-env create --rootfs-image` of an unmodified - `python:3.12-slim` produced the derived template, served `python3 --version` - from the image's own rootfs with the guest present under `/ate`, and deleted - cleanly. diff --git a/docs/rootfs-images/README.md b/docs/rootfs-images/README.md deleted file mode 100644 index 2843f20c..00000000 --- a/docs/rootfs-images/README.md +++ /dev/null @@ -1,144 +0,0 @@ -# Running rootfs images on ate-env - -Any OCI image can run as an ate-env environment without being rebuilt. The -`ate-env-guest` daemon that serves exec and file operations is mounted into the -image's container as a read-only image volume and started from there, so the -image's filesystem, interpreters and tools are exactly what was published. - -This page is the how-to. The reasoning behind it is in [DESIGN.md](DESIGN.md). - -## How it works - -An ActorTemplate normally names one container image, and for ate-env that image -has been the guest itself. A *rootfs-image template* splits the two things a -container sources: - -``` -rootfs = the image you name, unmodified, pinned by digest -process = /ate/ko-app/ate-env-guest, from a second image mounted read-only at /ate -``` - -Substrate materializes both at actor start. Everything else about the template -(worker selector, snapshot settings, sandbox class, resources, readiness probe) -is the same as for a guest-image template, and the guest behaves identically: -the environment's `shell`, process and file APIs, MCP, and suspend and resume -all work unchanged. - -## Prerequisites - -- A Substrate control plane whose ActorTemplate API supports image volumes - (`volumes[].image`). Every supported release does. -- Images pinned by digest (`repo@sha256:...`), both the rootfs image and the - guest. Substrate requires this because a changed image invalidates snapshots. -- Worker nodes that can pull the rootfs image's registry. -- A base template registered with `ate-env manifest template` (the quickstart - in the root README does this). Its guest image is the one injected. - -What the rootfs image needs: - -- A `sh` for `Env.shell()` and for the `shell` MCP tool. Images without one - (distroless) still work for `start_process` with absolute binaries and for - file transfer. -- Port 80 free inside the sandbox: the guest listens there and the readiness - probe hits `/readyz` on it. -- Nothing else. The guest is a static binary and runs as the image's default - user; a writable workspace is wherever `--workspace` points (the guest's - default is `/`). - -## Option A: register a rootfs-image template up front - -Good for a fixed set of images, or when you want the template to carry its own -name, resources or workspace. - -```bash -ate-env manifest template --template py312 \ - --rootfs-image docker.io/library/python@sha256: \ - --guest-image /ate-env-guest@sha256: \ - --workspace /workspace \ - --snapshots-bucket | kubectl-ate create actor-template -f - - -ate-env create job1 --template py312 -``` - -`--rootfs-image` makes the rootfs image the container and the guest an image -volume; `--workspace` is passed to the guest as `-workspace` in either mode. The -printed template is plain YAML you can edit before registering it (resources, -worker selector, sandbox class). - -## Option B: create from an image on demand - -Good when the set of images is open-ended or chosen by a caller, as in RL and -evaluation harnesses. - -```bash -ate-env create py1 --rootfs-image docker.io/library/python@sha256: -``` - -```python -env = await client.create("py1", rootfs_image="docker.io/library/python@sha256:") -(await env.info()).template.name # "default-template-<12 hex of the digest>" -``` - -```go -client.Create(ctx, &ateenvv1alpha.CreateEnvironmentRequest{ - Id: "py1", Image: "docker.io/library/python@sha256:", -}) -``` - -What happens on the server: - -1. The request's template (default `default-template`) is the **base**. It must - exist in the environment's atespace. -2. `ate-env-api` looks for a template named `-` next to the base. If it exists and runs that image, it is reused. -3. Otherwise the server derives it from the base and creates it: same worker - selector, snapshot and sandbox settings, env and readiness; the rootfs image as - the container; the base's guest image as a volume at `/ate`; the command - re-rooted to `/ate/ko-app/ate-env-guest`. -4. The environment is created from the derived template, which the response and - `GetEnvironment` report. - -The first environment on an image pays the image pull and a cold boot while the -template's golden snapshot bakes in the background; later ones restore from the -golden. Concurrent first creates on the same image are safe: whichever wins -creates the template, the others reuse it. - -Templates created this way are visible like any other: - -```bash -kubectl-ate get actor-template --atespace ate-env -``` - -A fresh template shows `Failed` in that listing for a few seconds while its -golden bakes, then `Ready`. Environments can be created in the meantime. - -## A second runtime - -The same mechanism carries more runtimes. `--layer` mounts another image, -`--sidecar` has the guest start a process from it, and `--sidecar-readyz` -folds its readiness into the actor's. See [RUNTIMES.md](RUNTIMES.md) and the -runnable example in [`examples/runtime-layers`](../../examples/runtime-layers). - -## Errors - -| Error | Meaning | -|---|---| -| `InvalidArgument: image ... is not pinned by digest` | Pass `repo@sha256:<64 hex>`; tags are rejected before anything is created | -| `FailedPrecondition: base actor template ... not found` | Register the base in the environment's atespace first (`ate-env manifest template`) | -| `FailedPrecondition: actor template ... exists but runs ...` | The derived name is taken by a template with a different image; delete it or use another base name | -| `InvalidArgument: ... cannot re-root command ...` | The base's command is not an absolute path to the guest (a shell wrapper); fix the base template | -| Environment never serves, derived template stays `Failed` | The rootfs image could not be pulled or does not start the guest (port 80 taken, no exec permission); check the worker pod logs | - -## Cleanup and limits - -- Derived templates are not deleted when environments are. They are small - control-plane records plus one golden snapshot in object storage each; remove - them with `kubectl-ate delete actor-template` when an image is retired. -- Bumping the guest means a new base template (templates are immutable), and a - new base name yields new derived names. Existing derived templates keep the - guest they were built with. -- Resources, sandbox class and worker selector come from the base; a rootfs image - that needs more is served by registering its own template (Option A) or a - second base. -- The rootfs image's existence is not checked at create time. A bad reference is - reported by the actor failing to start, not by `CreateEnvironment`. diff --git a/docs/rootfs-images/RUNTIMES.md b/docs/rootfs-images/RUNTIMES.md deleted file mode 100644 index cd50f31d..00000000 --- a/docs/rootfs-images/RUNTIMES.md +++ /dev/null @@ -1,149 +0,0 @@ -# Injecting a second runtime as a layer - -[README.md](README.md) explains how the ate-env guest is injected into an -unmodified rootfs image. The same mechanism carries any number of further -runtimes: a second daemon that an agent framework expects inside its sandbox, -such as OpenSandbox's `execd`, is one more read-only image volume plus a process -the guest starts next to itself. This page explains the model, the flags, and -what reaches the second runtime from where. - -## Model - -A runtime layer is three things: - -| part | where it is expressed | what it does | -|---|---|---| -| the binary | an image volume on the ActorTemplate (`--layer name=image@sha256:...=/mount`) | mounts the runtime's image read-only at a path of your choice, next to the guest at `/ate` | -| the process | a guest flag (`--sidecar '/mount/binary args...'`) | the guest starts it after it is listening, logs its output with a prefix, and restarts it with backoff if it exits | -| its readiness | a guest flag (`--sidecar-readyz http://127.0.0.1:PORT/health`) | the guest's `/readyz`, which is the actor's wakeup probe, answers 503 until every such URL has answered 2xx once | - -The guest is the container's only process by contract, which is why it is the -one that supervises the others. Sidecars run without a shell or `PATH` lookup, -so the binary is an absolute path into its mount, and the rootfs image needs -nothing for it. They inherit the container's user, environment and filesystem. - -Template derivation keeps all of it. A base with layers yields derived -templates with the same layers and the same guest flags, so environments -created on demand with `rootfs_image=` inherit the second runtime without any caller -knowing it exists. - -## Registering a layered template - -```bash -ate-env manifest template --template osb-base \ - --rootfs-image docker.io/library/python@sha256: \ - --guest-image /ate-env-guest@sha256: \ - --layer execd=/opensandbox/execd@sha256:=/opt/opensandbox \ - --sidecar '/opt/opensandbox/execd' \ - --sidecar-readyz http://127.0.0.1:44772/ready \ - --snapshots-bucket | kubectl-ate create actor-template -f - -``` - -The resulting template has two image volumes (`execd` at `/opt/opensandbox`, -`guest` at `/ate`), the rootfs image as the container, and the command - -``` -/ate/ko-app/ate-env-guest -sidecar "/opt/opensandbox/execd" -sidecar-readyz http://127.0.0.1:44772/ready -``` - -Everything else is as in a plain rootfs-image template. Environments are created -the usual way, from this template by name or from any other image with `--rootfs-image` -while naming it as the base: - -```bash -ate-env create py1 --template osb-base -ate-env create node1 --template osb-base --rootfs-image docker.io/library/node@sha256: -``` - -Both run the guest and `execd`; the second one from a derived template named -`osb-base-<12 hex of the node digest>`. - -### OpenSandbox `execd` specifics - -`execd` is a static binary that OpenSandbox injects into sandboxes without -modifying the base image, which is exactly the shape this page describes. Its -image carries the binary at `/execd`; mounting the image at `/opt/opensandbox` -puts it at `/opt/opensandbox/execd`, OpenSandbox's own default install path. It -serves HTTP on port 44772 by default, with `/ping` and `/ready` open and every -other endpoint behind the `X-EXECD-ACCESS-TOKEN` header when a token is -configured. Check the OpenSandbox documentation for the current flags and -environment variables for the port, token and workspace; pass them in the -`--sidecar` value, or wrap them in a script shipped in the layer if quoting is -needed. Its optional Jupyter code-interpreter mode starts a Jupyter server -inside the sandbox; for fleets of small environments use the plain command and -file endpoints. - -A verified stand-in for the mechanics, which needs nothing from any registry -but Docker Hub, is in [`examples/runtime-layers`](../../examples/runtime-layers): -the static `busybox:1.37-musl` mounted as the layer and its `httpd` as the -sidecar, run under two different rootfs images. - -## Who can reach the second runtime - -- **Processes inside the environment** reach it at `127.0.0.1:` right - away. An agent framework whose tool calls run inside the sandbox is served. -- **The guest's own APIs** are unaffected: exec, files and MCP keep going - through `ate-env-api` on port 80. -- **Clients outside the environment** reach it through `ate-env-api`. A plain - HTTP request to the API's listener that carries `X-Env-Id`, `X-Env-Port` and - optionally `X-Env-Atespace` goes to that port in that environment, whatever - its path: `ate-env-api` opens a CONNECT tunnel through the atenet router's - CONNECT listener (`--atenet-connect`, port 8081 on a stock install), the - router resumes the actor if needed and forwards to the port, and the request - and its streamed response pass through unchanged apart from the routing - headers, which are stripped. Idle tunnels are reused, one pool for each - environment and port. - - ```bash - curl -H 'X-Env-Id: py1' -H 'X-Env-Port: 44772' localhost:7777/ready - curl localhost:7777/v1alpha/envs/py1/ports/44772/ready # the same, for clients without headers - ``` - - The path form, `/v1alpha/envs/{id}/ports/{port}/...` (or - `/v1alpha/atespaces/{atespace}/envs/...`), is for clients that cannot set - headers: the prefix is stripped, so it works as a base URL. The tunnel - carries HTTP/1.1, not raw TCP. Streamed responses are flushed as they arrive - and WebSocket upgrades pass through; gRPC to the runtime is not supported - yet. The NeMo Gym - provider hands this out as its `endpoint(handle, port)` - ([integrations/nemo-gym](../../integrations/nemo-gym/README.md)). -- **The runtime must listen on `0.0.0.0`.** The router reaches the port on - the actor's address. A runtime bound to `127.0.0.1` serves only processes - inside the environment, and requests from outside fail with 502. -- **Nobody is authenticated on the way in.** `ate-env-api` forwards for any - caller that reaches its listener, to any port of any actor the router - serves. A runtime that executes commands or edits files must check its own - credential, as OpenSandbox's `execd` token and the OpenHands agent-server's - session key do. See the README section "Reaching a service inside an - environment" for the browser checks and the flags. - -## Suspend, resume and restarts - -A sidecar is part of the actor, so a suspend checkpoints its memory with the -guest's and a resume brings it back without a restart, open sockets included. If -a runtime does not survive that (it polls a clock it can no longer trust, say) -and exits, the guest restarts it: after one second, doubling up to thirty while -it keeps failing, and from one second again once it has run for a minute. -Readiness is gated once, at start; a sidecar that dies later is restarted but -does not take the actor out of service. - -## Limits - -- **A layer must be self-contained.** The rootfs image's libraries are whatever - that image ships, so a runtime binary that is dynamically linked works on - some images and not on others. Ship static binaries (the guest and `execd` - are) or put the runtime's libraries in the layer and point it at them. The - failure is easy to recognize: the guest log shows the sidecar exiting at - once with a loader error such as ``version `GLIBC_2.38' not found`` and - restarting with growing backoff, readiness never passes, and the actor - fails after the wakeup probe timeout. -- One container per actor: sidecars share the rootfs image's user, environment - variables and filesystem, and have no resource limits of their own. -- `--sidecar` values are split on whitespace with no quoting. Arguments that - need quotes or shell features go in a wrapper script inside the layer. -- Port 80 belongs to the guest. A runtime that insists on it needs a wrapper - that moves it elsewhere. -- Sidecar output goes to the guest's log, line by line, prefixed with the - binary's name. There is no separate log stream yet. -- Layers are read-only. A runtime that writes next to its binary needs a - writable working directory passed as an argument. diff --git a/examples/runtime-layers/README.md b/examples/runtime-layers/README.md deleted file mode 100644 index 79bba673..00000000 --- a/examples/runtime-layers/README.md +++ /dev/null @@ -1,60 +0,0 @@ -# Example: a second runtime as a layer - -Runs an unmodified rootfs image with two injected runtimes: the ate-env guest at -`/ate`, and a static busybox at `/opt/web` standing in for a real runtime -daemon such as OpenSandbox's `execd`. The guest starts `busybox httpd` on port 44772 as a -sidecar and holds the actor's readiness until it answers. The mechanics are the -ones [docs/rootfs-images/RUNTIMES.md](../../docs/rootfs-images/RUNTIMES.md) -describes; only the binary differs. - -`demo.sh` registers the template, creates an environment, fetches a file from -the second runtime from inside the environment, optionally creates another -environment from a different image on the same base to show the layer is -inherited, and deletes what it created. - -```bash -kubectl -n ate-env port-forward svc/ate-env-api 7777:7777 & - -GUEST_IMAGE=/ate-env-guest@sha256: \ -SNAPSHOTS_BUCKET= \ -ROOTFS_IMAGE=docker.io/library/python@sha256: \ -LAYER_IMAGE=docker.io/library/busybox@sha256: \ -OTHER_IMAGE=docker.io/library/node@sha256: \ -examples/runtime-layers/demo.sh -``` - -Digests for public images come from `crane digest --platform linux/amd64 -python:3.12-slim` or `docker manifest inspect`. The guest image must be built -from a revision of this repo that has the `-sidecar` flags. - -Use the `busybox:1.37-musl` tag, which is statically linked. The default -`busybox:1.37` is built against glibc and only runs on rootfs images whose glibc -is at least as new as its own: it works inside `python:3.12-slim` and fails -inside `debian:bookworm-slim` with ``GLIBC_2.38' not found``, which is the -textbook case of why a runtime layer has to be self-contained. - -What the template looks like, as `ate-env manifest template` prints it: - -```yaml -containers: -- command: - - /ate/ko-app/ate-env-guest - - -sidecar - - /opt/web/bin/busybox httpd -f -p 44772 -h / - - -sidecar-readyz - - http://127.0.0.1:44772/etc/hostname - image: docker.io/library/python@sha256: - name: guest - readyz: {httpGet: {path: /readyz, port: 80}} - volumeMounts: - - {name: web, mountPath: /opt/web} - - {name: guest, mountPath: /ate} -volumes: -- {name: web, image: {reference: docker.io/library/busybox@sha256:}} -- {name: guest, image: {reference: /ate-env-guest@sha256:}} -``` - -To swap in a real runtime, change the three values: the layer image, the -sidecar command, and the readiness URL. For `execd` that is its image mounted -at `/opt/opensandbox`, `/opt/opensandbox/execd` with its flags, and -`http://127.0.0.1:44772/ready`. diff --git a/examples/runtime-layers/demo.sh b/examples/runtime-layers/demo.sh deleted file mode 100755 index a49632d5..00000000 --- a/examples/runtime-layers/demo.sh +++ /dev/null @@ -1,107 +0,0 @@ -#!/usr/bin/env bash -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# A second runtime injected as a layer, end to end. -# -# busybox stands in for a real runtime daemon (OpenSandbox execd, say): its -# image is mounted read-only at /opt/web, the guest starts `busybox httpd` on -# port 44772 as a sidecar and waits for it before reporting the actor ready. -# The rootfs image is unmodified. The script registers the template, creates an -# environment from it, asks the second runtime for a file from inside the -# environment, creates a second environment from a different image on the -# same base to show the layer is inherited, and deletes both. -# -# Required: -# GUEST_IMAGE digest-pinned ate-env-guest image (built from this repo, -# at or after the -sidecar flags) -# SNAPSHOTS_BUCKET object-storage URL for actor snapshots -# ROOTFS_IMAGE digest-pinned rootfs image, e.g. docker.io/library/python@sha256:... -# LAYER_IMAGE digest-pinned STATIC busybox image (the 1.37-musl tag), -# e.g. docker.io/library/busybox@sha256:... -# Optional: -# OTHER_IMAGE a second digest-pinned image for the inheritance check -# TEMPLATE base template name (default layered-demo) -# ATESPACE atespace (default ate-env) -# SUBSTRATE_ENV_API ate-env-api address (default 127.0.0.1:7777; port-forward it) -# ATE_ENV, KUBECTL_ATE binaries (default: on PATH) -set -euo pipefail - -: "${GUEST_IMAGE:?set GUEST_IMAGE to the digest-pinned ate-env-guest image}" -: "${SNAPSHOTS_BUCKET:?set SNAPSHOTS_BUCKET to the snapshot object-storage URL}" -: "${ROOTFS_IMAGE:?set ROOTFS_IMAGE to a digest-pinned rootfs image}" -: "${LAYER_IMAGE:?set LAYER_IMAGE to a digest-pinned busybox image}" -TEMPLATE=${TEMPLATE:-layered-demo} -ATESPACE=${ATESPACE:-ate-env} -export SUBSTRATE_ENV_API=${SUBSTRATE_ENV_API:-127.0.0.1:7777} -ATE_ENV=${ATE_ENV:-ate-env} -KUBECTL_ATE=${KUBECTL_ATE:-kubectl-ate} - -PORT=44772 -SIDECAR="/opt/web/bin/busybox httpd -f -p ${PORT} -h /" -READYZ="http://127.0.0.1:${PORT}/etc/hostname" - -step() { printf '\n== %s\n' "$*"; } - -step "register ${ATESPACE}/${TEMPLATE}: ${ROOTFS_IMAGE} + guest + busybox layer" -"${ATE_ENV}" manifest template --template "${TEMPLATE}" --atespace "${ATESPACE}" \ - --rootfs-image "${ROOTFS_IMAGE}" --guest-image "${GUEST_IMAGE}" \ - --snapshots-bucket "${SNAPSHOTS_BUCKET}" \ - --layer "web=${LAYER_IMAGE}=/opt/web" \ - --sidecar "${SIDECAR}" \ - --sidecar-readyz "${READYZ}" | "${KUBECTL_ATE}" create actor-template -f - >/dev/null -"${KUBECTL_ATE}" get actor-template --atespace "${ATESPACE}" "${TEMPLATE}" - -ids=() -cleanup() { - for id in "${ids[@]:-}"; do - [ -n "${id}" ] && "${ATE_ENV}" delete "${id}" --atespace "${ATESPACE}" >/dev/null 2>&1 || true - done -} -trap cleanup EXIT - -# wait_serving : the guest answers only once the sidecar's readiness URL -# has answered, so the first successful shell means both runtimes are up. -wait_serving() { - local id=$1 i - for i in $(seq 1 60); do - if "${ATE_ENV}" "${id}" shell 'true' >/dev/null 2>&1; then return 0; fi - sleep 3 - done - echo "environment ${id} did not start serving" >&2 - return 1 -} - -ID="layered-$(date +%s)" -ids+=("${ID}") -step "create ${ID} from ${TEMPLATE}" -t0=$(date +%s) -"${ATE_ENV}" create "${ID}" --atespace "${ATESPACE}" --template "${TEMPLATE}" -wait_serving "${ID}" -echo "serving after $(( $(date +%s) - t0 ))s" - -step "ask the second runtime from inside the environment" -"${ATE_ENV}" "${ID}" shell "/opt/web/bin/busybox wget -qO- http://127.0.0.1:${PORT}/etc/os-release | head -2; echo guest: \$(ls /ate/ko-app); echo layer: \$(ls /opt/web/bin | head -1)" - -if [ -n "${OTHER_IMAGE:-}" ]; then - ID2="layered2-$(date +%s)" - ids+=("${ID2}") - step "create ${ID2} from ${OTHER_IMAGE} with ${TEMPLATE} as the base (layer inherited)" - "${ATE_ENV}" create "${ID2}" --atespace "${ATESPACE}" --template "${TEMPLATE}" --rootfs-image "${OTHER_IMAGE}" - "${KUBECTL_ATE}" get actor-template --atespace "${ATESPACE}" | grep "${TEMPLATE}-" || true - wait_serving "${ID2}" - "${ATE_ENV}" "${ID2}" shell "/opt/web/bin/busybox wget -qO- http://127.0.0.1:${PORT}/etc/os-release | head -1" -fi - -step "done; deleting" diff --git a/integrations/nemo-gym/.gitignore b/integrations/nemo-gym/.gitignore deleted file mode 100644 index e05d4abe..00000000 --- a/integrations/nemo-gym/.gitignore +++ /dev/null @@ -1,6 +0,0 @@ -*.egg-info/ -__pycache__/ -.pytest_cache/ -.venv/ -build/ -dist/ diff --git a/integrations/nemo-gym/README.md b/integrations/nemo-gym/README.md deleted file mode 100644 index 77cb0eaa..00000000 --- a/integrations/nemo-gym/README.md +++ /dev/null @@ -1,368 +0,0 @@ -# NeMo Gym sandbox provider for Agent Substrate - -Runs [NeMo Gym](https://github.com/NVIDIA-NeMo/Gym) rollout sandboxes as **Substrate actors**, -fronted by the `ate-env` API of this repo and reached through its Python client, -[`ate_env`](../../clients/python). -Registers as provider `substrate` next to the built-ins (Docker, Daytona, ECS Fargate, Enroot, -OpenShell, OpenSandbox, Apptainer) via the `nemo_gym.sandbox_providers` entry point — **no changes -to NeMo Gym required**, and no changes to your resource servers beyond the `sandbox:` config block. - -Why run rollouts on Substrate: - -- **Idle sandboxes hold no worker.** A sandbox that is waiting on model inference is suspended by - Substrate; the atenet router resumes it transparently on the next command. A rollout fleet is - mostly waiting, so effective density is bounded by the *active* set, not the fleet size. -- **Repeat sandboxes start from golden snapshots.** The first sandbox on an image pays the boot; - later ones restore a memory snapshot instead of booting. -- **Prepared states fork.** A sandbox that finished expensive task setup can be snapshotted and - branched into N independent rollouts (best-of-N from a mid-episode state). - -Verified against **nemo-gym 0.5.0**: the registry resolves this provider through the entry point -(`create_provider({"substrate": {...}})`) and the full unit suite passes against the real -`nemo_gym.sandbox.providers.base` types. - -## How it maps onto ate-env - -One gRPC endpoint, `ate-env-api`, carries everything (see the -[client README](../../clients/python/README.md#how-it-works)): - -| Provider call | `ate_env` call | Where it lands | -|---|---|---| -| `create(spec)` | `Client.create(id, template_name=…, template_atespace=…, rootfs_image=…)`, then `Env.shell("true")` until it answers, then `Env.write_file` for `spec.files` | `EnvironmentService` → Substrate control plane; the probe and the files go to the guest | -| `exec(handle, cmd, cwd=, env=, timeout_s=)` | `Env.start_process(["sh","-c",cmd], cwd, env, timeout=)` + `Process.output(follow=True)` | `ProcessService` in `ate-env-guest`, proxied through the atenet router | -| `upload_file` / `download_file` | `Env.write_file` / `Env.read_file_bytes` | `FileSystemService` in the guest | -| `status(handle)` | `Client.get(id)` → lifecycle status | `EnvironmentService`; does **not** wake a parked sandbox | -| `close(handle)` | `Client.delete(id)`, NOT_FOUND-safe | `EnvironmentService` | - -Environment ids are `gym-<10 hex>`; each handle remembers the atespace it was created in. - -## Prerequisites - -On the cluster: - -1. **Agent Substrate installed** (`ate-system` namespace healthy — api-server, controller, - atelet, atenet-router, valkey). See `hack/install-ate.sh` in the - [substrate repo](https://github.com/agent-substrate/substrate). The guest data - plane is gRPC through the router, which needs a substrate at or after - [substrate#1183](https://github.com/agent-substrate/substrate/pull/1183) (h2 on the ingress, - protocol mirrored to actors). Without it every exec and file call fails with - `server closed the stream without sending trailers`. -2. **The `ate-env` system deployed** from this repo at or after the `ateenv.v1alpha` proto - move: namespace, WorkerPool, ActorTemplate(s), and the `ate-env-api` service (see the - [root README](../../README.md#quickstart)). - - ```bash - ate-env manifest \ - --api-image \ - --worker-image | kubectl apply -f - - kubectl get pods -n ate-env # api + warm workers Running - # once ate-env-api is Running (it creates the ate-env atespace templates live in): - ate-env manifest template \ - --guest-image \ - --snapshots-bucket gs:///ate-env/ | kubectl-ate create actor-template -f - - ``` - - If your `ateapi` requires authentication (it does on any standard install), the - `ate-env-api` Deployment needs a projected ServiceAccount token with audience - `api.ate-system.svc` and the `-ateapi-token-file` flag pointing at it. -3. **A base ActorTemplate** in the `ate-env` atespace: `default-template`, which the root - README's quickstart registers. That is all a task image needs: `SandboxSpec.image` goes to - `ate-env-api` as `rootfs_image`, and the API derives a template from the base on first use - (the guest mounted into the image as a read-only volume) and reuses it afterwards. Register - extra templates only for images that need their own resources or workspace, and map them - under `create.image_templates`. - - > **Every template must run `ate-env-guest` inside the actor.** Readiness, exec, file - > transfer, even `create()` completing — all of it goes through the guest's gRPC services - > behind the router. Derived templates get the guest automatically; a hand-written template - > without it fails *completely*: the actor starts, but every guest call fails and `create()` - > times out with `SandboxCreateVerificationError`. Images must be digest-pinned - > (`repo@sha256:…`); the API rejects tags. -4. **Network path from wherever Gym's rollout workers run** to `ate-env-api`: in-cluster DNS - (`ate-env-api.ate-env:7777`) or, for a workstation, a port-forward: - - ```bash - kubectl port-forward -n ate-env svc/ate-env-api 7777:7777 - ``` - -Locally: - -5. **Python ≥ 3.10** and the `ate-env-client` package from this repo's `clients/python` (not on - PyPI yet). `nemo-gym` itself is only needed on the machine running the Gym resource servers — - this package's tests run without it (see Testing). - -## Task images - -Three ways to run a task image, from least to most work: - -**Pass the image (default).** Set `SandboxSpec.image` to a digest-pinned reference. `create()` -sends it as `rootfs_image` on the base template; the API derives `-<12 hex of the digest>` -once, bakes its golden snapshot, and later creates on the same image restore from it. Nothing to -register, nothing to rebuild. See [`docs/rootfs-images/README.md`](../../docs/rootfs-images/README.md). - -**Register a template per image** when an image needs its own resources, workspace or a second -runtime layer, then map it under `create.image_templates`: - -```bash -ate-env manifest template --template gym-py312 \ - --rootfs-image docker.io/library/python@sha256: \ - --guest-image /ate-env-guest@sha256: \ - --workspace /workspace \ - --snapshots-bucket | kubectl-ate create actor-template -f - -``` - -**Bake the guest into the image** (legacy, for control planes without image volumes): -`hack/bake-task-image.sh` builds the guest as a layer on top of any base image with `ko` and -prints the digest-pinned ref; the template's container then runs `/ko-app/ate-env-guest` directly. - -Digests for public images: `crane digest --platform linux/amd64 python:3.12-slim`. The task -image needs a `sh` for `exec()`. `timeout_s` no longer needs a `timeout` binary: the guest -enforces the deadline itself. - -## Installation - -```bash -# from a checkout of this repo: -pip install ./clients/python # ate-env-client (import name ate_env) -pip install ./integrations/nemo-gym # nemo-gym-substrate -# with NeMo Gym in the same environment: -pip install './integrations/nemo-gym[gym]' -``` - -The `gym` extra resolves only on Python ≥ 3.13.14, the floor the current `nemo-gym` release on PyPI -(0.7.0) declares. The provider, its tests and CI run on Python ≥ 3.10 without it: structural mirrors -of the `nemo_gym` base types stand in. - -The entry point makes the provider discoverable immediately: - -```python -from nemo_gym.sandbox.providers.registry import create_provider -provider = create_provider({"substrate": {"connection": {"api_url": "127.0.0.1:7777"}}}) -``` - -## How to use - -### From NeMo Gym config (the normal path) - -Add the `sandbox:` block to your resource-server config and select the `substrate` provider — -nothing else in the Gym setup changes: - -```yaml -sandbox: - default_metadata: - sandbox-api: substrate - substrate: - connection: - api_url: ate-env-api.ate-env:7777 # host:port or http(s):// URL; port-forward for dev - request_timeout_s: 30 # deadline per lifecycle RPC (create/get/delete) - create: - template: default-template # ActorTemplate when the spec names none - atespace: ate-env # atespace the templates live in and envs are created in - ready_timeout_s: 120 - ready_poll_interval_s: 1.0 - image_templates: # optional: pin an image to a registered template; others are derived - python:3.12-slim: gym-py312 -``` - -`namespace` is still accepted as an alias for `atespace` (the name this provider used before -ate-env adopted substrate's vocabulary); setting both is an error. `api_url` is plain gRPC -(h2c) like the Go client and CLI; an `https://` URL opens a TLS channel instead. - -Then run your resource server / `gym env start` as usual. Sandboxes created by Gym's resource -lifecycle now appear as Substrate actors; `kubectl get pods -n ate-env` shows the warm workers -hosting them. - -### Directly from Python (debugging, scripts) - -```python -import asyncio -from nemo_gym_substrate import SubstrateSandboxProvider -from nemo_gym.sandbox.providers.base import SandboxSpec # or nemo_gym_substrate._compat - -async def main(): - provider = SubstrateSandboxProvider({"connection": {"api_url": "127.0.0.1:7777"}}) - handle = await provider.create(SandboxSpec( - files={"/task/hello.sh": "echo hello from substrate"}, - workdir="/task", - env={"EPISODE": "1"}, - )) - result = await provider.exec(handle, "sh /task/hello.sh && echo episode=$EPISODE") - print(result.stdout, result.return_code) - await provider.close(handle) - await provider.aclose() - -asyncio.run(main()) -``` - -To share one `ate_env.Client` between the provider and other code, pass it in: -`SubstrateSandboxProvider(config, client=client)`. The provider then never closes it. - -### Per-sandbox `provider_options` (on `SandboxSpec`) - -| Option | Meaning | -|---|---| -| `template` | ActorTemplate for this sandbox (overrides `create.template` and `image_templates`) | -| `atespace` | Atespace of that template, and of the environment (`namespace` accepted as alias) | - -Unknown keys are rejected at `create()` time, as are unknown keys anywhere in the config block. - -### Services inside the sandbox: an injected daemon - -Some NeMo components talk to a service inside the sandbox instead of running commands through -`exec`: an exec server, a code-execution sandbox server, a tool daemon. The task images do not -contain it, and rebuilding hundreds of thousands of SWE-bench images to add it is not an option. -ate-env attaches such a daemon at runtime the same way it attaches its own guest: as a read-only -OCI image volume, a *runtime layer*, that the guest starts and supervises next to itself -([docs/rootfs-images/RUNTIMES.md](../../docs/rootfs-images/RUNTIMES.md)). The provider implements -nemo-gym 0.7's `SupportsSandboxEndpoint`, so Gym reaches the daemon from outside the cluster's -actor network. - -1. **Put the daemon on the base template, once.** Every environment created from this base with a - task image (`rootfs_image`) gets the daemon; the task images stay unmodified. - - ```bash - ate-env manifest template --template gym-base \ - --layer nemo=/nemo-daemon@sha256:=/opt/nemo \ - --sidecar '/opt/nemo/daemon --port 6000' \ - --sidecar-readyz http://127.0.0.1:6000/health \ - --snapshots-bucket gs:///ate-env/ | kubectl-ate create actor-template -f - - ``` - - The actor reports ready only once the daemon answers its readiness URL, so `create()` returns - a sandbox whose daemon is up. - -2. **Select the base and declare the port.** - - ```yaml - sandbox: - substrate: - connection: - api_url: ate-env-api.ate-env:7777 - # endpoint_url: http://ate-env-api.ate-env:7777 # default: api_url - create: - template: gym-base - ``` - - ```python - handle = await provider.create(SandboxSpec(image=task_image, ports=(6000,))) - ``` - -3. **Resolve the endpoint and call the daemon.** - - ```python - ep = await provider.endpoint(handle, 6000) - async with httpx.AsyncClient(base_url=ep.endpoint, headers=ep.headers) as http: - r = await http.post("/execute", json={"generated_code": "print(1)", "language": "python"}) - ``` - - `ep.endpoint` is ate-env-api, and `ep.headers` (`X-Env-Id`, `X-Env-Atespace`, `X-Env-Port`) - route each request through the Substrate router to port 6000 in this sandbox. `endpoint()` does - not touch the guest: the first request through the proxy resumes a parked sandbox, and the proxy - waits for the wake (ate-env-api's `--wake-budget`, two minutes by default, the same wait as a - guest call). It checks only that the sandbox exists, and raises `RuntimeError` at once for one - that is deleted or crashed. A port that `SandboxSpec.ports` did not declare is a `ValueError`, - as in Gym's Docker provider; `create()` rejects a spec whose ports are outside 1-65535 or name - the guest's port 80. Set `connection.endpoint_url` when the code that calls the daemon reaches - ate-env-api at another address than the provider does. - -Limits: - -- **HTTP/1.1 only.** The router tunnels HTTP to the port, not raw TCP. Streamed responses are - flushed as they arrive and WebSocket upgrades pass through; gRPC to the daemon is not supported - yet. -- **The daemon must be self-contained.** A layer cannot rely on the task image's libraries: ship a - static binary, or put the interpreter and its packages in the layer. NeMo Skills' sandbox server - (Flask on CPython, port 6000, `POST /execute`, `GET /health`) fits only with a relocatable Python - inside its layer; Gym's ECS exec server needs `python3` in the task image. -- **The daemon must listen on `0.0.0.0`.** The router reaches the port on the actor's address; a - daemon bound to `127.0.0.1` only answers processes inside the sandbox, and requests through the - endpoint fail with 502. -- **No authentication in front of the daemon.** Anyone who can reach ate-env-api can call any port - in any sandbox through it. Keep ate-env-api inside the cluster, and give the daemon its own - credential when the sandbox runs untrusted code. -- Port 80 belongs to the guest. The daemon shares the task image's user, environment and - filesystem. -- Each distinct task image yields one derived ActorTemplate on its base, created on first use and - reused after; the daemon adds no templates. - -## Testing - -```bash -cd integrations/nemo-gym -python3 -m venv .venv -.venv/bin/pip install ../../clients/python -e '.[test]' - -# Unit/contract tests — hermetic: an in-process fake ate-env-api (tests/fakes.py) -# behind the real ate_env client; no cluster, no nemo-gym needed (base types fall -# back to structural mirrors in _compat.py; with nemo-gym installed the same -# tests run against the real types): -.venv/bin/python -m pytest tests/test_provider.py -q # or: make nemo-gym-test (repo root) - -# End-to-end against a live cluster (full lifecycle on a real actor): -kubectl port-forward -n ate-env svc/ate-env-api 7777:7777 & -SUBSTRATE_E2E_API_URL=127.0.0.1:7777 .venv/bin/python -m pytest tests/test_e2e.py -q -# SUBSTRATE_E2E_TEMPLATE= picks the ActorTemplate (default: default-template) -``` - -## Benchmark - -`benchmarks/rollout_bench.py` shapes load like a NeMo-RL rollout batch: N parallel rollouts, -each *create → seed task file → T agent turns (exec + mocked model think time) → close*, on real -actors: - -```bash -python benchmarks/rollout_bench.py --api-url 127.0.0.1:7777 --rollouts 5 --turns 3 --think-s 2 -``` - -`benchmarks/driver_sim.py` is the start-time / time-to-first-exec probe over a list of templates -(one task per template per rollout), the same shape as the SWE-500 campaign's eval arm. - -First numbers on a small test cluster (5 warm workers, gVisor, cold creates; measured with the -earlier HTTP guest proxy, 2026-09-03 — rerun on the gRPC path before quoting): - -| Phase | p50 | max | Note | -|---|---|---|---| -| create | 4.95 s | 18.1 s | cold resume; max shows 5-way contention for 5 workers | -| exec (turn) | 0.14 s | 0.15 s | router + guest round trip | -| close | 3.06 s | 4.0 s | delete path | - -Turn overhead is already negligible against model think time; create/close dominate — which is -exactly what golden-snapshot starts and fork-from-golden creates are for. - -## Resource mapping and isolation - -- One Gym sandbox = one Substrate **actor**: a gVisor (or micro-VM) sandbox multiplexed onto a - warm WorkerPool worker. Isolation is the sandbox class; actors share worker nodes. -- `SandboxSpec.image` maps to a registered ActorTemplate via `image_templates`, or is sent as - `rootfs_image` and served from a derived template. The first create on an image pulls it and - bakes a golden snapshot; later creates restore from it. -- `spec.files` are seeded through the guest filesystem service after readiness. `spec.env` / - `spec.workdir` are applied per exec (`StartProcess` takes `env` and `cwd` natively). -- `exec(timeout_s=…)` is enforced guest-side and client-side. The guest kills the process group - with SIGKILL at the deadline, so the exit code is 137, and the provider appends - `substrate provider: exec timed out after …` to stderr. If the guest does not answer within - `timeout_s` plus a few seconds, the client kills the process and returns a `return_code=-1` - sentinel result rather than raising. Transport failures (env deleted mid-exec, router down) - raise `ate_env.EnvError` subclasses. -- `status()` reads the environment's lifecycle state from `ate-env-api` without touching the - guest, so polling it never wakes a parked sandbox. A **suspended or paused actor reports - RUNNING** — the router auto-resumes it on the next command, which is the behavior a rollout - loop wants. `RESUMING` → `STARTING`, `CRASHED` → `ERROR`, deleted or `DELETING` → `STOPPED`. - -Known limits (tracked in the RL-on-Substrate proposal): - -- `spec.ttl_s` is **not enforced** — substrate has no server-side TTL yet; clean up via - `close()`. A crashed harness leaks actors until deleted. -- `exec(user=…)` raises: commands run as the actor's configured user. - -## Troubleshooting - -| Symptom | Cause / fix | -|---|---| -| `SandboxCreateError: … rootfs_image … pinned by digest` (`InvalidArgument`) | `SandboxSpec.image` is a tag. Use `repo@sha256:` (`crane digest --platform linux/amd64 `), or map the image to a registered template under `create.image_templates` | -| `SandboxCreateError: … UNAVAILABLE` at once | `api_url` unreachable — port-forward died, or wrong kubectl context (`kubectl config current-context`) | -| Create times out (`SandboxCreateVerificationError: … UNAVAILABLE` / `INTERNAL`) | Check warm workers exist (`kubectl get pods -n ate-env`) and the template is Ready (`kubectl get actortemplate -n ate-env`). More parallel creates than warm workers will queue. If the message says `server closed the stream without sending trailers`, the router is not carrying gRPC to actors (see Prerequisites #1) | -| Create times out and the actor is `RUNNING` | The template's image does not run `ate-env-guest` — the actor is up but nothing serves the guest API. See "Task images" | -| Every exec fails with `NOT_FOUND` | The environment is gone (deleted by another party, or the `ate-env-api` restarted with a different atespace default). `status()` returns `STOPPED`; create a new sandbox | -| Every exec returns `UNAVAILABLE` with `CERTIFICATE_EXPIRED` | The atenet router's pod certificate expired and wasn't hot-reloaded (seen on routers running > ~1 day): `kubectl -n ate-system rollout restart deploy/atenet-router` | -| Actor stuck `RESUMING` forever (`status()` stays `STARTING`) | Restore landed on a CPU-incompatible node (mixed Intel/AMD pools) or raced a router restart. Delete the env and recreate; long-term fix is CPU-aware placement | diff --git a/integrations/nemo-gym/benchmarks/driver_sim.py b/integrations/nemo-gym/benchmarks/driver_sim.py deleted file mode 100644 index 5c5bef3f..00000000 --- a/integrations/nemo-gym/benchmarks/driver_sim.py +++ /dev/null @@ -1,161 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Substrate rollout-driver simulator: start-time / time-to-first-command bench. - -Mirrors the agent-sandbox `asrl-driver` bench convention (per task: create the -sandbox, deliver the environment, run a first `exec` probe — probe defaults to -`true`, isolating start latency). Here each task is a Substrate actor driven -through the substrate NeMo Gym provider (create -> wait ready -> exec -> delete). - -Phases measured per task: - * ready_s — create actor + wait until it serves commands (env delivery) - * first_cmd_s — the first exec (probe) round-trip - * ttfe_s — ready_s + first_cmd_s (time to first execution) - -Templates must be pre-provisioned (image + ate-env-guest, one per image). See -hack/bake-task-image.sh. Input file lists one substrate template name per line -(optionally `templatedisplay-image`); blank lines and #comments ignored. - -Usage: - kubectl port-forward -n ate-env svc/ate-env-api 7777:7777 & - python driver_sim.py --templates-file r2e_templates.txt \ - --rollouts 2 --concurrency 5 --probe true --out results.json -""" - -from __future__ import annotations - -import argparse -import asyncio -import json -import statistics -import time - -from nemo_gym_substrate._compat import SandboxSpec -from nemo_gym_substrate.provider import SubstrateSandboxProvider - - -def _pct(values: list[float], p: float) -> float | None: - if not values: - return None - if len(values) == 1: - return round(values[0], 3) - q = statistics.quantiles(values, n=100, method="inclusive") - return round(q[min(p, len(q)) - 1], 3) - - -def _summary(name: str, values: list[float]) -> dict: - if not values: - return {"phase": name, "n": 0} - return { - "phase": name, "n": len(values), - "p50": _pct(values, 50), "p90": _pct(values, 90), - "max": round(max(values), 3), "mean": round(statistics.fmean(values), 3), - } - - -async def _task(provider, template, namespace, rollout, probe, ready_timeout, sem, out): - async with sem: - opts = {"template": template} - if namespace: - opts["namespace"] = namespace - rec = {"template": template, "rollout": rollout, "ok": False} - handle = None - try: - t0 = time.monotonic() - handle = await provider.create( - SandboxSpec(provider_options=opts, ready_timeout_s=ready_timeout) - ) - rec["ready_s"] = time.monotonic() - t0 - t1 = time.monotonic() - res = await provider.exec(handle, probe) - rec["first_cmd_s"] = time.monotonic() - t1 - rec["ttfe_s"] = rec["ready_s"] + rec["first_cmd_s"] - rec["exit_code"] = res.return_code - rec["ok"] = res.return_code == 0 - except Exception as exc: # noqa: BLE001 - record, keep the sweep going - rec["error"] = f"{type(exc).__name__}: {str(exc)[:200]}" - finally: - if handle is not None: - try: - await provider.close(handle) - except Exception: # noqa: BLE001 - pass - status = "ok" if rec["ok"] else f"ERR {rec.get('error', rec.get('exit_code'))}" - print(f" [{template.split('/')[-1][:48]:48} r{rollout}] " - f"ttfe={rec.get('ttfe_s', float('nan')):.2f}s " - f"(ready={rec.get('ready_s', float('nan')):.2f} " - f"cmd={rec.get('first_cmd_s', float('nan')):.2f}) {status}", flush=True) - out.append(rec) - - -def _load_templates(path: str, limit: int) -> list[str]: - templates = [] - for line in open(path): - line = line.strip() - if not line or line.startswith("#"): - continue - templates.append(line.split("\t")[0].split()[0]) - return templates[:limit] if limit else templates - - -async def main() -> None: - ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument("--api-url", default="http://127.0.0.1:7777") - ap.add_argument("--templates-file", required=True) - ap.add_argument("--rollouts", type=int, default=1, help="tasks per template") - ap.add_argument("--concurrency", type=int, default=5) - ap.add_argument("--probe", default="true", help="first command (default no-op)") - ap.add_argument("--problems", type=int, default=0, help="limit templates (0=all)") - ap.add_argument("--ready-timeout-s", type=float, default=300.0) - ap.add_argument("--namespace", default=None) - ap.add_argument("--out", default=None) - a = ap.parse_args() - - templates = _load_templates(a.templates_file, a.problems) - tasks_plan = [(t, g) for t in templates for g in range(a.rollouts)] - print(f"driver: {len(templates)} templates x {a.rollouts} = {len(tasks_plan)} tasks; " - f"concurrency={a.concurrency} probe={a.probe!r}", flush=True) - - provider = SubstrateSandboxProvider({"connection": {"api_url": a.api_url}}) - sem = asyncio.Semaphore(a.concurrency) - out: list[dict] = [] - - wall0 = time.monotonic() - await asyncio.gather(*( - _task(provider, t, a.namespace, g, a.probe, a.ready_timeout_s, sem, out) - for t, g in tasks_plan - )) - wall = time.monotonic() - wall0 - await provider.aclose() - - ok = [r for r in out if r["ok"]] - ready = [r["ready_s"] for r in ok] - cmd = [r["first_cmd_s"] for r in ok] - ttfe = [r["ttfe_s"] for r in ok] - summary = { - "tasks": len(out), "ok": len(ok), "err": len(out) - len(ok), - "wall_s": round(wall, 1), - "throughput_tasks_per_s": round(len(out) / wall, 3) if wall else None, - "phases": [_summary("ready", ready), _summary("first_cmd", cmd), _summary("ttfe", ttfe)], - } - print("\n=== RESULT ===") - print(json.dumps(summary, indent=2)) - if a.out: - json.dump({"summary": summary, "tasks": out}, open(a.out, "w"), indent=2) - print(f"wrote {a.out}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/integrations/nemo-gym/benchmarks/rollout_bench.py b/integrations/nemo-gym/benchmarks/rollout_bench.py deleted file mode 100644 index 20bad5b2..00000000 --- a/integrations/nemo-gym/benchmarks/rollout_bench.py +++ /dev/null @@ -1,134 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Mock NeMo-RL rollout benchmark for the substrate sandbox provider. - -Shapes the load like a NeMo Gym rollout batch: N parallel rollouts, each -create -> seed task file -> T agent turns (exec + simulated model "think" -time) -> artifact download -> close. The LLM is mocked with asyncio.sleep; -the sandboxes are real Substrate actors. - -Usage: - kubectl port-forward -n ate-env svc/ate-env-api 7777:7777 & - python benchmarks/rollout_bench.py --rollouts 5 --turns 3 --think-s 2 -""" - -from __future__ import annotations - -import argparse -import asyncio -import statistics -import time - -from nemo_gym_substrate._compat import SandboxSpec -from nemo_gym_substrate.provider import SubstrateSandboxProvider - -# Pure-sh turn work: append the turn marker to the episode log, emit a -# checksum of the log so far. Runs in any guest with a POSIX shell. -TURN_CMD = 'echo "turn {turn} $(date +%s%N)" >> /task/episode.log && wc -c < /task/episode.log' - - -async def rollout( - provider: SubstrateSandboxProvider, - index: int, - turns: int, - think_s: float, - template: str | None, - metrics: dict[str, list[float]], - sem: asyncio.Semaphore, -) -> None: - async with sem: - opts = {"template": template} if template else {} - t0 = time.monotonic() - handle = await provider.create( - SandboxSpec( - files={"/task/episode.log": f"rollout {index}\n"}, - workdir="/task", - provider_options=opts, - ready_timeout_s=180, - ) - ) - metrics["create_s"].append(time.monotonic() - t0) - try: - for turn in range(turns): - t1 = time.monotonic() - result = await provider.exec(handle, TURN_CMD.format(turn=turn)) - metrics["exec_s"].append(time.monotonic() - t1) - if result.return_code != 0: - raise RuntimeError( - f"rollout {index} turn {turn} failed rc={result.return_code}: " - f"{result.stderr}" - ) - await asyncio.sleep(think_s) # mock model inference / eval - finally: - t2 = time.monotonic() - await provider.close(handle) - metrics["close_s"].append(time.monotonic() - t2) - - -def summarize(name: str, values: list[float]) -> str: - if not values: - return f"{name:>9}: n=0" - q = statistics.quantiles(values, n=20) if len(values) >= 2 else [values[0]] * 19 - return ( - f"{name:>9}: n={len(values):<4} p50={statistics.median(values):6.2f}s " - f"p95={q[18]:6.2f}s max={max(values):6.2f}s" - ) - - -async def main() -> None: - ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument("--api-url", default="http://127.0.0.1:7777") - ap.add_argument("--rollouts", type=int, default=5) - ap.add_argument("--turns", type=int, default=3) - ap.add_argument("--think-s", type=float, default=2.0) - ap.add_argument("--max-concurrency", type=int, default=5, - help="parallel rollouts (bound by warm workers on small clusters)") - ap.add_argument("--template", default=None) - args = ap.parse_args() - - provider = SubstrateSandboxProvider({"connection": {"api_url": args.api_url}}) - metrics: dict[str, list[float]] = {"create_s": [], "exec_s": [], "close_s": []} - sem = asyncio.Semaphore(args.max_concurrency) - - wall0 = time.monotonic() - results = await asyncio.gather( - *( - rollout(provider, i, args.turns, args.think_s, args.template, metrics, sem) - for i in range(args.rollouts) - ), - return_exceptions=True, - ) - wall = time.monotonic() - wall0 - await provider.aclose() - - failures = [r for r in results if isinstance(r, BaseException)] - execs = len(metrics["exec_s"]) - print(f"\nrollouts={args.rollouts} turns={args.turns} think={args.think_s}s " - f"concurrency={args.max_concurrency}") - print(summarize("create", metrics["create_s"])) - print(summarize("exec", metrics["exec_s"])) - print(summarize("close", metrics["close_s"])) - ideal = args.turns * args.think_s # pure think time per rollout - print(f"{'wall':>9}: {wall:6.2f}s execs/s={execs / wall:5.2f} " - f"think-only floor per rollout={ideal:.1f}s") - if failures: - print(f"FAILURES: {len(failures)}") - for f in failures[:3]: - print(" -", type(f).__name__, str(f)[:200]) - raise SystemExit(1) - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/integrations/nemo-gym/hack/bake-task-image.sh b/integrations/nemo-gym/hack/bake-task-image.sh deleted file mode 100755 index 24d35e03..00000000 --- a/integrations/nemo-gym/hack/bake-task-image.sh +++ /dev/null @@ -1,70 +0,0 @@ -#!/usr/bin/env bash -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# Bake ate-env-guest onto an arbitrary base image (Mode A: runtime baked in). -# -# ko builds the guest Go binary as a layer on top of BASE and pushes the result -# to REPO, printing the digest-pinned image ref (which substrate requires for -# ActorTemplates). The guest is a static binary, so BASE can be any distro. -# -# Usage: -# bake-task-image.sh [env-repo-path] -# -# Example: -# integrations/nemo-gym/hack/bake-task-image.sh python:3.12-slim gcr.io/PROJECT/py-task-guest -# -# Env: -# ATE_ENV_REPO path to an agent-substrate/env checkout (default: this repo) -# KO ko binary (default: ko on PATH, else ~/go/bin/ko) -set -o errexit -o nounset -o pipefail - -if [[ $# -lt 2 ]]; then - sed -n '18,32p' "$0" >&2 - exit 2 -fi - -BASE="$1" -REPO="$2" -ENV_REPO="${3:-${ATE_ENV_REPO:-$(cd "$(dirname "$0")/../../.." && pwd)}}" -KO="${KO:-$(command -v ko || echo "${HOME}/go/bin/ko")}" -GUEST_PKG="github.com/agent-substrate/env/cmd/ate-env-guest" - -[[ -x "${KO}" ]] || { echo "ko not found (set \$KO or install to ~/go/bin)" >&2; exit 1; } -[[ -d "${ENV_REPO}/cmd/ate-env-guest" ]] || { - echo "env repo not found at ${ENV_REPO} (set \$ATE_ENV_REPO)" >&2; exit 1; } - -# A ko config that overrides the guest's base image to BASE. This overrides the -# repo's own .ko.yaml (which pins bash:latest), so the guest lands on BASE's -# rootfs instead. -# ko infers the config format from the file extension, so the path must end in -# .yaml (macOS mktemp -t appends a random suffix, which breaks that) — use a -# temp dir holding a real .ko.yaml. -CFG_DIR="$(mktemp -d -t ko-bake-XXXX)" -CFG="${CFG_DIR}/.ko.yaml" -trap 'rm -rf "${CFG_DIR}"' EXIT -cat > "${CFG}" < ${REPO} ..." >&2 -cd "${ENV_REPO}" -IMAGE="$( - KO_DOCKER_REPO="${REPO}" \ - KO_CONFIG_PATH="${CFG}" \ - KO_DEFAULTPLATFORMS="${KO_DEFAULTPLATFORMS:-linux/amd64}" \ - "${KO}" build --bare ./cmd/ate-env-guest -)" -echo "${IMAGE}" diff --git a/integrations/nemo-gym/pyproject.toml b/integrations/nemo-gym/pyproject.toml deleted file mode 100644 index 29d93fdf..00000000 --- a/integrations/nemo-gym/pyproject.toml +++ /dev/null @@ -1,45 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -[build-system] -requires = ["setuptools>=68"] -build-backend = "setuptools.build_meta" - -[project] -name = "nemo-gym-substrate" -version = "0.2.0" -description = "NeMo Gym sandbox provider backed by Agent Substrate actors (via ate-env)" -readme = "README.md" -license = { text = "Apache-2.0" } -requires-python = ">=3.10" -# ate-env-client is this repo's Python client (clients/python); it is not on -# PyPI yet, so install it from the checkout first: pip install ./clients/python -dependencies = ["ate-env-client>=0.1.0", "grpcio>=1.84"] - -[project.optional-dependencies] -# The real deployment target. Without it the package falls back to local -# structural mirrors of the nemo_gym base types (tests only). -gym = ["nemo-gym; python_full_version >= '3.13.14'"] # what nemo-gym 0.7.0 on PyPI requires -test = ["pytest>=8"] - -[project.entry-points."nemo_gym.sandbox_providers"] -substrate = "nemo_gym_substrate.provider:SubstrateSandboxProvider" - -[tool.setuptools.packages.find] -where = ["src"] - -# ate-env-client is not published; resolve it from this checkout so -# `uv run --extra test pytest` works without a prior pip install. -[tool.uv.sources] -ate-env-client = { path = "../../clients/python", editable = true } diff --git a/integrations/nemo-gym/src/nemo_gym_substrate/__init__.py b/integrations/nemo-gym/src/nemo_gym_substrate/__init__.py deleted file mode 100644 index ef503383..00000000 --- a/integrations/nemo-gym/src/nemo_gym_substrate/__init__.py +++ /dev/null @@ -1,29 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""NeMo Gym sandbox provider for Agent Substrate.""" - -from .provider import ( - ConnectionConfig, - CreateConfig, - SubstrateProviderOptions, - SubstrateSandboxProvider, -) - -__all__ = [ - "ConnectionConfig", - "CreateConfig", - "SubstrateProviderOptions", - "SubstrateSandboxProvider", -] diff --git a/integrations/nemo-gym/src/nemo_gym_substrate/_compat.py b/integrations/nemo-gym/src/nemo_gym_substrate/_compat.py deleted file mode 100644 index c11eb635..00000000 --- a/integrations/nemo-gym/src/nemo_gym_substrate/_compat.py +++ /dev/null @@ -1,124 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""NeMo Gym base types, with hermetic fallbacks. - -When ``nemo_gym`` is installed (any real deployment), its types are used so -handles and results are the genuine articles. The fallback dataclasses below -are structural mirrors of ``nemo_gym.sandbox.providers.base`` that let this -demo's contract tests run inside the substrate repo without the nemo-gym -dependency. They must be kept shape-identical to the upstream types. -""" - -from __future__ import annotations - -try: # pragma: no cover - exercised only when nemo_gym is installed - from nemo_gym.sandbox.providers.base import ( # type: ignore - SandboxCreateError, - SandboxCreateVerificationError, - SandboxExecResult, - SandboxHandle, - SandboxSpec, - SandboxStatus, - ) - - NEMO_GYM_AVAILABLE = True -except ImportError: # pragma: no cover - the in-repo test path - from dataclasses import dataclass, field - from enum import Enum - from typing import Any - - NEMO_GYM_AVAILABLE = False - - class SandboxStatus(str, Enum): - STARTING = "starting" - RUNNING = "running" - STOPPED = "stopped" - ERROR = "error" - UNKNOWN = "unknown" - - @dataclass(frozen=True) - class SandboxHandle: - sandbox_id: str - provider_name: str - raw: Any - - @dataclass(frozen=True) - class SandboxExecResult: - stdout: str | None - stderr: str | None - return_code: int - error_type: str | None = None - - @dataclass(frozen=True) - class SandboxSpec: - image: str | None = None - ttl_s: int | float | None = None - ready_timeout_s: int | float | None = None - workdir: str | None = None - env: dict[str, str] = field(default_factory=dict) - files: dict[str, str] = field(default_factory=dict) - metadata: dict[str, str] = field(default_factory=dict) - resources: dict[str, Any] = field(default_factory=dict) - entrypoint: list[str] | None = None - provider_options: dict[str, Any] = field(default_factory=dict) - ports: tuple[int, ...] = () - - class SandboxCreateError(RuntimeError): - pass - - class SandboxCreateVerificationError(SandboxCreateError): - pass - - -# SandboxEndpoint is newer than the rest (nemo-gym 0.7): an installed 0.5 -# has the base module without it, so it falls back on its own. -try: # pragma: no cover - exercised only with nemo-gym >= 0.7 - from nemo_gym.sandbox.providers.base import SandboxEndpoint # type: ignore -except ImportError: # pragma: no cover - the in-repo test path and nemo-gym 0.5 - from collections.abc import Mapping as _Mapping - from dataclasses import dataclass as _dataclass - from dataclasses import field as _field - from urllib.parse import urlsplit as _urlsplit - - @_dataclass(frozen=True) - class SandboxEndpoint: # type: ignore[no-redef] - endpoint: str - headers: dict[str, str] = _field(default_factory=dict) - - def __post_init__(self) -> None: - if not isinstance(self.endpoint, str) or not self.endpoint.strip(): - raise ValueError("Sandbox endpoint must be a non-empty absolute URL") - endpoint = self.endpoint.strip() - parsed = _urlsplit(endpoint) - if not parsed.scheme or not parsed.netloc: - raise ValueError("Sandbox endpoint must be a non-empty absolute URL") - if not isinstance(self.headers, _Mapping): - raise TypeError("Sandbox endpoint headers must be a mapping") - object.__setattr__(self, "endpoint", endpoint) - object.__setattr__( - self, "headers", {str(k): str(v) for k, v in self.headers.items()} - ) - - -__all__ = [ - "NEMO_GYM_AVAILABLE", - "SandboxCreateError", - "SandboxCreateVerificationError", - "SandboxEndpoint", - "SandboxExecResult", - "SandboxHandle", - "SandboxSpec", - "SandboxStatus", -] diff --git a/integrations/nemo-gym/src/nemo_gym_substrate/provider.py b/integrations/nemo-gym/src/nemo_gym_substrate/provider.py deleted file mode 100644 index 89940311..00000000 --- a/integrations/nemo-gym/src/nemo_gym_substrate/provider.py +++ /dev/null @@ -1,567 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""NeMo Gym sandbox provider backed by Agent Substrate. - -Each sandbox is a Substrate actor fronted by the ``ate-env`` API -(https://github.com/agent-substrate/env), reached through this repo's Python -client, ``ate_env`` (``clients/python``). Environment lifecycle (create, get, delete) goes to -``ateenv.v1alpha.EnvironmentService`` on ``ate-env-api``; command execution and -file transfer go to the ``ProcessService`` / ``FileSystemService`` served by -the ``ate-env-guest`` daemon inside every actor, which ``ate-env-api`` proxies -through the atenet router. One gRPC endpoint carries both. Idle sandboxes can -be suspended by Substrate and are resumed transparently on the next guest -call, so a fleet of mostly-waiting rollout sandboxes holds no workers. - -Contract: https://docs.nvidia.com/nemo/gym/main/infrastructure/sandbox/adding-a-provider/ -""" - -from __future__ import annotations - -import asyncio -import contextlib -import logging -import time -import uuid -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Mapping -from urllib.parse import urlparse - -import grpc -import grpc.aio -from ate_env import Client, Env, EnvError, EnvironmentStatus, NotFoundError, Process - -from ._compat import ( - SandboxCreateError, - SandboxCreateVerificationError, - SandboxEndpoint, - SandboxExecResult, - SandboxHandle, - SandboxSpec, - SandboxStatus, -) - -logger = logging.getLogger(__name__) - -_PROVIDER_NAME = "substrate" - -# Client-side slack on top of exec(timeout_s=...): the guest enforces the -# deadline itself (see exec), so the client only has to outlive it. -_EXEC_GRACE_S = 5.0 -# Deadline for one readiness probe. A guest that is still booting may hold a -# call open; bounding it keeps the poll loop honest. -_PROBE_TIMEOUT_S = 10.0 -# Smallest guest deadline exec() asks for. The guest takes any positive -# duration, but zero selects its default (an hour) and a sub-second kill -# window is noise next to the RPC round trip. -_MIN_GUEST_TIMEOUT_S = 1.0 -# Exit code the guest reports for a process group it killed with SIGKILL. -_SIGKILL_EXIT = 137 - -# Routing headers of ate-env-api's port proxy (internal/portproxy): a plain -# HTTP request carrying them is forwarded to that port inside the environment. -_ENV_ID_HEADER = "X-Env-Id" -_ENV_ATESPACE_HEADER = "X-Env-Atespace" -_ENV_PORT_HEADER = "X-Env-Port" -# The ate-env guest's own port in every environment: not a daemon port, and -# the port proxy refuses it too. -_GUEST_PORT = 80 - - -def _validate_ports(ports: object) -> list[int]: - """``SandboxSpec.ports`` as a list of TCP ports a daemon can listen on. - - Checked before the actor starts, so a bad spec fails at once rather than - at ``endpoint()`` time. ``bool`` is excluded explicitly because it is an - ``int`` subclass in Python (``True == 1``).""" - out: list[int] = [] - for p in ports or (): - if isinstance(p, bool) or not isinstance(p, int): - raise TypeError(f"SandboxSpec.ports entries must be integers, got {p!r}") - if not 1 <= p <= 65535: - raise ValueError(f"SandboxSpec.ports entry {p} is not a TCP port (1-65535)") - if p == _GUEST_PORT: - raise ValueError( - f"SandboxSpec.ports entry {p} is the ate-env guest's port; a daemon listens on another port" - ) - out.append(p) - return out - -# Env lifecycle status -> Gym status. Parked (suspended / paused) actors report -# RUNNING: the router resumes them on the next guest call, which is what a -# rollout loop wants to see (see README, "Resource mapping"). -_STATUS_MAP = { - EnvironmentStatus.RESUMING: SandboxStatus.STARTING, - EnvironmentStatus.RUNNING: SandboxStatus.RUNNING, - EnvironmentStatus.SUSPENDING: SandboxStatus.RUNNING, - EnvironmentStatus.SUSPENDED: SandboxStatus.RUNNING, - EnvironmentStatus.PAUSING: SandboxStatus.RUNNING, - EnvironmentStatus.PAUSED: SandboxStatus.RUNNING, - EnvironmentStatus.CRASHED: SandboxStatus.ERROR, - EnvironmentStatus.DELETING: SandboxStatus.STOPPED, -} - - -def _require_keys(options: Mapping[str, Any], allowed: frozenset[str], where: str) -> None: - unknown = sorted(set(options) - allowed) - if unknown: - raise ValueError(f"{where}: unknown option(s) {unknown}; allowed: {sorted(allowed)}") - - -def _describe(exc: BaseException) -> str: - """``CODE: details`` for ate_env errors, so messages name the gRPC status.""" - code = getattr(exc, "code", None) - if isinstance(code, grpc.StatusCode): - return f"{code.name}: {exc}" - if isinstance(exc, EnvError): - return f"{type(exc).__name__}: {exc}" - return str(exc) or type(exc).__name__ - - -def _atespace_option(options: Mapping[str, Any], where: str) -> str | None: - """``atespace`` is the ate-env name; ``namespace`` is accepted as the older alias.""" - if "atespace" in options and "namespace" in options: - raise ValueError(f"{where}: set either atespace or namespace, not both") - value = options.get("atespace", options.get("namespace")) - return None if value is None else str(value) - - -@dataclass(frozen=True) -class ConnectionConfig: - """`connection:` block of the provider config.""" - - # host:port or an http(s):// URL of ate-env-api. Plain gRPC (h2c) unless https. - api_url: str = "http://127.0.0.1:7777" - # Deadline for each lifecycle RPC (create / get / delete). - request_timeout_s: float = 30.0 - # Base URL that endpoint() hands out for services inside a sandbox. Unset - # means api_url as an http(s) URL: ate-env-api serves its port proxy on the - # same listener. Set it when the code that calls the service reaches - # ate-env-api at another address than the provider does. - endpoint_url: str | None = None - - _ALLOWED = frozenset({"api_url", "request_timeout_s", "endpoint_url"}) - - @classmethod - def from_mapping(cls, options: Mapping[str, Any]) -> "ConnectionConfig": - _require_keys(options, cls._ALLOWED, "sandbox.substrate.connection") - endpoint_url = options.get("endpoint_url") - return cls( - api_url=str(options.get("api_url", cls.api_url)).rstrip("/"), - request_timeout_s=float(options.get("request_timeout_s", cls.request_timeout_s)), - endpoint_url=None if endpoint_url is None else str(endpoint_url).rstrip("/"), - ) - - def endpoint_base(self) -> str: - """The absolute URL of ate-env-api's port proxy: ``endpoint_url`` when - set, else ``api_url``; either gets ``http://`` when it names no scheme, - because ``SandboxEndpoint`` rejects a bare ``host:port``.""" - base = self.endpoint_url or self.api_url - return base if "://" in base else f"http://{base}" - - -@dataclass(frozen=True) -class CreateConfig: - """`create:` block of the provider config.""" - - # ActorTemplate when the spec names none; ate-env's own default. - template: str = "default-template" - # Atespace the templates live in and the environments are created in. - atespace: str = "ate-env" - ready_timeout_s: float = 120.0 - ready_poll_interval_s: float = 1.0 - # Optional mapping from SandboxSpec.image to an ActorTemplate name, for - # workloads that select sandboxes by image reference. Templates must be - # pre-provisioned on the cluster (see README). - image_templates: Mapping[str, str] = None # type: ignore[assignment] - - _ALLOWED = frozenset( - { - "template", - "atespace", - "namespace", - "ready_timeout_s", - "ready_poll_interval_s", - "image_templates", - } - ) - - @classmethod - def from_mapping(cls, options: Mapping[str, Any]) -> "CreateConfig": - _require_keys(options, cls._ALLOWED, "sandbox.substrate.create") - return cls( - template=str(options.get("template", cls.template)), - atespace=_atespace_option(options, "sandbox.substrate.create") or cls.atespace, - ready_timeout_s=float(options.get("ready_timeout_s", cls.ready_timeout_s)), - ready_poll_interval_s=float( - options.get("ready_poll_interval_s", cls.ready_poll_interval_s) - ), - image_templates=dict(options.get("image_templates") or {}), - ) - - -@dataclass(frozen=True) -class SubstrateProviderOptions: - """Per-sandbox options carried in ``SandboxSpec.provider_options``.""" - - template: str | None = None - atespace: str | None = None - - _ALLOWED = frozenset({"template", "atespace", "namespace"}) - - @classmethod - def from_mapping(cls, options: Mapping[str, Any]) -> "SubstrateProviderOptions": - _require_keys(options, cls._ALLOWED, "SandboxSpec.provider_options") - return cls( - template=options.get("template"), - atespace=_atespace_option(options, "SandboxSpec.provider_options"), - ) - - -class SubstrateSandboxProvider: - """NeMo Gym sandbox provider running sandboxes as Substrate actors.""" - - name = _PROVIDER_NAME - - def __init__( - self, - config: Mapping[str, Any] | None = None, - *, - client: Client | None = None, - **kwargs: Any, - ) -> None: - """``config`` is the ``sandbox.substrate`` block. Pass ``client`` to reuse a - caller-owned ``ate_env.Client`` (the provider then never closes it).""" - config = dict(config or {}) - config.update(kwargs) - _require_keys(config, frozenset({"connection", "create"}), "sandbox.substrate") - self._connection = ConnectionConfig.from_mapping(config.get("connection") or {}) - self._create = CreateConfig.from_mapping(config.get("create") or {}) - # The channel is opened lazily so constructing a provider never dials. - self._client: Client | None = client - self._owns_client = client is None - self._channel: grpc.aio.Channel | None = None - - # -- provider contract ------------------------------------------------- - - async def create(self, spec: SandboxSpec) -> SandboxHandle: - opts = SubstrateProviderOptions.from_mapping(spec.provider_options or {}) - # Ports a service inside the sandbox listens on (nemo-gym >= 0.7); - # endpoint() resolves only these. - ports = _validate_ports(getattr(spec, "ports", None)) - template, rootfs_image = self._resolve_template(spec, opts) - atespace = opts.atespace or self._create.atespace - env_id = f"gym-{uuid.uuid4().hex[:10]}" - - try: - env = await asyncio.wait_for( - self._ate().create( - env_id, - atespace=atespace, - template_name=template, - template_atespace=atespace, - rootfs_image=rootfs_image, - ), - timeout=self._connection.request_timeout_s, - ) - except EnvError as exc: - raise SandboxCreateError( - f"creating substrate env {env_id!r} (template {atespace}/{template}): " - f"{_describe(exc)}" - ) from exc - except asyncio.TimeoutError as exc: - raise SandboxCreateError( - f"creating substrate env {env_id!r}: no answer from " - f"{self._connection.api_url} within {self._connection.request_timeout_s}s" - ) from exc - - try: - await self._wait_ready(env, spec.ready_timeout_s or self._create.ready_timeout_s) - for target_path, content in (spec.files or {}).items(): - await env.write_file(target_path, content.encode(), mode=0o644) - except Exception: - await self._best_effort_delete(env_id, atespace) - raise - - return SandboxHandle( - sandbox_id=env_id, - provider_name=self.name, - raw={ - "env_id": env_id, - "atespace": atespace, - "workdir": spec.workdir, - "env": dict(spec.env or {}), - "ports": ports, - }, - ) - - async def exec( - self, - handle: SandboxHandle, - command: str, - *, - cwd: str | None = None, - env: dict[str, str] | None = None, - timeout_s: int | float | None = None, - user: str | int | None = None, - ) -> SandboxExecResult: - if user is not None: - raise ValueError( - "the substrate provider does not support per-exec `user`: " - "the guest daemon runs commands as the actor's configured user" - ) - raw = handle.raw or {} - merged_env = {**(raw.get("env") or {}), **(env or {})} - effective_cwd = cwd or raw.get("workdir") or "" - - argv = ["sh", "-c", command] - # The guest enforces the deadline itself and kills the process group - # with SIGKILL when it elapses, so a runaway process never outlives the - # call that started it and the image needs no `timeout` binary. `is not - # None` so an explicit 0 means "shortest deadline", not "unlimited". - guest_timeout = ( - None if timeout_s is None else max(_MIN_GUEST_TIMEOUT_S, float(timeout_s)) - ) - - sandbox = self._env(handle) - started = time.monotonic() - process = await sandbox.start_process( - argv, cwd=effective_cwd, env=merged_env, timeout=guest_timeout - ) - try: - result = await asyncio.wait_for( - self._collect(process), - timeout=None if timeout_s is None else timeout_s + _EXEC_GRACE_S, - ) - except asyncio.TimeoutError: - with contextlib.suppress(EnvError): - await process.kill() - return SandboxExecResult( - stdout=None, - stderr=f"substrate provider: exec timed out after {timeout_s}s", - return_code=-1, - ) - if ( - guest_timeout is not None - and result.return_code == _SIGKILL_EXIT - and time.monotonic() - started >= guest_timeout - ): - # The guest reports a deadline kill as a plain SIGKILL and nothing - # in the exit state says why; the elapsed time does. Name it: 137 - # with an empty stderr tells a rollout nothing. - note = f"substrate provider: exec timed out after {timeout_s}s (killed by the guest)" - return SandboxExecResult( - stdout=result.stdout, - stderr=f"{result.stderr}\n{note}" if result.stderr else note, - return_code=result.return_code, - ) - return result - - async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: - await self._env(handle).write_file( - target_path, Path(source_path).read_bytes(), mode=0o644 - ) - - async def download_file(self, handle: SandboxHandle, source_path: str, target_path: Path) -> None: - content = await self._env(handle).read_file_bytes(source_path) - target = Path(target_path) - target.parent.mkdir(parents=True, exist_ok=True) - target.write_bytes(content) - - async def status(self, handle: SandboxHandle) -> SandboxStatus: - # GetEnvironment reads the actor's state without touching the guest, - # so asking does not wake a parked sandbox. - raw = handle.raw or {} - try: - info = await asyncio.wait_for( - self._ate().get(handle.sandbox_id, atespace=self._atespace_of(raw)), - timeout=self._connection.request_timeout_s, - ) - except NotFoundError: - return SandboxStatus.STOPPED - except (EnvError, asyncio.TimeoutError): - return SandboxStatus.UNKNOWN - return _STATUS_MAP.get(info.status, SandboxStatus.UNKNOWN) - - async def endpoint(self, handle: SandboxHandle, port: int) -> SandboxEndpoint: - """Resolve a declared service port to a caller-reachable endpoint - (nemo-gym's ``SupportsSandboxEndpoint``). - - The service is a daemon inside the sandbox, typically one a runtime - layer on the base template mounts into the task image and the guest - starts next to itself (see README). The endpoint is ate-env-api's port - proxy; its headers route each plain HTTP request through the Substrate - router to ``port`` in this sandbox. Nothing here touches the guest: a - parked sandbox is resumed by the first request through the proxy, - which waits for the wake (ate-env-api's ``--wake-budget``). Only the - sandbox's existence is checked, from the control plane, so a deleted - or crashed sandbox raises ``RuntimeError`` here rather than on the - first request. - """ - raw = handle.raw or {} - declared = list(raw.get("ports") or ()) - if isinstance(port, bool) or port not in declared: - raise ValueError( - f"Sandbox port {port!r} was not declared in SandboxSpec.ports; " - f"declared ports: {declared!r}" - ) - atespace = self._atespace_of(raw) - try: - info = await asyncio.wait_for( - self._ate().get(handle.sandbox_id, atespace=atespace), - timeout=self._connection.request_timeout_s, - ) - except NotFoundError as exc: - raise RuntimeError(f"substrate env {handle.sandbox_id!r} does not exist") from exc - except EnvError as exc: - raise RuntimeError( - f"substrate env {handle.sandbox_id!r}: {_describe(exc)}" - ) from exc - except TimeoutError as exc: - raise RuntimeError( - f"substrate env {handle.sandbox_id!r}: no answer from " - f"{self._connection.api_url} within {self._connection.request_timeout_s}s" - ) from exc - if info.status == EnvironmentStatus.CRASHED: - raise RuntimeError(f"substrate env {handle.sandbox_id!r} crashed") - return SandboxEndpoint( - endpoint=self._connection.endpoint_base(), - headers={ - _ENV_ID_HEADER: handle.sandbox_id, - _ENV_ATESPACE_HEADER: atespace, - _ENV_PORT_HEADER: str(port), - }, - ) - - async def close(self, handle: SandboxHandle) -> None: - await self._best_effort_delete(handle.sandbox_id, self._atespace_of(handle.raw or {})) - - async def aclose(self) -> None: - if self._client is not None and self._owns_client: - await self._client.close() - if self._channel is not None: - await self._channel.close() - self._client = None if self._owns_client else self._client - self._channel = None - - # -- internals ---------------------------------------------------------- - - def _ate(self) -> Client: - if self._client is None: - self._client = self._open_client(self._connection.api_url) - return self._client - - def _open_client(self, api_url: str) -> Client: - url = urlparse(api_url if "://" in api_url else f"//{api_url}") - if url.scheme == "https": - # ate_env.Client only dials h2c itself; hand it a TLS channel. - target = f"{url.hostname}:{url.port or 443}" - self._channel = grpc.aio.secure_channel(target, grpc.ssl_channel_credentials()) - return Client(channel=self._channel) - return Client(api_url) - - def _atespace_of(self, raw: Mapping[str, Any]) -> str: - return str(raw.get("atespace") or self._create.atespace) - - def _env(self, handle: SandboxHandle) -> Env: - return self._ate().env(handle.sandbox_id, atespace=self._atespace_of(handle.raw or {})) - - def _resolve_template( - self, spec: SandboxSpec, opts: SubstrateProviderOptions - ) -> tuple[str, str | None]: - """Return ``(template, rootfs_image)`` for a spec. - - An explicit provider option or an ``image_templates`` entry names a - registered template. Any other image goes to ate-env-api as - ``rootfs_image``: the API derives a template from the base on first - use, with the guest mounted into the image as a read-only volume, and - reuses it afterwards. The image must be digest-pinned; the API rejects - tags. - """ - if opts.template: - return opts.template, None - if spec.image: - mapped = (self._create.image_templates or {}).get(spec.image) - if mapped: - return mapped, None - return self._create.template, spec.image - return self._create.template, None - - async def _wait_ready( - self, - env: Env, - timeout_s: float, - *, - error: type[Exception] = SandboxCreateVerificationError, - stop_on_not_found: bool = False, - ) -> None: - """Poll a trivial command until the guest answers it. A fresh actor may - still be booting, and a guest call made before it serves fails with a - transport-level error; both are retried until the deadline, then - ``error`` is raised. The first call also resumes a parked actor. With - ``stop_on_not_found``, a sandbox that no longer exists raises ``error`` - at once instead of after the deadline.""" - deadline = time.monotonic() + timeout_s - last_error = "no probe attempted" - while True: - try: - result = await asyncio.wait_for(env.shell("true"), timeout=_PROBE_TIMEOUT_S) - if result.exit_code == 0: - return - last_error = f"probe exited {result.exit_code}: {result.stderr.strip()[:200]}" - except NotFoundError as exc: - if stop_on_not_found: - raise error(f"substrate env {env.id!r} does not exist") from exc - last_error = _describe(exc) - except EnvError as exc: - last_error = _describe(exc) - except asyncio.TimeoutError: - last_error = f"probe did not answer within {_PROBE_TIMEOUT_S}s" - if time.monotonic() >= deadline: - break - await asyncio.sleep(self._create.ready_poll_interval_s) - raise error(f"substrate env {env.id!r} not ready after {timeout_s}s: {last_error}") - - @staticmethod - async def _collect(process: Process) -> SandboxExecResult: - stdout = bytearray() - stderr = bytearray() - final = None - async for chunk in process.output(follow=True): - if chunk.stdout is not None: - stdout.extend(chunk.stdout) - elif chunk.stderr is not None: - stderr.extend(chunk.stderr) - elif chunk.exit is not None: - final = chunk.exit - if final is None: - # The stream can end before the exit message; wait() polls for it. - final = await process.wait() - return SandboxExecResult( - stdout=stdout.decode("utf-8", errors="replace"), - stderr=stderr.decode("utf-8", errors="replace"), - return_code=final.exit_code, - ) - - async def _best_effort_delete(self, env_id: str, atespace: str) -> None: - try: - await asyncio.wait_for( - self._ate().delete(env_id, atespace=atespace), - timeout=self._connection.request_timeout_s, - ) - except NotFoundError: - pass # already gone: close() is idempotent - except (EnvError, asyncio.TimeoutError) as exc: - logger.warning("deleting substrate env %s: %s", env_id, _describe(exc)) diff --git a/integrations/nemo-gym/tests/__init__.py b/integrations/nemo-gym/tests/__init__.py deleted file mode 100644 index 58d482ea..00000000 --- a/integrations/nemo-gym/tests/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. diff --git a/integrations/nemo-gym/tests/fakes.py b/integrations/nemo-gym/tests/fakes.py deleted file mode 100644 index 32d87919..00000000 --- a/integrations/nemo-gym/tests/fakes.py +++ /dev/null @@ -1,272 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""In-process fake of ate-env-api speaking the current ateenv.v1alpha protocol. - -Mirrors just enough of EnvironmentService, ProcessService and FileSystemService -for the provider's contract tests; the knobs on FakeAteEnv script failures. -""" - -from __future__ import annotations - -import asyncio -import contextlib -import time -from dataclasses import dataclass, field - -import grpc -import grpc.aio -from ate_env import EnvironmentStatus, ProcessState -from ate_env._gen.ateenv.v1alpha import env_pb2, env_pb2_grpc, guest_pb2, guest_pb2_grpc - -DEFAULT_ATESPACE = "ate-env" -DEFAULT_TEMPLATE = "default-template" -SIGKILL = 9 - - -@dataclass -class FakeProc: - env_id: str - argv: list[str] - stdout: bytes = b"" - stderr: bytes = b"" - exit_code: int | None = None # None while running - killed: bool = False - timeout_s: float | None = None # guest-side deadline the caller asked for - started_at: float = field(default_factory=time.monotonic) - - @property - def state(self) -> int: - return int(ProcessState.RUNNING if self.exit_code is None else ProcessState.EXITED) - - def to_pb(self, process_id: str) -> guest_pb2.Process: - return guest_pb2.Process( - process_id=process_id, - command=self.argv, - pid=4242, - state=self.state, - # proto3 int32 has no unset value: the real guest serializes 0 for - # a running process too, and `state` is the discriminator. - exit_code=self.exit_code or 0, - ) - - -@dataclass -class FakeAteEnv: - """State shared by the fake services, plus the knobs tests script.""" - - envs: dict[str, dict] = field(default_factory=dict) - files: dict[tuple[str, str], bytes] = field(default_factory=dict) - modes: dict[tuple[str, str], int] = field(default_factory=dict) - starts: list[dict] = field(default_factory=list) - procs: dict[str, FakeProc] = field(default_factory=dict) - create_code: grpc.StatusCode | None = None - get_code: grpc.StatusCode | None = None - delete_code: grpc.StatusCode | None = None - ready_after_probes: int = 0 # fail this many `true` probes with UNAVAILABLE first - hang: set[str] = field(default_factory=set) # shell commands whose process never exits - _probes: int = 0 - _seq: int = 0 - - def env_pb(self, env_id: str) -> env_pb2.Environment: - e = self.envs[env_id] - return env_pb2.Environment( - id=env_id, - atespace=e["atespace"], - template=env_pb2.Template(name=e["template"], atespace=e["template_atespace"]), - status=e["status"], - ) - - def set_status(self, env_id: str, status: EnvironmentStatus) -> None: - self.envs[env_id]["status"] = int(status) - - -async def _env_from_metadata(state: FakeAteEnv, context) -> str: - md = dict(context.invocation_metadata()) - env_id = md.get("x-env-id", "") - if not env_id: - await context.abort(grpc.StatusCode.INVALID_ARGUMENT, "missing x-env-id metadata") - if env_id not in state.envs: - await context.abort(grpc.StatusCode.NOT_FOUND, f"environment {env_id!r} not found") - return env_id - - -class FakeEnvironmentService(env_pb2_grpc.EnvironmentServiceServicer): - def __init__(self, state: FakeAteEnv) -> None: - self.state = state - - async def CreateEnvironment(self, request, context): - s = self.state - if s.create_code is not None: - await context.abort(s.create_code, "create refused") - if request.id in s.envs: - await context.abort(grpc.StatusCode.ALREADY_EXISTS, "exists") - base = request.template.name or DEFAULT_TEMPLATE - # Like ate-env-api: an image on the request yields a template derived - # from the base, named after the base and the image's digest prefix. - template = base - if request.rootfs_image: - digest = request.rootfs_image.split("@sha256:", 1)[-1][:12] - template = f"{base}-{digest}" - s.envs[request.id] = { - "atespace": request.atespace or DEFAULT_ATESPACE, - "template": template, - "base_template": base, - "rootfs_image": request.rootfs_image, - "template_atespace": request.template.atespace or DEFAULT_ATESPACE, - "status": int(EnvironmentStatus.RUNNING), - } - return env_pb2.CreateEnvironmentResponse(environment=s.env_pb(request.id)) - - async def GetEnvironment(self, request, context): - s = self.state - if s.get_code is not None: - await context.abort(s.get_code, "get refused") - if request.id not in s.envs: - await context.abort(grpc.StatusCode.NOT_FOUND, "no such env") - return env_pb2.GetEnvironmentResponse(environment=s.env_pb(request.id)) - - async def SuspendEnvironment(self, request, context): - if request.id not in self.state.envs: - await context.abort(grpc.StatusCode.NOT_FOUND, "no such env") - self.state.set_status(request.id, EnvironmentStatus.SUSPENDED) - return env_pb2.SuspendEnvironmentResponse() - - async def DeleteEnvironment(self, request, context): - s = self.state - if s.delete_code is not None: - await context.abort(s.delete_code, "delete refused") - if s.envs.pop(request.id, None) is None: - await context.abort(grpc.StatusCode.NOT_FOUND, "no such env") - return env_pb2.DeleteEnvironmentResponse() - - -class FakeProcessService(guest_pb2_grpc.ProcessServiceServicer): - def __init__(self, state: FakeAteEnv) -> None: - self.state = state - - async def StartProcess(self, request, context): - s = self.state - env_id = await _env_from_metadata(s, context) - argv = list(request.command) - shell_cmd = argv[-1] if argv else "" - timeout_s = request.timeout.ToTimedelta().total_seconds() if request.HasField("timeout") else None - s.starts.append( - { - "env_id": env_id, - "command": argv, - "cwd": request.cwd, - "env": dict(request.env), - "timeout": timeout_s, - } - ) - if shell_cmd == "true" and s._probes < s.ready_after_probes: - s._probes += 1 - await context.abort(grpc.StatusCode.UNAVAILABLE, "guest not up yet") - s._seq += 1 - pid = f"p{s._seq}" - proc = FakeProc(env_id=env_id, argv=argv, timeout_s=timeout_s) - if shell_cmd in s.hang: - proc.exit_code = None - elif "exit 7" in shell_cmd: - proc.stderr, proc.exit_code = b"boom", 7 - else: - proc.stdout, proc.exit_code = b"ok", 0 - s.procs[pid] = proc - return proc.to_pb(pid) - - async def _lookup(self, process_id: str, context) -> FakeProc: - await _env_from_metadata(self.state, context) - proc = self.state.procs.get(process_id) - if proc is None: - await context.abort(grpc.StatusCode.NOT_FOUND, f"process {process_id!r} not found") - return proc - - async def GetProcess(self, request, context): - proc = await self._lookup(request.process_id, context) - return proc.to_pb(request.process_id) - - async def StreamProcessOutput(self, request, context): - proc = await self._lookup(request.process_id, context) - out = proc.stdout[request.stdout_offset :] - err = proc.stderr[request.stderr_offset :] - if out: - yield guest_pb2.ProcessOutput(stdout=out) - if err: - yield guest_pb2.ProcessOutput(stderr=err) - if request.follow: - while proc.exit_code is None: - if proc.timeout_s is not None and time.monotonic() - proc.started_at >= proc.timeout_s: - # The guest's watchdog: SIGKILL to the group at the deadline. - proc.exit_code, proc.killed = 128 + SIGKILL, True - break - await asyncio.sleep(0.01) - if proc.exit_code is not None: - yield guest_pb2.ProcessOutput(exit=proc.to_pb(request.process_id)) - - async def WriteProcessInput(self, request_iterator, context): - await _env_from_metadata(self.state, context) - n = 0 - async for req in request_iterator: - n += len(req.data) - return guest_pb2.WriteProcessInputResponse(bytes_written=n) - - async def SignalProcess(self, request, context): - proc = await self._lookup(request.process_id, context) - if proc.exit_code is None and request.signal in (SIGKILL, 15): - proc.exit_code, proc.killed = 128 + int(request.signal), True - return proc.to_pb(request.process_id) - - -class FakeFileSystemService(guest_pb2_grpc.FileSystemServiceServicer): - def __init__(self, state: FakeAteEnv) -> None: - self.state = state - - async def ReadFile(self, request, context): - env_id = await _env_from_metadata(self.state, context) - data = self.state.files.get((env_id, request.path)) - if data is None: - await context.abort(grpc.StatusCode.NOT_FOUND, f"{request.path}: no such file") - # Small chunks so the client's streaming reassembly is exercised. - for i in range(0, len(data), 4): - yield guest_pb2.ReadFileResponse(chunk=data[i : i + 4]) - - async def WriteFile(self, request_iterator, context): - env_id = await _env_from_metadata(self.state, context) - path, mode, buf = None, 0, bytearray() - async for req in request_iterator: - if path is None: - path, mode = req.path, req.mode - buf.extend(req.chunk) - if not path: - await context.abort(grpc.StatusCode.INVALID_ARGUMENT, "path is required") - self.state.files[(env_id, path)] = bytes(buf) - self.state.modes[(env_id, path)] = mode - return guest_pb2.WriteFileResponse(bytes_written=len(buf)) - - -@contextlib.asynccontextmanager -async def serve(state: FakeAteEnv | None = None): - """Run the fake ate-env-api on a loopback port; yields ``(target, state)``.""" - state = state or FakeAteEnv() - server = grpc.aio.server() - env_pb2_grpc.add_EnvironmentServiceServicer_to_server(FakeEnvironmentService(state), server) - guest_pb2_grpc.add_ProcessServiceServicer_to_server(FakeProcessService(state), server) - guest_pb2_grpc.add_FileSystemServiceServicer_to_server(FakeFileSystemService(state), server) - port = server.add_insecure_port("127.0.0.1:0") - await server.start() - try: - yield f"127.0.0.1:{port}", state - finally: - await server.stop(None) diff --git a/integrations/nemo-gym/tests/test_e2e.py b/integrations/nemo-gym/tests/test_e2e.py deleted file mode 100644 index d5605366..00000000 --- a/integrations/nemo-gym/tests/test_e2e.py +++ /dev/null @@ -1,88 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""End-to-end test against a live ate-env deployment. - -Skipped unless SUBSTRATE_E2E_API_URL points at a reachable ate-env-api (e.g. -`kubectl port-forward -n ate-env svc/ate-env-api 7777:7777` then -`SUBSTRATE_E2E_API_URL=127.0.0.1:7777 pytest tests/test_e2e.py`; an http:// -URL works too). SUBSTRATE_E2E_TEMPLATE picks the ActorTemplate (default -`default-template`, ate-env's own default). -Exercises the full provider lifecycle on a real Substrate actor: -create → exec (incl. nonzero exit, env/cwd) → file round-trip → status → -close → status STOPPED. -""" - -from __future__ import annotations - -import asyncio -import os -from pathlib import Path - -import pytest - -from nemo_gym_substrate._compat import SandboxSpec, SandboxStatus -from nemo_gym_substrate.provider import SubstrateSandboxProvider - -API_URL = os.environ.get("SUBSTRATE_E2E_API_URL") -TEMPLATE = os.environ.get("SUBSTRATE_E2E_TEMPLATE", "default-template") - -pytestmark = pytest.mark.skipif( - not API_URL, reason="SUBSTRATE_E2E_API_URL not set; e2e needs a live ate-env-api" -) - - -def test_full_lifecycle(tmp_path: Path) -> None: - async def scenario() -> None: - provider = SubstrateSandboxProvider({"connection": {"api_url": API_URL}}) - handle = None - try: - handle = await provider.create( - SandboxSpec( - files={"/task/hello.txt": "hello from nemo-gym-substrate e2e\n"}, - workdir="/task", - env={"E2E_MARK": "42"}, - provider_options={"template": TEMPLATE}, - ready_timeout_s=180, - ) - ) - - # exec: env + cwd honored, output captured. - result = await provider.exec(handle, "echo -n $E2E_MARK; pwd") - assert result.return_code == 0 - assert "42" in (result.stdout or "") - assert "/task" in (result.stdout or "") - - # nonzero exit is a result, not an exception. - result = await provider.exec(handle, "exit 7") - assert result.return_code == 7 - - # file round-trip: the seeded file comes back, an uploaded file too. - out = tmp_path / "hello.txt" - await provider.download_file(handle, "/task/hello.txt", out) - assert out.read_text().startswith("hello from nemo-gym-substrate") - src = tmp_path / "up.bin" - src.write_bytes(bytes(range(64))) - await provider.upload_file(handle, src, "/task/up.bin") - result = await provider.exec(handle, "wc -c < /task/up.bin") - assert (result.stdout or "").strip() == "64" - - assert await provider.status(handle) == SandboxStatus.RUNNING - finally: - if handle is not None: - await provider.close(handle) - assert await provider.status(handle) == SandboxStatus.STOPPED - await provider.aclose() - - asyncio.run(scenario()) diff --git a/integrations/nemo-gym/tests/test_provider.py b/integrations/nemo-gym/tests/test_provider.py deleted file mode 100644 index 134d0968..00000000 --- a/integrations/nemo-gym/tests/test_provider.py +++ /dev/null @@ -1,589 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Contract tests for the substrate NeMo Gym provider. - -Runs against an in-process fake of ate-env-api (tests/fakes.py) reached -through the real ``ate_env`` client over loopback gRPC, asserting the rules -from the adding-a-provider contract: create returns only when the sandbox -executes commands, exec never raises on nonzero exits, close is -cleanup-safe, and provider_options are validated strictly. -""" - -from __future__ import annotations - -import asyncio -import contextlib -import logging -import time -from pathlib import Path - -import grpc -import pytest -from ate_env import Client, EnvironmentStatus, NotFoundError - -from nemo_gym_substrate import provider as provider_mod -from nemo_gym_substrate._compat import ( - SandboxCreateError, - SandboxCreateVerificationError, - SandboxEndpoint, - SandboxSpec, - SandboxStatus, -) -from nemo_gym_substrate.provider import SubstrateSandboxProvider - -from .fakes import FakeAteEnv, serve - - -@contextlib.asynccontextmanager -async def stack( - create: dict | None = None, - *, - state: FakeAteEnv | None = None, - connection: dict | None = None, -): - """A provider dialing a fake ate-env-api; yields ``(provider, fake)``.""" - async with serve(state) as (target, fake): - config = { - "connection": {"api_url": target, "request_timeout_s": 5.0, **(connection or {})}, - "create": { - "ready_poll_interval_s": 0.01, - "image_templates": {"img:1": "tpl-img1"}, - **(create or {}), - }, - } - provider = SubstrateSandboxProvider(config) - try: - yield provider, fake - finally: - await provider.aclose() - - -def run(coro): - return asyncio.run(coro) - - -def _probe_count(fake: FakeAteEnv) -> int: - return sum(1 for s in fake.starts if s["command"][-1:] == ["true"]) - - -def test_create_returns_ready_handle_and_seeds_files(): - async def t(): - async with stack() as (provider, fake): - handle = await provider.create( - SandboxSpec(files={"/task/input.json": '{"n": 1}'}, workdir="/task") - ) - assert handle.provider_name == "substrate" - assert handle.sandbox_id.startswith("gym-") - env = fake.envs[handle.sandbox_id] - assert (env["template"], env["template_atespace"], env["atespace"]) == ( - "default-template", - "ate-env", - "ate-env", - ) - assert fake.files[(handle.sandbox_id, "/task/input.json")] == b'{"n": 1}' - assert fake.modes[(handle.sandbox_id, "/task/input.json")] == 0o644 - - run(t()) - - -def test_create_waits_for_readiness(): - async def t(): - async with stack() as (provider, fake): - fake.ready_after_probes = 2 - handle = await provider.create(SandboxSpec()) - assert handle.sandbox_id in fake.envs - assert _probe_count(fake) == 3 # two UNAVAILABLE + one success - - run(t()) - - -def test_create_failure_raises_create_error(): - async def t(): - async with stack() as (provider, fake): - fake.create_code = grpc.StatusCode.RESOURCE_EXHAUSTED - with pytest.raises(SandboxCreateError, match="RESOURCE_EXHAUSTED"): - await provider.create(SandboxSpec()) - assert fake.envs == {} - - run(t()) - - -def test_create_readiness_timeout_cleans_up(): - async def t(): - async with stack() as (provider, fake): - fake.ready_after_probes = 10_000 - with pytest.raises(SandboxCreateVerificationError, match="UNAVAILABLE"): - await provider.create(SandboxSpec(ready_timeout_s=0.05)) - assert fake.envs == {} # the half-created env was deleted - - run(t()) - - -def test_image_resolves_via_mapping_or_derives(): - async def t(): - async with stack() as (provider, fake): - # A mapped image names a registered template; no rootfs_image is sent. - handle = await provider.create(SandboxSpec(image="img:1")) - assert fake.envs[handle.sandbox_id]["template"] == "tpl-img1" - assert fake.envs[handle.sandbox_id]["rootfs_image"] == "" - # Any other image goes to the API as rootfs_image on the base - # template, which derives (and later reuses) a template for it. - image = "docker.io/library/python@sha256:" + "ab" * 32 - handle = await provider.create(SandboxSpec(image=image)) - env = fake.envs[handle.sandbox_id] - assert env["rootfs_image"] == image - assert env["base_template"] == "default-template" - assert env["template"] == "default-template-" + "ab" * 6 - - run(t()) - - -def test_exec_nonzero_exit_returns_result_not_raise(): - async def t(): - async with stack() as (provider, _): - handle = await provider.create(SandboxSpec()) - result = await provider.exec(handle, "exit 7") - assert (result.return_code, result.stderr, result.stdout) == (7, "boom", "") - - run(t()) - - -def test_exec_merges_spec_env_and_cwd(): - async def t(): - async with stack() as (provider, fake): - handle = await provider.create(SandboxSpec(workdir="/w", env={"A": "1", "B": "spec"})) - await provider.exec(handle, "echo hi", env={"B": "call"}) - sent = fake.starts[-1] - assert sent["command"] == ["sh", "-c", "echo hi"] - assert sent["env"] == {"A": "1", "B": "call"} - assert sent["cwd"] == "/w" - await provider.exec(handle, "echo hi", cwd="/other") - assert fake.starts[-1]["cwd"] == "/other" - - run(t()) - - -@pytest.mark.parametrize( - ("timeout_s", "guest_deadline"), - [(3, 3.0), (0.5, 1.0), (0, 1.0), (2.2, 2.2)], -) -def test_exec_timeout_is_enforced_guest_side(timeout_s, guest_deadline): - # The deadline travels as the process's own timeout, not as a `timeout` - # wrapper, so the image needs no coreutils. Zero and sub-second values - # floor to 1 s: a zero duration would select the guest default instead. - async def t(): - async with stack() as (provider, fake): - handle = await provider.create(SandboxSpec()) - await provider.exec(handle, "sleep 100", timeout_s=timeout_s) - assert fake.starts[-1]["command"] == ["sh", "-c", "sleep 100"] - assert fake.starts[-1]["timeout"] == pytest.approx(guest_deadline) - - run(t()) - - -def test_exec_without_timeout_has_no_guest_deadline(): - async def t(): - async with stack() as (provider, fake): - handle = await provider.create(SandboxSpec()) - await provider.exec(handle, "sleep 1") - assert fake.starts[-1]["command"] == ["sh", "-c", "sleep 1"] - assert fake.starts[-1]["timeout"] is None - - run(t()) - - -def test_exec_rejects_user(): - async def t(): - async with stack() as (provider, _): - handle = await provider.create(SandboxSpec()) - with pytest.raises(ValueError, match="user"): - await provider.exec(handle, "id", user="root") - - run(t()) - - -def test_exec_client_timeout_returns_sentinel_and_kills(monkeypatch): - monkeypatch.setattr(provider_mod, "_EXEC_GRACE_S", 0.05) - - async def t(): - async with stack() as (provider, fake): - handle = await provider.create(SandboxSpec()) - fake.hang.add("sleepy") - result = await provider.exec(handle, "sleepy", timeout_s=0.01) - assert result.return_code == -1 - assert result.stdout is None - assert "timed out" in (result.stderr or "") - hung = [p for p in fake.procs.values() if p.argv[-1] == "sleepy"] - assert hung and hung[0].killed - - run(t()) - - -def test_exec_guest_deadline_reports_timeout(monkeypatch): - # The guest kills the process group at the deadline and reports a plain - # SIGKILL; the result keeps that exit code and says why on stderr. - monkeypatch.setattr(provider_mod, "_MIN_GUEST_TIMEOUT_S", 0.05) - - async def t(): - async with stack() as (provider, fake): - handle = await provider.create(SandboxSpec()) - fake.hang.add("sleepy") - result = await provider.exec(handle, "sleepy", timeout_s=0.05) - assert result.return_code == 137 - assert "timed out" in result.stderr and "guest" in result.stderr - assert fake.starts[-1]["timeout"] == pytest.approx(0.05) - - run(t()) - - -def test_file_roundtrip(tmp_path: Path): - async def t(): - async with stack() as (provider, _): - handle = await provider.create(SandboxSpec()) - src = tmp_path / "up.bin" - src.write_bytes(b"\x00\x01payload-longer-than-one-chunk") - await provider.upload_file(handle, src, "/data/up.bin") - dst = tmp_path / "down" / "up.bin" - await provider.download_file(handle, "/data/up.bin", dst) - assert dst.read_bytes() == b"\x00\x01payload-longer-than-one-chunk" - - run(t()) - - -def test_download_missing_file_raises(tmp_path: Path): - async def t(): - async with stack() as (provider, _): - handle = await provider.create(SandboxSpec()) - with pytest.raises(NotFoundError): - await provider.download_file(handle, "/nope", tmp_path / "out") - - run(t()) - - -def test_status_running_stopped_unknown(): - async def t(): - async with stack() as (provider, fake): - handle = await provider.create(SandboxSpec()) - assert await provider.status(handle) == SandboxStatus.RUNNING - fake.get_code = grpc.StatusCode.UNAVAILABLE - assert await provider.status(handle) == SandboxStatus.UNKNOWN - fake.get_code = None - await provider.close(handle) - assert await provider.status(handle) == SandboxStatus.STOPPED - - run(t()) - - -@pytest.mark.parametrize( - ("env_status", "expected"), - [ - (EnvironmentStatus.RESUMING, SandboxStatus.STARTING), - (EnvironmentStatus.SUSPENDED, SandboxStatus.RUNNING), # parked is transparent - (EnvironmentStatus.PAUSED, SandboxStatus.RUNNING), - (EnvironmentStatus.CRASHED, SandboxStatus.ERROR), - (EnvironmentStatus.DELETING, SandboxStatus.STOPPED), - (EnvironmentStatus.UNSPECIFIED, SandboxStatus.UNKNOWN), - ], -) -def test_status_maps_env_lifecycle(env_status, expected): - async def t(): - async with stack() as (provider, fake): - handle = await provider.create(SandboxSpec()) - fake.set_status(handle.sandbox_id, env_status) - assert await provider.status(handle) == expected - - run(t()) - - -def test_status_does_not_touch_the_guest(): - async def t(): - async with stack() as (provider, fake): - handle = await provider.create(SandboxSpec()) - before = len(fake.starts) - await provider.status(handle) - assert len(fake.starts) == before # no probe exec, so no wake-up - - run(t()) - - -def test_close_is_idempotent(): - async def t(): - async with stack() as (provider, fake): - handle = await provider.create(SandboxSpec()) - await provider.close(handle) - await provider.close(handle) # second delete hits NOT_FOUND; must not raise - assert handle.sandbox_id not in fake.envs - - run(t()) - - -def test_close_logs_but_does_not_raise_on_delete_failure(caplog): - async def t(): - async with stack() as (provider, fake): - handle = await provider.create(SandboxSpec()) - fake.delete_code = grpc.StatusCode.INTERNAL - with caplog.at_level(logging.WARNING, logger="nemo_gym_substrate.provider"): - await provider.close(handle) # best-effort: never raises - assert any( - "deleting substrate env" in r.message and "INTERNAL" in r.message - for r in caplog.records - ) - - run(t()) - - -def test_provider_options_override_template_and_atespace(): - async def t(): - async with stack() as (provider, fake): - handle = await provider.create( - SandboxSpec(provider_options={"template": "tpl-x", "atespace": "as-x"}) - ) - env = fake.envs[handle.sandbox_id] - assert (env["template"], env["template_atespace"], env["atespace"]) == ( - "tpl-x", - "as-x", - "as-x", - ) - # The handle remembers its atespace, so later calls route there. - assert await provider.status(handle) == SandboxStatus.RUNNING - await provider.close(handle) - assert handle.sandbox_id not in fake.envs - - run(t()) - - -def test_namespace_is_accepted_as_alias_for_atespace(): - async def t(): - async with stack({"namespace": "legacy"}) as (provider, fake): - handle = await provider.create(SandboxSpec(provider_options={"namespace": "per-sb"})) - assert fake.envs[handle.sandbox_id]["atespace"] == "per-sb" - handle2 = await provider.create(SandboxSpec()) - assert fake.envs[handle2.sandbox_id]["atespace"] == "legacy" - - run(t()) - - -def test_atespace_and_namespace_together_rejected(): - with pytest.raises(ValueError, match="either atespace or namespace"): - SubstrateSandboxProvider({"create": {"atespace": "a", "namespace": "b"}}) - - -def test_unknown_provider_option_rejected(): - async def t(): - async with stack() as (provider, _): - with pytest.raises(ValueError, match="unknown option"): - await provider.create(SandboxSpec(provider_options={"tempalte": "oops"})) - - run(t()) - - -def test_unknown_config_key_rejected(): - with pytest.raises(ValueError, match="unknown option"): - SubstrateSandboxProvider({"connection": {"api_urll": "typo"}}) - - -def test_create_timeout_on_unreachable_api(): - async def t(): - # A closed port refuses at once (UNAVAILABLE); either way it must not hang. - provider = SubstrateSandboxProvider( - {"connection": {"api_url": "127.0.0.1:9", "request_timeout_s": 0.5}} - ) - try: - with pytest.raises(SandboxCreateError): - await provider.create(SandboxSpec()) - finally: - await provider.aclose() - - run(t()) - - -def test_aclose_is_safe_to_repeat_and_drops_the_client(): - async def t(): - async with stack() as (provider, _): - await provider.create(SandboxSpec()) - assert provider._client is not None # noqa: SLF001 - test seam - await provider.aclose() - await provider.aclose() - assert provider._client is None # noqa: SLF001 - - run(t()) - - -def test_caller_owned_client_is_not_closed(): - async def t(): - async with serve() as (target, _): - client = Client(target) - try: - provider = SubstrateSandboxProvider(client=client) - handle = await provider.create(SandboxSpec()) - await provider.aclose() - # Still usable: aclose() left the caller's client open. - info = await client.get(handle.sandbox_id) - assert info.id == handle.sandbox_id - finally: - await client.close() - - run(t()) - - -def test_https_api_url_uses_a_tls_channel(): - async def t(): - provider = SubstrateSandboxProvider( - {"connection": {"api_url": "https://ate-env.example:443"}} - ) - try: - assert provider._ate() is not None # noqa: SLF001 - opens the channel, never dials - assert provider._channel is not None # noqa: SLF001 - finally: - await provider.aclose() - assert provider._channel is None # noqa: SLF001 - - run(t()) - - -# -- endpoint(): services inside the sandbox --------------------------------- - - -def test_endpoint_routes_a_declared_port_through_the_api(): - async def t(): - async with stack() as (provider, fake): - handle = await provider.create(SandboxSpec(ports=(6000,))) - assert handle.raw["ports"] == [6000] - probes = _probe_count(fake) - ep = await provider.endpoint(handle, 6000) - assert isinstance(ep, SandboxEndpoint) - # The port proxy shares ate-env-api's listener. - assert ep.endpoint == provider._connection.endpoint_base() - assert ep.endpoint.startswith("http://") and ep.endpoint.endswith( - provider._connection.api_url.split("://")[-1] - ) - assert ep.headers == { - "X-Env-Id": handle.sandbox_id, - "X-Env-Atespace": "ate-env", - "X-Env-Port": "6000", - } - # Nothing touches the guest: the first request through the proxy - # wakes a parked sandbox, and the proxy waits for it. - assert _probe_count(fake) == probes - - run(t()) - - -def test_endpoint_rejects_an_undeclared_port(): - async def t(): - async with stack() as (provider, _fake): - handle = await provider.create(SandboxSpec(ports=(6000,))) - with pytest.raises(ValueError, match="not declared in SandboxSpec.ports"): - await provider.endpoint(handle, 8080) - bare = await provider.create(SandboxSpec()) - with pytest.raises(ValueError, match=r"declared ports: \[\]"): - await provider.endpoint(bare, 6000) - - run(t()) - - -def test_endpoint_uses_the_configured_url_and_the_sandbox_atespace(): - async def t(): - async with stack( - {"atespace": "team-a"}, - connection={"endpoint_url": "https://gym-proxy.example.com/"}, - ) as (provider, _fake): - handle = await provider.create(SandboxSpec(ports=[6000])) - ep = await provider.endpoint(handle, 6000) - assert ep.endpoint == "https://gym-proxy.example.com" - assert ep.headers["X-Env-Atespace"] == "team-a" - - run(t()) - - -def test_endpoint_raises_for_a_crashed_sandbox(): - async def t(): - async with stack() as (provider, fake): - handle = await provider.create(SandboxSpec(ports=(6000,))) - fake.set_status(handle.sandbox_id, EnvironmentStatus.CRASHED) - with pytest.raises(RuntimeError, match="crashed"): - await provider.endpoint(handle, 6000) - - run(t()) - - -def test_endpoint_does_not_wake_a_parked_sandbox(): - async def t(): - async with stack() as (provider, fake): - handle = await provider.create(SandboxSpec(ports=(6000,))) - fake.set_status(handle.sandbox_id, EnvironmentStatus.SUSPENDED) - before = len(fake.starts) - ep = await provider.endpoint(handle, 6000) - assert ep.headers["X-Env-Port"] == "6000" - assert len(fake.starts) == before # no probe exec, so no wake-up - - run(t()) - - -def test_create_rejects_invalid_ports(): - async def t(): - async with stack() as (provider, fake): - cases = [ - ((0,), ValueError, "not a TCP port"), - ((70000,), ValueError, "not a TCP port"), - ((80,), ValueError, "guest"), - ((True,), TypeError, "integers"), # bool is an int subclass - (("6000",), TypeError, "integers"), - ] - for ports, error, why in cases: - before = len(fake.envs) - with pytest.raises(error, match=why): - await provider.create(SandboxSpec(ports=ports)) - assert len(fake.envs) == before, "a bad spec must fail before the actor starts" - - run(t()) - - -def test_endpoint_rejects_a_bool_port(): - async def t(): - async with stack() as (provider, _fake): - handle = await provider.create(SandboxSpec(ports=(1,))) - with pytest.raises(ValueError, match="not declared"): - await provider.endpoint(handle, True) # True == 1 in Python - - run(t()) - - -def test_endpoint_stops_at_once_for_a_deleted_sandbox(): - async def t(): - async with stack({"ready_timeout_s": 30}) as (provider, _fake): - handle = await provider.create(SandboxSpec(ports=(6000,))) - await provider.close(handle) - started = time.monotonic() - with pytest.raises(RuntimeError, match="does not exist"): - await provider.endpoint(handle, 6000) - assert time.monotonic() - started < 5, "endpoint() waited out the ready timeout" - - run(t()) - - -def test_endpoint_base_defaults_to_the_api_url(): - cfg = provider_mod.ConnectionConfig.from_mapping({"api_url": "ate-env-api.ate-env:7777"}) - assert cfg.endpoint_base() == "http://ate-env-api.ate-env:7777" - cfg = provider_mod.ConnectionConfig.from_mapping({"api_url": "https://env.example.com/"}) - assert cfg.endpoint_base() == "https://env.example.com" - # A bare host:port endpoint_url gets a scheme too; SandboxEndpoint rejects it otherwise. - cfg = provider_mod.ConnectionConfig.from_mapping( - {"api_url": "ate-env-api.ate-env:7777", "endpoint_url": "gym-proxy.example.com:8080"} - ) - assert cfg.endpoint_base() == "http://gym-proxy.example.com:8080" - SandboxEndpoint(endpoint=cfg.endpoint_base(), headers={}) - with pytest.raises(ValueError, match="unknown option"): - provider_mod.ConnectionConfig.from_mapping({"endpoint": "x"}) diff --git a/internal/apiservice/imagetemplate.go b/internal/apiservice/imagetemplate.go index af4d66c2..a90fde73 100644 --- a/internal/apiservice/imagetemplate.go +++ b/internal/apiservice/imagetemplate.go @@ -40,19 +40,8 @@ const defaultGuestBinary = "/ko-app/ate-env-guest" // guest's own path moves under the mount, never a shell or another binary. const guestBinaryName = "ate-env-guest" -// imageTemplateDigestLen is how many hex digits of the digest go into a -// derived template's name: enough to make collisions a non-concern and short -// enough to keep the name a valid k8s short name next to any base name. -const imageTemplateDigestLen = 12 - var digestRef = regexp.MustCompile(`@sha256:([0-9a-f]{64})$`) -// shortName is the shape substrate accepts for a resource name (k8s short -// name): lowercase alphanumerics and '-', at most 63 characters. -var shortName = regexp.MustCompile(`^[a-z0-9]([-a-z0-9]*[a-z0-9])?$`) - -const maxShortNameLen = 63 - // ImageDigest returns the hex digest of a digest-pinned image reference, or an // error when the reference is not pinned. Substrate requires pinned images in // ActorTemplates because changing the image invalidates snapshots. @@ -64,24 +53,6 @@ func ImageDigest(image string) (string, error) { return m[1], nil } -// ImageTemplateName is the name of the ActorTemplate derived from base for -// image: "-". -func ImageTemplateName(base, image string) (string, error) { - digest, err := ImageDigest(image) - if err != nil { - return "", err - } - name := base + "-" + digest[:imageTemplateDigestLen] - if len(name) > maxShortNameLen { - return "", fmt.Errorf("derived template name %q is %d characters, the limit is %d; use a shorter base template name", - name, len(name), maxShortNameLen) - } - if !shortName.MatchString(name) { - return "", fmt.Errorf("derived template name %q is not a valid resource name (lowercase alphanumerics and '-')", name) - } - return name, nil -} - // DeriveImageTemplate returns a copy of base, named name, whose container runs // image unmodified with the ate-env-guest mounted into it as a read-only image // volume at GuestMountPath. @@ -92,8 +63,7 @@ func ImageTemplateName(base, image string) (string, error) { // mount; or one that already mounts the guest as an image volume (named // GuestVolumeName), in which case only the container image changes. // Everything else carries over: worker selector, snapshots, sandbox config, -// resources, env, readiness, and any other volumes, such as extra runtime -// layers, together with the command-line flags that start them. +// resources, env, readiness, and any other volumes. func DeriveImageTemplate(base *ateapipb.ActorTemplate, name, image string) (*ateapipb.ActorTemplate, error) { if base == nil { return nil, errors.New("base template is required") @@ -142,8 +112,7 @@ func DeriveImageTemplate(base *ateapipb.ActorTemplate, name, image string) (*ate } // guestVolume returns the template's guest image volume, or nil when the -// guest is the container image. Other image volumes (extra runtime layers) -// do not count: they can be present in either kind of base. +// guest is the container image. func guestVolume(tmpl *ateapipb.ActorTemplate) *ateapipb.Volume { for _, v := range tmpl.GetVolumes() { if v.GetName() == GuestVolumeName && v.GetImage() != nil { diff --git a/internal/apiservice/imagetemplate_test.go b/internal/apiservice/imagetemplate_test.go index d7050e23..f7444eea 100644 --- a/internal/apiservice/imagetemplate_test.go +++ b/internal/apiservice/imagetemplate_test.go @@ -60,29 +60,15 @@ func TestDeriveImageTemplateKeepsResources(t *testing.T) { } } -func TestImageDigestAndTemplateName(t *testing.T) { - name, err := ImageTemplateName("default-template", testRootfsImage) - if err != nil { - t.Fatal(err) - } - if name != "default-template-0123456789ab" { - t.Errorf("name = %q", name) +func TestImageDigest(t *testing.T) { + if d, err := ImageDigest(testRootfsImage); err != nil || d != testRootfsImage[len(testRootfsImage)-64:] { + t.Errorf("ImageDigest(%q) = %q, %v", testRootfsImage, d, err) } for _, bad := range []string{"python:3.12", "python@sha256:abc", "python@sha512:" + strings.Repeat("a", 128), ""} { if _, err := ImageDigest(bad); err == nil { t.Errorf("ImageDigest(%q) accepted an unpinned reference", bad) } } - // The derived name must stay a valid resource name. - if _, err := ImageTemplateName(strings.Repeat("b", 51), testRootfsImage); err == nil { - t.Error("a 64-character derived name was accepted") - } - if name, err := ImageTemplateName(strings.Repeat("b", 50), testRootfsImage); err != nil || len(name) != 63 { - t.Errorf("a 63-character derived name should be accepted: %q, %v", name, err) - } - if _, err := ImageTemplateName("Default", testRootfsImage); err == nil { - t.Error("an uppercase base name was accepted") - } } func TestDeriveImageTemplateFromGuestImageBase(t *testing.T) { @@ -192,38 +178,3 @@ func TestDeriveImageTemplateRejects(t *testing.T) { } } } - -func TestDeriveImageTemplateKeepsRuntimeLayers(t *testing.T) { - // A guest-image base that also carries a second runtime as an image - // volume and starts it as a guest sidecar: the layer and the flags - // survive derivation, and the guest volume is still added. - layerImage := "example.com/execd@sha256:" + strings.Repeat("c", 64) - base := modeABase() - base.Volumes = []*ateapipb.Volume{{Name: "execd", Image: &ateapipb.ImageVolumeSource{Reference: layerImage}}} - base.Containers[0].VolumeMounts = []*ateapipb.VolumeMount{{Name: "execd", MountPath: "/opt/opensandbox"}} - base.Containers[0].Command = []string{"/ko-app/ate-env-guest", "-sidecar", "/opt/opensandbox/execd --port 44772", "-sidecar-readyz", "http://127.0.0.1:44772/ready"} - - got, err := DeriveImageTemplate(base, "x", testRootfsImage) - if err != nil { - t.Fatal(err) - } - if len(got.GetVolumes()) != 2 || got.GetVolumes()[0].GetName() != "execd" || got.GetVolumes()[1].GetName() != GuestVolumeName { - t.Errorf("volumes = %v, want the layer followed by the guest volume", got.GetVolumes()) - } - c := got.GetContainers()[0] - if len(c.GetVolumeMounts()) != 2 || c.GetVolumeMounts()[1].GetMountPath() != GuestMountPath { - t.Errorf("mounts = %v", c.GetVolumeMounts()) - } - want := "/ate/ko-app/ate-env-guest -sidecar /opt/opensandbox/execd --port 44772 -sidecar-readyz http://127.0.0.1:44772/ready" - if got := strings.Join(c.GetCommand(), " "); got != want { - t.Errorf("command = %q, want %q", got, want) - } - // Deriving again from the result (an already-injected base) changes only the image. - again, err := DeriveImageTemplate(got, "y", "example.com/other@sha256:"+strings.Repeat("d", 64)) - if err != nil { - t.Fatal(err) - } - if len(again.GetVolumes()) != 2 || strings.Join(again.GetContainers()[0].GetCommand(), " ") != want { - t.Errorf("second derivation altered layers or command: %v", again) - } -} diff --git a/internal/apiservice/server.go b/internal/apiservice/server.go index f3466f60..aa911b1c 100644 --- a/internal/apiservice/server.go +++ b/internal/apiservice/server.go @@ -96,24 +96,15 @@ func (s *Server) Close() {} // and to create the actor; a template.atespace that differs from a non-empty // request atespace is ignored. // -// When rootfs_image is set, the template acts as the base and the environment is -// created from a template derived from it, named -// "-" in the same atespace. -// The derived template keeps the base's settings, runs the rootfs image as -// the container image and mounts the base's guest as a read-only image volume, -// so the image runs unmodified (see ensureImageTemplate and -// DeriveImageTemplate). It is created on first use and reused by later -// creates on the same image, and the response reports it in -// Environment.template. When rootfs_image is empty the template is used as is. -// -// Errors: InvalidArgument for a missing id, a rootfs image not pinned by digest, a -// derived name that is too long or not a valid resource name, or a base that -// cannot be derived from (no containers, unpinned guest image, or a command -// that cannot be re-rooted under the guest volume). FailedPrecondition when -// the base template does not exist in the atespace, or when the derived name -// is already held by a template running a different image. Other control -// plane failures are mapped by toGRPCError. +// Errors: InvalidArgument for a missing id or for a request carrying fields +// this server does not know, such as the removed rootfs_image (field 4) that +// older clients still send; silently ignoring it would create the environment +// from the wrong template. Other control plane failures are mapped by +// toGRPCError. func (s *Server) CreateEnvironment(ctx context.Context, req *ateenvv1alpha.CreateEnvironmentRequest) (*ateenvv1alpha.CreateEnvironmentResponse, error) { + if len(req.ProtoReflect().GetUnknown()) > 0 { + return nil, status.Error(codes.InvalidArgument, "request has unknown fields; rootfs_image is no longer supported, register an image template with `ate-env manifest template --rootfs-image` and pass its name in template.name") + } if req.GetId() == "" { return nil, status.Error(codes.InvalidArgument, "id is required") } @@ -130,14 +121,6 @@ func (s *Server) CreateEnvironment(ctx context.Context, req *ateenvv1alpha.Creat atespace = DefaultAtespace } - if image := req.GetRootfsImage(); image != "" { - name, err := s.ensureImageTemplate(ctx, atespace, templateName, image) - if err != nil { - return nil, err - } - templateName = name - } - opts := ate.CreateOptions{ ID: req.GetId(), Template: templateName, @@ -160,51 +143,6 @@ func (s *Server) CreateEnvironment(ctx context.Context, req *ateenvv1alpha.Creat }, nil } -// ensureImageTemplate returns the name of the ActorTemplate that runs image -// on top of base, creating it from base on first use. The derived template -// lives in the environment's atespace, next to base. -func (s *Server) ensureImageTemplate(ctx context.Context, atespace, base, image string) (string, error) { - name, err := ImageTemplateName(base, image) - if err != nil { - return "", status.Error(codes.InvalidArgument, err.Error()) - } - existing, err := s.client.GetActorTemplate(ctx, atespace, name) - switch { - case err == nil: - if got := firstContainerImage(existing); got != image { - return "", status.Errorf(codes.FailedPrecondition, - "actor template %q exists but runs %q, not %q; delete it or use another base template", name, got, image) - } - return name, nil - case !errors.Is(err, ate.ErrNotFound): - return "", toGRPCError(err) - } - baseTmpl, err := s.client.GetActorTemplate(ctx, atespace, base) - if err != nil { - if errors.Is(err, ate.ErrNotFound) { - return "", status.Errorf(codes.FailedPrecondition, - "base actor template %q not found in atespace %q; register it first (ate-env manifest template)", base, atespace) - } - return "", toGRPCError(err) - } - derived, err := DeriveImageTemplate(baseTmpl, name, image) - if err != nil { - return "", status.Error(codes.InvalidArgument, err.Error()) - } - // A concurrent create may have won the race; it built the same template. - if _, err := s.client.CreateActorTemplate(ctx, derived); err != nil && !errors.Is(err, ate.ErrAlreadyExists) { - return "", toGRPCError(err) - } - return name, nil -} - -func firstContainerImage(t *ateapipb.ActorTemplate) string { - if len(t.GetContainers()) == 0 { - return "" - } - return t.GetContainers()[0].GetImage() -} - // GetEnvironment retrieves the status and configuration of an existing environment. func (s *Server) GetEnvironment(ctx context.Context, req *ateenvv1alpha.GetEnvironmentRequest) (*ateenvv1alpha.GetEnvironmentResponse, error) { if req.GetId() == "" { diff --git a/internal/apiservice/server_test.go b/internal/apiservice/server_test.go index 9ecccd5c..6969bc83 100644 --- a/internal/apiservice/server_test.go +++ b/internal/apiservice/server_test.go @@ -19,6 +19,7 @@ import ( "io" "math" "net" + "strings" "testing" "github.com/agent-substrate/env/guest" @@ -33,6 +34,7 @@ import ( "google.golang.org/grpc/credentials/insecure" "google.golang.org/grpc/metadata" "google.golang.org/grpc/status" + "google.golang.org/protobuf/encoding/protowire" ) type testEnv struct { @@ -224,6 +226,35 @@ func TestCreateValidation(t *testing.T) { } } +// TestCreateRejectsUnknownFields sends the removed rootfs_image field (4) the +// way an old client does. proto3 keeps it as unknown bytes instead of failing +// to unmarshal, so the server has to reject it explicitly rather than quietly +// creating the environment from the base template. +func TestCreateRejectsUnknownFields(t *testing.T) { + client, control := newTestEnv(t) + ctx := context.Background() + + req := &ateenvv1alpha.CreateEnvironmentRequest{Id: "old-client"} + var raw []byte + raw = protowire.AppendTag(raw, 4, protowire.BytesType) + raw = protowire.AppendString(raw, "ghcr.io/example/app@sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef") + req.ProtoReflect().SetUnknown(raw) + + _, err := client.CreateEnvironment(ctx, req) + if err == nil { + t.Fatal("expected error for request with rootfs_image (field 4)") + } + if status.Code(err) != codes.InvalidArgument { + t.Errorf("got code %v, want %v", status.Code(err), codes.InvalidArgument) + } + if msg := status.Convert(err).Message(); !strings.Contains(msg, "ate-env manifest template --rootfs-image") { + t.Errorf("error %q does not name the replacement command", msg) + } + if st := control.Status("old-client"); st != ateapipb.ActorState_ACTOR_STATE_UNSPECIFIED { + t.Errorf("environment was created despite the unknown field (state %v)", st) + } +} + func TestGetNotFound(t *testing.T) { client, _ := newTestEnv(t) ctx := context.Background() @@ -579,131 +610,3 @@ func TestActorStatusToEnvStatus(t *testing.T) { } } } - -const ( - testGuestImage = "example.com/ate-env-guest@sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" - testRootfsImage = "docker.io/library/python@sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef" -) - -// seedGuestTemplate registers the template `ate-env manifest template` -// writes: the guest is the container image. -func seedGuestTemplate(control *fakecontrol.Server, atespace, name string) { - control.AddTemplate(&ateapipb.ActorTemplate{ - Metadata: &ateapipb.ResourceMetadata{Atespace: atespace, Name: name}, - WorkerSelector: &ateapipb.Selector{MatchLabels: map[string]string{"workload": "default-env"}}, - Containers: []*ateapipb.Container{{ - Name: "guest", - Image: testGuestImage, - Command: []string{"/ko-app/ate-env-guest"}, - Env: []*ateapipb.EnvVar{{Name: "PORT", Value: "80"}}, - WakeupProbe: &ateapipb.WakeupProbe{HttpGet: &ateapipb.HttpGetAction{Path: "/readyz", Port: 80}}, - }}, - SnapshotConfig: &ateapipb.SnapshotConfig{StorageLocation: "gs://bucket/ate-env/"}, - SandboxConfig: &ateapipb.SandboxConfig{SandboxClass: "gvisor", ConfigName: "gvisor-default"}, - }) -} - -func TestCreateFromImageDerivesTemplate(t *testing.T) { - envClient, control := newTestEnv(t) - seedGuestTemplate(control, "ate-env", "default-template") - ctx := context.Background() - - resp, err := envClient.CreateEnvironment(ctx, &ateenvv1alpha.CreateEnvironmentRequest{Id: "img1", RootfsImage: testRootfsImage}) - if err != nil { - t.Fatal(err) - } - const want = "default-template-0123456789ab" - if got := resp.GetEnvironment().GetTemplate().GetName(); got != want { - t.Errorf("response template = %q, want %q", got, want) - } - - tmpl := control.Template("ate-env", want) - if tmpl == nil { - t.Fatalf("derived template %q was not created", want) - } - c := tmpl.GetContainers()[0] - if c.GetImage() != testRootfsImage { - t.Errorf("derived container image = %q", c.GetImage()) - } - if len(tmpl.GetVolumes()) != 1 || tmpl.GetVolumes()[0].GetImage().GetReference() != testGuestImage { - t.Errorf("derived volumes = %v, want the guest image volume", tmpl.GetVolumes()) - } - if len(c.GetVolumeMounts()) != 1 || c.GetVolumeMounts()[0].GetMountPath() != apiservice.GuestMountPath { - t.Errorf("derived mounts = %v", c.GetVolumeMounts()) - } - if len(c.GetCommand()) != 1 || c.GetCommand()[0] != "/ate/ko-app/ate-env-guest" { - t.Errorf("derived command = %v", c.GetCommand()) - } - if tmpl.GetSnapshotConfig().GetStorageLocation() != "gs://bucket/ate-env/" || tmpl.GetWorkerSelector() == nil { - t.Error("snapshots config or worker selector did not carry over") - } - - // The actor references the derived template. - got, err := envClient.GetEnvironment(ctx, &ateenvv1alpha.GetEnvironmentRequest{Id: "img1"}) - if err != nil { - t.Fatal(err) - } - if got.GetEnvironment().GetTemplate().GetName() != want { - t.Errorf("actor template = %q, want %q", got.GetEnvironment().GetTemplate().GetName(), want) - } - - // A second environment on the same image reuses the template. - if _, err := envClient.CreateEnvironment(ctx, &ateenvv1alpha.CreateEnvironmentRequest{Id: "img2", RootfsImage: testRootfsImage}); err != nil { - t.Fatal(err) - } - if n := control.TemplateCount(); n != 2 { - t.Errorf("template count = %d, want 2 (base + one derived)", n) - } -} - -func TestCreateFromImageUsesNamedBaseAndAtespace(t *testing.T) { - envClient, control := newTestEnv(t) - seedGuestTemplate(control, "team-a", "py-base") - - resp, err := envClient.CreateEnvironment(context.Background(), &ateenvv1alpha.CreateEnvironmentRequest{ - Id: "img1", - Atespace: "team-a", - Template: &ateenvv1alpha.Template{Name: "py-base"}, - RootfsImage: testRootfsImage, - }) - if err != nil { - t.Fatal(err) - } - if got := resp.GetEnvironment().GetTemplate().GetName(); got != "py-base-0123456789ab" { - t.Errorf("template = %q", got) - } - if control.Template("team-a", "py-base-0123456789ab") == nil { - t.Error("derived template not created in the environment's atespace") - } -} - -func TestCreateFromImageRejects(t *testing.T) { - envClient, control := newTestEnv(t) - seedGuestTemplate(control, "ate-env", "default-template") - ctx := context.Background() - - cases := []struct { - name string - req *ateenvv1alpha.CreateEnvironmentRequest - code codes.Code - }{ - {"unpinned image", &ateenvv1alpha.CreateEnvironmentRequest{Id: "a", RootfsImage: "python:3.12"}, codes.InvalidArgument}, - {"missing base", &ateenvv1alpha.CreateEnvironmentRequest{Id: "b", Template: &ateenvv1alpha.Template{Name: "nope"}, RootfsImage: testRootfsImage}, codes.FailedPrecondition}, - } - for _, tc := range cases { - _, err := envClient.CreateEnvironment(ctx, tc.req) - if status.Code(err) != tc.code { - t.Errorf("%s: code = %v (%v), want %v", tc.name, status.Code(err), err, tc.code) - } - } - - // A derived name already taken by a template running something else. - control.AddTemplate(&ateapipb.ActorTemplate{ - Metadata: &ateapipb.ResourceMetadata{Atespace: "ate-env", Name: "default-template-0123456789ab"}, - Containers: []*ateapipb.Container{{Name: "guest", Image: "example.com/other@sha256:" + testGuestImage[len(testGuestImage)-64:]}}, - }) - _, err := envClient.CreateEnvironment(ctx, &ateenvv1alpha.CreateEnvironmentRequest{Id: "c", RootfsImage: testRootfsImage}) - if status.Code(err) != codes.FailedPrecondition { - t.Errorf("conflicting derived template: code = %v (%v), want FailedPrecondition", status.Code(err), err) - } -} diff --git a/internal/ate/client.go b/internal/ate/client.go index 344e953e..d5079156 100644 --- a/internal/ate/client.go +++ b/internal/ate/client.go @@ -60,9 +60,6 @@ const DefaultHostSuffix = "actors.resources.substrate.ate.dev" const ( DefaultControlAddr = "api.ate-system.svc.cluster.local:443" DefaultRouterAddr = "atenet-router.ate-system.svc.cluster.local:80" - // DefaultRouterConnectAddr is the router's plaintext CONNECT listener, which - // tunnels to a port on an actor other than 80. - DefaultRouterConnectAddr = "atenet-router.ate-system.svc.cluster.local:8081" ) // DefaultAtespace is the Substrate atespace every environment actor lives in. @@ -71,9 +68,6 @@ const DefaultAtespace = "ate-env" // ErrNotFound is returned when an env, file, or directory does not exist. var ErrNotFound = errors.New("not found") -// ErrAlreadyExists is returned when a resource with the same name exists. -var ErrAlreadyExists = errors.New("already exists") - // Options configures a Client. type Options struct { // ControlAddr is the ateapi gRPC endpoint, e.g. "localhost:8080" @@ -319,39 +313,6 @@ func (c *Client) Delete(ctx context.Context, atespace, id string) error { // ref returns the ObjectRef identifying the actor backing actor id in atespace. // If atespace is empty, DefaultAtespace is used. -// GetActorTemplate reads an ActorTemplate. A missing template is ErrNotFound. -func (c *Client) GetActorTemplate(ctx context.Context, atespace, name string) (*ateapipb.ActorTemplate, error) { - if name == "" { - return nil, errors.New("ate: template name is required") - } - tmpl, err := c.control.GetActorTemplate(ctx, &ateapipb.GetActorTemplateRequest{ActorTemplate: c.ref(atespace, name)}) - if err != nil { - return nil, fmt.Errorf("ate: getting template %q: %w", name, wrapGRPCError(err)) - } - return tmpl, nil -} - -// CreateActorTemplate creates an ActorTemplate in its metadata's atespace. A -// name collision is ErrAlreadyExists, so callers racing to create the same -// derived template can treat it as success. -func (c *Client) CreateActorTemplate(ctx context.Context, tmpl *ateapipb.ActorTemplate) (*ateapipb.ActorTemplate, error) { - name := tmpl.GetMetadata().GetName() - if name == "" { - return nil, errors.New("ate: template metadata.name is required") - } - if tmpl.GetMetadata().GetAtespace() == "" { - tmpl.Metadata.Atespace = DefaultAtespace - } - created, err := c.control.CreateActorTemplate(ctx, &ateapipb.CreateActorTemplateRequest{ActorTemplate: tmpl}) - if err != nil { - if status.Code(err) == codes.AlreadyExists { - return nil, fmt.Errorf("ate: creating template %q: %w: %s", name, ErrAlreadyExists, status.Convert(err).Message()) - } - return nil, fmt.Errorf("ate: creating template %q: %w", name, wrapGRPCError(err)) - } - return created, nil -} - func (c *Client) ref(atespace, id string) *ateapipb.ObjectRef { if atespace == "" { atespace = DefaultAtespace diff --git a/internal/benchmarking/swebench/backend.go b/internal/benchmarking/swebench/backend.go deleted file mode 100644 index 9ec190ba..00000000 --- a/internal/benchmarking/swebench/backend.go +++ /dev/null @@ -1,247 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package swebench - -import ( - "bytes" - "context" - "errors" - "fmt" - "io" - "io/fs" - "sync/atomic" - "time" - - env "github.com/agent-substrate/env/clients/go" - ateenvv1alpha "github.com/agent-substrate/env/proto/ateenv/v1alpha" - "google.golang.org/grpc/codes" - "google.golang.org/grpc/status" - "google.golang.org/protobuf/types/known/durationpb" -) - -// Result is what a Run step returned. -type Result struct { - ExitCode int - Stdout string - Stderr string -} - -// Sandbox is the subset of an environment the loader drives. -type Sandbox interface { - Run(ctx context.Context, command string, timeout time.Duration) (Result, error) - WriteFile(ctx context.Context, path string, content []byte, mode fs.FileMode) error - Suspend(ctx context.Context) error - Delete(ctx context.Context) error -} - -// Backend creates sandboxes. The loader is written against this interface so -// tests run it against an in-memory fake and a second SDK or backend can be -// compared with the same workload. -type Backend interface { - Create(ctx context.Context, id, image string) (Sandbox, error) - // Delete removes an environment by name without a Sandbox handle, for - // one whose create call failed on the client side but may have - // succeeded on the server. - Delete(ctx context.Context, id string) error -} - -// EnvBackend is the Backend over the ate-env API through the Go SDK. -// -// It holds several connections and hands them out round-robin per -// environment: a Kubernetes Service balances per connection, so one -// connection would put every task on one API replica, and gRPC caps the -// concurrent streams on a connection. -type EnvBackend struct { - clients []*env.Client - next atomic.Uint64 - atespace string - baseTemplate string -} - -// NewEnvBackend dials the ate-env API over connections connections (at -// least one). Environments are created in atespace from baseTemplate with -// the task image as rootfs; the API derives and bakes one template per -// distinct image on first use. -func NewEnvBackend(endpoint, atespace, baseTemplate string, connections int) (*EnvBackend, error) { - if connections < 1 { - connections = 1 - } - b := &EnvBackend{atespace: atespace, baseTemplate: baseTemplate} - for i := 0; i < connections; i++ { - c, err := env.NewClient(env.ClientOptions{Endpoint: endpoint}) - if err != nil { - b.Close() - return nil, err - } - b.clients = append(b.clients, c) - } - return b, nil -} - -// Close releases the API connections. -func (b *EnvBackend) Close() error { - var first error - for _, c := range b.clients { - if err := c.Close(); err != nil && first == nil { - first = err - } - } - return first -} - -// pick returns the next connection round-robin. -func (b *EnvBackend) pick() *env.Client { - return b.clients[int(b.next.Add(1)-1)%len(b.clients)] -} - -// Delete implements Backend. -func (b *EnvBackend) Delete(ctx context.Context, id string) error { - return b.pick().Delete(ctx, b.atespace, id) -} - -// Create implements Backend. The environment keeps the connection it was -// created on for the rest of its life. -func (b *EnvBackend) Create(ctx context.Context, id, image string) (Sandbox, error) { - e, err := b.pick().Create(ctx, &ateenvv1alpha.CreateEnvironmentRequest{ - Id: id, - Atespace: b.atespace, - Template: &ateenvv1alpha.Template{Name: b.baseTemplate, Atespace: b.atespace}, - RootfsImage: image, - }) - if err != nil { - return nil, fmt.Errorf("create %s: %w", id, err) - } - return envSandbox{e: e}, nil -} - -type envSandbox struct { - e *env.Env -} - -// Polling bounds for a running process: short at first so quick commands -// return quickly, backing off so a long test run does not hammer the API. -// Variables so tests can shorten them. -var ( - pollMin = 50 * time.Millisecond - pollMax = 2 * time.Second -) - -const ( - // pollRetries is how many consecutive transient poll errors are - // tolerated before the step fails; a wake or an API pod restart costs a - // few. - pollRetries = 10 - // killTimeout bounds the kill of a process Run abandons. - killTimeout = 10 * time.Second -) - -// Run starts the command, polls until it exits and then reads its output. -// -// The SDK's streaming Shell holds one request open through the router for -// the whole command, and the router's route timeout (five minutes by -// default) cuts it off; SWE-bench test runs take up to twelve minutes. -// Start, poll and read are each short calls, so no single request can -// outlive it. This belongs in the SDK as an option of Env.Run, since every -// SDK user with long commands meets the same timeout; it lives here until -// that change lands so the benchmark does not wait on it. -// -// If Run returns before the process has exited, it kills the process, so a -// retry of the same command never runs beside the first copy. -func (s envSandbox) Run(ctx context.Context, command string, timeout time.Duration) (res Result, err error) { - start := &ateenvv1alpha.StartProcessRequest{Command: []string{"sh", "-c", command}} - if timeout > 0 { - start.Timeout = durationpb.New(timeout) - } - proc, err := s.e.StartProcess(ctx, start) - if err != nil { - return Result{}, err - } - exited := false - defer func() { - if exited { - return - } - kctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), killTimeout) - defer cancel() - if _, kerr := proc.Kill(kctx); kerr != nil && !errors.Is(kerr, env.ErrProcessExited) && err != nil { - err = fmt.Errorf("%w (and killing the process failed: %v)", err, kerr) - } - }() - - wait := pollMin - failures := 0 - for { - select { - case <-ctx.Done(): - return Result{}, ctx.Err() - case <-time.After(wait): - } - p, perr := s.e.FindProcess(ctx, proc.ID()) - switch { - case perr == nil: - failures = 0 - if p.GetState() == ateenvv1alpha.ProcessState_PROCESS_STATE_EXITED { - exited = true - return s.collect(ctx, proc, int(p.GetExitCode())) - } - case status.Code(perr) == codes.Unavailable || status.Code(perr) == codes.DeadlineExceeded: - failures++ - if failures > pollRetries { - return Result{}, fmt.Errorf("poll %s: %w", proc.ID(), perr) - } - default: - return Result{}, perr - } - if wait < pollMax { - wait = min(wait*2, pollMax) - } - } -} - -// collect reads the exited process's buffered output in one non-following -// call. -func (s envSandbox) collect(ctx context.Context, proc *env.Process, exitCode int) (Result, error) { - stream, err := proc.Output(ctx, &ateenvv1alpha.StreamProcessOutputRequest{}) - if err != nil { - return Result{}, err - } - var stdout, stderr bytes.Buffer - for { - msg, err := stream.Recv() - if errors.Is(err, io.EOF) { - break - } - if err != nil { - return Result{}, err - } - switch out := msg.GetOutput().(type) { - case *ateenvv1alpha.ProcessOutput_Stdout: - stdout.Write(out.Stdout) - case *ateenvv1alpha.ProcessOutput_Stderr: - stderr.Write(out.Stderr) - case *ateenvv1alpha.ProcessOutput_Exit: - exitCode = int(out.Exit.GetExitCode()) - } - } - return Result{ExitCode: exitCode, Stdout: stdout.String(), Stderr: stderr.String()}, nil -} - -func (s envSandbox) WriteFile(ctx context.Context, path string, content []byte, mode fs.FileMode) error { - return s.e.WriteFile(ctx, path, bytes.NewReader(content), mode) -} - -func (s envSandbox) Suspend(ctx context.Context) error { return s.e.Suspend(ctx) } - -func (s envSandbox) Delete(ctx context.Context) error { return s.e.Delete(ctx) } diff --git a/internal/benchmarking/swebench/backend_test.go b/internal/benchmarking/swebench/backend_test.go deleted file mode 100644 index f0cc88d4..00000000 --- a/internal/benchmarking/swebench/backend_test.go +++ /dev/null @@ -1,270 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package swebench - -import ( - "context" - "fmt" - "net" - "strings" - "sync" - "testing" - "time" - - env "github.com/agent-substrate/env/clients/go" - ateenvv1alpha "github.com/agent-substrate/env/proto/ateenv/v1alpha" - "google.golang.org/grpc" - "google.golang.org/grpc/codes" - "google.golang.org/grpc/credentials/insecure" - "google.golang.org/grpc/status" - "google.golang.org/grpc/test/bufconn" -) - -// fakeAPI is an in-process ate-env API: enough of the environment, process -// and file services for the backend's start, poll and read path. -type fakeAPI struct { - ateenvv1alpha.UnimplementedEnvironmentServiceServer - ateenvv1alpha.UnimplementedProcessServiceServer - ateenvv1alpha.UnimplementedFileSystemServiceServer - - mu sync.Mutex - // pollsUntilExit is how many GetProcess calls a process stays running; - // unavailablePolls is how many GetProcess calls answer Unavailable first. - pollsUntilExit int - unavailablePolls int - exitCode int32 - stdout, stderr string - - nextID int - polls map[string]int - exited map[string]bool - killed []string - created []string - deleted []string - commands []string -} - -func (f *fakeAPI) CreateEnvironment(_ context.Context, req *ateenvv1alpha.CreateEnvironmentRequest) (*ateenvv1alpha.CreateEnvironmentResponse, error) { - f.mu.Lock() - defer f.mu.Unlock() - f.created = append(f.created, req.GetId()) - return &ateenvv1alpha.CreateEnvironmentResponse{Environment: &ateenvv1alpha.Environment{Id: req.GetId(), Atespace: req.GetAtespace()}}, nil -} - -func (f *fakeAPI) DeleteEnvironment(_ context.Context, req *ateenvv1alpha.DeleteEnvironmentRequest) (*ateenvv1alpha.DeleteEnvironmentResponse, error) { - f.mu.Lock() - defer f.mu.Unlock() - f.deleted = append(f.deleted, req.GetId()) - return &ateenvv1alpha.DeleteEnvironmentResponse{}, nil -} - -func (f *fakeAPI) StartProcess(_ context.Context, req *ateenvv1alpha.StartProcessRequest) (*ateenvv1alpha.Process, error) { - f.mu.Lock() - defer f.mu.Unlock() - f.nextID++ - id := fmt.Sprintf("p%d", f.nextID) - f.commands = append(f.commands, strings.Join(req.GetCommand(), " ")) - return &ateenvv1alpha.Process{ProcessId: id, State: ateenvv1alpha.ProcessState_PROCESS_STATE_RUNNING}, nil -} - -func (f *fakeAPI) GetProcess(_ context.Context, req *ateenvv1alpha.GetProcessRequest) (*ateenvv1alpha.Process, error) { - f.mu.Lock() - defer f.mu.Unlock() - if f.unavailablePolls > 0 { - f.unavailablePolls-- - return nil, status.Error(codes.Unavailable, "router: no healthy upstream") - } - f.polls[req.GetProcessId()]++ - p := &ateenvv1alpha.Process{ProcessId: req.GetProcessId(), State: ateenvv1alpha.ProcessState_PROCESS_STATE_RUNNING} - if f.polls[req.GetProcessId()] > f.pollsUntilExit { - f.exited[req.GetProcessId()] = true - p.State = ateenvv1alpha.ProcessState_PROCESS_STATE_EXITED - p.ExitCode = f.exitCode - } - return p, nil -} - -func (f *fakeAPI) StreamProcessOutput(req *ateenvv1alpha.StreamProcessOutputRequest, stream grpc.ServerStreamingServer[ateenvv1alpha.ProcessOutput]) error { - f.mu.Lock() - stdout, stderr, exited, code := f.stdout, f.stderr, f.exited[req.GetProcessId()], f.exitCode - f.mu.Unlock() - if stdout != "" { - if err := stream.Send(&ateenvv1alpha.ProcessOutput{Output: &ateenvv1alpha.ProcessOutput_Stdout{Stdout: []byte(stdout)}}); err != nil { - return err - } - } - if stderr != "" { - if err := stream.Send(&ateenvv1alpha.ProcessOutput{Output: &ateenvv1alpha.ProcessOutput_Stderr{Stderr: []byte(stderr)}}); err != nil { - return err - } - } - if exited { - return stream.Send(&ateenvv1alpha.ProcessOutput{Output: &ateenvv1alpha.ProcessOutput_Exit{Exit: &ateenvv1alpha.Process{ - ProcessId: req.GetProcessId(), State: ateenvv1alpha.ProcessState_PROCESS_STATE_EXITED, ExitCode: code}}}) - } - return nil -} - -func (f *fakeAPI) SignalProcess(_ context.Context, req *ateenvv1alpha.SignalProcessRequest) (*ateenvv1alpha.Process, error) { - f.mu.Lock() - defer f.mu.Unlock() - f.killed = append(f.killed, req.GetProcessId()) - f.exited[req.GetProcessId()] = true - return &ateenvv1alpha.Process{ProcessId: req.GetProcessId(), State: ateenvv1alpha.ProcessState_PROCESS_STATE_EXITED, ExitCode: 137}, nil -} - -// newFakeBackend serves fakeAPI over an in-memory connection and returns an -// EnvBackend on it, the same path the binary uses with a real address. -func newFakeBackend(t *testing.T, api *fakeAPI) *EnvBackend { - t.Helper() - api.polls = map[string]int{} - api.exited = map[string]bool{} - lis := bufconn.Listen(1 << 20) - srv := grpc.NewServer() - ateenvv1alpha.RegisterEnvironmentServiceServer(srv, api) - ateenvv1alpha.RegisterProcessServiceServer(srv, api) - ateenvv1alpha.RegisterFileSystemServiceServer(srv, api) - go func() { _ = srv.Serve(lis) }() - conn, err := grpc.NewClient("passthrough:///bufconn", - grpc.WithContextDialer(func(ctx context.Context, _ string) (net.Conn, error) { return lis.DialContext(ctx) }), - grpc.WithTransportCredentials(insecure.NewCredentials())) - if err != nil { - t.Fatal(err) - } - c, err := env.NewClient(env.ClientOptions{GRPCConn: conn}) - if err != nil { - t.Fatal(err) - } - t.Cleanup(func() { _ = c.Close(); _ = conn.Close(); srv.Stop() }) - return &EnvBackend{clients: []*env.Client{c}, atespace: "ate-env", baseTemplate: "default-template"} -} - -func fastPolls(t *testing.T) { - t.Helper() - oldMin, oldMax := pollMin, pollMax - pollMin, pollMax = time.Millisecond, 5*time.Millisecond - t.Cleanup(func() { pollMin, pollMax = oldMin, oldMax }) -} - -func TestEnvBackendRunCollectsExitAndOutput(t *testing.T) { - fastPolls(t) - api := &fakeAPI{pollsUntilExit: 3, exitCode: 4, stdout: "out", stderr: "err"} - b := newFakeBackend(t, api) - sb, err := b.Create(context.Background(), "e1", "reg/img@"+goodDigest) - if err != nil { - t.Fatal(err) - } - res, err := sb.Run(context.Background(), "echo hi", time.Minute) - if err != nil { - t.Fatal(err) - } - if res.ExitCode != 4 || res.Stdout != "out" || res.Stderr != "err" { - t.Fatalf("result %+v", res) - } - if len(api.killed) != 0 { - t.Fatalf("a process that exited must not be killed: %v", api.killed) - } - if len(api.created) != 1 || api.created[0] != "e1" { - t.Fatalf("created %v", api.created) - } - if !strings.HasSuffix(api.commands[0], "echo hi") { - t.Fatalf("command %q", api.commands[0]) - } -} - -func TestEnvBackendRunToleratesTransientPolls(t *testing.T) { - fastPolls(t) - api := &fakeAPI{pollsUntilExit: 1, unavailablePolls: pollRetries} - b := newFakeBackend(t, api) - sb, _ := b.Create(context.Background(), "e1", "reg/img@"+goodDigest) - if _, err := sb.Run(context.Background(), "true", time.Minute); err != nil { - t.Fatalf("%d transient poll errors must be tolerated: %v", pollRetries, err) - } -} - -func TestEnvBackendRunGivesUpAndKills(t *testing.T) { - fastPolls(t) - api := &fakeAPI{pollsUntilExit: 100, unavailablePolls: pollRetries + 1} - b := newFakeBackend(t, api) - sb, _ := b.Create(context.Background(), "e1", "reg/img@"+goodDigest) - _, err := sb.Run(context.Background(), "sleep 100", time.Minute) - if err == nil || status.Code(err) != codes.Unavailable { - t.Fatalf("expected the poll error after the retry limit, got %v", err) - } - if len(api.killed) != 1 { - t.Fatalf("an abandoned process must be killed once, got %v", api.killed) - } -} - -func TestEnvBackendRunKillsOnCancel(t *testing.T) { - fastPolls(t) - api := &fakeAPI{pollsUntilExit: 1 << 30} - b := newFakeBackend(t, api) - sb, _ := b.Create(context.Background(), "e1", "reg/img@"+goodDigest) - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond) - defer cancel() - if _, err := sb.Run(ctx, "sleep 100", time.Minute); err == nil { - t.Fatal("a cancelled Run must return an error") - } - if len(api.killed) != 1 { - t.Fatalf("the process must be killed after a cancel, got %v", api.killed) - } -} - -func TestEnvBackendDeleteByName(t *testing.T) { - api := &fakeAPI{} - b := newFakeBackend(t, api) - if err := b.Delete(context.Background(), "gone"); err != nil { - t.Fatal(err) - } - if len(api.deleted) != 1 || api.deleted[0] != "gone" { - t.Fatalf("deleted %v", api.deleted) - } -} - -func TestEnvBackendSpreadsConnections(t *testing.T) { - // grpc.NewClient does not dial until the first call, so a backend over - // an unreachable address is enough to check the pool and the rotation. - b, err := NewEnvBackend("127.0.0.1:1", "ate-env", "default-template", 3) - if err != nil { - t.Fatal(err) - } - defer b.Close() - if len(b.clients) != 3 { - t.Fatalf("%d clients, want 3", len(b.clients)) - } - seen := map[int]int{} - for i := 0; i < 9; i++ { - c := b.pick() - for j, cc := range b.clients { - if cc == c { - seen[j]++ - } - } - } - for j := 0; j < 3; j++ { - if seen[j] != 3 { - t.Fatalf("connection %d picked %d times, want 3: %v", j, seen[j], seen) - } - } - one, err := NewEnvBackend("127.0.0.1:1", "ate-env", "default-template", 0) - if err != nil { - t.Fatal(err) - } - defer one.Close() - if len(one.clients) != 1 { - t.Fatalf("connections 0 must mean 1, got %d", len(one.clients)) - } -} diff --git a/internal/benchmarking/swebench/dataset.go b/internal/benchmarking/swebench/dataset.go deleted file mode 100644 index 32016462..00000000 --- a/internal/benchmarking/swebench/dataset.go +++ /dev/null @@ -1,216 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package swebench - -import ( - "bufio" - "bytes" - "encoding/json" - "errors" - "fmt" - "io" - "net/http" - "os" - "regexp" - "strings" - "time" -) - -// Instance is one SWE-bench task: the digest-pinned task image and the -// material a trajectory needs to evaluate it. The manifest is produced by -// benchmarking/swebench/make-dataset.py, one JSON object per line. -type Instance struct { - InstanceID string `json:"instance_id"` - Repo string `json:"repo"` - // Image is the task image, pinned by digest (repo@sha256:...). Substrate - // requires pinned images because changing the image invalidates snapshots. - Image string `json:"image"` - // Patch is the gold patch; applying it makes FailToPass pass. - Patch string `json:"patch"` - // TestPatch adds or changes the tests that check the fix. - TestPatch string `json:"test_patch"` - // EvalScript is the SWE-bench harness eval script for this instance: it - // applies TestPatch, runs the repo's test command and restores the test - // files. Exit code 0 means every selected test passed. - EvalScript string `json:"eval_script"` - FailToPass []string `json:"fail_to_pass"` - PassToPass []string `json:"pass_to_pass"` - // LogParser names the SWE-bench harness parser for this repo's test - // output (parse_log_pytest, parse_log_django, ...); EvalType is - // pass_and_fail or fail_only. Both come from the dataset row. - LogParser string `json:"log_parser"` - EvalType string `json:"eval_type"` -} - -var digestPinned = regexp.MustCompile(`@sha256:[0-9a-f]{64}$`) - -// Validate reports the first thing that would make the instance unusable. -func (in Instance) Validate() error { - switch { - case in.InstanceID == "": - return errors.New("instance_id is required") - case in.Image == "": - return fmt.Errorf("%s: image is required", in.InstanceID) - case !digestPinned.MatchString(in.Image): - return fmt.Errorf("%s: image %q is not digest-pinned", in.InstanceID, in.Image) - case in.EvalScript == "": - return fmt.Errorf("%s: eval_script is required", in.InstanceID) - case in.LogParser == "": - return fmt.Errorf("%s: log_parser is required", in.InstanceID) - } - return nil -} - -// datasetFetchTimeout bounds a manifest download; a stalled server must not -// hang the loader before it starts. -const datasetFetchTimeout = 2 * time.Minute - -// LoadDataset reads a JSONL manifest from a local path or an https URL. The -// manifest carries scripts that run inside the sandboxes, so a remote source -// must be https. -func LoadDataset(source string) ([]Instance, error) { - var r io.ReadCloser - if strings.HasPrefix(source, "http://") { - return nil, fmt.Errorf("dataset %s: remote manifests must use https", source) - } - if strings.HasPrefix(source, "https://") { - client := &http.Client{Timeout: datasetFetchTimeout} - resp, err := client.Get(source) - if err != nil { - return nil, fmt.Errorf("fetch dataset: %w", err) - } - if resp.StatusCode != http.StatusOK { - resp.Body.Close() - return nil, fmt.Errorf("fetch dataset: HTTP %d", resp.StatusCode) - } - r = resp.Body - } else { - f, err := os.Open(source) - if err != nil { - return nil, fmt.Errorf("open dataset: %w", err) - } - r = f - } - defer r.Close() - return ParseDataset(r) -} - -// ParseDataset decodes a JSONL manifest. Blank lines are skipped; every -// record must validate and instance ids must be unique. -func ParseDataset(r io.Reader) ([]Instance, error) { - sc := bufio.NewScanner(r) - // Eval scripts and patches run to hundreds of KiB for some repos. - sc.Buffer(make([]byte, 0, 1<<20), 64<<20) - var out []Instance - seen := map[string]bool{} - line := 0 - for sc.Scan() { - line++ - b := bytes.TrimSpace(sc.Bytes()) - if len(b) == 0 { - continue - } - var in Instance - if err := json.Unmarshal(b, &in); err != nil { - return nil, fmt.Errorf("dataset line %d: %w", line, err) - } - if err := in.Validate(); err != nil { - return nil, fmt.Errorf("dataset line %d: %w", line, err) - } - if seen[in.InstanceID] { - return nil, fmt.Errorf("dataset line %d: duplicate instance %s", line, in.InstanceID) - } - seen[in.InstanceID] = true - out = append(out, in) - } - if err := sc.Err(); err != nil { - return nil, fmt.Errorf("read dataset: %w", err) - } - if len(out) == 0 { - return nil, errors.New("dataset is empty") - } - return out, nil -} - -// Select narrows a dataset: exclude drops the named instances first (a -// known-bad list; unknown names are ignored), ids then keeps only the named -// instances (in manifest order) and errors on an unknown or excluded id, and -// limit > 0 keeps the first limit of what is left. -func Select(all []Instance, ids, exclude []string, limit int) ([]Instance, error) { - out := all - drop := map[string]bool{} - if len(exclude) > 0 { - for _, id := range exclude { - drop[id] = true - } - out = nil - for _, in := range all { - if !drop[in.InstanceID] { - out = append(out, in) - } - } - all = out - } - if len(ids) > 0 { - for _, id := range ids { - if drop[id] { - return nil, fmt.Errorf("instance %s is on the exclude list", id) - } - } - want := map[string]bool{} - for _, id := range ids { - want[id] = true - } - out = nil - for _, in := range all { - if want[in.InstanceID] { - out = append(out, in) - delete(want, in.InstanceID) - } - } - if len(want) > 0 { - missing := make([]string, 0, len(want)) - for id := range want { - missing = append(missing, id) - } - return nil, fmt.Errorf("instances not in dataset: %s", strings.Join(missing, ", ")) - } - } - if limit > 0 && limit < len(out) { - out = out[:limit] - } - return out, nil -} - -// ReadIDList reads instance ids from a file, one per line. Blank lines and -// lines starting with '#' are skipped, and anything after the first -// whitespace on a line is a comment, so a list can say why each id is there. -func ReadIDList(path string) ([]string, error) { - f, err := os.Open(path) - if err != nil { - return nil, err - } - defer f.Close() - var ids []string - sc := bufio.NewScanner(f) - for sc.Scan() { - line := strings.TrimSpace(sc.Text()) - if line == "" || strings.HasPrefix(line, "#") { - continue - } - ids = append(ids, strings.Fields(line)[0]) - } - return ids, sc.Err() -} diff --git a/internal/benchmarking/swebench/dataset_test.go b/internal/benchmarking/swebench/dataset_test.go deleted file mode 100644 index f2955dea..00000000 --- a/internal/benchmarking/swebench/dataset_test.go +++ /dev/null @@ -1,110 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package swebench - -import ( - "os" - "path/filepath" - "strings" - "testing" -) - -const goodDigest = "sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef" - -func manifest(ids ...string) string { - var b strings.Builder - for _, id := range ids { - b.WriteString(`{"instance_id":"` + id + `","repo":"r","image":"reg/img@` + goodDigest + `","patch":"p","eval_script":"e","log_parser":"parse_log_pytest"}` + "\n") - } - return b.String() -} - -func TestParseDataset(t *testing.T) { - got, err := ParseDataset(strings.NewReader("\n" + manifest("a__a-1", "b__b-2") + "\n")) - if err != nil { - t.Fatal(err) - } - if len(got) != 2 || got[0].InstanceID != "a__a-1" || got[1].InstanceID != "b__b-2" { - t.Fatalf("got %+v", got) - } -} - -func TestParseDatasetRejects(t *testing.T) { - cases := map[string]string{ - "empty": "", - "bad json": "{not json}\n", - "no id": `{"image":"reg/img@` + goodDigest + `"}` + "\n", - "no image": `{"instance_id":"x"}` + "\n", - "tag not pin": `{"instance_id":"x","image":"reg/img:latest"}` + "\n", - "duplicate": manifest("a", "a"), - "no eval": `{"instance_id":"x","image":"reg/img@` + goodDigest + `","log_parser":"parse_log_pytest"}` + "\n", - "no parser": `{"instance_id":"x","image":"reg/img@` + goodDigest + `","eval_script":"e"}` + "\n", - } - for name, in := range cases { - if _, err := ParseDataset(strings.NewReader(in)); err == nil { - t.Errorf("%s: expected an error", name) - } - } -} - -func TestSelect(t *testing.T) { - all, err := ParseDataset(strings.NewReader(manifest("a", "b", "c"))) - if err != nil { - t.Fatal(err) - } - got, err := Select(all, nil, nil, 2) - if err != nil || len(got) != 2 || got[1].InstanceID != "b" { - t.Fatalf("limit: got %+v, %v", got, err) - } - got, err = Select(all, []string{"c", "a"}, nil, 0) - if err != nil || len(got) != 2 || got[0].InstanceID != "a" || got[1].InstanceID != "c" { - t.Fatalf("ids keep manifest order: got %+v, %v", got, err) - } - got, err = Select(all, []string{"c", "a"}, nil, 1) - if err != nil || len(got) != 1 || got[0].InstanceID != "a" { - t.Fatalf("ids then limit: got %+v, %v", got, err) - } - if _, err := Select(all, []string{"zz"}, nil, 0); err == nil { - t.Fatal("unknown id must error") - } -} - -func TestSelectExcludeAndIDList(t *testing.T) { - all, err := ParseDataset(strings.NewReader(manifest("a", "b", "c"))) - if err != nil { - t.Fatal(err) - } - got, err := Select(all, nil, []string{"b", "unknown"}, 0) - if err != nil || len(got) != 2 || got[0].InstanceID != "a" || got[1].InstanceID != "c" { - t.Fatalf("exclude: got %+v, %v", got, err) - } - if _, err := Select(all, []string{"b"}, []string{"b"}, 0); err == nil || !strings.Contains(err.Error(), "exclude list") { - t.Fatalf("an excluded id asked for by name must say so: %v", err) - } - path := filepath.Join(t.TempDir(), "ids.txt") - if err := os.WriteFile(path, []byte("# known network-dependent\nb needs pypi\n\n c\n"), 0o644); err != nil { - t.Fatal(err) - } - ids, err := ReadIDList(path) - if err != nil || len(ids) != 2 || ids[0] != "b" || ids[1] != "c" { - t.Fatalf("ReadIDList: %v, %v", ids, err) - } -} - -func TestLoadDatasetRefusesPlainHTTP(t *testing.T) { - if _, err := LoadDataset("http://example.invalid/manifest.jsonl"); err == nil || !strings.Contains(err.Error(), "https") { - t.Fatalf("plain http must be refused: %v", err) - } -} diff --git a/internal/benchmarking/swebench/doc.go b/internal/benchmarking/swebench/doc.go deleted file mode 100644 index ae8ef3ef..00000000 --- a/internal/benchmarking/swebench/doc.go +++ /dev/null @@ -1,35 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Package swebench holds the data types of the Fast Substrate RL benchmark on -// SWE-bench Verified, and the runner that drives them through the ate-env API. -// -// - Instance and LoadDataset: the JSONL manifest, one record per SWE-bench -// instance (digest-pinned task image, gold patch, harness eval script, -// the tests that decide the verdict, the repo's log parser). -// - Trajectory: the scripted rollout a task replays (run a command or write -// a file, with think time, timeout, expected exit and one reward step), -// with setup steps that run once per environment and reset steps that -// run before each later step. The built-in swebench-gold trajectory -// applies the gold patch, runs the harness eval script and grades the -// log with swebench-grade.py, a port of the SWE-bench harness's rules. -// - Event and Recorder: one JSON line per operation of a run. -// - Summarize and Build: the report over those events, with percentiles, -// retry counts, a per-step table and the unresolved instances in -// exclude-file format. -// -// The layout mirrors substrate's internal/benchmarking/boomer so a user class -// can be lifted into that harness later; it is not imported because Go does -// not allow importing another module's internal packages. -package swebench diff --git a/internal/benchmarking/swebench/events.go b/internal/benchmarking/swebench/events.go deleted file mode 100644 index 78d51c14..00000000 --- a/internal/benchmarking/swebench/events.go +++ /dev/null @@ -1,111 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package swebench - -import ( - "encoding/json" - "io" - "sync" - "time" -) - -// Event names. Every timed operation of a run is one of these. -const ( - EventCreate = "create" // CreateEnvironment call - EventTTFI = "ttfi" // create start to the first command's return - EventWake = "wake" // first command after a park, including the wake it triggers - EventStep = "step" // one trajectory step (name is the step name) - EventTask = "task" // one trajectory replay end to end - EventSuspend = "suspend" // SuspendEnvironment call - EventDelete = "delete" // DeleteEnvironment call - EventBarrier = "step_barrier" // steps mode: first task start to last task end in a training step - EventPark = "park" // steps mode: parking the whole fleet between steps - EventRun = "run" // whole run -) - -// Event is one JSONL record. Dur is wall time as seen by the loader; it -// includes the network and the API server, which is what a trainer sees. -type Event struct { - Time time.Time `json:"ts"` - RunID string `json:"run"` - Event string `json:"event"` - Name string `json:"name,omitempty"` - Env string `json:"env,omitempty"` - Instance string `json:"instance,omitempty"` - // Step is the training step in steps mode (1-based); 0 in eval mode. - Step int `json:"step,omitempty"` - // Task numbers the replays of one instance (eval: task index; steps: - // rollout index). - Task int `json:"task,omitempty"` - DurMS float64 `json:"dur_ms"` - OK bool `json:"ok"` - Exit *int `json:"exit,omitempty"` - Reward *int `json:"reward,omitempty"` - Err string `json:"err,omitempty"` - // Count is the number of tasks or environments an aggregate event covers. - Count int `json:"count,omitempty"` - // Tail is the last bytes of a failed or reward step's output. - Tail string `json:"tail,omitempty"` - // Attempts is how many tries a ttfi or wake took; above 1 means the - // environment was not serving when first asked. - Attempts int `json:"attempts,omitempty"` -} - -// Recorder writes events as JSON lines. It is safe for concurrent use. -type Recorder struct { - mu sync.Mutex - w io.Writer - enc *json.Encoder - n int - err error -} - -// NewRecorder writes to w; a nil w discards events. -func NewRecorder(w io.Writer) *Recorder { - if w == nil { - w = io.Discard - } - return &Recorder{w: w, enc: json.NewEncoder(w)} -} - -// Record writes one event, stamping the time if unset. -func (r *Recorder) Record(e Event) { - if e.Time.IsZero() { - e.Time = time.Now() - } - r.mu.Lock() - defer r.mu.Unlock() - if err := r.enc.Encode(e); err != nil && r.err == nil { - r.err = err - } - r.n++ -} - -// Len is the number of events written so far. -func (r *Recorder) Len() int { - r.mu.Lock() - defer r.mu.Unlock() - return r.n -} - -// Err is the first write error, if any; a run whose events were lost must -// not report success. -func (r *Recorder) Err() error { - r.mu.Lock() - defer r.mu.Unlock() - return r.err -} - -func intp(v int) *int { return &v } diff --git a/internal/benchmarking/swebench/report.go b/internal/benchmarking/swebench/report.go deleted file mode 100644 index b78b8e13..00000000 --- a/internal/benchmarking/swebench/report.go +++ /dev/null @@ -1,317 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package swebench - -import ( - "bufio" - "bytes" - "encoding/json" - "fmt" - "io" - "math" - "sort" - "strings" - "text/tabwriter" -) - -// Stat summarizes one kind of event. -type Stat struct { - Event string `json:"event"` - Name string `json:"name,omitempty"` - Count int `json:"count"` - Failed int `json:"failed"` - P50MS float64 `json:"p50_ms"` - P95MS float64 `json:"p95_ms"` - MaxMS float64 `json:"max_ms"` - // Retried counts ttfi and wake events that took more than one attempt: - // the environment was not serving when first asked. - Retried int `json:"retried,omitempty"` -} - -// Unresolved is an instance whose reward step scored 0, with the grader's -// verdict when the step's output ended in one. -type Unresolved struct { - Instance string `json:"instance"` - Reason string `json:"reason"` -} - -// StepStat is one training step in steps mode. -type StepStat struct { - Step int `json:"step"` - Tasks int `json:"tasks"` - Failed int `json:"failed"` - Rewarded int `json:"rewarded"` - // Live is how many environments were still in the fleet at this step; - // steps mode drops an environment after its first failure, so a - // shrinking Live makes later steps' counts incomparable. - Live int `json:"live,omitempty"` - BarrierMS float64 `json:"barrier_ms"` - ParkMS float64 `json:"park_ms"` - WakeP50MS float64 `json:"wake_p50_ms"` - WakeP95MS float64 `json:"wake_p95_ms"` -} - -// Report is the summary of one run's JSONL, stable enough to diff between -// runs. -type Report struct { - RunID string `json:"run"` - Mode string `json:"mode,omitempty"` - RunMS float64 `json:"run_ms,omitempty"` - Tasks int `json:"tasks"` - Failed int `json:"failed"` - Rewarded int `json:"rewarded"` - RewardRate float64 `json:"reward_rate"` - Stats []Stat `json:"stats"` - Steps []StepStat `json:"steps,omitempty"` - // Errors counts distinct error strings, most frequent first. - Errors []ErrorCount `json:"errors,omitempty"` - // Unresolved lists the instances whose reward was 0, once each. - Unresolved []Unresolved `json:"unresolved,omitempty"` -} - -// ErrorCount is one error string and how often it occurred. -type ErrorCount struct { - Err string `json:"err"` - Count int `json:"count"` -} - -// Summarize reads JSONL events and builds the report. -func Summarize(r io.Reader) (*Report, error) { - sc := bufio.NewScanner(r) - sc.Buffer(make([]byte, 0, 1<<20), 16<<20) - var events []Event - for sc.Scan() { - b := bytes.TrimSpace(sc.Bytes()) - if len(b) == 0 { - continue - } - var e Event - if err := json.Unmarshal(b, &e); err != nil { - return nil, fmt.Errorf("parse event: %w", err) - } - events = append(events, e) - } - if err := sc.Err(); err != nil { - return nil, err - } - return Build(events), nil -} - -// Build computes the report from events in memory. -func Build(events []Event) *Report { - rep := &Report{} - type key struct{ event, name string } - durs := map[key][]float64{} - failed := map[key]int{} - retried := map[key]int{} - unresolved := map[string]string{} - var order []key - steps := map[int]*StepStat{} - wakes := map[int][]float64{} - errs := map[string]int{} - - for _, e := range events { - if rep.RunID == "" { - rep.RunID = e.RunID - } - k := key{e.Event, ""} - if e.Event == EventStep || e.Event == EventPark { - k.name = e.Name - } - if _, ok := durs[k]; !ok { - order = append(order, k) - durs[k] = nil - } - if e.Attempts > 1 { - retried[k]++ - } - if e.Event == EventStep && e.Reward != nil && *e.Reward == 0 { - if _, seen := unresolved[e.Instance]; !seen { - unresolved[e.Instance] = verdict(e.Tail) - } - } - if e.OK { - durs[k] = append(durs[k], e.DurMS) - } else { - failed[k]++ - if e.Err != "" { - errs[e.Err]++ - } - } - switch e.Event { - case EventRun: - rep.Mode = e.Name - rep.RunMS = e.DurMS - case EventTask: - rep.Tasks++ - if !e.OK { - rep.Failed++ - } - if e.Reward != nil && *e.Reward == 1 { - rep.Rewarded++ - } - if e.Step > 0 { - s := stepStat(steps, e.Step) - s.Tasks++ - if !e.OK { - s.Failed++ - } - if e.Reward != nil && *e.Reward == 1 { - s.Rewarded++ - } - } - case EventBarrier: - stepStat(steps, e.Step).BarrierMS = e.DurMS - stepStat(steps, e.Step).Live = e.Count - case EventPark: - stepStat(steps, e.Step).ParkMS = e.DurMS - case EventWake: - if e.OK { - wakes[e.Step] = append(wakes[e.Step], e.DurMS) - } - } - } - if rep.Tasks > 0 { - rep.RewardRate = float64(rep.Rewarded) / float64(rep.Tasks) - } - for _, k := range order { - d := durs[k] - sort.Float64s(d) - rep.Stats = append(rep.Stats, Stat{Event: k.event, Name: k.name, Count: len(d) + failed[k], Failed: failed[k], - P50MS: percentile(d, 0.50), P95MS: percentile(d, 0.95), MaxMS: percentile(d, 1), Retried: retried[k]}) - } - for step, s := range steps { - w := wakes[step] - sort.Float64s(w) - s.WakeP50MS = percentile(w, 0.5) - s.WakeP95MS = percentile(w, 0.95) - rep.Steps = append(rep.Steps, *s) - } - sort.Slice(rep.Steps, func(i, j int) bool { return rep.Steps[i].Step < rep.Steps[j].Step }) - for msg, n := range errs { - rep.Errors = append(rep.Errors, ErrorCount{Err: msg, Count: n}) - } - for inst, why := range unresolved { - rep.Unresolved = append(rep.Unresolved, Unresolved{Instance: inst, Reason: why}) - } - sort.Slice(rep.Unresolved, func(i, j int) bool { return rep.Unresolved[i].Instance < rep.Unresolved[j].Instance }) - sort.Slice(rep.Errors, func(i, j int) bool { - if rep.Errors[i].Count != rep.Errors[j].Count { - return rep.Errors[i].Count > rep.Errors[j].Count - } - return rep.Errors[i].Err < rep.Errors[j].Err - }) - return rep -} - -// verdict extracts the grader's verdict from the end of a reward step's -// output: "invalid run", "f2p A/B p2p C/D", or "reward 0" when the output -// does not end in the grader's JSON line. -func verdict(tail string) string { - i := strings.LastIndex(tail, `{"valid"`) - if i < 0 { - return "reward 0" - } - var g struct { - Valid bool `json:"valid"` - F2P string `json:"f2p"` - P2P string `json:"p2p"` - } - if err := json.Unmarshal([]byte(strings.TrimSpace(tail[i:])), &g); err != nil { - return "reward 0" - } - if !g.Valid { - return "invalid run" - } - return fmt.Sprintf("f2p %s p2p %s", g.F2P, g.P2P) -} - -// WriteExclusions writes the unresolved instances in the exclude-file format: -// one id per line followed by the reason as a comment. -func (rep *Report) WriteExclusions(w io.Writer, header string) { - fmt.Fprintf(w, "# %s\n# Generated from run %s: %d of %d tasks unresolved.\n", header, rep.RunID, len(rep.Unresolved), rep.Tasks) - for _, u := range rep.Unresolved { - fmt.Fprintf(w, "%s %s\n", u.Instance, u.Reason) - } -} - -func stepStat(m map[int]*StepStat, step int) *StepStat { - s, ok := m[step] - if !ok { - s = &StepStat{Step: step} - m[step] = s - } - return s -} - -// percentile of a sorted slice by the nearest-rank method, ceil(p*n)-1; 0 -// for an empty slice. -func percentile(sorted []float64, p float64) float64 { - n := len(sorted) - if n == 0 { - return 0 - } - if p >= 1 { - return sorted[n-1] - } - i := int(math.Ceil(p*float64(n))) - 1 - if i < 0 { - i = 0 - } - if i >= n { - i = n - 1 - } - return sorted[i] -} - -// WriteText prints the report as tables. -func (rep *Report) WriteText(w io.Writer) { - fmt.Fprintf(w, "run %s mode=%s tasks=%d failed=%d rewarded=%d reward_rate=%.3f run=%.1fs\n\n", - rep.RunID, rep.Mode, rep.Tasks, rep.Failed, rep.Rewarded, rep.RewardRate, rep.RunMS/1000) - tw := tabwriter.NewWriter(w, 0, 0, 2, ' ', 0) - fmt.Fprintln(tw, "event\tname\tcount\tfailed\tretried\tp50_ms\tp95_ms\tmax_ms") - for _, s := range rep.Stats { - fmt.Fprintf(tw, "%s\t%s\t%d\t%d\t%d\t%.0f\t%.0f\t%.0f\n", s.Event, s.Name, s.Count, s.Failed, s.Retried, s.P50MS, s.P95MS, s.MaxMS) - } - tw.Flush() - if len(rep.Steps) > 0 { - fmt.Fprintln(w) - tw = tabwriter.NewWriter(w, 0, 0, 2, ' ', 0) - fmt.Fprintln(tw, "step\tlive\ttasks\tfailed\trewarded\tbarrier_s\tpark_s\twake_p50_ms\twake_p95_ms") - for _, s := range rep.Steps { - fmt.Fprintf(tw, "%d\t%d\t%d\t%d\t%d\t%.1f\t%.1f\t%.0f\t%.0f\n", s.Step, s.Live, s.Tasks, s.Failed, s.Rewarded, - s.BarrierMS/1000, s.ParkMS/1000, s.WakeP50MS, s.WakeP95MS) - } - tw.Flush() - } - if len(rep.Unresolved) > 0 { - fmt.Fprintf(w, "\nunresolved: %d instances (report -unresolved FILE writes them in exclude-file format)\n", len(rep.Unresolved)) - } - if len(rep.Errors) > 0 { - fmt.Fprintln(w) - fmt.Fprintln(w, "errors:") - for i, e := range rep.Errors { - if i == 10 { - fmt.Fprintf(w, " ... %d more\n", len(rep.Errors)-10) - break - } - msg := strings.ReplaceAll(e.Err, "\n", " ") - if len(msg) > 160 { - msg = msg[:160] + "..." - } - fmt.Fprintf(w, " %5d %s\n", e.Count, msg) - } - } -} diff --git a/internal/benchmarking/swebench/report_test.go b/internal/benchmarking/swebench/report_test.go deleted file mode 100644 index dadf87a4..00000000 --- a/internal/benchmarking/swebench/report_test.go +++ /dev/null @@ -1,139 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package swebench - -import ( - "bytes" - "errors" - "strings" - "testing" -) - -func TestPercentile(t *testing.T) { - s := []float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10} - if p := percentile(s, 0.5); p != 5 { - t.Errorf("p50 = %v", p) - } - if p := percentile(s, 0.95); p != 10 { - t.Errorf("p95 = %v", p) - } - if p := percentile(s, 1); p != 10 { - t.Errorf("max = %v", p) - } - if p := percentile(nil, 0.5); p != 0 { - t.Errorf("empty = %v", p) - } - thirteen := make([]float64, 13) - for i := range thirteen { - thirteen[i] = float64(i + 1) - } - if p := percentile(thirteen, 0.95); p != 13 { - t.Errorf("nearest rank p95 of 13 samples = %v, want the 13th", p) - } -} - -func TestSummarizeRoundTrip(t *testing.T) { - var buf bytes.Buffer - rec := NewRecorder(&buf) - rec.Record(Event{RunID: "r", Event: EventCreate, Env: "e1", DurMS: 100, OK: true}) - rec.Record(Event{RunID: "r", Event: EventCreate, Env: "e2", DurMS: 300, OK: false, Err: "boom"}) - rec.Record(Event{RunID: "r", Event: EventStep, Name: "run_tests", DurMS: 50, OK: true, Exit: intp(0)}) - rec.Record(Event{RunID: "r", Event: EventStep, Name: "explore", DurMS: 5, OK: true, Exit: intp(0)}) - rec.Record(Event{RunID: "r", Event: EventTask, Step: 1, DurMS: 500, OK: true, Reward: intp(1)}) - rec.Record(Event{RunID: "r", Event: EventTask, Step: 1, DurMS: 600, OK: false}) - rec.Record(Event{RunID: "r", Event: EventWake, Step: 2, DurMS: 40, OK: true, Attempts: 3}) - rec.Record(Event{RunID: "r", Event: EventStep, Name: "grade", Instance: "repo__x-1", DurMS: 9, OK: true, Reward: intp(0), - Tail: `{"valid": true, "resolved": false, "f2p": "1/1", "p2p": "3/4"}`}) - rec.Record(Event{RunID: "r", Event: EventStep, Name: "grade", Instance: "repo__x-2", DurMS: 9, OK: true, Reward: intp(0), Tail: "no json"}) - rec.Record(Event{RunID: "r", Event: EventStep, Name: "grade", Instance: "repo__x-1", DurMS: 9, OK: true, Reward: intp(0), - Tail: `{"valid": false, "resolved": false}`}) - rec.Record(Event{RunID: "r", Event: EventTask, Step: 2, DurMS: 700, OK: true, Reward: intp(0)}) - rec.Record(Event{RunID: "r", Event: EventBarrier, Step: 1, DurMS: 650, OK: true, Count: 2}) - rec.Record(Event{RunID: "r", Event: EventPark, Name: "suspend", Step: 1, DurMS: 90, OK: true, Count: 1}) - rec.Record(Event{RunID: "r", Event: EventRun, Name: "steps", DurMS: 2000, OK: false, Count: 3}) - if rec.Len() != 14 { - t.Fatalf("recorded %d", rec.Len()) - } - - rep, err := Summarize(strings.NewReader(buf.String())) - if err != nil { - t.Fatal(err) - } - if rep.RunID != "r" || rep.Mode != "steps" || rep.RunMS != 2000 { - t.Fatalf("header: %+v", rep) - } - if rep.Tasks != 3 || rep.Failed != 1 || rep.Rewarded != 1 { - t.Fatalf("tasks: %+v", rep) - } - find := func(event, name string) Stat { - for _, s := range rep.Stats { - if s.Event == event && s.Name == name { - return s - } - } - t.Fatalf("no stat %s/%s in %+v", event, name, rep.Stats) - return Stat{} - } - if c := find(EventCreate, ""); c.Count != 2 || c.Failed != 1 || c.P50MS != 100 || c.MaxMS != 100 { - t.Fatalf("create: %+v", c) - } - if s := find(EventStep, "run_tests"); s.Count != 1 || s.P50MS != 50 { - t.Fatalf("step: %+v", s) - } - if len(rep.Steps) != 2 || rep.Steps[0].BarrierMS != 650 || rep.Steps[0].ParkMS != 90 || rep.Steps[0].Live != 2 || rep.Steps[1].WakeP50MS != 40 { - t.Fatalf("steps: %+v", rep.Steps) - } - if len(rep.Errors) != 1 || rep.Errors[0].Err != "boom" || rep.Errors[0].Count != 1 { - t.Fatalf("errors: %+v", rep.Errors) - } - if w := find(EventWake, ""); w.Retried != 1 { - t.Fatalf("wake retried: %+v", w) - } - if len(rep.Unresolved) != 2 || rep.Unresolved[0].Instance != "repo__x-1" || rep.Unresolved[0].Reason != "f2p 1/1 p2p 3/4" || - rep.Unresolved[1].Reason != "reward 0" { - t.Fatalf("unresolved: %+v", rep.Unresolved) - } - var excl bytes.Buffer - rep.WriteExclusions(&excl, "test list") - if !strings.Contains(excl.String(), "# test list\n") || !strings.Contains(excl.String(), "repo__x-1 f2p 1/1 p2p 3/4\n") { - t.Fatalf("exclusions:\n%s", excl.String()) - } - if v := verdict(`{"valid": false}`); v != "invalid run" { - t.Fatalf("verdict: %q", v) - } - var text bytes.Buffer - rep.WriteText(&text) - for _, want := range []string{"run r", "create", "run_tests", "barrier_s", "boom"} { - if !strings.Contains(text.String(), want) { - t.Errorf("text report lacks %q:\n%s", want, text.String()) - } - } -} - -type failingWriter struct{} - -func (failingWriter) Write([]byte) (int, error) { return 0, errors.New("disk full") } - -func TestRecorderKeepsFirstWriteError(t *testing.T) { - rec := NewRecorder(failingWriter{}) - rec.Record(Event{RunID: "r", Event: EventCreate}) - rec.Record(Event{RunID: "r", Event: EventDelete}) - if err := rec.Err(); err == nil || err.Error() != "disk full" { - t.Fatalf("Err = %v", err) - } - if NewRecorder(nil).Err() != nil { - t.Fatal("a discarding recorder has no error") - } -} diff --git a/internal/benchmarking/swebench/run.go b/internal/benchmarking/swebench/run.go deleted file mode 100644 index ebd713af..00000000 --- a/internal/benchmarking/swebench/run.go +++ /dev/null @@ -1,637 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package swebench - -import ( - "context" - "errors" - "fmt" - "hash/fnv" - "log/slog" - "regexp" - "strings" - "sync" - "time" - - env "github.com/agent-substrate/env/clients/go" - "google.golang.org/grpc/codes" - "google.golang.org/grpc/status" -) - -// Mode selects the workload shape. -type Mode string - -const ( - // ModeEval runs every (image, task) once: create, replay, delete. - ModeEval Mode = "eval" - // ModeSteps keeps a fleet across training steps and parks it in between. - ModeSteps Mode = "steps" -) - -// Park is what happens to the fleet between training steps. -type Park string - -const ( - ParkNone Park = "none" // environments stay running - ParkSuspend Park = "suspend" // SuspendEnvironment; the next command wakes it - ParkDelete Park = "delete" // DeleteEnvironment; the next step creates again -) - -// Config is one run. Every field that shapes the load is a parameter so the -// same binary serves the smoke test and the 500-wide run. -type Config struct { - RunID string - Backend Backend - Instances []Instance - Trajectory *Trajectory - Recorder *Recorder - Logger *slog.Logger - - Mode Mode - // Width is the most tasks in flight at once (the "500 wide"). - Width int - // TasksPerImage is the eval mode fan-out per image. - TasksPerImage int - // Rollouts and Steps shape steps mode: Instances x Rollouts environments - // driven through Steps training steps. - Rollouts int - Steps int - Park Park - - // CreateTimeout bounds one create call; StepTimeout is the default for a - // run step without its own timeout. - CreateTimeout time.Duration - StepTimeout time.Duration - // WakeBudget is how long the first command after a create or a park is - // retried on transient errors: a wake that outlasts the router's parking - // budget answers Unavailable, and a first create whose image is still - // being golden-baked answers Unknown until the bake is done. WakeBackoff - // is the initial pause between attempts; it doubles up to five times - // that. - WakeBudget time.Duration - WakeBackoff time.Duration - // Keep leaves the environments in place when the run ends. - Keep bool -} - -// Summary is what a run produced, for the exit status and the log. -type Summary struct { - Tasks int - Failed int - Rewarded int - Duration time.Duration -} - -func (c *Config) validate() error { - switch { - case c.RunID == "": - return errors.New("run id is required") - case !runIDPattern.MatchString(c.RunID): - return fmt.Errorf("run id %q must be lowercase letters, digits and '-'", c.RunID) - case len(c.RunID) > maxRunIDLen: - return fmt.Errorf("run id %q is longer than %d characters; it prefixes every environment name", c.RunID, maxRunIDLen) - case c.Backend == nil: - return errors.New("backend is required") - case len(c.Instances) == 0: - return errors.New("no instances") - case c.Trajectory == nil: - return errors.New("trajectory is required") - case c.Width < 1: - return errors.New("width must be >= 1") - } - switch c.Mode { - case ModeEval: - if c.TasksPerImage < 1 { - return errors.New("tasks per image must be >= 1") - } - case ModeSteps: - if c.Rollouts < 1 || c.Steps < 1 { - return errors.New("rollouts and steps must be >= 1") - } - switch c.Park { - case ParkNone, ParkSuspend, ParkDelete: - default: - return fmt.Errorf("unknown park mode %q", c.Park) - } - default: - return fmt.Errorf("unknown mode %q", c.Mode) - } - return nil -} - -func (c *Config) defaults() { - if c.Recorder == nil { - c.Recorder = NewRecorder(nil) - } - if c.Logger == nil { - c.Logger = slog.Default() - } - if c.CreateTimeout <= 0 { - c.CreateTimeout = 5 * time.Minute - } - if c.StepTimeout <= 0 { - c.StepTimeout = 10 * time.Minute - } - if c.WakeBudget <= 0 { - c.WakeBudget = 3 * time.Minute - } - if c.WakeBackoff <= 0 { - c.WakeBackoff = DefaultWakeBackoff - } -} - -// DefaultWakeBackoff is the initial pause between first-command retries. -const DefaultWakeBackoff = 2 * time.Second - -// runIDPattern is Substrate's name rule (a k8s short name); the run id is a -// prefix of every environment name. maxRunIDLen keeps the longest name the -// loader builds (run id, 36 characters of instance id, a replay index and a -// step suffix) under the 63-character limit. -var runIDPattern = regexp.MustCompile(`^[a-z0-9]([-a-z0-9]*[a-z0-9])?$`) - -const maxRunIDLen = 20 - -// cleanupTimeout bounds the deletes that run after a cancelled or finished -// run; they use a context that survives the cancellation so nothing leaks. -const cleanupTimeout = 10 * time.Minute - -func ms(d time.Duration) float64 { return float64(d) / float64(time.Millisecond) } - -// tail returns at most n trailing bytes of s. -func tail(s string, n int) string { - if len(s) <= n { - return s - } - return s[len(s)-n:] -} - -// environment is one sandbox and its progress. One goroutine owns it at a -// time; phases are separated by barriers, so it needs no lock. -type environment struct { - // id is the current actor name; base is the name of the first one. A - // recreate after a delete park gets a fresh name, since Substrate - // deletes asynchronously and the old name is not free immediately. - id string - base string - inst Instance - task int - traj *Trajectory - - sb Sandbox - createStart time.Time - // failed environments are skipped for the rest of a steps run. - failed bool -} - -type runner struct { - cfg *Config - rec *Recorder - log *slog.Logger - - mu sync.Mutex - tasks int - failed int - rewarded int -} - -// Run executes the configured workload and returns its summary. A non-nil -// error means the run could not be carried out; task failures are counted in -// the Summary and recorded as events, not returned. -func Run(ctx context.Context, cfg Config) (*Summary, error) { - if err := cfg.validate(); err != nil { - return nil, err - } - cfg.defaults() - r := &runner{cfg: &cfg, rec: cfg.Recorder, log: cfg.Logger} - - envs, err := r.plan() - if err != nil { - return nil, err - } - start := time.Now() - switch cfg.Mode { - case ModeEval: - r.runEval(ctx, envs) - case ModeSteps: - r.runSteps(ctx, envs) - } - dur := time.Since(start) - s := &Summary{Tasks: r.tasks, Failed: r.failed, Rewarded: r.rewarded, Duration: dur} - r.rec.Record(Event{RunID: cfg.RunID, Event: EventRun, Name: string(cfg.Mode), DurMS: ms(dur), - OK: s.Failed == 0, Count: s.Tasks}) - return s, ctx.Err() -} - -// plan renders the trajectory per instance and names every environment. -func (r *runner) plan() ([]*environment, error) { - fanout := r.cfg.TasksPerImage - if r.cfg.Mode == ModeSteps { - fanout = r.cfg.Rollouts - } - envs := make([]*environment, 0, len(r.cfg.Instances)*fanout) - for _, in := range r.cfg.Instances { - traj, err := r.cfg.Trajectory.Render(in) - if err != nil { - return nil, err - } - for k := 0; k < fanout; k++ { - id := EnvID(r.cfg.RunID, in.InstanceID, k) - envs = append(envs, &environment{ - id: id, - base: id, - inst: in, - task: k, - traj: traj, - }) - } - } - return envs, nil -} - -// EnvID derives a Substrate-valid name (lowercase alphanumerics and '-', at -// most 63 characters) from the run, the instance and the replay index. A -// truncated instance id gets a hash suffix so distinct instances stay -// distinct. -func EnvID(runID, instanceID string, k int) string { - const maxInstance = 36 - b := strings.Builder{} - for _, c := range strings.ToLower(instanceID) { - switch { - case c >= 'a' && c <= 'z', c >= '0' && c <= '9': - b.WriteRune(c) - default: - b.WriteByte('-') - } - } - inst := strings.Trim(b.String(), "-") - if len(inst) > maxInstance { - h := fnv.New32a() - _, _ = h.Write([]byte(instanceID)) - inst = fmt.Sprintf("%s-%08x", strings.Trim(inst[:maxInstance-9], "-"), h.Sum32()) - } - return fmt.Sprintf("%s-%s-%d", runID, inst, k) -} - -// forEach runs fn over every index with at most width in flight and waits -// for all of them. It stops launching when ctx is done. -func forEach(ctx context.Context, n, width int, fn func(i int)) { - sem := make(chan struct{}, width) - var wg sync.WaitGroup - for i := 0; i < n; i++ { - select { - case <-ctx.Done(): - wg.Wait() - return - case sem <- struct{}{}: - } - wg.Add(1) - go func(i int) { - defer wg.Done() - defer func() { <-sem }() - fn(i) - }(i) - } - wg.Wait() -} - -// runEval is the evaluation shape: every environment is created, replayed -// once and deleted, Width at a time. -func (r *runner) runEval(ctx context.Context, envs []*environment) { - var done int - forEach(ctx, len(envs), r.cfg.Width, func(i int) { - e := envs[i] - if r.create(ctx, e, 0) { - r.replay(ctx, e, 0) - } - if !r.cfg.Keep { - // Delete with a context that survives a cancelled run, or a - // Ctrl-C at width 500 leaves 500 environments behind. - cleanup, cancel := context.WithTimeout(context.WithoutCancel(ctx), cleanupTimeout) - r.remove(cleanup, e, 0) - cancel() - } - r.mu.Lock() - done++ - if done%100 == 0 || done == len(envs) { - r.log.Info("eval progress", slog.Int("done", done), slog.Int("total", len(envs)), - slog.Int("failed", r.failed), slog.Int("rewarded", r.rewarded)) - } - r.mu.Unlock() - }) -} - -// runSteps is the training shape: the fleet is created once, every step -// replays the trajectory on all of it behind a barrier, and the fleet is -// parked between steps. -func (r *runner) runSteps(ctx context.Context, envs []*environment) { - forEach(ctx, len(envs), r.cfg.Width, func(i int) { r.create(ctx, envs[i], 0) }) - r.log.Info("fleet created", slog.Int("environments", len(envs)), slog.Int("failed", r.countFailed(envs))) - - for step := 1; step <= r.cfg.Steps; step++ { - if ctx.Err() != nil { - break - } - if step > 1 && r.cfg.Park == ParkDelete { - forEach(ctx, len(envs), r.cfg.Width, func(i int) { - if !envs[i].failed { - r.create(ctx, envs[i], step) - } - }) - } - barrier := time.Now() - forEach(ctx, len(envs), r.cfg.Width, func(i int) { - if !envs[i].failed { - r.replay(ctx, envs[i], step) - } - }) - live := len(envs) - r.countFailed(envs) - r.rec.Record(Event{RunID: r.cfg.RunID, Event: EventBarrier, Step: step, DurMS: ms(time.Since(barrier)), - OK: ctx.Err() == nil, Count: live}) - r.log.Info("step done", slog.Int("step", step), slog.Duration("barrier", time.Since(barrier)), - slog.Int("live", live), slog.Int("failed", r.failed), slog.Int("rewarded", r.rewarded)) - - if step == r.cfg.Steps || r.cfg.Park == ParkNone || ctx.Err() != nil { - continue - } - parkStart := time.Now() - forEach(ctx, len(envs), r.cfg.Width, func(i int) { - e := envs[i] - if e.failed { - return - } - switch r.cfg.Park { - case ParkSuspend: - r.suspend(ctx, e, step) - case ParkDelete: - r.remove(ctx, e, step) - } - }) - r.rec.Record(Event{RunID: r.cfg.RunID, Event: EventPark, Name: string(r.cfg.Park), Step: step, - DurMS: ms(time.Since(parkStart)), OK: ctx.Err() == nil, Count: live}) - } - - if r.cfg.Keep { - return - } - // Tear down with a context that survives a cancelled run so nothing - // leaks on Ctrl-C. - cleanup, cancel := context.WithTimeout(context.WithoutCancel(ctx), cleanupTimeout) - defer cancel() - forEach(cleanup, len(envs), r.cfg.Width, func(i int) { - if envs[i].sb != nil { - r.remove(cleanup, envs[i], r.cfg.Steps) - } - }) -} - -func (r *runner) countFailed(envs []*environment) int { - n := 0 - for _, e := range envs { - if e.failed { - n++ - } - } - return n -} - -// create makes the environment and records the call. It returns false and -// marks the environment failed when the create did not succeed. -func (r *runner) create(ctx context.Context, e *environment, step int) bool { - if step > 1 { - e.id = fmt.Sprintf("%s-s%d", e.base, step) - } - cctx, cancel := context.WithTimeout(ctx, r.cfg.CreateTimeout) - defer cancel() - e.createStart = time.Now() - sb, err := r.cfg.Backend.Create(cctx, e.id, e.inst.Image) - ev := r.event(EventCreate, e, step, time.Since(e.createStart), err) - r.rec.Record(ev) - if err != nil { - e.failed = true - // A task that never got an environment is still a failed task, in - // the report as well as in the summary. - task := r.event(EventTask, e, step, 0, err) - r.rec.Record(task) - r.countTask(e, step, false, nil) - r.log.Warn("create failed", slog.String("env", e.id), slog.String("err", err.Error())) - // The server may still finish a create whose call failed or timed - // out on our side; delete by name so it cannot leak or collide. - cleanup, cancel := context.WithTimeout(context.WithoutCancel(ctx), cleanupTimeout) - defer cancel() - if derr := r.cfg.Backend.Delete(cleanup, e.id); derr != nil && !isNotFound(derr) { - r.log.Warn("delete after failed create", slog.String("env", e.id), slog.String("err", derr.Error())) - } - return false - } - e.sb = sb - return true -} - -// replay runs the trajectory once on e and records one task event. A fresh -// environment (its first step, or any step after a delete park) runs the -// Setup steps first; one that survived a park runs the Reset steps first. -func (r *runner) replay(ctx context.Context, e *environment, step int) { - fresh := step <= 1 || r.cfg.Park == ParkDelete - prefix := e.traj.Reset - if fresh { - prefix = e.traj.Setup - } - steps := append(append([]Step{}, prefix...), e.traj.Steps...) - // The first instruction of a task either follows a create (eval, or a - // delete park) or a park (suspend); its wall time is the TTFI or the - // wake, and it is the one that may find the environment not yet - // serving. - first := true - firstEvent := EventTTFI - firstStart := e.createStart - if step > 1 && r.cfg.Park != ParkDelete { - firstEvent = EventWake - firstStart = time.Time{} - } - - taskStart := time.Now() - ok := true - var reward *int - for _, s := range steps { - // No think time before the first command: it would land in the - // ttfi or wake measurement, and the "model" has nothing to think - // about before it has seen the environment. - if s.Think > 0 && !first { - select { - case <-ctx.Done(): - ok = false - case <-time.After(time.Duration(s.Think)): - } - if !ok { - break - } - } - stepStart := time.Now() - if firstStart.IsZero() { - firstStart = stepStart - } - res, attempts, err := r.doStep(ctx, e, s, first) - dur := time.Since(stepStart) - if first { - fe := r.event(firstEvent, e, step, time.Since(firstStart), err) - fe.Attempts = attempts - r.rec.Record(fe) - first = false - } - ev := r.event(EventStep, e, step, dur, err) - ev.Name = s.Name - if s.Run != "" && err == nil { - ev.Exit = intp(res.ExitCode) - } - switch { - case err != nil: - ok = false - case s.Reward: - reward = intp(0) - if res.ExitCode == 0 { - *reward = 1 - } - ev.Reward = reward - // The grader's verdict is the last line of a reward step's - // output; keep enough of it for the report to read it back. - ev.Tail = tail(res.Stdout+res.Stderr, 4000) - case s.ExpectExit != nil && res.ExitCode != *s.ExpectExit: - ok = false - ev.OK = false - ev.Err = fmt.Sprintf("exit %d, expected %d", res.ExitCode, *s.ExpectExit) - ev.Tail = tail(res.Stdout+res.Stderr, 2000) - } - r.rec.Record(ev) - if !ok { - r.log.Warn("step failed", slog.String("env", e.id), slog.String("step", s.Name), slog.String("err", ev.Err)) - break - } - } - if ok && reward == nil && e.traj.HasReward() { - ok = false - } - ev := r.event(EventTask, e, step, time.Since(taskStart), nil) - ev.OK = ok - ev.Reward = reward - r.rec.Record(ev) - r.countTask(e, step, ok, reward) - if !ok && r.cfg.Mode == ModeSteps { - e.failed = true - } -} - -// doStep performs one step and reports how many attempts it took. The -// first step of a task retries transient errors within WakeBudget: the -// environment may still be waking, or its image may still be baking. -func (r *runner) doStep(ctx context.Context, e *environment, s Step, first bool) (Result, int, error) { - deadline := time.Now().Add(r.cfg.WakeBudget) - wait := r.cfg.WakeBackoff - var ( - res Result - err error - ) - for attempt := 1; ; attempt++ { - if s.Write != nil { - mode, _ := s.Write.FileMode() - err = e.sb.WriteFile(ctx, s.Write.Path, []byte(s.Write.Content), mode) - } else { - timeout := time.Duration(s.Timeout) - if timeout <= 0 { - timeout = r.cfg.StepTimeout - } - sctx, cancel := context.WithTimeout(ctx, timeout+time.Minute) - res, err = e.sb.Run(sctx, s.Run, timeout) - cancel() - } - if err == nil || !first || !isTransient(err) || ctx.Err() != nil || time.Now().Add(wait).After(deadline) { - return res, attempt, err - } - select { - case <-ctx.Done(): - return res, attempt, err - case <-time.After(wait): - } - wait = min(wait*2, 5*r.cfg.WakeBackoff) - } -} - -// isTransient reports whether a first-command error is worth another try. -// Unavailable is the router giving up on a wake; Unknown, Internal and -// Aborted are what a resume that races the golden bake or a worker restart -// surfaces as. Argument and precondition errors are not retried. -// isNotFound reports whether a delete found nothing to delete. -func isNotFound(err error) bool { - if errors.Is(err, env.ErrNotFound) { - return true - } - st, ok := status.FromError(err) - return ok && st.Code() == codes.NotFound -} - -func isTransient(err error) bool { - st, ok := status.FromError(err) - if !ok { - return false - } - switch st.Code() { - case codes.Unavailable, codes.DeadlineExceeded, codes.Unknown, codes.Internal, codes.Aborted: - return true - } - return false -} - -func (r *runner) suspend(ctx context.Context, e *environment, step int) { - start := time.Now() - err := e.sb.Suspend(ctx) - r.rec.Record(r.event(EventSuspend, e, step, time.Since(start), err)) - if err != nil { - r.log.Warn("suspend failed", slog.String("env", e.id), slog.String("err", err.Error())) - } -} - -func (r *runner) remove(ctx context.Context, e *environment, step int) { - if e.sb == nil { - return - } - start := time.Now() - err := e.sb.Delete(ctx) - r.rec.Record(r.event(EventDelete, e, step, time.Since(start), err)) - if err != nil { - r.log.Warn("delete failed", slog.String("env", e.id), slog.String("err", err.Error())) - return - } - e.sb = nil -} - -func (r *runner) event(kind string, e *environment, step int, d time.Duration, err error) Event { - ev := Event{RunID: r.cfg.RunID, Event: kind, Env: e.id, Instance: e.inst.InstanceID, - Step: step, Task: e.task, DurMS: ms(d), OK: err == nil} - if err != nil { - ev.Err = err.Error() - } - return ev -} - -func (r *runner) countTask(e *environment, step int, ok bool, reward *int) { - r.mu.Lock() - defer r.mu.Unlock() - r.tasks++ - if !ok { - r.failed++ - } - if reward != nil && *reward == 1 { - r.rewarded++ - } -} diff --git a/internal/benchmarking/swebench/run_test.go b/internal/benchmarking/swebench/run_test.go deleted file mode 100644 index 908221d7..00000000 --- a/internal/benchmarking/swebench/run_test.go +++ /dev/null @@ -1,588 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package swebench - -import ( - "bytes" - "context" - "encoding/json" - "errors" - "io" - "io/fs" - "log/slog" - "strings" - "sync" - "sync/atomic" - "testing" - "time" - - "google.golang.org/grpc/codes" - "google.golang.org/grpc/status" -) - -// fakeBackend is an in-memory Backend that records calls and can fail on -// demand. Its sandboxes answer every command with exit code 0 unless the -// command names a nonzero exit (e.g. "exit 1"). -type fakeBackend struct { - mu sync.Mutex - created map[string]int - suspended map[string]int - deleted map[string]int - deletedByID map[string]int - // onCreate, when set, runs after every successful create; tests use it - // to cancel a run at a known point. - onCreate func(id string) - commands []string - files map[string]string - failCreate map[string]bool - unavailable int32 // first-command Unavailable errors still to return - - inflight atomic.Int32 - maxInflight atomic.Int32 -} - -func newFake() *fakeBackend { - return &fakeBackend{created: map[string]int{}, suspended: map[string]int{}, deleted: map[string]int{}, - deletedByID: map[string]int{}, files: map[string]string{}, failCreate: map[string]bool{}} -} - -func (f *fakeBackend) Delete(ctx context.Context, id string) error { - f.mu.Lock() - defer f.mu.Unlock() - f.deletedByID[id]++ - if f.created[id] == 0 { - return status.Error(codes.NotFound, "no such environment") - } - return nil -} - -func (f *fakeBackend) Create(ctx context.Context, id, image string) (Sandbox, error) { - f.mu.Lock() - defer f.mu.Unlock() - f.created[id]++ - if f.failCreate[id] { - return nil, errors.New("create refused") - } - if f.onCreate != nil { - f.onCreate(id) - } - return &fakeSandbox{b: f, id: id}, nil -} - -type fakeSandbox struct { - b *fakeBackend - id string -} - -func (s *fakeSandbox) enter() func() { - n := s.b.inflight.Add(1) - for { - m := s.b.maxInflight.Load() - if n <= m || s.b.maxInflight.CompareAndSwap(m, n) { - break - } - } - return func() { s.b.inflight.Add(-1) } -} - -func (s *fakeSandbox) Run(ctx context.Context, command string, timeout time.Duration) (Result, error) { - defer s.enter()() - if atomic.AddInt32(&s.b.unavailable, -1) >= 0 { - return Result{}, status.Error(codes.Unavailable, "waking") - } - atomic.AddInt32(&s.b.unavailable, 1) - time.Sleep(time.Millisecond) - s.b.mu.Lock() - s.b.commands = append(s.b.commands, s.id+": "+command) - s.b.mu.Unlock() - if strings.HasPrefix(command, "exit ") { - return Result{ExitCode: int(command[5] - '0'), Stderr: "failed"}, nil - } - return Result{Stdout: "ok"}, nil -} - -func (s *fakeSandbox) WriteFile(ctx context.Context, path string, content []byte, mode fs.FileMode) error { - defer s.enter()() - s.b.mu.Lock() - defer s.b.mu.Unlock() - s.b.files[s.id+":"+path] = string(content) - return nil -} - -func (s *fakeSandbox) Suspend(ctx context.Context) error { - s.b.mu.Lock() - defer s.b.mu.Unlock() - s.b.suspended[s.id]++ - return nil -} - -// Delete fails on a done context, as the real API would, so a cleanup that -// runs on a cancelled run context is visible as a leak. -func (s *fakeSandbox) Delete(ctx context.Context) error { - if err := ctx.Err(); err != nil { - return err - } - s.b.mu.Lock() - defer s.b.mu.Unlock() - s.b.deleted[s.id]++ - return nil -} - -func instances(n int) []Instance { - out := make([]Instance, n) - for i := range out { - out[i] = Instance{InstanceID: "repo__repo-" + string(rune('a'+i)), Image: "reg/img@" + goodDigest, - Patch: "PATCH", EvalScript: "EVAL"} - } - return out -} - -func testTrajectory(t *testing.T, rewardCmd string) *Trajectory { - t.Helper() - tr, err := ParseTrajectory([]byte(` -name: test -setup: - - {name: baseline, run: "git commit baseline", expect_exit: 0} -reset: - - {name: reset, run: "git checkout", expect_exit: 0} -steps: - - {name: write, write: {path: /tmp/p, content: "{{.Patch}}"}} - - {name: apply, run: "git apply", expect_exit: 0} - - {name: tests, run: "` + rewardCmd + `", reward: true} -`)) - if err != nil { - t.Fatal(err) - } - return tr -} - -func baseConfig(t *testing.T, b Backend, out *bytes.Buffer) Config { - t.Helper() - var w io.Writer - if out != nil { - w = out - } - return Config{ - RunID: "t", - Backend: b, - Instances: instances(3), - Trajectory: testTrajectory(t, "true"), - Recorder: NewRecorder(w), - Logger: slog.New(slog.NewTextHandler(&bytes.Buffer{}, nil)), - Width: 2, - WakeBackoff: time.Millisecond, - } -} - -func TestEvalCreatesReplaysDeletes(t *testing.T) { - f := newFake() - var out bytes.Buffer - cfg := baseConfig(t, f, &out) - cfg.Mode = ModeEval - cfg.TasksPerImage = 2 - - sum, err := Run(context.Background(), cfg) - if err != nil { - t.Fatal(err) - } - if sum.Tasks != 6 || sum.Failed != 0 || sum.Rewarded != 6 { - t.Fatalf("summary %+v", sum) - } - if len(f.created) != 6 || len(f.deleted) != 6 { - t.Fatalf("created %d deleted %d", len(f.created), len(f.deleted)) - } - for id, n := range f.created { - if n != 1 || f.deleted[id] != 1 { - t.Fatalf("env %s created %d deleted %d", id, n, f.deleted[id]) - } - } - if got := f.maxInflight.Load(); got > 2 { - t.Fatalf("max in flight %d > width 2", got) - } - if f.files["t-repo--repo-a-0:/tmp/p"] != "PATCH" { - t.Fatalf("patch not written: %v", f.files) - } - resets, setups := countCommands(f, "git checkout", "git commit baseline") - if resets != 0 || setups != 6 { - t.Fatalf("eval mode: %d resets (want 0), %d setups (want 6, once per environment)", resets, setups) - } - rep := Build(parseEvents(t, out.String())) - if rep.Tasks != 6 || rep.RewardRate != 1 { - t.Fatalf("report %+v", rep) - } - for _, ev := range []string{EventCreate, EventTTFI, EventTask, EventDelete, EventRun} { - if !hasStat(rep, ev) { - t.Errorf("no %s events", ev) - } - } -} - -func TestEvalCleanupSurvivesCancel(t *testing.T) { - f := newFake() - var out bytes.Buffer - cfg := baseConfig(t, f, &out) - cfg.Mode = ModeEval - cfg.TasksPerImage = 1 - cfg.Instances = instances(6) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - // Cancel the run while its first environments are in flight, so later - // creates never start and the in-flight ones must still be deleted. - var once sync.Once - f.onCreate = func(string) { once.Do(cancel) } - if _, err := Run(ctx, cfg); err == nil { - t.Fatal("a cancelled run must report the cancellation") - } - f.mu.Lock() - defer f.mu.Unlock() - if len(f.created) == 0 { - t.Fatal("nothing was created before the cancel; the test proves nothing") - } - for id := range f.created { - if f.deleted[id] != 1 { - t.Fatalf("environment %s leaked after cancel: deleted %d", id, f.deleted[id]) - } - } -} - -func TestFailedCreateIsDeletedByName(t *testing.T) { - f := newFake() - bad := EnvID("t", "repo__repo-a", 0) - f.failCreate[bad] = true - var out bytes.Buffer - cfg := baseConfig(t, f, &out) - cfg.Mode = ModeEval - cfg.TasksPerImage = 1 - if _, err := Run(context.Background(), cfg); err != nil { - t.Fatal(err) - } - if f.deletedByID[bad] != 1 { - t.Fatalf("the failed create must be deleted by name once, got %d", f.deletedByID[bad]) - } - for id, n := range f.deletedByID { - if id != bad && n != 0 { - t.Fatalf("%s deleted by name although its create succeeded", id) - } - } -} - -func TestNoThinkBeforeFirstCommand(t *testing.T) { - f := newFake() - var out bytes.Buffer - cfg := baseConfig(t, f, &out) - cfg.Mode = ModeEval - cfg.TasksPerImage = 1 - cfg.Instances = instances(1) - tr, err := ParseTrajectory([]byte("name: t\nsteps:\n - {name: a, run: \"true\", think: 5s, reward: true}\n")) - if err != nil { - t.Fatal(err) - } - cfg.Trajectory = tr - start := time.Now() - if _, err := Run(context.Background(), cfg); err != nil { - t.Fatal(err) - } - if time.Since(start) > time.Second { - t.Fatal("think time before the first command was not skipped") - } -} - -func TestEvalCountsFailures(t *testing.T) { - f := newFake() - f.failCreate[EnvID("t", "repo__repo-a", 0)] = true - var out bytes.Buffer - cfg := baseConfig(t, f, &out) - cfg.Mode = ModeEval - cfg.TasksPerImage = 1 - cfg.Trajectory = testTrajectory(t, "exit 1") - - sum, err := Run(context.Background(), cfg) - if err != nil { - t.Fatal(err) - } - // One create failure; the other two run to the reward step and get 0. - if sum.Tasks != 3 || sum.Failed != 1 || sum.Rewarded != 0 { - t.Fatalf("summary %+v", sum) - } - rep := Build(parseEvents(t, out.String())) - if rep.RewardRate != 0 || rep.Failed != 1 { - t.Fatalf("report %+v", rep) - } -} - -func TestExpectExitFailsTask(t *testing.T) { - f := newFake() - var out bytes.Buffer - cfg := baseConfig(t, f, &out) - cfg.Mode = ModeEval - cfg.TasksPerImage = 1 - tr, err := ParseTrajectory([]byte("name: t\nsteps:\n - {name: a, run: \"exit 2\", expect_exit: 0}\n - {name: b, run: \"true\"}\n")) - if err != nil { - t.Fatal(err) - } - cfg.Trajectory = tr - sum, err := Run(context.Background(), cfg) - if err != nil { - t.Fatal(err) - } - if sum.Failed != 3 { - t.Fatalf("summary %+v", sum) - } - for _, c := range f.commands { - if strings.HasSuffix(c, ": true") { - t.Fatal("steps after a failed expectation must not run") - } - } -} - -func TestStepsSuspendPark(t *testing.T) { - f := newFake() - var out bytes.Buffer - cfg := baseConfig(t, f, &out) - cfg.Mode = ModeSteps - cfg.Rollouts = 2 - cfg.Steps = 3 - cfg.Park = ParkSuspend - - sum, err := Run(context.Background(), cfg) - if err != nil { - t.Fatal(err) - } - if sum.Tasks != 18 || sum.Failed != 0 || sum.Rewarded != 18 { - t.Fatalf("summary %+v", sum) - } - if len(f.created) != 6 { - t.Fatalf("created %d envs, want 6 created once", len(f.created)) - } - for id, n := range f.created { - if n != 1 || f.suspended[id] != 2 || f.deleted[id] != 1 { - t.Fatalf("env %s created %d suspended %d deleted %d", id, n, f.suspended[id], f.deleted[id]) - } - } - resets, setups := countCommands(f, "git checkout", "git commit baseline") - if resets != 12 || setups != 6 { - t.Fatalf("suspend park: %d resets (want 6 envs x 2 later steps), %d setups (want 6, once per environment)", resets, setups) - } - rep := Build(parseEvents(t, out.String())) - if len(rep.Steps) != 3 || rep.Steps[0].Tasks != 6 || rep.Steps[0].ParkMS == 0 || rep.Steps[2].ParkMS != 0 { - t.Fatalf("steps %+v", rep.Steps) - } - if s := stat(rep, EventWake); s.Count != 12 { - t.Fatalf("wake events %+v, want 12 (6 envs x steps 2 and 3)", s) - } - if s := stat(rep, EventTTFI); s.Count != 6 { - t.Fatalf("ttfi events %+v, want 6", s) - } -} - -func TestStepsDeleteParkRecreates(t *testing.T) { - f := newFake() - var out bytes.Buffer - cfg := baseConfig(t, f, &out) - cfg.Mode = ModeSteps - cfg.Rollouts = 1 - cfg.Steps = 2 - cfg.Park = ParkDelete - - if _, err := Run(context.Background(), cfg); err != nil { - t.Fatal(err) - } - // Each step creates under a fresh name: 3 environments x 2 steps. - if len(f.created) != 6 { - t.Fatalf("created %d distinct names, want 6: %v", len(f.created), f.created) - } - fresh := 0 - for id, n := range f.created { - if n != 1 || f.deleted[id] != 1 || f.suspended[id] != 0 { - t.Fatalf("env %s created %d deleted %d suspended %d", id, n, f.deleted[id], f.suspended[id]) - } - if strings.HasSuffix(id, "-s2") { - fresh++ - } - } - if fresh != 3 { - t.Fatalf("%d step-2 names, want 3", fresh) - } - rep := Build(parseEvents(t, out.String())) - if s := stat(rep, EventTTFI); s.Count != 6 { - t.Fatalf("ttfi events %+v, want 3 envs x 2 creates", s) - } - if hasStat(rep, EventWake) { - t.Fatal("delete park has no wake events") - } - resets, setups := countCommands(f, "git checkout", "git commit baseline") - if resets != 0 || setups != 6 { - t.Fatalf("delete park: %d resets (want 0, every environment is new), %d setups (want 3 envs x 2 creates)", resets, setups) - } -} - -// countCommands counts the fake's recorded commands containing each needle. -func countCommands(f *fakeBackend, resetNeedle, setupNeedle string) (resets, setups int) { - f.mu.Lock() - defer f.mu.Unlock() - for _, c := range f.commands { - if strings.Contains(c, resetNeedle) { - resets++ - } - if strings.Contains(c, setupNeedle) { - setups++ - } - } - return resets, setups -} - -func TestStepsKeepLeavesFleet(t *testing.T) { - f := newFake() - var out bytes.Buffer - cfg := baseConfig(t, f, &out) - cfg.Mode = ModeSteps - cfg.Rollouts = 1 - cfg.Steps = 1 - cfg.Park = ParkNone - cfg.Keep = true - if _, err := Run(context.Background(), cfg); err != nil { - t.Fatal(err) - } - if len(f.deleted) != 0 || len(f.suspended) != 0 { - t.Fatalf("keep: deleted %v suspended %v", f.deleted, f.suspended) - } -} - -func TestFirstCommandRetriesUnavailable(t *testing.T) { - f := newFake() - f.unavailable = 2 - var out bytes.Buffer - cfg := baseConfig(t, f, &out) - cfg.Mode = ModeEval - cfg.TasksPerImage = 1 - cfg.Instances = instances(1) - cfg.WakeBudget = time.Second - tr, err := ParseTrajectory([]byte("name: t\nsteps:\n - {name: a, run: \"true\", reward: true}\n")) - if err != nil { - t.Fatal(err) - } - cfg.Trajectory = tr - sum, err := Run(context.Background(), cfg) - if err != nil { - t.Fatal(err) - } - if sum.Failed != 0 { - t.Fatalf("summary %+v", sum) - } - - rep := Build(parseEvents(t, out.String())) - if s := stat(rep, EventTTFI); s.Count != 1 { - t.Fatalf("ttfi %+v", s) - } - for _, e := range parseEvents(t, out.String()) { - if e.Event == EventTTFI && e.Attempts != 3 { - t.Fatalf("ttfi attempts = %d, want 3", e.Attempts) - } - } - - f = newFake() - f.unavailable = 1000 - cfg.Backend = f - cfg.WakeBudget = 5 * time.Millisecond - cfg.WakeBackoff = 2 * time.Millisecond - sum, err = Run(context.Background(), cfg) - if err != nil { - t.Fatal(err) - } - if sum.Failed != 1 { - t.Fatalf("budget exhausted must fail the task: %+v", sum) - } -} - -func TestIsTransient(t *testing.T) { - for code, want := range map[codes.Code]bool{ - codes.Unavailable: true, codes.Unknown: true, codes.Internal: true, codes.Aborted: true, codes.DeadlineExceeded: true, - codes.InvalidArgument: false, codes.NotFound: false, codes.FailedPrecondition: false, codes.AlreadyExists: false, - } { - if got := isTransient(status.Error(code, "x")); got != want { - t.Errorf("%v: transient = %v, want %v", code, got, want) - } - } - if isTransient(errors.New("plain")) { - t.Error("a non-status error is not transient") - } -} - -func TestConfigValidation(t *testing.T) { - f := newFake() - good := baseConfig(t, f, nil) - good.Mode = ModeEval - good.TasksPerImage = 1 - bad := []func(c *Config){ - func(c *Config) { c.RunID = "" }, - func(c *Config) { c.RunID = "Swe-1" }, - func(c *Config) { c.RunID = "swe_1" }, - func(c *Config) { c.RunID = strings.Repeat("a", 21) }, - func(c *Config) { c.Backend = nil }, - func(c *Config) { c.Instances = nil }, - func(c *Config) { c.Trajectory = nil }, - func(c *Config) { c.Width = 0 }, - func(c *Config) { c.TasksPerImage = 0 }, - func(c *Config) { c.Mode = "bogus" }, - func(c *Config) { c.Mode = ModeSteps; c.Rollouts = 0; c.Steps = 1 }, - func(c *Config) { c.Mode = ModeSteps; c.Rollouts = 1; c.Steps = 1; c.Park = "bogus" }, - } - for i, mutate := range bad { - c := good - mutate(&c) - if _, err := Run(context.Background(), c); err == nil { - t.Errorf("case %d: expected a validation error", i) - } - } -} - -func TestEnvID(t *testing.T) { - if got := EnvID("swe-1", "astropy__astropy-7336", 2); got != "swe-1-astropy--astropy-7336-2" { - t.Fatalf("got %q", got) - } - long := EnvID("r", strings.Repeat("x", 80)+"__y", 0) - if len(long) > 63 { - t.Fatalf("too long: %q", long) - } - if long == EnvID("r", strings.Repeat("x", 80)+"__z", 0) { - t.Fatal("truncated ids must stay distinct") - } -} - -func parseEvents(t *testing.T, jsonl string) []Event { - t.Helper() - var events []Event - for _, line := range strings.Split(strings.TrimSpace(jsonl), "\n") { - var e Event - if err := json.Unmarshal([]byte(line), &e); err != nil { - t.Fatal(err) - } - events = append(events, e) - } - return events -} - -func stat(rep *Report, event string) Stat { - for _, s := range rep.Stats { - if s.Event == event { - return s - } - } - return Stat{} -} - -func hasStat(rep *Report, event string) bool { - return stat(rep, event).Count > 0 -} diff --git a/internal/benchmarking/swebench/trajectories/swebench-gold.yaml b/internal/benchmarking/swebench/trajectories/swebench-gold.yaml deleted file mode 100644 index c6b6102a..00000000 --- a/internal/benchmarking/swebench/trajectories/swebench-gold.yaml +++ /dev/null @@ -1,86 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# The default trajectory: what a SWE agent does on a SWE-bench task, with the -# gold patch standing in for the model's edit so the reward is known in -# advance. A healthy run rewards every task; a drop in reward rate is a -# correctness regression, not a model result. -# -# /testbed and /opt/miniconda3/envs/testbed are where SWE-bench task images -# keep the repository and its interpreter. Think times are the simulated -# policy latency between commands. -name: swebench-gold -# The image's working tree is the state a rollout starts from, and it is not -# the git HEAD: SWE-bench edits files when it builds the image (sphinx's -# tox.ini gains the per-test output flag the grader reads) and editable -# installs leave metadata in untracked directories. Setup commits that state -# as a local baseline once per environment; reset returns to it. -setup: - - name: baseline - run: >- - cd /testbed && - git -c user.email=bench@localhost -c user.name=bench add -A && - git -c user.email=bench@localhost -c user.name=bench commit -q --allow-empty -m "benchmark baseline" - expect_exit: 0 -reset: - - name: reset_repo - run: cd /testbed && git reset -q --hard && git clean -qfd - expect_exit: 0 -steps: - - name: explore - run: cd /testbed && git status --short | head -n 20; ls - think: 1s - expect_exit: 0 - - name: read_files - run: cd /testbed && git diff --stat HEAD && git log -1 --format=%H - think: 2s - expect_exit: 0 - - name: write_patch - write: - path: /tmp/gold.diff - content: "{{.Patch}}" - - name: apply_patch - run: cd /testbed && git apply -v /tmp/gold.diff - think: 3s - expect_exit: 0 - - name: write_eval - write: - path: /tmp/eval.sh - content: "{{.EvalScript}}" - mode: "0755" - # The harness script's exit code is not the verdict: it runs without -e, - # ends with a git checkout that fails when the test patch added files, and - # the test command itself exits nonzero when unrelated tests in the same - # files fail. The verdict is what the SWE-bench harness computes: parse the - # log with the repo's parser and check the FAIL_TO_PASS and PASS_TO_PASS - # tests. swebench-grade.py does that inside the sandbox from the instance - # record; its exit code is the reward. - - name: run_tests - run: >- - /bin/bash /tmp/eval.sh > /tmp/eval.log 2>&1; - tail -c 1500 /tmp/eval.log; - grep -o '>>>>> Test Exit Code: [0-9]*' /tmp/eval.log | tail -n 1 - think: 1s - timeout: 20m - - name: write_grader - write: - path: /tmp/grade.py - file: swebench-grade.py - - name: write_instance - write: - path: /tmp/instance.json - content: "{{json .}}" - - name: grade - run: /opt/miniconda3/bin/python /tmp/grade.py /tmp/instance.json /tmp/eval.log - reward: true diff --git a/internal/benchmarking/swebench/trajectories/swebench-grade.py b/internal/benchmarking/swebench/trajectories/swebench-grade.py deleted file mode 100644 index e5d7db53..00000000 --- a/internal/benchmarking/swebench/trajectories/swebench-grade.py +++ /dev/null @@ -1,317 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Grade a SWE-bench eval log the way the SWE-bench harness does. - - python swebench-grade.py INSTANCE.json EVAL.log - -INSTANCE.json is the dataset record (log_parser, eval_type, fail_to_pass, -pass_to_pass). The log is parsed with the repo's parser into a test -> status -map; the instance is resolved when every FAIL_TO_PASS test passed and every -PASS_TO_PASS test passed or was skipped. Exit 0 when resolved, 1 otherwise, -2 when the log shows the run itself was invalid (patch did not apply, tests -errored or timed out). One JSON line with the counts goes to stdout. - -Standard library only: it runs inside the task image with whatever Python -the image carries. -""" - -import json -import re -import sys - -PASSED, FAILED, SKIPPED, ERROR, XFAIL = "PASSED", "FAILED", "SKIPPED", "ERROR", "XFAIL" -STATUSES = (FAILED, PASSED, SKIPPED, ERROR, XFAIL) -PASSING = {PASSED, XFAIL} - -START_OUTPUT = ">>>>> Start Test Output" -END_OUTPUT = ">>>>> End Test Output" -BAD_MARKERS = (">>>>> Patch Apply Failed", ">>>>> Reset Failed", ">>>>> Tests Errored", ">>>>> Tests Timed Out") -EXIT_CODE = re.compile(r">>>>> Test Exit Code:\s*(-?\d+)") -# Evidence that a suite ran at all, for logs whose parser records only -# failures; mirrors the harness's SUITE_RAN. -SUITE_RAN = re.compile( - r"Executed [1-9]\d* of \d+" - r"|TOTAL: [1-9]\d* (?:SUCCESS|FAILED)" - r"|[1-9]\d* passing" - r"|Tests:\s+[1-9]\d*" - r"|Test Suites:\s+(?:\d+ \w+, )*[1-9]\d* total" - r"|^# tests [1-9]\d*" - r"|[1-9]\d* specs?, \d+ failures?" - r"|': ok$", - re.M, -) -SKIP_SUMMARY = re.compile(r"^\[\d+\]$") -CONTROL = str.maketrans("", "", "".join(chr(c) for c in range(1, 32))) - - -def _status_line(line, status_map, rewrite=None): - """pytest -rA short summary: ' [ - reason]'.""" - if not line.startswith(STATUSES): - return - if line.startswith(FAILED): - line = line.replace(" - ", " ") - parts = line.split() - if len(parts) < 2: - return - status, name = parts[0], parts[1] - if status == SKIPPED and SKIP_SUMMARY.match(name): - return - if rewrite: - name = rewrite(name) - status_map[name] = status - - -def parse_pytest(log): - status_map = {} - for line in log.split("\n"): - _status_line(line, status_map) - return status_map - - -_OPTION = re.compile(r"(.*?)\[(.*)\]") - - -def parse_pytest_options(log): - def rewrite(name): - m = _OPTION.search(name) - if not m: - return name - main, option = m.groups() - if option.startswith("/") and not option.startswith("//") and "*" not in option: - option = "/" + option.split("/")[-1] - return f"{main}[{option}]" - - status_map = {} - for line in log.split("\n"): - _status_line(line, status_map, rewrite) - return status_map - - -def parse_pytest_v2(log): - status_map = {} - for line in log.split("\n"): - line = re.sub(r"\[(\d+)m", "", line).translate(CONTROL) - if line.startswith(STATUSES): - if line.startswith(FAILED): - line = line.split(" - ", 1)[0] - parts = line.split() - if len(parts) >= 2 and not (parts[0] == SKIPPED and SKIP_SUMMARY.match(parts[1])): - status_map[" ".join(parts[1:])] = parts[0] - elif line.endswith(STATUSES): - parts = line.split() - if len(parts) >= 2: - status_map[" ".join(parts[:-1])] = parts[-1] - return status_map - - -def parse_matplotlib(log): - log = log.replace("MouseButton.LEFT", "1").replace("MouseButton.RIGHT", "3") - return parse_pytest(log) - - -def parse_seaborn(log): - status_map = {} - for line in log.split("\n"): - parts = line.split() - if len(parts) < 2: - continue - if line.startswith(FAILED): - status_map[parts[1]] = FAILED - elif f" {PASSED} " in line and parts[1] == PASSED: - status_map[parts[0]] = PASSED - elif line.startswith(PASSED): - status_map[parts[1]] = PASSED - return status_map - - -def parse_sympy(log): - status_map = {} - for m in re.findall(r"(_*) (.*)\.py:(.*) (_*)", log): - status_map[f"{m[1]}.py:{m[2]}"] = FAILED - for line in log.split("\n"): - line = line.strip() - if not line.startswith("test_"): - continue - name = line.split()[0] - if line.endswith(" E"): - status_map[name] = ERROR - elif line.endswith(" F"): - status_map[name] = FAILED - elif line.endswith(" ok"): - status_map[name] = PASSED - return status_map - - -def parse_django(log): - status_map = {} - prev_test = None - for raw in log.split("\n"): - line = raw.strip() - if "--version is equivalent to version" in line: - status_map["--version is equivalent to version"] = PASSED - if " ... " in line: - prev_test = line.split(" ... ")[0] - for suffix in (" ... ok", " ... OK", " ... OK"): - if line.endswith(suffix): - if line.startswith("Applying sites.0002_alter_domain_unique...test_no_migrations"): - line = line.split("...", 1)[-1].strip() - status_map[line.rsplit(suffix, 1)[0]] = PASSED - break - if " ... skipped" in line: - status_map[line.split(" ... skipped")[0]] = SKIPPED - if line.endswith(" ... FAIL"): - status_map[line.split(" ... FAIL")[0]] = FAILED - if line.startswith("FAIL:"): - status_map[line.split()[1].strip()] = FAILED - if line.endswith(" ... ERROR"): - status_map[line.split(" ... ERROR")[0]] = ERROR - if line.startswith("ERROR:"): - status_map[line.split()[1].strip()] = ERROR - if line.startswith("ok") and prev_test is not None: - status_map[prev_test] = PASSED - for pattern in ( - r"^(.*?)\s\.\.\.\sTesting\ against\ Django\ installed\ in\ ((?s:.*?))\ silenced\)\.\nok$", - r"^(.*?)\s\.\.\.\sInternal\ Server\ Error:\ \/(.*)\/\nok$", - r"^(.*?)\s\.\.\.\sSystem check identified no issues \(0 silenced\)\nok$", - ): - for m in re.finditer(pattern, log, re.MULTILINE): - status_map[m.group(1)] = PASSED - return status_map - - -PARSERS = { - "parse_log_pytest": parse_pytest, - "parse_log_astroid": parse_pytest, - "parse_log_flask": parse_pytest, - "parse_log_marshmallow": parse_pytest, - "parse_log_pvlib": parse_pytest, - "parse_log_pyvista": parse_pytest, - "parse_log_sqlfluff": parse_pytest, - "parse_log_xarray": parse_pytest, - "parse_log_pytest_options": parse_pytest_options, - "parse_log_pydicom": parse_pytest_options, - "parse_log_requests": parse_pytest_options, - "parse_log_pylint": parse_pytest_options, - "parse_log_pytest_v2": parse_pytest_v2, - "parse_log_astropy": parse_pytest_v2, - "parse_log_scikit": parse_pytest_v2, - "parse_log_sphinx": parse_pytest_v2, - "parse_log_matplotlib": parse_matplotlib, - "parse_log_seaborn": parse_seaborn, - "parse_log_sympy": parse_sympy, - "parse_log_django": parse_django, -} - - -def resolve_case(case, status_map): - """Exact match, or a prefix match for ids truncated mid-parameter when the - candidates agree on the outcome (676 Verified ids are truncated).""" - if case in status_map: - return case - if case.count("[") > case.count("]"): - matches = [k for k in status_map if k.startswith(case)] - if matches and len({status_map[k] in PASSING for k in matches}) == 1: - return matches[0] - return None - - -def passed(case, status_map): - key = resolve_case(case, status_map) - return key is not None and status_map[key] in PASSING - - -def maintained(case, status_map): - key = resolve_case(case, status_map) - return passed(case, status_map) or (key is not None and status_map[key] == SKIPPED) - - -def failed(case, status_map): - """Missing, FAILED, ERROR or SKIPPED, as the harness's test_failed.""" - key = resolve_case(case, status_map) - return key is None or status_map[key] in (FAILED, ERROR, SKIPPED) - - -def status_map_from_log(parser, content): - """The harness's validity rules, then the parser over the marked region.""" - if any(marker in content for marker in BAD_MARKERS): - return None - start, end = content.find(START_OUTPUT), content.find(END_OUTPUT) - if start < 0 or end < 0 or end < start: - return None - sliced = content[start + len(START_OUTPUT):end] - status_map = parser(sliced) or parser(content) - if not status_map and not SUITE_RAN.search(content): - return None - m = EXIT_CODE.search(content) - exit_code = int(m.group(1)) if m else None - if exit_code not in (None, 0) and status_map and not any(s in (FAILED, ERROR) for s in status_map.values()): - return None - return status_map - - -def grade(instance, content): - parser = PARSERS.get(instance.get("log_parser", "")) - if parser is None: - raise SystemExit(f"unknown log_parser {instance.get('log_parser')!r}") - status_map = status_map_from_log(parser, content) - if status_map is None: - return {"valid": False, "resolved": False} - fail_only = instance.get("eval_type", "pass_and_fail").lower() == "fail_only" - f2p = instance.get("fail_to_pass", []) - p2p = instance.get("pass_to_pass", []) - - # As in the harness, every test lands in a success list, a failure list, - # or neither (an unknown status counts for nothing); the instance is - # resolved when both failure lists are empty. - def split(tests, ok, bad): - success = [t for t in tests if ok(t)] - failure = [t for t in tests if not ok(t) and bad(t)] - return success, failure - - if fail_only: - bad = lambda t: status_map.get(t) == FAILED # noqa: E731 - f2p_ok, f2p_bad = split(f2p, lambda t: not bad(t), bad) - p2p_ok, p2p_bad = split(p2p, lambda t: not bad(t), bad) - else: - f2p_ok, f2p_bad = split(f2p, lambda t: passed(t, status_map), lambda t: failed(t, status_map)) - p2p_ok, p2p_bad = split(p2p, lambda t: maintained(t, status_map), lambda t: failed(t, status_map)) - resolved = not f2p_bad and not p2p_bad - return { - "valid": True, - "resolved": resolved, - "f2p": f"{len(f2p_ok)}/{len(f2p)}", - "p2p": f"{len(p2p_ok)}/{len(p2p)}", - "f2p_failed": f2p_bad[:5], - "p2p_failed": p2p_bad[:5], - } - - -def main(argv): - if len(argv) != 3: - print(__doc__, file=sys.stderr) - return 2 - with open(argv[1]) as f: - instance = json.load(f) - with open(argv[2], errors="replace") as f: - content = f.read() - report = grade(instance, content) - print(json.dumps(report)) - if not report["valid"]: - return 2 - return 0 if report["resolved"] else 1 - - -if __name__ == "__main__": - sys.exit(main(sys.argv)) diff --git a/internal/benchmarking/swebench/trajectories/test_swebench_grade.py b/internal/benchmarking/swebench/trajectories/test_swebench_grade.py deleted file mode 100644 index 7e906af2..00000000 --- a/internal/benchmarking/swebench/trajectories/test_swebench_grade.py +++ /dev/null @@ -1,154 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Tests for the in-sandbox grader next to it: - - python3 -m pytest internal/benchmarking/swebench/trajectories -q -""" - -import importlib.util -import pathlib - -GRADER = pathlib.Path(__file__).resolve().with_name("swebench-grade.py") - -spec = importlib.util.spec_from_file_location("grader", GRADER) -grader = importlib.util.module_from_spec(spec) -spec.loader.exec_module(grader) - - -def wrap(body, exit_code=0): - return f"+ : '>>>>> Start Test Output'\n{body}\n+ : '>>>>> End Test Output'\n>>>>> Test Exit Code: {exit_code}\n" - - -PYTEST_LOG = """\ -============================= test session starts ============================== -PASSED astropy/units/tests/test_a.py::test_one -PASSED astropy/units/tests/test_a.py::test_param[log(photon s-1)] -FAILED astropy/units/tests/test_a.py::test_unrelated - AssertionError: boom -XFAIL astropy/units/tests/test_a.py::test_known -SKIPPED [3] astropy/units/tests/test_a.py:12: no network -========================= 1 failed, 2 passed in 0.42s ========================== -""" - - -def test_pytest_v2_resolves_despite_unrelated_failure(): - inst = { - "log_parser": "parse_log_astropy", - "eval_type": "pass_and_fail", - "fail_to_pass": ["astropy/units/tests/test_a.py::test_one"], - "pass_to_pass": ["astropy/units/tests/test_a.py::test_known"], - } - rep = grader.grade(inst, wrap(PYTEST_LOG, exit_code=1)) - assert rep["valid"] and rep["resolved"], rep - - -def test_pytest_truncated_parametrized_id_prefix_matches(): - inst = { - "log_parser": "parse_log_pytest", - "fail_to_pass": ["astropy/units/tests/test_a.py::test_param[log(photon"], - "pass_to_pass": [], - } - assert grader.grade(inst, wrap(PYTEST_LOG))["resolved"] - - -def test_f2p_failed_is_not_resolved(): - inst = { - "log_parser": "parse_log_pytest", - "fail_to_pass": ["astropy/units/tests/test_a.py::test_unrelated"], - "pass_to_pass": [], - } - rep = grader.grade(inst, wrap(PYTEST_LOG, exit_code=1)) - assert rep["valid"] and not rep["resolved"] and rep["f2p"] == "0/1" - - -def test_skipped_f2p_is_not_resolved_but_skipped_p2p_is_maintained(): - log = "SKIPPED tests/test_x.py::test_a\nPASSED tests/test_x.py::test_b\n" - inst = {"log_parser": "parse_log_pytest", "fail_to_pass": ["tests/test_x.py::test_a"], "pass_to_pass": []} - assert not grader.grade(inst, wrap(log))["resolved"] - inst = {"log_parser": "parse_log_pytest", "fail_to_pass": ["tests/test_x.py::test_b"], "pass_to_pass": ["tests/test_x.py::test_a"]} - assert grader.grade(inst, wrap(log))["resolved"] - - -DJANGO_LOG = """\ -test_loader (migrations.test_loader.LoaderTests) ... ok -test_namespace (migrations.test_loader.LoaderTests) -Some docstring here ... ok -test_broken (migrations.test_loader.LoaderTests) ... FAIL -test_skipped (migrations.test_loader.LoaderTests) ... skipped 'no db' -test_err (migrations.test_loader.LoaderTests) ... ERROR ----------------------------------------------------------------------- -Ran 5 tests in 0.1s -FAILED (failures=1, errors=1) -""" - - -def test_django_parser_statuses(): - sm = grader.parse_django(DJANGO_LOG) - assert sm["test_loader (migrations.test_loader.LoaderTests)"] == "PASSED" - # A docstring on the next line is what the harness records; the port keeps - # that behavior so verdicts match the harness, brittle as it is. - assert sm["Some docstring here"] == "PASSED" - assert sm["test_broken (migrations.test_loader.LoaderTests)"] == "FAILED" - assert sm["test_skipped (migrations.test_loader.LoaderTests)"] == "SKIPPED" - assert sm["test_err (migrations.test_loader.LoaderTests)"] == "ERROR" - - -def test_django_resolution(): - inst = { - "log_parser": "parse_log_django", - "fail_to_pass": ["test_loader (migrations.test_loader.LoaderTests)"], - "pass_to_pass": ["test_skipped (migrations.test_loader.LoaderTests)"], - } - assert grader.grade(inst, wrap(DJANGO_LOG, exit_code=1))["resolved"] - - -def test_sympy_parser(): - log = "test_one ok\ntest_two F\ntest_three E\n________ sympy/core/tests/test_x.py:test_four ________\n" - sm = grader.parse_sympy(log) - assert sm == {"test_one": "PASSED", "test_two": "FAILED", "test_three": "ERROR", - "sympy/core/tests/test_x.py:test_four": "FAILED"} - - -def test_pytest_options_normalizes_paths(): - log = "PASSED tests/test_a.py::test_file[/tmp/xyz/data.txt]\n" - sm = grader.parse_pytest_options(log) - assert sm == {"tests/test_a.py::test_file[/data.txt]": "PASSED"} - - -def test_invalid_runs(): - inst = {"log_parser": "parse_log_pytest", "fail_to_pass": ["t"], "pass_to_pass": []} - assert not grader.grade(inst, ">>>>> Patch Apply Failed\n" + wrap("PASSED t"))["valid"] - assert not grader.grade(inst, "PASSED t\n")["valid"] - # Nonzero exit with a log that shows no failure means the log is not the run. - assert not grader.grade(inst, wrap("PASSED t", exit_code=1))["valid"] - # Markers in the wrong order, or no parsed result and no sign the suite ran. - assert not grader.grade(inst, "+ : '>>>>> End Test Output'\nPASSED t\n+ : '>>>>> Start Test Output'\n")["valid"] - assert not grader.grade(inst, wrap("nothing here"))["valid"] - assert grader.grade(inst, wrap("Executed 3 of 3 specs"))["valid"] - - -def test_unknown_status_counts_for_nothing(): - # A recorded status outside the known set is neither a pass nor a failure - # in the harness, so it must not block resolution. - log = "PASSEDX tests/test_x.py::test_a\nPASSED tests/test_x.py::test_b\n" - inst = {"log_parser": "parse_log_pytest", "fail_to_pass": ["tests/test_x.py::test_b"], - "pass_to_pass": ["tests/test_x.py::test_a"]} - rep = grader.grade(inst, wrap(log)) - assert rep["resolved"] and rep["p2p"] == "0/1" and rep["p2p_failed"] == [], rep - - -def test_unknown_parser_exits(): - import pytest - with pytest.raises(SystemExit): - grader.grade({"log_parser": "nope"}, wrap("")) diff --git a/internal/benchmarking/swebench/trajectory.go b/internal/benchmarking/swebench/trajectory.go deleted file mode 100644 index 80227089..00000000 --- a/internal/benchmarking/swebench/trajectory.go +++ /dev/null @@ -1,332 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package swebench - -import ( - "bytes" - "embed" - "encoding/json" - "errors" - "fmt" - "io/fs" - "os" - "path/filepath" - "strconv" - "strings" - "text/template" - "time" - - "sigs.k8s.io/yaml" -) - -//go:embed trajectories/* -var trajectoryFS embed.FS - -// DefaultTrajectory is the built-in trajectory used when none is given. -const DefaultTrajectory = "swebench-gold" - -// Duration is a time.Duration that decodes from a string such as "1s". -type Duration time.Duration - -// UnmarshalJSON accepts a quoted Go duration string. -func (d *Duration) UnmarshalJSON(b []byte) error { - s, err := strconv.Unquote(string(b)) - if err != nil { - return fmt.Errorf("duration must be a quoted string like \"500ms\": %s", string(b)) - } - v, err := time.ParseDuration(s) - if err != nil { - return err - } - *d = Duration(v) - return nil -} - -// MarshalJSON writes the duration back as a string. -func (d Duration) MarshalJSON() ([]byte, error) { - return []byte(strconv.Quote(time.Duration(d).String())), nil -} - -// WriteSpec puts a file into the sandbox. Content is a template over the -// Instance, so a step can write the gold patch or the eval script. -type WriteSpec struct { - Path string `json:"path"` - Content string `json:"content"` - // File names a file next to the trajectory whose contents become - // Content when the trajectory is loaded: a built-in asset for built-in - // trajectories, a path relative to the YAML file otherwise. File - // contents are written as they are; only inline Content is a template. - File string `json:"file,omitempty"` - // Mode is an octal file mode string; empty means 0644. - Mode string `json:"mode,omitempty"` - - // raw marks Content as loaded from File, so Render leaves it alone. - raw bool -} - -// FileMode parses Mode. -func (w WriteSpec) FileMode() (fs.FileMode, error) { - if w.Mode == "" { - return 0o644, nil - } - v, err := strconv.ParseUint(w.Mode, 8, 32) - if err != nil { - return 0, fmt.Errorf("mode %q: %w", w.Mode, err) - } - return fs.FileMode(v), nil -} - -// Step is one action of a trajectory: exactly one of Run or Write. Run and -// Write.Content are text/template strings evaluated against the Instance. -type Step struct { - Name string `json:"name"` - Run string `json:"run,omitempty"` - Write *WriteSpec `json:"write,omitempty"` - // Think is the simulated model latency before the step (an RL rollout - // waits for the policy between commands). - Think Duration `json:"think,omitempty"` - // Timeout bounds a Run step; zero means the runner's default. - Timeout Duration `json:"timeout,omitempty"` - // ExpectExit, when set, fails the task if the exit code differs. A - // Reward step never fails the task on exit code: its exit code is the - // reward. - ExpectExit *int `json:"expect_exit,omitempty"` - // Reward marks the step whose exit code decides the task reward: - // 1 when it exits 0, else 0. At most one step per trajectory. - Reward bool `json:"reward,omitempty"` -} - -// Trajectory is the scripted rollout replayed in every environment. -type Trajectory struct { - Name string `json:"name"` - // Setup runs once per environment, before its first rollout, for - // example to record the image's working state as the baseline that - // Reset returns to. - Setup []Step `json:"setup,omitempty"` - // Reset runs before Steps on every later step an environment survives - // into, so a rollout starts from the image's state again (the gold patch - // cannot be applied twice). - Reset []Step `json:"reset,omitempty"` - Steps []Step `json:"steps"` -} - -// Names lists the built-in trajectories. -func Names() []string { - entries, err := fs.ReadDir(trajectoryFS, "trajectories") - if err != nil { - return nil - } - var names []string - for _, e := range entries { - if n, ok := strings.CutSuffix(e.Name(), ".yaml"); ok { - names = append(names, n) - } - } - return names -} - -// LoadTrajectory returns a built-in trajectory by name or one read from a -// file when source names an existing path. Empty means DefaultTrajectory. -func LoadTrajectory(source string) (*Trajectory, error) { - if source == "" { - source = DefaultTrajectory - } - if data, err := os.ReadFile(source); err == nil { - dir := filepath.Dir(source) - return parseWithAssets(data, func(name string) ([]byte, error) { - return os.ReadFile(filepath.Join(dir, name)) - }) - } - data, err := trajectoryFS.ReadFile("trajectories/" + source + ".yaml") - if err != nil { - return nil, fmt.Errorf("no file or built-in trajectory %q (built-in: %s)", source, strings.Join(Names(), ", ")) - } - return parseWithAssets(data, func(name string) ([]byte, error) { - return trajectoryFS.ReadFile("trajectories/" + name) - }) -} - -// parseWithAssets parses the trajectory and resolves every write.file -// through read, so the loaded trajectory carries its contents. -func parseWithAssets(data []byte, read func(name string) ([]byte, error)) (*Trajectory, error) { - t, err := ParseTrajectory(data) - if err != nil { - return nil, err - } - for _, steps := range [][]Step{t.Setup, t.Reset, t.Steps} { - for i := range steps { - w := steps[i].Write - if w == nil || w.File == "" { - continue - } - if strings.Contains(w.File, "/") || strings.Contains(w.File, "..") { - return nil, fmt.Errorf("trajectory %s step %s: write.file must be a plain file name", t.Name, steps[i].Name) - } - b, err := read(w.File) - if err != nil { - return nil, fmt.Errorf("trajectory %s step %s: %w", t.Name, steps[i].Name, err) - } - w.Content = string(b) - w.File = "" - w.raw = true - } - } - return t, nil -} - -// ParseTrajectory decodes and validates a YAML trajectory. -func ParseTrajectory(data []byte) (*Trajectory, error) { - var t Trajectory - if err := yaml.UnmarshalStrict(data, &t); err != nil { - return nil, fmt.Errorf("decode trajectory: %w", err) - } - if err := t.Validate(); err != nil { - return nil, err - } - return &t, nil -} - -// Validate checks the trajectory's shape without rendering it. -func (t *Trajectory) Validate() error { - if t.Name == "" { - return errors.New("trajectory: name is required") - } - if len(t.Steps) == 0 { - return fmt.Errorf("trajectory %s: at least one step is required", t.Name) - } - rewards := 0 - for _, s := range t.Steps { - if err := s.validate(); err != nil { - return fmt.Errorf("trajectory %s: %w", t.Name, err) - } - if s.Reward { - rewards++ - } - } - if rewards > 1 { - return fmt.Errorf("trajectory %s: more than one reward step", t.Name) - } - for section, steps := range map[string][]Step{"setup": t.Setup, "reset": t.Reset} { - for _, s := range steps { - if err := s.validate(); err != nil { - return fmt.Errorf("trajectory %s %s: %w", t.Name, section, err) - } - if s.Reward { - return fmt.Errorf("trajectory %s %s: step %s cannot be a reward step", t.Name, section, s.Name) - } - } - } - return nil -} - -func (s Step) validate() error { - if s.Name == "" { - return errors.New("step without a name") - } - switch { - case s.Run != "" && s.Write != nil: - return fmt.Errorf("step %s: run and write are exclusive", s.Name) - case s.Run == "" && s.Write == nil: - return fmt.Errorf("step %s: one of run or write is required", s.Name) - case s.Write != nil && s.Write.Path == "": - return fmt.Errorf("step %s: write.path is required", s.Name) - case s.Write != nil && s.Write.File != "" && s.Write.Content != "": - return fmt.Errorf("step %s: write.file and write.content are exclusive", s.Name) - case s.Write != nil && (s.Reward || s.ExpectExit != nil): - return fmt.Errorf("step %s: reward and expect_exit apply to run steps only", s.Name) - } - if s.Write != nil { - if _, err := s.Write.FileMode(); err != nil { - return fmt.Errorf("step %s: %w", s.Name, err) - } - } - if s.Think < 0 || s.Timeout < 0 { - return fmt.Errorf("step %s: think and timeout cannot be negative", s.Name) - } - return nil -} - -// HasReward reports whether a step decides the task reward. -func (t *Trajectory) HasReward() bool { - for _, s := range t.Steps { - if s.Reward { - return true - } - } - return false -} - -// Render evaluates every template in the trajectory against in and returns -// the concrete copy to replay. The original is not modified. -func (t *Trajectory) Render(in Instance) (*Trajectory, error) { - out := &Trajectory{Name: t.Name} - var err error - if out.Setup, err = renderSteps(t.Name, t.Setup, in); err != nil { - return nil, err - } - if out.Reset, err = renderSteps(t.Name, t.Reset, in); err != nil { - return nil, err - } - if out.Steps, err = renderSteps(t.Name, t.Steps, in); err != nil { - return nil, err - } - return out, nil -} - -func renderSteps(tname string, steps []Step, in Instance) ([]Step, error) { - out := make([]Step, len(steps)) - for i, s := range steps { - r := s - var err error - if r.Run, err = renderText(s.Run, in); err != nil { - return nil, fmt.Errorf("trajectory %s step %s run: %w", tname, s.Name, err) - } - if s.Write != nil { - w := *s.Write - if !w.raw { - if w.Content, err = renderText(s.Write.Content, in); err != nil { - return nil, fmt.Errorf("trajectory %s step %s write: %w", tname, s.Name, err) - } - } - r.Write = &w - } - out[i] = r - } - return out, nil -} - -// templateFuncs are available in step templates. json renders a value as -// JSON, for handing an instance's fields to a script in the sandbox. -var templateFuncs = template.FuncMap{ - "json": func(v any) (string, error) { - b, err := json.Marshal(v) - return string(b), err - }, -} - -func renderText(text string, in Instance) (string, error) { - if !strings.Contains(text, "{{") { - return text, nil - } - tmpl, err := template.New("step").Option("missingkey=error").Funcs(templateFuncs).Parse(text) - if err != nil { - return "", err - } - var buf bytes.Buffer - if err := tmpl.Execute(&buf, in); err != nil { - return "", err - } - return buf.String(), nil -} diff --git a/internal/benchmarking/swebench/trajectory_test.go b/internal/benchmarking/swebench/trajectory_test.go deleted file mode 100644 index ab656052..00000000 --- a/internal/benchmarking/swebench/trajectory_test.go +++ /dev/null @@ -1,185 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package swebench - -import ( - "encoding/json" - "os" - "path/filepath" - "strings" - "testing" - "time" -) - -func TestBuiltinTrajectoriesLoadAndRender(t *testing.T) { - names := Names() - if len(names) == 0 { - t.Fatal("no built-in trajectories") - } - in := Instance{InstanceID: "x", Image: "reg/img@" + goodDigest, Patch: "PATCH", EvalScript: "EVAL"} - for _, n := range names { - tr, err := LoadTrajectory(n) - if err != nil { - t.Fatalf("%s: %v", n, err) - } - r, err := tr.Render(in) - if err != nil { - t.Fatalf("%s render: %v", n, err) - } - if len(r.Steps) != len(tr.Steps) { - t.Fatalf("%s: rendered %d steps, want %d", n, len(r.Steps), len(tr.Steps)) - } - for _, s := range r.Steps { - if s.Write != nil && (s.Write.Content == "" || s.Write.File != "") { - t.Fatalf("%s step %s: write.file not resolved into content", n, s.Name) - } - } - if len(r.Setup) != len(tr.Setup) || len(r.Reset) != len(tr.Reset) { - t.Fatalf("%s: setup/reset not rendered", n) - } - } - def, err := LoadTrajectory("") - if err != nil || def.Name != DefaultTrajectory { - t.Fatalf("default: %v %v", def, err) - } - if !def.HasReward() { - t.Fatal("default trajectory has no reward step") - } -} - -func TestRenderSubstitutesInstanceFields(t *testing.T) { - tr, err := ParseTrajectory([]byte(` -name: t -steps: - - name: w - write: {path: /tmp/p, content: "{{.Patch}}", mode: "0755"} - - name: r - run: "echo {{.InstanceID}}" - think: 250ms - timeout: 1m - reward: true -`)) - if err != nil { - t.Fatal(err) - } - r, err := tr.Render(Instance{InstanceID: "id-1", Patch: "line1\nline2"}) - if err != nil { - t.Fatal(err) - } - if r.Steps[0].Write.Content != "line1\nline2" || r.Steps[1].Run != "echo id-1" { - t.Fatalf("render: %+v", r.Steps) - } - if tr.Steps[1].Run != "echo {{.InstanceID}}" { - t.Fatal("Render must not modify the original") - } - if time.Duration(r.Steps[1].Think) != 250*time.Millisecond || time.Duration(r.Steps[1].Timeout) != time.Minute { - t.Fatalf("durations: %+v", r.Steps[1]) - } - if m, _ := r.Steps[0].Write.FileMode(); m != 0o755 { - t.Fatalf("mode: %v", m) - } -} - -func TestRenderUnknownFieldFails(t *testing.T) { - tr, err := ParseTrajectory([]byte("name: t\nsteps:\n - name: r\n run: \"{{.Nope}}\"\n")) - if err != nil { - t.Fatal(err) - } - if _, err := tr.Render(Instance{}); err == nil { - t.Fatal("unknown field must fail at render") - } -} - -func TestParseTrajectoryRejects(t *testing.T) { - cases := map[string]string{ - "no name": "steps:\n - name: a\n run: true\n", - "no steps": "name: t\n", - "run and write": "name: t\nsteps:\n - name: a\n run: x\n write: {path: /p}\n", - "neither": "name: t\nsteps:\n - name: a\n", - "two rewards": "name: t\nsteps:\n - {name: a, run: x, reward: true}\n - {name: b, run: y, reward: true}\n", - "reward reset": "name: t\nreset:\n - {name: a, run: x, reward: true}\nsteps:\n - {name: b, run: y}\n", - "reward setup": "name: t\nsetup:\n - {name: a, run: x, reward: true}\nsteps:\n - {name: b, run: y}\n", - "write reward": "name: t\nsteps:\n - name: a\n write: {path: /p}\n reward: true\n", - "bad mode": "name: t\nsteps:\n - name: a\n write: {path: /p, mode: \"zz\"}\n", - "bad think": "name: t\nsteps:\n - name: a\n run: x\n think: soon\n", - "neg think": "name: t\nsteps:\n - name: a\n run: x\n think: -1s\n", - "neg timeout": "name: t\nsteps:\n - name: a\n run: x\n timeout: -1m\n", - "unknown field": "name: t\nsteps:\n - name: a\n run: x\n bogus: 1\n", - } - for name, in := range cases { - if _, err := ParseTrajectory([]byte(in)); err == nil { - t.Errorf("%s: expected an error", name) - } else if strings.Contains(err.Error(), "panic") { - t.Errorf("%s: %v", name, err) - } - } -} - -func TestWriteFileResolvesNextToTrajectory(t *testing.T) { - dir := t.TempDir() - // The asset contains template syntax on purpose: file contents are not - // rendered. - if err := os.WriteFile(filepath.Join(dir, "asset.py"), []byte("print('{{not a template}}')\n"), 0o644); err != nil { - t.Fatal(err) - } - yamlPath := filepath.Join(dir, "t.yaml") - if err := os.WriteFile(yamlPath, []byte("name: t\nsteps:\n - name: w\n write: {path: /tmp/a.py, file: asset.py}\n - name: r\n run: \"true\"\n"), 0o644); err != nil { - t.Fatal(err) - } - tr, err := LoadTrajectory(yamlPath) - if err != nil { - t.Fatal(err) - } - if got := tr.Steps[0].Write.Content; got != "print('{{not a template}}')\n" { - t.Fatalf("content %q", got) - } - r, err := tr.Render(Instance{InstanceID: "x"}) - if err != nil { - t.Fatalf("rendering a trajectory with a file asset: %v", err) - } - if r.Steps[0].Write.Content != tr.Steps[0].Write.Content { - t.Fatal("file asset contents must not be rendered") - } - for _, bad := range []string{"missing.py", "../asset.py", "sub/asset.py"} { - if err := os.WriteFile(yamlPath, []byte("name: t\nsteps:\n - name: w\n write: {path: /tmp/a, file: "+bad+"}\n"), 0o644); err != nil { - t.Fatal(err) - } - if _, err := LoadTrajectory(yamlPath); err == nil { - t.Errorf("write.file %q must fail", bad) - } - } - if _, err := ParseTrajectory([]byte("name: t\nsteps:\n - name: w\n write: {path: /p, file: a, content: b}\n")); err == nil { - t.Error("file and content together must fail") - } -} - -func TestJSONTemplateFunc(t *testing.T) { - tr, err := ParseTrajectory([]byte("name: t\nsteps:\n - name: w\n write: {path: /tmp/i.json, content: \"{{json .}}\"}\n")) - if err != nil { - t.Fatal(err) - } - in := Instance{InstanceID: "x", Image: "reg/img@" + goodDigest, FailToPass: []string{"a::b[1]"}, LogParser: "parse_log_pytest"} - r, err := tr.Render(in) - if err != nil { - t.Fatal(err) - } - var back Instance - if err := json.Unmarshal([]byte(r.Steps[0].Write.Content), &back); err != nil { - t.Fatalf("not JSON: %v: %s", err, r.Steps[0].Write.Content) - } - if back.InstanceID != "x" || back.LogParser != "parse_log_pytest" || len(back.FailToPass) != 1 { - t.Fatalf("round trip: %+v", back) - } -} diff --git a/internal/internaltest/fakecontrol/fakecontrol.go b/internal/internaltest/fakecontrol/fakecontrol.go index 1a87acc1..15a3ab3f 100644 --- a/internal/internaltest/fakecontrol/fakecontrol.go +++ b/internal/internaltest/fakecontrol/fakecontrol.go @@ -50,7 +50,6 @@ type Server struct { // DenyCreateAtespace makes CreateAtespace fail with PermissionDenied, as // Substrate does for an identity allowed to use but not create atespaces. DenyCreateAtespace bool - templates map[string]*ateapipb.ActorTemplate } // New returns an empty fake control server. @@ -58,77 +57,9 @@ func New() *Server { return &Server{ actors: make(map[string]*ateapipb.Actor), atespaces: make(map[string]*ateapipb.Atespace), - templates: make(map[string]*ateapipb.ActorTemplate), } } -// AddTemplate seeds an ActorTemplate, as an operator would have registered it. -func (s *Server) AddTemplate(tmpl *ateapipb.ActorTemplate) { - s.mu.Lock() - defer s.mu.Unlock() - s.templates[key(tmpl.GetMetadata().GetAtespace(), tmpl.GetMetadata().GetName())] = proto.Clone(tmpl).(*ateapipb.ActorTemplate) -} - -// Template returns a copy of a stored ActorTemplate, or nil. -func (s *Server) Template(atespace, name string) *ateapipb.ActorTemplate { - s.mu.Lock() - defer s.mu.Unlock() - t, ok := s.templates[key(atespace, name)] - if !ok { - return nil - } - return proto.Clone(t).(*ateapipb.ActorTemplate) -} - -// TemplateCount is how many ActorTemplates are stored. -func (s *Server) TemplateCount() int { - s.mu.Lock() - defer s.mu.Unlock() - return len(s.templates) -} - -func (s *Server) GetActorTemplate(ctx context.Context, req *ateapipb.GetActorTemplateRequest) (*ateapipb.ActorTemplate, error) { - s.mu.Lock() - defer s.mu.Unlock() - ref := req.GetActorTemplate() - t, ok := s.templates[key(ref.GetAtespace(), ref.GetName())] - if !ok { - return nil, status.Errorf(codes.NotFound, "actor template %q not found", ref.GetName()) - } - return proto.Clone(t).(*ateapipb.ActorTemplate), nil -} - -func (s *Server) CreateActorTemplate(ctx context.Context, req *ateapipb.CreateActorTemplateRequest) (*ateapipb.ActorTemplate, error) { - s.mu.Lock() - defer s.mu.Unlock() - t := req.GetActorTemplate() - md := t.GetMetadata() - if md.GetName() == "" { - return nil, status.Error(codes.InvalidArgument, "metadata.name is required") - } - k := key(md.GetAtespace(), md.GetName()) - if _, exists := s.templates[k]; exists { - return nil, status.Errorf(codes.AlreadyExists, "actor template %q already exists", md.GetName()) - } - // Mirror the server's volume_mounts validation so a bad derivation fails here too. - declared := map[string]bool{} - for _, v := range t.GetVolumes() { - declared[v.GetName()] = true - } - for _, c := range t.GetContainers() { - for _, m := range c.GetVolumeMounts() { - if !declared[m.GetName()] { - return nil, status.Errorf(codes.InvalidArgument, "container %q mounts undeclared volume %q", c.GetName(), m.GetName()) - } - } - } - stored := proto.Clone(t).(*ateapipb.ActorTemplate) - stored.Metadata.Uid = "tmpl-" + md.GetName() - stored.Metadata.Version = 1 - s.templates[k] = stored - return proto.Clone(stored).(*ateapipb.ActorTemplate), nil -} - // Serve starts the fake on a random localhost port and returns its // address and a shutdown function. Like the real ateapi, it serves TLS // with a certificate no client can verify (clients set SkipVerify). diff --git a/internal/internaltest/fakerouter/fakerouter.go b/internal/internaltest/fakerouter/fakerouter.go index 8a3ba9f9..bf0b45ad 100644 --- a/internal/internaltest/fakerouter/fakerouter.go +++ b/internal/internaltest/fakerouter/fakerouter.go @@ -16,17 +16,12 @@ // resolves the target environment from the request's ate-target-actor header // (or, failing that, its Host header) and forwards to that environment's guest // handler, returning 404 for an invalid or duplicated actor reference and 503 -// when the environment is not running. Like the real router's CONNECT -// listener, it also tunnels an HTTP CONNECT whose authority names a port to -// the handler registered for that environment and port. +// when the environment is not running. package fakerouter import ( - "bufio" - "net" "net/http" "net/http/httptest" - "strconv" "strings" "sync" @@ -41,26 +36,12 @@ type Router struct { mu sync.Mutex guests map[string]http.Handler - ports map[portKey]http.Handler -} - -type portKey struct { - id string - port int } // New returns a Router that considers every actor running unless a Running // callback is set. func New() *Router { - return &Router{guests: make(map[string]http.Handler), ports: make(map[portKey]http.Handler)} -} - -// RegisterPort installs the handler that serves CONNECT tunnels to port on -// environment id, standing in for a daemon listening there. -func (r *Router) RegisterPort(id string, port int, h http.Handler) { - r.mu.Lock() - defer r.mu.Unlock() - r.ports[portKey{id, port}] = h + return &Router{guests: make(map[string]http.Handler)} } // Register installs the guest handler serving a environment ID. @@ -76,10 +57,6 @@ func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) { http.Error(w, "invalid actor reference", http.StatusNotFound) return } - if req.Method == http.MethodConnect { - r.serveConnect(w, req, id) - return - } r.mu.Lock() guest := r.guests[id] r.mu.Unlock() @@ -125,80 +102,3 @@ func (r *Router) Serve() (addr string, stop func()) { srv.Start() return strings.TrimPrefix(srv.URL, "http://"), srv.Close } - -// serveConnect answers a CONNECT like the real router: the actor comes from -// ate-target-actor, the port from the authority. On success the hijacked -// connection carries HTTP to the registered port handler. -func (r *Router) serveConnect(w http.ResponseWriter, req *http.Request, id string) { - _, portStr, err := net.SplitHostPort(req.Host) - port, perr := strconv.Atoi(portStr) - if err != nil || perr != nil { - http.Error(w, "CONNECT authority names no port", http.StatusBadRequest) - return - } - r.mu.Lock() - h := r.ports[portKey{id, port}] - r.mu.Unlock() - if h == nil { - http.Error(w, "no actor "+id+" listening on "+portStr, http.StatusNotFound) - return - } - if r.Running != nil && !r.Running(id) { - http.Error(w, "actor "+id+" is not running", http.StatusServiceUnavailable) - return - } - conn, rw, err := http.NewResponseController(w).Hijack() - if err != nil { - http.Error(w, "hijack: "+err.Error(), http.StatusInternalServerError) - return - } - if _, err := conn.Write([]byte("HTTP/1.1 200 OK\r\n\r\n")); err != nil { - conn.Close() - return - } - l := &oneConnListener{conn: &tunnelConn{Conn: conn, r: rw.Reader}, done: make(chan struct{})} - _ = (&http.Server{Handler: h}).Serve(l) -} - -// tunnelConn reads what the hijacked reader buffered before the connection. -type tunnelConn struct { - net.Conn - r *bufio.Reader - closeFn func() -} - -func (c *tunnelConn) Read(b []byte) (int, error) { return c.r.Read(b) } - -func (c *tunnelConn) Close() error { - if c.closeFn != nil { - c.closeFn() - } - return c.Conn.Close() -} - -// oneConnListener hands out one connection, then blocks until it is closed. -type oneConnListener struct { - conn *tunnelConn - once sync.Once - closed sync.Once - done chan struct{} -} - -func (l *oneConnListener) Accept() (net.Conn, error) { - var c net.Conn - l.once.Do(func() { - l.conn.closeFn = func() { l.closed.Do(func() { close(l.done) }) } - c = l.conn - }) - if c != nil { - return c, nil - } - <-l.done - return nil, net.ErrClosed -} - -func (l *oneConnListener) Close() error { - l.closed.Do(func() { close(l.done) }) - return nil -} -func (l *oneConnListener) Addr() net.Addr { return l.conn.LocalAddr() } diff --git a/internal/portproxy/portproxy.go b/internal/portproxy/portproxy.go deleted file mode 100644 index 759ec158..00000000 --- a/internal/portproxy/portproxy.go +++ /dev/null @@ -1,503 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Package portproxy forwards HTTP requests from callers outside the actor -// network to a TCP port inside an environment, such as a daemon that a runtime -// layer starts next to the guest (docs/rootfs-images/RUNTIMES.md). -// -// A request names its target with three headers, EnvIDHeader, -// EnvAtespaceHeader (optional) and EnvPortHeader, or, for clients that cannot -// set headers (WebSocket clients, browsers), with a path prefix that the proxy -// strips: /v1alpha/envs/{id}/ports/{port}/... in the default atespace, or -// /v1alpha/atespaces/{atespace}/envs/{id}/ports/{port}/... . The proxy opens -// an HTTP CONNECT tunnel through the atenet router's CONNECT listener, naming -// the actor in ate-target-actor and the port in the CONNECT authority. The -// router resumes a suspended actor before it answers, then forwards the tunnel -// to that port. The proxy speaks HTTP/1.1 to the daemon through the tunnel and -// keeps idle tunnels for reuse, one pool for each environment and port. -// -// Trust model: the proxy does not authenticate callers. Whoever can reach the -// listener it is served on can reach any TCP port of any actor the router -// serves, in any atespace, including actors that ate-env did not create. A -// daemon behind it must authenticate its own clients. The path form refuses -// requests that a browser sends from another origin (see Options.TrustedOrigins), -// so a web page cannot drive a daemon through a port-forwarded listener. -package portproxy - -import ( - "bufio" - "context" - "errors" - "fmt" - "io" - "log" - "net" - "net/http" - "net/http/httputil" - "net/url" - "regexp" - "strconv" - "strings" - "time" - - "github.com/agent-substrate/env/internal/ate" -) - -// The routing headers. They match the x-env-id and x-env-atespace metadata -// keys of the gRPC API, so a caller uses one vocabulary for both. -const ( - EnvIDHeader = "X-Env-Id" - EnvAtespaceHeader = "X-Env-Atespace" - EnvPortHeader = "X-Env-Port" -) - -// DefaultHandshakeTimeout bounds dialing the router and waiting for its -// answer to CONNECT, which includes resuming a suspended actor. -const DefaultHandshakeTimeout = 2 * time.Minute - -// DefaultWakeBudget bounds how long a request waits for a suspended -// environment to wake. The router's CONNECT is what resumes a suspended -// actor, and the router holds it only for its own parking budget, a few -// seconds; a wake that pulls an image or reads a snapshot across a region -// takes longer. The same two minutes as ate-env-api's guest calls -// (apiservice.DefaultWakeBudget). -const DefaultWakeBudget = 2 * time.Minute - -// Idle tunnels: each one is a TCP stream from the API to an actor, held open -// by the router and the worker. Keep few, and not for long, so an idle -// environment is not held awake and a large fleet does not pin thousands -// of sockets in the API pod. -const ( - idleTunnelTimeout = 10 * time.Second - idleTunnelsPerPort = 4 - idleTunnelsTotal = 1024 -) - -const ( - wakeRetryInitialBackoff = 200 * time.Millisecond - wakeRetryMaxBackoff = 2 * time.Second -) - -// guestPort is where the ate-env guest listens in every environment. It is -// not a daemon port: the API already serves it, and a daemon cannot bind it. -const guestPort = 80 - -// Options configures a Proxy. -type Options struct { - // ConnectAddr is the router's plaintext CONNECT listener, host:port. - ConnectAddr string - // HostSuffix is the router's actor host suffix; the CONNECT authority is - // .:. The router reads only the port from it. - HostSuffix string - // DefaultAtespace applies when a request carries no EnvAtespaceHeader. - DefaultAtespace string - // HandshakeTimeout bounds the CONNECT handshake. Zero means - // DefaultHandshakeTimeout. - HandshakeTimeout time.Duration - // ResponseHeaderTimeout bounds the wait for the daemon's response headers - // after the request is sent. Zero means no limit: a daemon that answers a - // long command only when it finishes is served, and the caller's own - // deadline ends a wait on a daemon that never answers. - ResponseHeaderTimeout time.Duration - // TrustedOrigins are browser origins (scheme://host[:port]) allowed to use - // the path form from a page they serve, next to the listener's own origin. - TrustedOrigins []string - // WakeBudget bounds how long a request waits while the router answers - // 503 for an environment that is suspended or still waking. Zero means - // DefaultWakeBudget. - WakeBudget time.Duration -} - -// Proxy is an http.Handler that forwards each request to the environment port -// its headers name. -type Proxy struct { - opts Options - rp *httputil.ReverseProxy - tr *http.Transport - cop *http.CrossOriginProtection -} - -// New returns a Proxy. It fails only on a malformed trusted origin. -func New(opts Options) (*Proxy, error) { - if opts.HostSuffix == "" { - opts.HostSuffix = ate.DefaultHostSuffix - } - if opts.HandshakeTimeout <= 0 { - opts.HandshakeTimeout = DefaultHandshakeTimeout - } - if opts.WakeBudget <= 0 { - opts.WakeBudget = DefaultWakeBudget - } - p := &Proxy{opts: opts, cop: http.NewCrossOriginProtection()} - for _, o := range opts.TrustedOrigins { - if err := p.cop.AddTrustedOrigin(o); err != nil { - return nil, fmt.Errorf("trusted origin: %w", err) - } - } - p.tr = &http.Transport{ - DialContext: p.dial, - MaxIdleConnsPerHost: idleTunnelsPerPort, - MaxIdleConns: idleTunnelsTotal, - IdleConnTimeout: idleTunnelTimeout, - ResponseHeaderTimeout: opts.ResponseHeaderTimeout, - // Pass encodings through untouched; the daemon and the caller - // negotiate them. - DisableCompression: true, - } - p.rp = &httputil.ReverseProxy{ - Rewrite: p.rewrite, - Transport: p.tr, - // Flush each write at once, so streamed responses (server-sent events, - // chunked output) reach the caller as the daemon produces them. - FlushInterval: -1, - ErrorHandler: p.handleError, - } - return p, nil -} - -// Wants reports whether r uses the header form: it carries EnvPortHeader and -// is not a gRPC call. Callers of the gRPC API never send the header, and the -// tunnel carries HTTP/1.1 only. -// Close drops the idle tunnels. Requests in flight are not affected. -func (p *Proxy) Close() { p.tr.CloseIdleConnections() } - -func Wants(r *http.Request) bool { - return r.Header.Get(EnvPortHeader) != "" && - !strings.HasPrefix(r.Header.Get("Content-Type"), "application/grpc") -} - -// Dispatch returns a handler that sends header-form requests to p.ServeHTTP, -// path-form requests to p.ServePath, and all others to next. The path form is -// matched here rather than on a ServeMux, so the daemon gets the path exactly -// as sent: a mux would redirect a missing trailing slash, "//" and "..", and -// WebSocket clients do not follow redirects. -func Dispatch(p *Proxy, next http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - switch { - case Wants(r): - p.ServeHTTP(w, r) - case isPathForm(r.URL.EscapedPath()): - p.ServePath(w, r) - default: - next.ServeHTTP(w, r) - } - }) -} - -type target struct { - atespace string - id string - port int - // path and rawPath, when set, replace the request's path upstream (the - // path form strips its routing prefix). - path, rawPath string -} - -type targetKey struct{} - -// label is a DNS-1123 label, the shape of environment ids and atespaces. The -// values go into a CONNECT request line and header, so nothing else may pass. -var label = regexp.MustCompile(`^[a-z0-9]([-a-z0-9]{0,61}[a-z0-9])?$`) - -// canonicalPort is a TCP port in decimal: no sign, no leading zero. -var canonicalPort = regexp.MustCompile(`^[1-9][0-9]{0,4}$`) - -func parseTarget(r *http.Request, defaultAtespace string) (target, error) { - for _, h := range []string{EnvIDHeader, EnvAtespaceHeader, EnvPortHeader} { - if len(r.Header.Values(h)) > 1 { - return target{}, fmt.Errorf("%s appears more than once", h) - } - } - if r.Header.Get(EnvIDHeader) == "" { - return target{}, fmt.Errorf("%s is required next to %s", EnvIDHeader, EnvPortHeader) - } - return makeTarget(r.Header.Get(EnvIDHeader), r.Header.Get(EnvAtespaceHeader), r.Header.Get(EnvPortHeader), defaultAtespace) -} - -// makeTarget validates a target from either form. -func makeTarget(id, atespace, portStr, defaultAtespace string) (target, error) { - t := target{id: id, atespace: atespace} - if !label.MatchString(t.id) { - return t, fmt.Errorf("environment id %q is not valid (lowercase alphanumerics and '-', at most 63)", t.id) - } - if t.atespace == "" { - t.atespace = defaultAtespace - } - if !label.MatchString(t.atespace) { - return t, fmt.Errorf("atespace %q is not valid", t.atespace) - } - n, err := strconv.Atoi(portStr) - if !canonicalPort.MatchString(portStr) || err != nil || n > 65535 { - return t, fmt.Errorf("port %q is not a TCP port (1-65535)", portStr) - } - if n == guestPort { - return t, fmt.Errorf("port %d is the ate-env guest, which ate-env-api serves itself; a daemon listens on another port", guestPort) - } - t.port = n - return t, nil -} - -// pathPrefix splits an escaped path of the path form into its routing -// segments and the rest. n is 0 when the path does not have that shape. The -// split is by segment on the escaped path, the way a ServeMux matches, so an -// escaped routing segment (%65nv1 for env1) still ends the prefix where it -// should. -func pathPrefix(escaped string) (id, atespace, portStr, rest string, n int) { - segs := strings.Split(escaped, "/") - switch { - case len(segs) >= 6 && segs[0] == "" && segs[1] == "v1alpha" && segs[2] == "envs" && segs[4] == "ports": - id, portStr, n = segs[3], segs[5], 6 - case len(segs) >= 8 && segs[0] == "" && segs[1] == "v1alpha" && segs[2] == "atespaces" && segs[4] == "envs" && segs[6] == "ports": - atespace, id, portStr, n = segs[3], segs[5], segs[7], 8 - default: - return "", "", "", "", 0 - } - return id, atespace, portStr, "/" + strings.Join(segs[n:], "/"), n -} - -func isPathForm(escaped string) bool { - _, _, _, _, n := pathPrefix(escaped) - return n > 0 -} - -func parsePath(escaped, defaultAtespace string) (target, error) { - id, atespace, portStr, rest, n := pathPrefix(escaped) - if n == 0 { - return target{}, errors.New("path is not /v1alpha/[atespaces/{atespace}/]envs/{id}/ports/{port}[/rest]") - } - // Validation admits only labels and digits, thus an escaped routing - // segment is unescaped first and the value that passes reads the same both - // ways. - for _, s := range []*string{&id, &atespace, &portStr} { - v, err := url.PathUnescape(*s) - if err != nil { - return target{}, fmt.Errorf("bad path escaping: %w", err) - } - *s = v - } - t, err := makeTarget(id, atespace, portStr, defaultAtespace) - if err != nil { - return t, err - } - // The rest keeps its own escaping. The bare prefix, with or without a - // trailing slash, is the daemon's root. - t.rawPath = rest - if t.path, err = url.PathUnescape(rest); err != nil { - return t, fmt.Errorf("bad path escaping: %w", err) - } - if t.rawPath == t.path { - t.rawPath = "" - } - return t, nil -} - -// ServePath forwards a path-form request to the environment port its path -// names, after it strips the routing prefix. It refuses a request that a -// browser sends from another origin: a state-changing method, or a WebSocket -// handshake, which is a GET that opens a channel both ways. The header form -// needs no such check, because a browser sends custom headers to another -// origin only after a CORS preflight, which the proxy never approves. -func (p *Proxy) ServePath(w http.ResponseWriter, r *http.Request) { - t, err := parsePath(r.URL.EscapedPath(), p.opts.DefaultAtespace) - if err != nil { - http.Error(w, err.Error(), http.StatusBadRequest) - return - } - check := r - if r.Header.Get("Upgrade") != "" { - // CrossOriginProtection lets safe methods through; judge a handshake - // as the state-changing request it is. - check = r.WithContext(r.Context()) - check.Method = http.MethodPost - } - if err := p.cop.Check(check); err != nil { - http.Error(w, "cross-origin request to an environment port refused; "+ - "add the page's origin to the port proxy's trusted origins", http.StatusForbidden) - return - } - p.rp.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), targetKey{}, t))) -} - -// ServeHTTP forwards r to the environment port its headers name. -func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { - t, err := parseTarget(r, p.opts.DefaultAtespace) - if err != nil { - http.Error(w, err.Error(), http.StatusBadRequest) - return - } - p.rp.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), targetKey{}, t))) -} - -func (p *Proxy) rewrite(pr *httputil.ProxyRequest) { - t := pr.In.Context().Value(targetKey{}).(target) - pr.SetXForwarded() - // The URL host is the transport's pool key, not an address: one pool for - // each environment and port. dial reads the target back from it. - pr.Out.URL.Scheme = "http" - pr.Out.URL.Host = net.JoinHostPort(t.id+"."+t.atespace, strconv.Itoa(t.port)) - // What a client inside the environment would send. - pr.Out.Host = net.JoinHostPort("127.0.0.1", strconv.Itoa(t.port)) - if t.path != "" { - pr.Out.URL.Path, pr.Out.URL.RawPath = t.path, t.rawPath - } - for _, h := range []string{EnvIDHeader, EnvAtespaceHeader, EnvPortHeader} { - pr.Out.Header.Del(h) - } -} - -// routerError is a CONNECT the router refused. -type routerError struct { - status int - msg string -} - -func (e *routerError) Error() string { - return fmt.Sprintf("router answered CONNECT with %d: %s", e.status, e.msg) -} - -// dial opens a CONNECT tunnel to the environment port encoded in addr (see -// rewrite). While the router answers 503, which is what a suspended actor -// and one still waking both look like, it tries again with backoff until -// WakeBudget is spent: the first CONNECT is what resumes the actor, and the -// router holds it only for its own parking budget. -func (p *Proxy) dial(ctx context.Context, _, addr string) (net.Conn, error) { - host, portStr, err := net.SplitHostPort(addr) - if err != nil { - return nil, err - } - id, atespace, _ := strings.Cut(host, ".") - deadline := time.Now().Add(p.opts.WakeBudget) - backoff := wakeRetryInitialBackoff - for attempt := 1; ; attempt++ { - conn, err := p.connect(ctx, id, atespace, portStr) - var re *routerError - if !errors.As(err, &re) || re.status != http.StatusServiceUnavailable { - return conn, err - } - remaining := time.Until(deadline) - if remaining <= 0 { - // Logged here: handleError keeps router 503s out of the log, and - // a spent wake budget is the one 503 an operator wants to see, - // as the close of the "waiting" line from the first attempt. - err = fmt.Errorf("environment %s/%s did not wake within %s (%d attempts): %w", atespace, id, p.opts.WakeBudget, attempt, err) - log.Printf("port proxy: %v", err) - return nil, err - } - if attempt == 1 { - log.Printf("port proxy: waiting for environment %s/%s to wake (wake budget %s)", atespace, id, p.opts.WakeBudget) - } - select { - case <-ctx.Done(): - return nil, ctx.Err() - case <-time.After(min(backoff, remaining)): - } - backoff = min(backoff*2, wakeRetryMaxBackoff) - } -} - -// connect sends one CONNECT for the environment port and returns the tunnel. -func (p *Proxy) connect(ctx context.Context, id, atespace, portStr string) (net.Conn, error) { - ctx, cancel := context.WithTimeout(ctx, p.opts.HandshakeTimeout) - defer cancel() - var d net.Dialer - conn, err := d.DialContext(ctx, "tcp", p.opts.ConnectAddr) - if err != nil { - return nil, fmt.Errorf("dialing the router's CONNECT listener %s: %w", p.opts.ConnectAddr, err) - } - // A context that ends mid-handshake unblocks the write or read below. - stop := context.AfterFunc(ctx, func() { _ = conn.SetDeadline(time.Unix(1, 0)) }) - - authority := net.JoinHostPort(id+"."+p.opts.HostSuffix, portStr) - req := &http.Request{ - Method: http.MethodConnect, - URL: &url.URL{Opaque: authority}, - Host: authority, - Header: http.Header{ate.TargetActorHeader: {atespace + "/" + id}}, - } - br := bufio.NewReader(conn) - resp, err := func() (*http.Response, error) { - if err := req.Write(conn); err != nil { - return nil, err - } - return http.ReadResponse(br, req) - }() - if !stop() { - err = errors.Join(ctx.Err(), err) - } - if err != nil { - conn.Close() - return nil, fmt.Errorf("CONNECT to %s/%s port %s: %w", atespace, id, portStr, err) - } - if resp.StatusCode != http.StatusOK { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 512)) - conn.Close() - return nil, &routerError{status: resp.StatusCode, msg: strings.TrimSpace(string(body))} - } - if br.Buffered() > 0 { - return &bufferedConn{Conn: conn, r: br}, nil - } - return conn, nil -} - -// bufferedConn reads what the handshake reader buffered past the CONNECT -// response before it reads the connection again. -type bufferedConn struct { - net.Conn - r *bufio.Reader -} - -func (c *bufferedConn) Read(b []byte) (int, error) { return c.r.Read(b) } - -// CloseWrite half-closes the tunnel when the connection can: after an -// upgrade, httputil.ReverseProxy half-closes the backend once the client -// side is done, and embedding net.Conn would otherwise hide the method. -func (c *bufferedConn) CloseWrite() error { - if cw, ok := c.Conn.(interface{ CloseWrite() error }); ok { - return cw.CloseWrite() - } - return nil -} - -// handleError answers a failed forward. The body names the target and a -// short cause; the full error, which can name the router's address, goes to -// the log only. -func (p *Proxy) handleError(w http.ResponseWriter, r *http.Request, err error) { - t, _ := r.Context().Value(targetKey{}).(target) - var re *routerError - var ne net.Error - status, cause := http.StatusBadGateway, "is unreachable; the ate-env-api log has the cause" - switch { - case errors.Is(err, context.Canceled): - return // the caller went away; nobody reads an answer - case errors.As(err, &re): - // 404 (no such environment) and 503 (not running) mean the same to - // the caller as to the router; anything else is a gateway failure. - switch re.status { - case http.StatusNotFound: - status, cause = re.status, "was not found" - case http.StatusServiceUnavailable: - status, cause = re.status, "is not available (not running, or still waking)" - case http.StatusBadGateway: - cause = "refused the connection; a daemon must listen on 0.0.0.0, not only on 127.0.0.1" - } - case errors.Is(err, context.DeadlineExceeded), errors.As(err, &ne) && ne.Timeout(): - status, cause = http.StatusGatewayTimeout, "did not answer in time" - } - // A 404 or 503 from the router is the caller's answer, not an operator's - // problem; a log line for each would flood the log when a fleet wakes. - if status != http.StatusNotFound && status != http.StatusServiceUnavailable { - log.Printf("port proxy: %s/%s port %d: %v", t.atespace, t.id, t.port, err) - } - http.Error(w, fmt.Sprintf("environment %s/%s port %d %s", t.atespace, t.id, t.port, cause), status) -} diff --git a/internal/portproxy/portproxy_test.go b/internal/portproxy/portproxy_test.go deleted file mode 100644 index 3d2b7c47..00000000 --- a/internal/portproxy/portproxy_test.go +++ /dev/null @@ -1,651 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package portproxy_test - -import ( - "bufio" - "fmt" - "io" - "net" - "net/http" - "net/http/httptest" - "strings" - "sync" - "testing" - "time" - - "github.com/agent-substrate/env/internal/ate" - "github.com/agent-substrate/env/internal/internaltest/fakerouter" - "github.com/agent-substrate/env/internal/portproxy" -) - -// recordingRouter wraps the fake router and records each CONNECT's actor -// reference and authority. -type recordingRouter struct { - *fakerouter.Router - mu sync.Mutex - connects []string -} - -func (r *recordingRouter) ServeHTTP(w http.ResponseWriter, req *http.Request) { - if req.Method == http.MethodConnect { - r.mu.Lock() - r.connects = append(r.connects, req.Header.Get(ate.TargetActorHeader)+" "+req.Host) - r.mu.Unlock() - } - r.Router.ServeHTTP(w, req) -} - -func (r *recordingRouter) connectCount() int { - r.mu.Lock() - defer r.mu.Unlock() - return len(r.connects) -} - -type setup struct { - router *recordingRouter - routerAddr string - proxy *httptest.Server -} - -func newSetup(t *testing.T) *setup { - t.Helper() - rr := &recordingRouter{Router: fakerouter.New()} - srv := httptest.NewServer(rr) - t.Cleanup(srv.Close) - addr := strings.TrimPrefix(srv.URL, "http://") - p := mustNew(t, portproxy.Options{ - ConnectAddr: addr, - DefaultAtespace: "ate-env", - HandshakeTimeout: 5 * time.Second, - WakeBudget: 2 * time.Second, - }) - proxy := httptest.NewServer(p) - t.Cleanup(proxy.Close) - return &setup{router: rr, routerAddr: addr, proxy: proxy} -} - -// mustNew builds a proxy and drops its idle tunnels when the test ends, so -// the fake router's tunnel goroutines end with the test. -func mustNew(t *testing.T, opts portproxy.Options) *portproxy.Proxy { - t.Helper() - p, err := portproxy.New(opts) - if err != nil { - t.Fatal(err) - } - t.Cleanup(p.Close) - return p -} - -func (s *setup) do(t *testing.T, method, path, body string, headers map[string]string) (*http.Response, string) { - t.Helper() - req, err := http.NewRequest(method, s.proxy.URL+path, strings.NewReader(body)) - if err != nil { - t.Fatal(err) - } - for k, v := range headers { - req.Header.Set(k, v) - } - resp, err := http.DefaultClient.Do(req) - if err != nil { - t.Fatalf("%s %s: %v", method, path, err) - } - defer resp.Body.Close() - b, _ := io.ReadAll(resp.Body) - return resp, string(b) -} - -func target(id string, port int) map[string]string { - return map[string]string{portproxy.EnvIDHeader: id, portproxy.EnvPortHeader: fmt.Sprint(port)} -} - -func TestForwardsToEnvironmentPort(t *testing.T) { - s := newSetup(t) - var seen *http.Request - var seenBody string - s.router.RegisterPort("env1", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - seen = r - b, _ := io.ReadAll(r.Body) - seenBody = string(b) - w.Header().Set("X-Daemon", "yes") - w.WriteHeader(http.StatusCreated) - io.WriteString(w, "daemon says hi") - })) - - h := target("env1", 8000) - h["X-Custom"] = "kept" - resp, body := s.do(t, http.MethodPost, "/execute?lang=python", `{"code":"1+1"}`, h) - - if resp.StatusCode != http.StatusCreated || body != "daemon says hi" || resp.Header.Get("X-Daemon") != "yes" { - t.Fatalf("got %d %q (X-Daemon %q), want 201 from the daemon", resp.StatusCode, body, resp.Header.Get("X-Daemon")) - } - if seen.Method != http.MethodPost || seen.URL.RequestURI() != "/execute?lang=python" || seenBody != `{"code":"1+1"}` { - t.Errorf("daemon saw %s %s %q", seen.Method, seen.URL.RequestURI(), seenBody) - } - if seen.Host != "127.0.0.1:8000" { - t.Errorf("daemon saw Host %q, want 127.0.0.1:8000 (what an in-sandbox client sends)", seen.Host) - } - for _, hdr := range []string{portproxy.EnvIDHeader, portproxy.EnvAtespaceHeader, portproxy.EnvPortHeader} { - if v := seen.Header.Get(hdr); v != "" { - t.Errorf("daemon saw routing header %s=%q; the proxy must strip it", hdr, v) - } - } - if seen.Header.Get("X-Custom") != "kept" { - t.Error("other request headers must reach the daemon") - } - if got := s.router.connects; len(got) != 1 || got[0] != "ate-env/env1 env1."+ate.DefaultHostSuffix+":8000" { - t.Errorf("router saw CONNECTs %q, want one for ate-env/env1 on port 8000", got) - } -} - -func TestAtespaceHeaderNamesTheActor(t *testing.T) { - s := newSetup(t) - s.router.RegisterPort("env1", 9000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) - h := target("env1", 9000) - h[portproxy.EnvAtespaceHeader] = "team-a" - if resp, body := s.do(t, http.MethodGet, "/", "", h); resp.StatusCode != http.StatusOK { - t.Fatalf("got %d %q", resp.StatusCode, body) - } - if got := s.router.connects; len(got) != 1 || !strings.HasPrefix(got[0], "team-a/env1 ") { - t.Errorf("router saw %q, want the actor team-a/env1", got) - } -} - -func TestStreamsResponses(t *testing.T) { - s := newSetup(t) - release := make(chan struct{}) - s.router.RegisterPort("env1", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - io.WriteString(w, "first\n") - w.(http.Flusher).Flush() - <-release - io.WriteString(w, "second\n") - })) - req, _ := http.NewRequest(http.MethodGet, s.proxy.URL+"/events", nil) - for k, v := range target("env1", 8000) { - req.Header.Set(k, v) - } - resp, err := http.DefaultClient.Do(req) - if err != nil { - t.Fatal(err) - } - defer resp.Body.Close() - buf := make([]byte, len("first\n")) - if _, err := io.ReadFull(resp.Body, buf); err != nil || string(buf) != "first\n" { - t.Fatalf("first chunk %q, %v: the proxy must flush before the daemon finishes", buf, err) - } - close(release) - rest, _ := io.ReadAll(resp.Body) - if string(rest) != "second\n" { - t.Errorf("rest %q", rest) - } -} - -func TestRejectsBadTargets(t *testing.T) { - s := newSetup(t) - for name, h := range map[string]map[string]string{ - "no id": {portproxy.EnvPortHeader: "8000"}, - "port not number": {portproxy.EnvIDHeader: "env1", portproxy.EnvPortHeader: "http"}, - "port zero": {portproxy.EnvIDHeader: "env1", portproxy.EnvPortHeader: "0"}, - "port too large": {portproxy.EnvIDHeader: "env1", portproxy.EnvPortHeader: "70000"}, - "id uppercase": {portproxy.EnvIDHeader: "Env1", portproxy.EnvPortHeader: "8000"}, - "id with slash": {portproxy.EnvIDHeader: "a/b", portproxy.EnvPortHeader: "8000"}, - "port with sign": {portproxy.EnvIDHeader: "env1", portproxy.EnvPortHeader: "+8000"}, - "port zero-led": {portproxy.EnvIDHeader: "env1", portproxy.EnvPortHeader: "08000"}, - "bad atespace": {portproxy.EnvIDHeader: "env1", portproxy.EnvPortHeader: "8000", portproxy.EnvAtespaceHeader: "x y"}, - "guest port": {portproxy.EnvIDHeader: "env1", portproxy.EnvPortHeader: "80"}, - } { - if resp, body := s.do(t, http.MethodGet, "/", "", h); resp.StatusCode != http.StatusBadRequest { - t.Errorf("%s: got %d %q, want 400", name, resp.StatusCode, body) - } - } - // A routing header sent twice is ambiguous. - req, _ := http.NewRequest(http.MethodGet, s.proxy.URL+"/", nil) - req.Header.Add(portproxy.EnvIDHeader, "env1") - req.Header.Add(portproxy.EnvIDHeader, "env2") - req.Header.Set(portproxy.EnvPortHeader, "8000") - resp, err := http.DefaultClient.Do(req) - if err != nil { - t.Fatal(err) - } - resp.Body.Close() - if resp.StatusCode != http.StatusBadRequest { - t.Errorf("duplicate %s: got %d, want 400", portproxy.EnvIDHeader, resp.StatusCode) - } - if n := s.router.connectCount(); n != 0 { - t.Errorf("router saw %d CONNECTs for rejected requests, want 0", n) - } -} - -func TestPassesRouterRefusals(t *testing.T) { - s := newSetup(t) - s.router.RegisterPort("asleep", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) - s.router.Running = func(id string) bool { return id != "asleep" } - - if resp, body := s.do(t, http.MethodGet, "/", "", target("missing", 8000)); resp.StatusCode != http.StatusNotFound { - t.Errorf("unknown environment: got %d %q, want 404", resp.StatusCode, body) - } - if resp, body := s.do(t, http.MethodGet, "/", "", target("asleep", 8000)); resp.StatusCode != http.StatusServiceUnavailable || !strings.Contains(body, "ate-env/asleep port 8000") { - t.Errorf("not running: got %d %q, want 503 naming the target", resp.StatusCode, body) - } -} - -// A suspended environment answers 503 at the router until it is back; the -// proxy tries again within its wake budget instead of passing the first 503 -// on, the way ate-env-api's guest calls do. -func TestRetriesWhileTheEnvironmentWakes(t *testing.T) { - s := newSetup(t) - s.router.RegisterPort("waking", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, "awake") })) - var mu sync.Mutex - refusals := 0 - s.router.Running = func(id string) bool { - mu.Lock() - defer mu.Unlock() - if refusals < 2 { - refusals++ - return false - } - return true - } - resp, body := s.do(t, http.MethodGet, "/", "", target("waking", 8000)) - if resp.StatusCode != http.StatusOK || body != "awake" { - t.Fatalf("got %d %q, want the daemon's answer after the wake", resp.StatusCode, body) - } - if n := s.router.connectCount(); n != 3 { - t.Errorf("router saw %d CONNECTs, want 3 (two refused while waking, one accepted)", n) - } -} - -func TestWakeBudgetBoundsTheWait(t *testing.T) { - s := newSetup(t) - p := mustNew(t, portproxy.Options{ConnectAddr: s.routerAddr, DefaultAtespace: "ate-env", HandshakeTimeout: 5 * time.Second, WakeBudget: 300 * time.Millisecond}) - srv := httptest.NewServer(p) - t.Cleanup(srv.Close) - s.router.RegisterPort("asleep", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) - s.router.Running = func(string) bool { return false } - req, _ := http.NewRequest(http.MethodGet, srv.URL+"/", nil) - for k, v := range target("asleep", 8000) { - req.Header.Set(k, v) - } - start := time.Now() - resp, err := http.DefaultClient.Do(req) - if err != nil { - t.Fatal(err) - } - b, _ := io.ReadAll(resp.Body) - resp.Body.Close() - if resp.StatusCode != http.StatusServiceUnavailable || !strings.Contains(string(b), "not available") { - t.Errorf("got %d %q, want 503 once the wake budget is spent", resp.StatusCode, b) - } - if took := time.Since(start); took < 300*time.Millisecond || took > 3*time.Second { - t.Errorf("waited %s, want about the 300ms budget", took) - } - if n := s.router.connectCount(); n < 2 { - t.Errorf("router saw %d CONNECTs, want at least 2 (the budget allows retries)", n) - } -} - -func TestPoolsConnectionsPerTarget(t *testing.T) { - s := newSetup(t) - for _, id := range []string{"env1", "env2"} { - id := id - s.router.RegisterPort(id, 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, id) })) - } - for i := 0; i < 3; i++ { - for _, id := range []string{"env1", "env2"} { - if _, body := s.do(t, http.MethodGet, "/", "", target(id, 8000)); body != id { - t.Fatalf("request for %s answered by %q: tunnels must not be shared across environments", id, body) - } - } - } - if n := s.router.connectCount(); n != 2 { - t.Errorf("router saw %d CONNECTs for 6 sequential requests to 2 targets, want 2 (idle tunnels reused)", n) - } -} - -func TestDispatch(t *testing.T) { - s := newSetup(t) - s.router.RegisterPort("env1", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, "proxied") })) - next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, "api") }) - h := portproxy.Dispatch(mustNew(t, portproxy.Options{ConnectAddr: s.routerAddr, DefaultAtespace: "ate-env"}), next) - - rec := httptest.NewRecorder() - h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/healthz", nil)) - if rec.Body.String() != "api" { - t.Errorf("request without %s went to %q, want the API", portproxy.EnvPortHeader, rec.Body.String()) - } - srv := httptest.NewServer(h) - defer srv.Close() - req, _ := http.NewRequest(http.MethodGet, srv.URL+"/healthz", nil) - for k, v := range target("env1", 8000) { - req.Header.Set(k, v) - } - resp, err := http.DefaultClient.Do(req) - if err != nil { - t.Fatal(err) - } - defer resp.Body.Close() - if b, _ := io.ReadAll(resp.Body); string(b) != "proxied" { - t.Errorf("request with %s got %q, want the daemon's answer", portproxy.EnvPortHeader, b) - } - - // A gRPC call is never proxied, even with the header: the tunnel carries - // HTTP/1.1 only. - rec = httptest.NewRecorder() - req = httptest.NewRequest(http.MethodPost, "/ateenv.v1alpha.EnvironmentService/Get", nil) - req.Header.Set("Content-Type", "application/grpc") - req.Header.Set(portproxy.EnvPortHeader, "8000") - h.ServeHTTP(rec, req) - if rec.Body.String() != "api" { - t.Errorf("gRPC request with %s went to %q, want the API", portproxy.EnvPortHeader, rec.Body.String()) - } -} - -// pathSetup serves the proxy the way ate-env-api does: both forms dispatched -// in front of the API's mux. -func pathSetup(t *testing.T, origins ...string) (*setup, string) { - t.Helper() - s := newSetup(t) - p := mustNew(t, portproxy.Options{ConnectAddr: s.routerAddr, DefaultAtespace: "ate-env", HandshakeTimeout: 5 * time.Second, TrustedOrigins: origins}) - mux := http.NewServeMux() - mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, "api") }) - srv := httptest.NewServer(portproxy.Dispatch(p, mux)) - t.Cleanup(srv.Close) - return s, srv.URL -} - -func TestPathFormStripsPrefix(t *testing.T) { - s, base := pathSetup(t) - var seen *http.Request - s.router.RegisterPort("env1", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - seen = r - io.WriteString(w, "daemon") - })) - resp, err := http.Get(base + "/v1alpha/envs/env1/ports/8000/api/file/download?path=/etc/os-release") - if err != nil { - t.Fatal(err) - } - b, _ := io.ReadAll(resp.Body) - resp.Body.Close() - if string(b) != "daemon" { - t.Fatalf("got %d %q", resp.StatusCode, b) - } - if got := seen.URL.RequestURI(); got != "/api/file/download?path=/etc/os-release" { - t.Errorf("daemon saw %q, want the path after the routing prefix with its query", got) - } - if got := s.router.connects; len(got) != 1 || !strings.HasPrefix(got[0], "ate-env/env1 ") { - t.Errorf("router saw %q", got) - } - - // Escaping after the prefix is kept as sent. - resp, err = http.Get(base + "/v1alpha/envs/env1/ports/8000/files/a%2Fb") - if err != nil { - t.Fatal(err) - } - resp.Body.Close() - if got := seen.URL.EscapedPath(); got != "/files/a%2Fb" { - t.Errorf("daemon saw escaped path %q, want /files/a%%2Fb", got) - } - - // The bare prefix with a trailing slash is the daemon's root. - resp, err = http.Get(base + "/v1alpha/envs/env1/ports/8000/") - if err != nil { - t.Fatal(err) - } - resp.Body.Close() - if seen.URL.Path != "/" { - t.Errorf("daemon saw %q for the bare prefix, want /", seen.URL.Path) - } -} - -func TestPathFormAtespaceAndValidation(t *testing.T) { - s, base := pathSetup(t) - s.router.RegisterPort("env1", 9000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, r.URL.Path) })) - resp, err := http.Get(base + "/v1alpha/atespaces/team-a/envs/env1/ports/9000/x") - if err != nil { - t.Fatal(err) - } - b, _ := io.ReadAll(resp.Body) - resp.Body.Close() - if string(b) != "/x" { - t.Errorf("got %d %q, want the daemon's /x", resp.StatusCode, b) - } - if got := s.router.connects; len(got) != 1 || !strings.HasPrefix(got[0], "team-a/env1 ") { - t.Errorf("router saw %q, want team-a/env1", got) - } - for _, path := range []string{ - "/v1alpha/envs/Env1/ports/9000/x", - "/v1alpha/envs/env1/ports/0/x", - "/v1alpha/envs/env1/ports/http/x", - "/v1alpha/envs/env1/ports/09000/x", - "/v1alpha/envs//ports/9000/x", - "/v1alpha/envs/a%2Fb/ports/9000/x", - "/v1alpha/atespaces/A_B/envs/env1/ports/9000/x", - } { - resp, err := http.Get(base + path) - if err != nil { - t.Fatal(err) - } - resp.Body.Close() - if resp.StatusCode != http.StatusBadRequest { - t.Errorf("%s: got %d, want 400", path, resp.StatusCode) - } - } - // Other paths still reach the API. - resp, err = http.Get(base + "/v1alpha/envs/env1/mcp") - if err != nil { - t.Fatal(err) - } - b, _ = io.ReadAll(resp.Body) - resp.Body.Close() - if string(b) != "api" { - t.Errorf("/v1alpha/envs/env1/mcp went to %q, want the API", b) - } -} - -// An HTTP upgrade (what a WebSocket handshake is) passes through the proxy and -// the tunnel, after which bytes flow both ways. -func TestUpgradePassesThrough(t *testing.T) { - s, base := pathSetup(t) - s.router.RegisterPort("env1", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.Header.Get("Upgrade") != "echo" { - http.Error(w, "want Upgrade: echo", http.StatusBadRequest) - return - } - conn, rw, err := http.NewResponseController(w).Hijack() - if err != nil { - return - } - defer conn.Close() - rw.WriteString("HTTP/1.1 101 Switching Protocols\r\nConnection: Upgrade\r\nUpgrade: echo\r\n\r\n") - rw.Flush() - buf := make([]byte, 4) - if _, err := io.ReadFull(rw, buf); err == nil { - rw.Write(append([]byte("re:"), buf...)) - rw.Flush() - } - })) - addr := strings.TrimPrefix(base, "http://") - for name, req := range map[string]string{ - "path form": "GET /v1alpha/envs/env1/ports/8000/sockets/events/c1 HTTP/1.1\r\nHost: " + addr + "\r\n", - "header form": "GET /sockets/events/c1 HTTP/1.1\r\nHost: " + addr + "\r\nX-Env-Id: env1\r\nX-Env-Port: 8000\r\n", - } { - conn, err := net.Dial("tcp", addr) - if err != nil { - t.Fatal(err) - } - conn.SetDeadline(time.Now().Add(10 * time.Second)) - fmt.Fprint(conn, req+"Connection: Upgrade\r\nUpgrade: echo\r\n\r\n") - br := bufio.NewReader(conn) - resp, err := http.ReadResponse(br, nil) - if err != nil || resp.StatusCode != http.StatusSwitchingProtocols { - t.Fatalf("%s: upgrade answered %v, %v", name, resp, err) - } - conn.Write([]byte("ping")) - got := make([]byte, 7) - if _, err := io.ReadFull(br, got); err != nil || string(got) != "re:ping" { - t.Errorf("%s: after upgrade got %q, %v; want re:ping", name, got, err) - } - conn.Close() - } -} - -// The path form reaches the daemon exactly as sent: an escaped routing segment -// still ends the prefix, and nothing is redirected or cleaned. -func TestPathFormKeepsPathAsSent(t *testing.T) { - s, base := pathSetup(t) - var seen *http.Request - s.router.RegisterPort("env1", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - seen = r - io.WriteString(w, "daemon") - })) - noRedirect := &http.Client{CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }} - for _, c := range []struct{ method, path, want string }{ - {http.MethodGet, "/v1alpha/envs/%65nv1/ports/8000/a", "/a"}, - {http.MethodGet, "/v1alpha/envs/env1/ports/%38000/a", "/a"}, - {http.MethodGet, "/v1alpha/envs/env1/ports/8000", "/"}, - {http.MethodPost, "/v1alpha/envs/env1/ports/8000", "/"}, - {http.MethodPost, "/v1alpha/atespaces/ate-env/envs/env1/ports/8000", "/"}, - {http.MethodGet, "/v1alpha/envs/env1/ports/8000/a//b/../c?x=1", "/a//b/../c?x=1"}, - } { - seen = nil - req, _ := http.NewRequest(c.method, base+c.path, nil) - resp, err := noRedirect.Do(req) - if err != nil { - t.Fatal(err) - } - b, _ := io.ReadAll(resp.Body) - resp.Body.Close() - if resp.StatusCode != http.StatusOK || string(b) != "daemon" || seen == nil { - t.Errorf("%s %s: got %d %q, want the daemon's 200", c.method, c.path, resp.StatusCode, b) - continue - } - if got := seen.URL.RequestURI(); got != c.want { - t.Errorf("%s %s: daemon saw %q, want %q", c.method, c.path, got, c.want) - } - } -} - -// A browser page from another origin cannot drive a daemon through the path -// form: no state-changing request, no WebSocket handshake. Reads, non-browser -// clients, the listener's own origin and trusted origins pass. -func TestPathFormRefusesCrossOrigin(t *testing.T) { - s, base := pathSetup(t, "https://ui.example") - s.router.RegisterPort("env1", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.Header.Get("Upgrade") != "" { - w.Header().Set("Connection", "Upgrade") - w.Header().Set("Upgrade", "echo") - w.WriteHeader(http.StatusSwitchingProtocols) - return - } - io.WriteString(w, "daemon") - })) - self := base // the listener's own origin, as a browser on it sends - path := base + "/v1alpha/envs/env1/ports/8000/run" - for _, c := range []struct { - name string - method string - headers map[string]string - want int - }{ - {"cross-site POST", http.MethodPost, map[string]string{"Sec-Fetch-Site": "cross-site", "Origin": "https://evil.example"}, http.StatusForbidden}, - {"foreign Origin POST", http.MethodPost, map[string]string{"Origin": "https://evil.example"}, http.StatusForbidden}, - {"foreign Origin upgrade", http.MethodGet, map[string]string{"Origin": "https://evil.example", "Connection": "Upgrade", "Upgrade": "echo"}, http.StatusForbidden}, - {"cross-site upgrade", http.MethodGet, map[string]string{"Sec-Fetch-Site": "cross-site", "Connection": "Upgrade", "Upgrade": "echo"}, http.StatusForbidden}, - {"cross-site GET", http.MethodGet, map[string]string{"Sec-Fetch-Site": "cross-site", "Origin": "https://evil.example"}, http.StatusOK}, - {"no browser headers", http.MethodPost, nil, http.StatusOK}, - {"own origin POST", http.MethodPost, map[string]string{"Origin": self}, http.StatusOK}, - {"same-origin upgrade", http.MethodGet, map[string]string{"Sec-Fetch-Site": "same-origin", "Connection": "Upgrade", "Upgrade": "echo"}, http.StatusSwitchingProtocols}, - {"trusted origin POST", http.MethodPost, map[string]string{"Sec-Fetch-Site": "cross-site", "Origin": "https://ui.example"}, http.StatusOK}, - } { - req, _ := http.NewRequest(c.method, path, nil) - for k, v := range c.headers { - req.Header.Set(k, v) - } - resp, err := http.DefaultClient.Do(req) - if err != nil { - t.Fatalf("%s: %v", c.name, err) - } - resp.Body.Close() - if resp.StatusCode != c.want { - t.Errorf("%s: got %d, want %d", c.name, resp.StatusCode, c.want) - } - } - if _, err := portproxy.New(portproxy.Options{ConnectAddr: s.routerAddr, TrustedOrigins: []string{"not an origin"}}); err == nil { - t.Error("New accepted a malformed trusted origin") - } -} - -// An error body names the target and a short cause. The router's message and -// address stay in the log. -func TestErrorBodiesNameCauseOnly(t *testing.T) { - refuser := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - http.Error(w, "dial tcp 10.1.2.3:8000: connect: connection refused", http.StatusBadGateway) - })) - defer refuser.Close() - p := mustNew(t, portproxy.Options{ConnectAddr: strings.TrimPrefix(refuser.URL, "http://"), DefaultAtespace: "ate-env"}) - srv := httptest.NewServer(p) - defer srv.Close() - req, _ := http.NewRequest(http.MethodGet, srv.URL+"/", nil) - for k, v := range target("env1", 8000) { - req.Header.Set(k, v) - } - resp, err := http.DefaultClient.Do(req) - if err != nil { - t.Fatal(err) - } - b, _ := io.ReadAll(resp.Body) - resp.Body.Close() - body := string(b) - if resp.StatusCode != http.StatusBadGateway || !strings.Contains(body, "ate-env/env1 port 8000") || !strings.Contains(body, "0.0.0.0") { - t.Errorf("got %d %q, want 502 naming the target and the listen rule", resp.StatusCode, body) - } - if strings.Contains(body, "10.1.2.3") || strings.Contains(body, refuser.URL[len("http://"):]) { - t.Errorf("body %q leaks an internal address", body) - } -} - -func TestResponseHeaderTimeout(t *testing.T) { - rr := &recordingRouter{Router: fakerouter.New()} - router := httptest.NewServer(rr) - defer router.Close() - release := make(chan struct{}) - defer close(release) - rr.RegisterPort("env1", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - select { - case <-release: - case <-r.Context().Done(): - } - })) - p := mustNew(t, portproxy.Options{ - ConnectAddr: strings.TrimPrefix(router.URL, "http://"), - DefaultAtespace: "ate-env", - ResponseHeaderTimeout: 200 * time.Millisecond, - }) - srv := httptest.NewServer(p) - defer srv.Close() - req, _ := http.NewRequest(http.MethodGet, srv.URL+"/", nil) - for k, v := range target("env1", 8000) { - req.Header.Set(k, v) - } - client := &http.Client{Timeout: 10 * time.Second} // fail, not hang, without the limit - resp, err := client.Do(req) - if err != nil { - t.Fatal(err) - } - resp.Body.Close() - if resp.StatusCode != http.StatusGatewayTimeout { - t.Errorf("silent daemon: got %d, want 504", resp.StatusCode) - } -} diff --git a/proto/ateenv/v1alpha/env.pb.go b/proto/ateenv/v1alpha/env.pb.go index 98c0a5db..69e7c34f 100644 --- a/proto/ateenv/v1alpha/env.pb.go +++ b/proto/ateenv/v1alpha/env.pb.go @@ -252,18 +252,7 @@ type CreateEnvironmentRequest struct { // Substrate atespace for the environment (defaults to "ate-env"). Atespace string `protobuf:"bytes,2,opt,name=atespace,proto3" json:"atespace,omitempty"` // ActorTemplate configuration to instantiate (defaults to name "default-template" in atespace "ate-env"). - Template *Template `protobuf:"bytes,3,opt,name=template,proto3" json:"template,omitempty"` - // Rootfs image to run, pinned by digest (repo@sha256:...). Optional. When - // set, the environment is created from an ActorTemplate derived from - // `template`, which acts as the base: its worker selector, snapshot, - // sandbox and resource settings are kept, its container image is replaced - // by the rootfs image, and the ate-env-guest from the base is mounted into - // it as a read-only image volume (a runtime layer, not part of the task - // image), so the rootfs image runs unmodified. The derived template is named - // "-" in the base's atespace; it - // is created on first use and reused afterwards, and the response reports - // it. - RootfsImage string `protobuf:"bytes,4,opt,name=rootfs_image,json=rootfsImage,proto3" json:"rootfs_image,omitempty"` + Template *Template `protobuf:"bytes,3,opt,name=template,proto3" json:"template,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -319,13 +308,6 @@ func (x *CreateEnvironmentRequest) GetTemplate() *Template { return nil } -func (x *CreateEnvironmentRequest) GetRootfsImage() string { - if x != nil { - return x.RootfsImage - } - return "" -} - // Response returned after creating an environment. type CreateEnvironmentResponse struct { state protoimpl.MessageState `protogen:"open.v1"` @@ -669,12 +651,11 @@ const file_proto_ateenv_v1alpha_env_proto_rawDesc = "" + "\x02id\x18\x01 \x01(\tR\x02id\x12\x1a\n" + "\batespace\x18\x02 \x01(\tR\batespace\x124\n" + "\btemplate\x18\x03 \x01(\v2\x18.ateenv.v1alpha.TemplateR\btemplate\x129\n" + - "\x06status\x18\x04 \x01(\x0e2!.ateenv.v1alpha.EnvironmentStatusR\x06status\"\x9f\x01\n" + + "\x06status\x18\x04 \x01(\x0e2!.ateenv.v1alpha.EnvironmentStatusR\x06status\"\x90\x01\n" + "\x18CreateEnvironmentRequest\x12\x0e\n" + "\x02id\x18\x01 \x01(\tR\x02id\x12\x1a\n" + "\batespace\x18\x02 \x01(\tR\batespace\x124\n" + - "\btemplate\x18\x03 \x01(\v2\x18.ateenv.v1alpha.TemplateR\btemplate\x12!\n" + - "\frootfs_image\x18\x04 \x01(\tR\vrootfsImage\"Z\n" + + "\btemplate\x18\x03 \x01(\v2\x18.ateenv.v1alpha.TemplateR\btemplateJ\x04\b\x04\x10\x05R\frootfs_image\"Z\n" + "\x19CreateEnvironmentResponse\x12=\n" + "\venvironment\x18\x01 \x01(\v2\x1b.ateenv.v1alpha.EnvironmentR\venvironment\"C\n" + "\x15GetEnvironmentRequest\x12\x0e\n" + diff --git a/proto/ateenv/v1alpha/env.proto b/proto/ateenv/v1alpha/env.proto index 034b6841..ae68c808 100644 --- a/proto/ateenv/v1alpha/env.proto +++ b/proto/ateenv/v1alpha/env.proto @@ -102,17 +102,8 @@ message CreateEnvironmentRequest { string atespace = 2; // ActorTemplate configuration to instantiate (defaults to name "default-template" in atespace "ate-env"). Template template = 3; - // Rootfs image to run, pinned by digest (repo@sha256:...). Optional. When - // set, the environment is created from an ActorTemplate derived from - // `template`, which acts as the base: its worker selector, snapshot, - // sandbox and resource settings are kept, its container image is replaced - // by the rootfs image, and the ate-env-guest from the base is mounted into - // it as a read-only image volume (a runtime layer, not part of the task - // image), so the rootfs image runs unmodified. The derived template is named - // "-" in the base's atespace; it - // is created on first use and reused afterwards, and the response reports - // it. - string rootfs_image = 4; + reserved 4; + reserved "rootfs_image"; } // Response returned after creating an environment.