From 6ee7e085d5a029713faafc7656a4a88bccf42ddc Mon Sep 17 00:00:00 2001 From: Jaana Dogan Date: Sun, 18 Oct 2026 23:50:00 -0700 Subject: [PATCH] Remove runtime layers, sidecars, port proxy, on-demand rootfs images This removes a set of features that widened what ate-env-api and the guest can be made to do, without review in proportion to that reach. Our goal for this project is to keep the data plane APIs to a minimum and see if we can provide the guest services from Substrate actors only via config in the long-term. Removed: - Runtime layers and sidecars: `ate-env manifest template --layer`, `--sidecar` and `--sidecar-readyz`, and the guest's `-sidecar` and `-sidecar-readyz` support (cmd/ate-env-guest/sidecar.go). /readyz answers once the guest is listening. - Port proxy (internal/portproxy) and its ate-env-api flags `--atenet-connect`, `--port-proxy-trusted-origins` and `--port-proxy-response-header-timeout`, plus the CONNECT support in the fake router. - On-demand rootfs images: CreateEnvironmentRequest.rootfs_image (field 4 and its name are now reserved), the server-side derivation of `-` ActorTemplates, `ate-env create --rootfs-image`, and `rootfs_image=` in the Python client. `ate-env manifest template --rootfs-image` stays: operators can still register image templates themselves. - The NeMo Gym integration (integrations/nemo-gym), which depended on the port proxy and on runtime layers. - SWE-bench benchmarking (cmd/benchmarking, internal/benchmarking, benchmarking/), which depended on on-demand rootfs images. - Docs for all of the above: docs/rootfs-images, examples/runtime-layers, and the related README sections. - CreateEnvironment now returns InvalidArgument when the request has unknown fields. The message names the replacement: register an image template with `ate-env manifest template --rootfs-image` and pass its name in template.name. Process: Changes of this size to a project like this, which adds new network surface, new privileges for the API service, new wire-protocol fields, and new integrations, need substantial design discussion and security review before they are merged. They should not land incrementally through feature PRs. Any of these capabilities can come back through a design proposal with a threat model and explicit maintainer sign-off. Compatibility: - Clients that set rootfs_image get an environment created from the base template, because the server drops the unknown field. Update them. - ate-env-api deployments that pass the removed port-proxy flags will fail to start; the generated manifest never set them. - Templates derived by earlier versions (`-<12 hex>`) are unused and can be deleted with `kubectl-ate delete actor-template`. --- .github/dependabot.yml | 4 +- .github/workflows/ci.yaml | 6 +- .gitignore | 1 - Makefile | 6 - README.md | 104 +-- benchmarking/swebench/README.md | 150 ---- benchmarking/swebench/make-dataset.py | 184 ----- benchmarking/swebench/mirror-images.sh | 64 -- benchmarking/swebench/offline-exclusions.txt | 57 -- benchmarking/swebench/run-in-cluster.sh | 165 ----- clients/python/README.md | 16 +- .../ate_env/_gen/ateenv/v1alpha/env_pb2.py | 40 +- .../ate_env/_gen/ateenv/v1alpha/env_pb2.pyi | 6 +- clients/python/src/ate_env/client.py | 10 +- clients/python/tests/e2e/test_full_stack.py | 84 --- clients/python/tests/fakes.py | 7 - clients/python/tests/test_client.py | 24 - clients/python/tests/test_config.py | 2 +- cmd/ate-env-api/main.go | 36 +- cmd/ate-env-guest/main.go | 52 -- cmd/ate-env-guest/sidecar.go | 259 ------- cmd/ate-env-guest/sidecar_test.go | 159 ----- cmd/ate-env/main.go | 7 +- cmd/ate-env/main_test.go | 17 +- cmd/ate-env/manifest.go | 91 --- cmd/ate-env/manifest_test.go | 58 -- cmd/benchmarking/swebench-loader/main.go | 246 ------- docs/rootfs-images/DESIGN.md | 224 ------ docs/rootfs-images/README.md | 144 ---- docs/rootfs-images/RUNTIMES.md | 149 ---- examples/runtime-layers/README.md | 60 -- examples/runtime-layers/demo.sh | 107 --- integrations/nemo-gym/.gitignore | 6 - integrations/nemo-gym/README.md | 368 ---------- .../nemo-gym/benchmarks/driver_sim.py | 161 ----- .../nemo-gym/benchmarks/rollout_bench.py | 134 ---- integrations/nemo-gym/hack/bake-task-image.sh | 70 -- integrations/nemo-gym/pyproject.toml | 45 -- .../src/nemo_gym_substrate/__init__.py | 29 - .../src/nemo_gym_substrate/_compat.py | 124 ---- .../src/nemo_gym_substrate/provider.py | 567 --------------- integrations/nemo-gym/tests/__init__.py | 13 - integrations/nemo-gym/tests/fakes.py | 272 -------- integrations/nemo-gym/tests/test_e2e.py | 88 --- integrations/nemo-gym/tests/test_provider.py | 589 ---------------- internal/apiservice/imagetemplate.go | 35 +- internal/apiservice/imagetemplate_test.go | 55 +- internal/apiservice/server.go | 78 +-- internal/apiservice/server_test.go | 159 +---- internal/ate/client.go | 39 -- internal/benchmarking/swebench/backend.go | 247 ------- .../benchmarking/swebench/backend_test.go | 270 -------- internal/benchmarking/swebench/dataset.go | 216 ------ .../benchmarking/swebench/dataset_test.go | 110 --- internal/benchmarking/swebench/doc.go | 35 - internal/benchmarking/swebench/events.go | 111 --- internal/benchmarking/swebench/report.go | 317 --------- internal/benchmarking/swebench/report_test.go | 139 ---- internal/benchmarking/swebench/run.go | 637 ----------------- internal/benchmarking/swebench/run_test.go | 588 ---------------- .../swebench/trajectories/swebench-gold.yaml | 86 --- .../swebench/trajectories/swebench-grade.py | 317 --------- .../trajectories/test_swebench_grade.py | 154 ----- internal/benchmarking/swebench/trajectory.go | 332 --------- .../benchmarking/swebench/trajectory_test.go | 185 ----- .../internaltest/fakecontrol/fakecontrol.go | 69 -- .../internaltest/fakerouter/fakerouter.go | 104 +-- internal/portproxy/portproxy.go | 503 -------------- internal/portproxy/portproxy_test.go | 651 ------------------ proto/ateenv/v1alpha/env.pb.go | 25 +- proto/ateenv/v1alpha/env.proto | 13 +- 71 files changed, 96 insertions(+), 10384 deletions(-) delete mode 100644 benchmarking/swebench/README.md delete mode 100755 benchmarking/swebench/make-dataset.py delete mode 100755 benchmarking/swebench/mirror-images.sh delete mode 100644 benchmarking/swebench/offline-exclusions.txt delete mode 100755 benchmarking/swebench/run-in-cluster.sh delete mode 100644 cmd/ate-env-guest/sidecar.go delete mode 100644 cmd/ate-env-guest/sidecar_test.go delete mode 100644 cmd/benchmarking/swebench-loader/main.go delete mode 100644 docs/rootfs-images/DESIGN.md delete mode 100644 docs/rootfs-images/README.md delete mode 100644 docs/rootfs-images/RUNTIMES.md delete mode 100644 examples/runtime-layers/README.md delete mode 100755 examples/runtime-layers/demo.sh delete mode 100644 integrations/nemo-gym/.gitignore delete mode 100644 integrations/nemo-gym/README.md delete mode 100644 integrations/nemo-gym/benchmarks/driver_sim.py delete mode 100644 integrations/nemo-gym/benchmarks/rollout_bench.py delete mode 100755 integrations/nemo-gym/hack/bake-task-image.sh delete mode 100644 integrations/nemo-gym/pyproject.toml delete mode 100644 integrations/nemo-gym/src/nemo_gym_substrate/__init__.py delete mode 100644 integrations/nemo-gym/src/nemo_gym_substrate/_compat.py delete mode 100644 integrations/nemo-gym/src/nemo_gym_substrate/provider.py delete mode 100644 integrations/nemo-gym/tests/__init__.py delete mode 100644 integrations/nemo-gym/tests/fakes.py delete mode 100644 integrations/nemo-gym/tests/test_e2e.py delete mode 100644 integrations/nemo-gym/tests/test_provider.py delete mode 100644 internal/benchmarking/swebench/backend.go delete mode 100644 internal/benchmarking/swebench/backend_test.go delete mode 100644 internal/benchmarking/swebench/dataset.go delete mode 100644 internal/benchmarking/swebench/dataset_test.go delete mode 100644 internal/benchmarking/swebench/doc.go delete mode 100644 internal/benchmarking/swebench/events.go delete mode 100644 internal/benchmarking/swebench/report.go delete mode 100644 internal/benchmarking/swebench/report_test.go delete mode 100644 internal/benchmarking/swebench/run.go delete mode 100644 internal/benchmarking/swebench/run_test.go delete mode 100644 internal/benchmarking/swebench/trajectories/swebench-gold.yaml delete mode 100644 internal/benchmarking/swebench/trajectories/swebench-grade.py delete mode 100644 internal/benchmarking/swebench/trajectories/test_swebench_grade.py delete mode 100644 internal/benchmarking/swebench/trajectory.go delete mode 100644 internal/benchmarking/swebench/trajectory_test.go delete mode 100644 internal/portproxy/portproxy.go delete mode 100644 internal/portproxy/portproxy_test.go 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.