From 9c9420fefbbdbc686108126d85a0f472967f6baf Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Thu, 9 Jul 2026 11:22:02 -0700 Subject: [PATCH 1/8] Add Google Cloud Run sandbox integration Cloud Run sandboxes (public preview) provide fast, strongly isolated environments for executing untrusted code, but expose no standalone REST API: the sandbox CLI only exists inside a Cloud Run service deployed with --sandbox-launcher. This adds a "cloudrun" sandbox flavor to the GCP integration that talks to a small bridge service deployed on such a service, following the bridge pattern established by the Cloudflare sandbox flavor. The client authenticates with Google-signed ID tokens (Cloud Run IAM invoker flow) derived from the standard GCP credential chain (service connector, SA key file, or ADC). Unlike the Cloudflare flavor it supports per-exec env vars (native `sandbox exec -e`) and filesystem snapshots to GCS via `sandbox tar`, giving create_snapshot()/restore() parity with Modal. The deployable bridge (sync stdlib HTTP server shelling out to the sandbox CLI) ships as examples/cloudrun_sandbox_bridge/. Design notes in CLOUDRUN_SANDBOX_SPEC.md. Co-Authored-By: Claude Fable 5 --- CLOUDRUN_SANDBOX_SPEC.md | 270 ++++ .../component-guide/sandboxes/cloudrun.md | 113 ++ docs/book/component-guide/toc.md | 1 + examples/cloudrun_sandbox_bridge/Dockerfile | 8 + examples/cloudrun_sandbox_bridge/README.md | 93 ++ examples/cloudrun_sandbox_bridge/main.py | 476 +++++++ .../cloudrun_sandbox_bridge/requirements.txt | 1 + src/zenml/integrations/gcp/__init__.py | 3 + .../integrations/gcp/flavors/__init__.py | 8 + .../gcp/flavors/cloudrun_sandbox_flavor.py | 249 ++++ .../integrations/gcp/sandboxes/__init__.py | 26 + .../gcp/sandboxes/cloudrun_sandbox.py | 1206 +++++++++++++++++ .../integrations/gcp/sandboxes/__init__.py | 13 + .../gcp/sandboxes/test_cloudrun_sandbox.py | 463 +++++++ 14 files changed, 2930 insertions(+) create mode 100644 CLOUDRUN_SANDBOX_SPEC.md create mode 100644 docs/book/component-guide/sandboxes/cloudrun.md create mode 100644 examples/cloudrun_sandbox_bridge/Dockerfile create mode 100644 examples/cloudrun_sandbox_bridge/README.md create mode 100644 examples/cloudrun_sandbox_bridge/main.py create mode 100644 examples/cloudrun_sandbox_bridge/requirements.txt create mode 100644 src/zenml/integrations/gcp/flavors/cloudrun_sandbox_flavor.py create mode 100644 src/zenml/integrations/gcp/sandboxes/__init__.py create mode 100644 src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py create mode 100644 tests/unit/integrations/gcp/sandboxes/__init__.py create mode 100644 tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py diff --git a/CLOUDRUN_SANDBOX_SPEC.md b/CLOUDRUN_SANDBOX_SPEC.md new file mode 100644 index 00000000000..a6ca288ddc9 --- /dev/null +++ b/CLOUDRUN_SANDBOX_SPEC.md @@ -0,0 +1,270 @@ +# Spec: Google Cloud Run Sandbox integration + +Add a `cloudrun` sandbox flavor to the existing GCP integration, backed by +[Cloud Run sandboxes](https://docs.cloud.google.com/run/docs/code-execution) +(public preview, July 2026). + +## 1. Background + +### 1.1 What Cloud Run sandboxes are + +Cloud Run sandboxes provide fast (<500 ms), strongly isolated environments for +executing untrusted code — LLM/agent-generated code being the headline use +case. Unlike Vercel/Cloudflare/Modal sandboxes there is **no standalone REST +API**: sandboxes are created *inside* a Cloud Run service instance that was +deployed with the sandbox launcher enabled: + +```bash +gcloud beta run deploy SERVICE --image IMAGE_URL --sandbox-launcher +``` + +Inside such an instance a CLI at `/usr/local/gcp/bin/sandbox` is the control +surface: + +| Command | Purpose | +|---|---| +| `sandbox do -- CMD` | ephemeral sandbox: create, run, destroy | +| `sandbox run [CMD] --detach` | create a persistent sandbox | +| `sandbox exec CMD [-e K=V] [-w DIR]` | run a command in a persistent sandbox | +| `sandbox delete [--force]` | destroy a sandbox | +| `sandbox tar --file F` | export the writable overlay as a tarball | +| `sandbox fork ` | clone a running sandbox | + +Key flags: `--env`, `--mount type=bind,source=…,destination=…`, `--write`, +`--allow-egress` (no egress by default), `--import-tar`/`--export-tar`, +`--rootfs`, `--workdir`. Sandboxes share the host instance's CPU/memory and +see the instance's filesystem read-only (plus a writable tmpfs overlay). + +### 1.2 What ZenML's sandbox abstraction expects + +`src/zenml/sandboxes/` (on `develop`) defines the `SANDBOX` stack component +type: + +- `BaseSandbox.create_session(settings, destroy_on_exit=False) -> SandboxSession` + (required), `attach(session_id)` and `restore(snapshot)` (optional). +- `SandboxSession` hooks: `_exec` (required), `_close` (required), + `_upload_file` / `_download_file` / `_create_snapshot` / `_destroy` / + `_get_dashboard_url` (optional). +- `SandboxProcess`: `stdout()` / `stderr()` line iterators, `wait(timeout)`, + `kill()`, `exit_code`; `collect()` is provided by the base class. +- `SandboxSnapshot(sandbox_id: UUID, ref: str, metadata: dict)`. + +Reference implementations: `local`, `modal` (only flavor with snapshots), +`kubernetes`, and `cloudflare` (branch `feature/sandbox-cloudflare`, PR #4907) +which established the **bridge pattern**: a small HTTP service deployed on the +provider mediates between the ZenML client and the provider-local sandbox +primitive. Cloud Run's launcher architecture forces the same pattern. + +## 2. Design + +### 2.1 Architecture + +``` +ZenML client (step code) Google Cloud +┌──────────────────────────┐ HTTPS + ID ┌──────────────────────────────┐ +│ CloudRunSandbox │ token (IAM) │ Cloud Run service │ +│ └ CloudRunSandboxSession├────────────────►│ "sandbox bridge" │ +│ └ CloudRunSandbox- │ SSE streams │ (--sandbox-launcher) │ +│ Process │◄────────────────┤ ┌────────┐ ┌────────┐ │ +└──────────────────────────┘ │ │sandbox1│ │sandbox2│ ... │ + │ └────────┘ └────────┘ │ + └───────────┬──────────────────┘ + │ snapshots (tar) + ▼ + GCS bucket (optional) +``` + +The user deploys the provided **bridge service** (a small synchronous Python +HTTP server that shells out to the `sandbox` CLI) to Cloud Run with +`--sandbox-launcher`. The ZenML `cloudrun` sandbox component talks to it over +HTTPS, authenticating with Google-signed **ID tokens** (standard Cloud Run +IAM invoker auth) minted from the component's Google credentials. + +The bridge speaks the **same wire protocol v1 as the Cloudflare bridge** +(`POST /v1/sandbox`, `POST /v1/sandbox/:id/exec` streaming SSE frames +`stdout`/`stderr` (base64) + `exit`/`error`, `PUT|GET /v1/sandbox/:id/file/*`, +`GET /v1/sandbox/:id/running`, `DELETE /v1/sandbox/:id`), with two +Cloud-Run-specific extensions: + +1. **Per-exec env**: exec body accepts `"env": {K: V}` (the `sandbox` CLI + supports `-e` natively), removing the Cloudflare limitation where per-exec + env raises `NotImplementedError`. Session-level `sandbox_environment` is + merged into every exec by the client — no bridge-side session objects. +2. **Snapshots**: `POST /v1/sandbox/:id/snapshot {"gcs_uri": "gs://…"}` runs + `sandbox tar` and uploads the tarball to GCS; `POST /v1/sandbox` accepts + `"import_tar_uri": "gs://…"` to boot from a snapshot. This gives Cloud Run + parity with Modal (`create_snapshot`/`restore`), which Cloudflare lacks. + +### 2.2 Deployment constraints (documented, not enforced) + +- Persistent sandboxes are **instance-local**. The bridge must be deployed + with `--max-instances 1` (or session affinity for advanced setups) so + `exec` calls for a sandbox land on the instance that owns it. +- `--no-cpu-throttling` and `--min-instances 1` keep detached sandboxes alive + between requests; otherwise Cloud Run may throttle/scale-to-zero and kill + them. +- Cloud Run request timeout caps a single SSE exec stream at 60 min; the + per-exec `timeout_ms` setting must stay below the service timeout. +- Sandbox CPU/memory come out of the bridge instance's allocation — size the + service accordingly. + +### 2.3 Placement: extend the `gcp` integration + +The GCP integration already ships six flavors and depends on +`google-cloud-run` and `google-cloud-storage`; the flavor rides on the +existing `GoogleCredentialsConfigMixin`/`GoogleCredentialsMixin` (service +connector → SA key file → ADC) and the `gcp` service connector +(`resource_type="gcp-generic"`). Flavor name: **`cloudrun`** (precedent: +Vertex components use the service name `vertex`, not `gcp`). + +### 2.4 New classes + +`src/zenml/integrations/gcp/flavors/cloudrun_sandbox_flavor.py` (no google +imports — flavor files must import without the integration installed): + +```python +class CloudRunSandboxSettings(BaseSandboxSettings): + timeout_ms: int = 120_000 # per-exec cap, passed to the bridge + cwd: Optional[str] = None # default workdir inside the sandbox + allow_egress: bool = False # sandbox outbound network (CR default: off) + +class CloudRunSandboxConfig( + BaseSandboxConfig, GoogleCredentialsConfigMixin, CloudRunSandboxSettings +): + service_url: str # https URL of the bridge service + audience: Optional[str] = None # ID-token audience; default service_url + allow_unauthenticated: bool = False # skip ID tokens (public bridge; discouraged) + snapshot_uri_prefix: Optional[str] = None # gs://bucket/prefix enabling snapshots + # + project / service_account_path from GoogleCredentialsConfigMixin + is_remote -> True + +class CloudRunSandboxFlavor(BaseSandboxFlavor): + name = "cloudrun" # GCP_CLOUDRUN_SANDBOX_FLAVOR + service_connector_requirements = ServiceConnectorRequirements( + connector_type=GCP_CONNECTOR_TYPE, resource_type=GCP_RESOURCE_TYPE + ) +``` + +`service_url` gets the same https-or-localhost validator as the Cloudflare +`worker_url`. + +`src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py` (google imports +allowed here): + +- `_CloudRunBridgeClient` — httpx wrapper (httpx is a core dependency), + adapted from `_CloudflareBridgeClient`: same retry policy (idempotent + methods only; 429/502/503/504), same SSE parsing, plus `env` on exec, + `import_tar_uri` on create, and a `snapshot(sandbox_id, gcs_uri)` call. + Auth is a pluggable header callback instead of a static bearer token. +- `CloudRunSandboxProcess` — SSE demux into stdout/stderr line buffers on a + pump thread (same design as `CloudflareSandboxProcess`). +- `CloudRunSandboxSession` — implements `_exec` (with per-exec env support), + `_upload_file`, `_download_file`, `_close`, `_destroy`, `_create_snapshot` + (only when `snapshot_uri_prefix` is configured; `ref` = the GCS tarball URI). +- `CloudRunSandbox(BaseSandbox, GoogleCredentialsMixin)` — implements + `create_session` (threading `destroy_on_exit`, per develop's signature), + `attach`, and `restore` (creates a sandbox with `import_tar_uri=snapshot.ref` + after `_validate_snapshot`). + +**ID-token minting** (`_get_id_token_credentials`): start from +`self._get_authentication()` (mixin), then: + +1. `service_account.Credentials` → `service_account.IDTokenCredentials` + (same signer, `target_audience=audience`). +2. `impersonated_credentials.Credentials` → + `impersonated_credentials.IDTokenCredentials(..., include_email=True)`. +3. Otherwise (plain ADC) → `google.oauth2.id_token.fetch_id_token`. +4. Failure → actionable error naming the three supported paths. + +Tokens are cached and refreshed via `google.auth.transport.requests.Request` +before expiry; the bearer header is injected per request. + +### 2.5 Bridge service (deployable example) + +Following the Cloudflare precedent the bridge is not part of the `zenml` +package. It ships as `examples/cloudrun_sandbox_bridge/`: + +- `main.py` — synchronous stdlib/`ThreadingHTTPServer` implementation of the + wire protocol, shelling out to `/usr/local/gcp/bin/sandbox`: + - create → `sandbox run --detach [--allow-egress] [--import-tar …] + --mount type=bind,source=/tmp/zenml-share/,destination=/mnt/zenml + --write -- sleep infinity` (the bind mount backs file transfer) + - exec → `sandbox exec -e K=V … -w CWD -- argv…`, stdout/stderr piped + and re-emitted as base64 SSE frames, exit code in the `exit` frame + - file PUT/GET → write/read `/tmp/zenml-share//…` on the host plus a + `cp` exec inside the sandbox + - snapshot → `sandbox tar --file /tmp/….tar` then GCS upload + (`google-cloud-storage`) +- `Dockerfile` + `README.md` with the exact deploy command: + + ```bash + gcloud beta run deploy zenml-sandbox-bridge --source . \ + --sandbox-launcher --no-allow-unauthenticated \ + --max-instances 1 --min-instances 1 --no-cpu-throttling + ``` + +The client treats the bridge URL as opaque, so users can substitute their own +implementation of the protocol. + +### 2.6 Registration & metadata + +- `src/zenml/integrations/gcp/__init__.py`: add + `GCP_CLOUDRUN_SANDBOX_FLAVOR = "cloudrun"`; register + `CloudRunSandboxFlavor` in `flavors()`. +- `src/zenml/integrations/gcp/flavors/__init__.py`: export config + flavor. +- Logo: `https://public-flavor-logos.s3.eu-central-1.amazonaws.com/sandbox/cloudrun.png` + (asset upload tracked outside this repo, same as other new flavors). + +## 3. Testing + +`tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py`, mirroring +the Cloudflare suite (`httpx.MockTransport`, no real bridge): + +- flavor metadata / config validation (https enforcement, snapshot prefix + validation) +- bridge client: create/delete/running/exec/file/snapshot request shapes, + retry policy, error surfacing +- SSE parsing: interleaved stdout/stderr, exit frames, error frames, + truncated streams +- process demux: concurrent stdout/stderr draining, `wait` timeout, `kill` +- session lifecycle: per-exec env merging, `destroy_on_exit`, closed-handle + errors, upload/download size guards +- sandbox: `create_session`/`attach`/`restore` wiring, snapshot gating on + `snapshot_uri_prefix`, ID-token path selection (mocked google.auth) + +Live end-to-end testing against a real deployed bridge is manual, per the +repo's policy for external-service integrations. + +## 4. Documentation + +- `docs/book/component-guide/sandboxes/cloudrun.md` — deploy-the-bridge + walkthrough, IAM setup (`roles/run.invoker` for the caller; + GCS access for the bridge SA when snapshots are enabled), registration: + + ```bash + zenml sandbox register my_cloudrun_sandbox --flavor cloudrun \ + --service_url=https://zenml-sandbox-bridge-….run.app + zenml stack update my_stack --sandbox my_cloudrun_sandbox + ``` + +- Add the page to `docs/book/component-guide/toc.md` under Sandboxes. + +## 5. Out of scope (future work) + +- `sandbox fork` → no abstraction hook today (a `fork()` on `SandboxSession` + would benefit Modal too). +- Auto-deploying the bridge from the client via `google-cloud-run` (the + Deployer already owns "deploy a Cloud Run service"; a convenience CLI could + reuse it). +- Multi-instance bridges with session affinity routing. +- `aexec` async execution. + +## 6. Risks + +- **Preview API**: the `sandbox` CLI is pre-GA; flags may change. The CLI + surface is isolated in the bridge (one file), not in the ZenML package. +- **Instance lifetime**: Cloud Run may recycle the bridge instance, killing + persistent sandboxes. `attach()` reports this cleanly via + `GET /running`; docs recommend snapshot checkpointing for long-lived state. +- **Protocol drift** with the Cloudflare bridge: mitigated by versioned + paths (`/v1/…`) and per-flavor unit suites. diff --git a/docs/book/component-guide/sandboxes/cloudrun.md b/docs/book/component-guide/sandboxes/cloudrun.md new file mode 100644 index 00000000000..85d2bca8e28 --- /dev/null +++ b/docs/book/component-guide/sandboxes/cloudrun.md @@ -0,0 +1,113 @@ +--- +description: Running agent code in Google Cloud Run sandboxes. +--- + +# Google Cloud Run Sandbox + +[Cloud Run sandboxes](https://docs.cloud.google.com/run/docs/code-execution) give you fast (sub-second), strongly isolated environments for executing untrusted code — most commonly code written by an LLM agent. The `cloudrun` sandbox flavor wraps them in ZenML's [Sandbox](README.md) interface, so an agent inside a step can `exec()` commands, stream output, transfer files, snapshot the filesystem to Cloud Storage, and restore from those snapshots. + +### How it works + +Cloud Run sandboxes have no standalone REST API: the `sandbox` CLI that creates and drives them only exists *inside* a Cloud Run service deployed with the sandbox launcher enabled. ZenML therefore talks to a small **bridge service** that you deploy once to Cloud Run — see [`examples/cloudrun_sandbox_bridge`](https://github.com/zenml-io/zenml/tree/main/examples/cloudrun_sandbox_bridge) — and authenticates to it with Google-signed ID tokens, the standard Cloud Run IAM invoker flow. + +``` +step code ──HTTPS + ID token──▶ bridge service (--sandbox-launcher) ──sandbox CLI──▶ isolated sandboxes +``` + +### When to use it + +Use the Cloud Run sandbox flavor when: + +- You're already running on Google Cloud and want sandbox isolation without adding another vendor. +- You need no-egress-by-default isolation for LLM-generated code (network access is opt-in per sandbox). +- You want filesystem snapshots stored durably in Cloud Storage that can boot new sandboxes. + +### How to deploy it + +1. Deploy the bridge service (from `examples/cloudrun_sandbox_bridge`): + + ```bash + gcloud beta run deploy zenml-sandbox-bridge \ + --source . \ + --region europe-west1 \ + --sandbox-launcher \ + --no-allow-unauthenticated \ + --max-instances 1 --min-instances 1 --no-cpu-throttling + ``` + + `--max-instances 1` matters: persistent sandboxes live on a single instance, and a scaled-out bridge would route execs to instances that don't own the sandbox. `--min-instances 1 --no-cpu-throttling` keep detached sandboxes running between requests. + +2. Grant the identity that runs your pipelines `roles/run.invoker` on the service. For snapshots, grant the *bridge's* runtime service account write access to your snapshot bucket. + +3. Install the GCP integration: + + ```bash + zenml integration install gcp + ``` + +### How to register the component + +```bash +zenml sandbox register cloudrun_sandbox \ + --flavor=cloudrun \ + --service_url=https://zenml-sandbox-bridge-abc123-ew.a.run.app \ + --snapshot_uri_prefix=gs://my-bucket/zenml-sandbox-snapshots + +zenml stack update --sandbox cloudrun_sandbox +``` + +Authentication follows the usual GCP component rules: a linked [GCP Service Connector](https://docs.zenml.io/how-to/infrastructure-deployment/auth-management/gcp-service-connector) takes priority, then `service_account_path`, then Application Default Credentials. Note that ID tokens (which Cloud Run IAM requires) can only be minted from service-account-based credentials — user credentials from `gcloud auth application-default login` won't work. + +```bash +zenml service-connector register gcp_connector --type gcp --auth-method service-account \ + --service_account_json=@path/to/key.json +zenml sandbox connect cloudrun_sandbox --connector gcp_connector +``` + +### How to use it + +```python +from zenml import step +from zenml.client import Client + + +@step +def run_agent_code(code: str) -> str: + sandbox = Client().active_stack.sandbox + with sandbox.create_session(destroy_on_exit=True) as session: + session.upload_file("requirements.txt", "/tmp/requirements.txt") + output = session.exec(["python3", "-c", code]).collect() + if output.exit_code != 0: + raise RuntimeError(output.stderr) + return output.stdout +``` + +Snapshots capture the sandbox's writable filesystem overlay as a tarball in Cloud Storage and can seed new sandboxes — including after the original bridge instance is gone: + +```python +snapshot = session.create_snapshot() # -> tarball at snapshot_uri_prefix +... +restored = sandbox.restore(snapshot) # new sandbox from the tarball +``` + +### Settings reference + +`CloudRunSandboxSettings` (override on individual `@step` decorations): + +| Field | Purpose | +|---|---| +| `sandbox_environment` | Env vars applied to every command in the session. Cloud Run sandboxes inherit nothing from the host, so anything the code needs must be set here (or per-exec via `session.exec(..., env=...)`). | +| `timeout_ms` | Per-exec timeout in milliseconds (default 120000). Must stay below the bridge service's Cloud Run request timeout. | +| `cwd` | Default working directory for commands. | +| `allow_egress` | Give sandboxes outbound network access (off by default). | + +Config-only fields on the component: `service_url`, `audience`, `allow_unauthenticated`, `snapshot_uri_prefix`, plus the standard GCP `project` / `service_account_path`. + +### Limitations + +- Persistent sandboxes are **instance-local**: if Cloud Run recycles the bridge instance, running sandboxes are lost (`attach()` reports this cleanly). Snapshot anything you need to keep. +- File transfer through the bridge is capped at 32 MiB per request; move larger data through Cloud Storage. +- A single exec is bounded by the Cloud Run request timeout (max 60 minutes). +- Cloud Run sandboxes are in public preview; the underlying CLI may change. + +For the full config surface, see the [SDK docs](https://sdkdocs.zenml.io). diff --git a/docs/book/component-guide/toc.md b/docs/book/component-guide/toc.md index 91ac16c6b36..06d606a2771 100644 --- a/docs/book/component-guide/toc.md +++ b/docs/book/component-guide/toc.md @@ -53,6 +53,7 @@ * [Develop a Custom Log Store](log-stores/custom.md) * [Sandboxes](sandboxes/README.md) * [Local](sandboxes/local.md) + * [Google Cloud Run](sandboxes/cloudrun.md) * [Kubernetes](sandboxes/kubernetes.md) * [Modal](sandboxes/modal.md) * [Step Operators](step-operators/README.md) diff --git a/examples/cloudrun_sandbox_bridge/Dockerfile b/examples/cloudrun_sandbox_bridge/Dockerfile new file mode 100644 index 00000000000..1089b3476a8 --- /dev/null +++ b/examples/cloudrun_sandbox_bridge/Dockerfile @@ -0,0 +1,8 @@ +FROM python:3.12-slim + +WORKDIR /app +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt +COPY main.py . + +CMD ["python", "main.py"] diff --git a/examples/cloudrun_sandbox_bridge/README.md b/examples/cloudrun_sandbox_bridge/README.md new file mode 100644 index 00000000000..e4693e8109f --- /dev/null +++ b/examples/cloudrun_sandbox_bridge/README.md @@ -0,0 +1,93 @@ +# ZenML sandbox bridge for Google Cloud Run + +Cloud Run sandboxes have no standalone REST API: the `sandbox` CLI only +exists inside a Cloud Run service deployed with the sandbox launcher. This +directory contains the small HTTP service — the *bridge* — that the ZenML +`cloudrun` sandbox flavor talks to. + +## What it does + +The bridge translates the ZenML sandbox bridge protocol (v1) into +`/usr/local/gcp/bin/sandbox` CLI calls: + +| Endpoint | CLI | +|---|---| +| `POST /v1/sandbox` | `sandbox run --detach … -- sleep infinity` | +| `POST /v1/sandbox//exec` | `sandbox exec [-e K=V] [-w DIR] -- argv…` (streamed as SSE) | +| `PUT/GET /v1/sandbox//file/` | staged through a bind-mounted share directory | +| `POST /v1/sandbox//snapshot` | `sandbox tar ` + upload to Cloud Storage | +| `GET /v1/sandbox//running` | `sandbox exec /bin/true` | +| `DELETE /v1/sandbox/` | `sandbox delete --force` | + +Authentication is delegated to Cloud Run IAM — the bridge itself contains no +auth code. Deploy with `--no-allow-unauthenticated` and Cloud Run verifies +the caller's Google ID token before the request reaches the container. + +## Deploy + +```bash +gcloud beta run deploy zenml-sandbox-bridge \ + --source . \ + --region europe-west1 \ + --sandbox-launcher \ + --no-allow-unauthenticated \ + --max-instances 1 \ + --min-instances 1 \ + --no-cpu-throttling \ + --memory 2Gi --cpu 2 +``` + +Why these flags: + +- `--sandbox-launcher` — enables the `sandbox` CLI inside the instance + (requires the second-generation execution environment). +- `--max-instances 1` — persistent sandboxes are instance-local; a second + instance would not see sandboxes created on the first. +- `--min-instances 1 --no-cpu-throttling` — keeps detached sandboxes alive + and scheduled between requests. +- Sandboxes share the instance's CPU/memory: size `--memory`/`--cpu` for + your workload. + +Grant callers invoke rights: + +```bash +gcloud run services add-iam-policy-binding zenml-sandbox-bridge \ + --region europe-west1 \ + --member "serviceAccount:YOUR_CALLER_SA" \ + --role roles/run.invoker +``` + +For snapshots, also grant the *service's* runtime service account +`roles/storage.objectAdmin` on the snapshot bucket. + +## Register the ZenML component + +```bash +zenml integration install gcp +zenml sandbox register cloudrun_sandbox --flavor cloudrun \ + --service_url="$(gcloud run services describe zenml-sandbox-bridge \ + --region europe-west1 --format 'value(status.url)')" \ + --snapshot_uri_prefix="gs://my-bucket/zenml-sandbox-snapshots" +zenml stack update my_stack --sandbox cloudrun_sandbox +``` + +Then, inside a step: + +```python +from zenml.client import Client + +sandbox = Client().active_stack.sandbox +with sandbox.create_session(destroy_on_exit=True) as session: + output = session.exec("python3 -c 'print(21 * 2)'").collect() + assert output.stdout.strip() == "42" +``` + +## Notes and limits + +- File transfer is capped at 32 MiB per request (a Cloud Run request-body + limit); route bigger payloads through GCS and fetch them from inside the + sandbox (`allow_egress` or a mounted path). +- A single exec stream is bounded by the Cloud Run request timeout (max 60 + minutes) — keep `timeout_ms` below the service timeout. +- The `sandbox` CLI is in public preview; flag names may change. All CLI + interaction is contained in `main.py`. diff --git a/examples/cloudrun_sandbox_bridge/main.py b/examples/cloudrun_sandbox_bridge/main.py new file mode 100644 index 00000000000..a38d729c925 --- /dev/null +++ b/examples/cloudrun_sandbox_bridge/main.py @@ -0,0 +1,476 @@ +"""ZenML sandbox bridge for Google Cloud Run sandboxes. + +A minimal synchronous HTTP service that exposes the ZenML sandbox bridge +protocol (v1) on top of the Cloud Run ``sandbox`` CLI. Deploy it to a Cloud +Run service with ``--sandbox-launcher`` enabled; the ZenML ``cloudrun`` +sandbox flavor talks to it over HTTPS. + +Authentication is delegated to Cloud Run IAM: deploy the service with +``--no-allow-unauthenticated`` and grant callers ``roles/run.invoker``. +Requests that reach this process have already presented a valid Google ID +token. + +Endpoints: + POST /v1/sandbox create a sandbox + GET /v1/sandbox//running liveness probe + POST /v1/sandbox//exec run a command, stream SSE + PUT /v1/sandbox//file/ upload a file + GET /v1/sandbox//file/ download a file + POST /v1/sandbox//snapshot export overlay tarball to GCS + DELETE /v1/sandbox/ delete a sandbox +""" + +import base64 +import json +import os +import posixpath +import re +import shlex +import shutil +import subprocess +import tempfile +import threading +import urllib.parse +import uuid +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Any, Dict, List, Optional, Tuple + +SANDBOX_BIN = os.environ.get("SANDBOX_BIN", "/usr/local/gcp/bin/sandbox") +SHARE_ROOT = os.environ.get("SHARE_ROOT", "/tmp/zenml-share") +# Mount point of the per-sandbox shared directory inside each sandbox, +# used to move file payloads across the isolation boundary. +SHARE_MOUNT = "/mnt/zenml" +MAX_FILE_BYTES = 32 * 1024 * 1024 +DEFAULT_EXEC_TIMEOUT_MS = 120_000 + +_SANDBOX_ID_RE = re.compile(r"^sb-[0-9a-f]{12}$") + +# Sandboxes created by this instance. Persistent sandboxes are +# instance-local, so an id missing here is gone (or belongs to a recycled +# instance) and reads as 404. +_sandboxes: Dict[str, bool] = {} +_sandboxes_lock = threading.Lock() + + +def _run_cli( + args: List[str], timeout: Optional[float] = None +) -> "subprocess.CompletedProcess[bytes]": + """Run the sandbox CLI and capture output.""" + return subprocess.run( + [SANDBOX_BIN, *args], capture_output=True, timeout=timeout + ) + + +def _share_dir(sandbox_id: str) -> str: + return os.path.join(SHARE_ROOT, sandbox_id) + + +def _create_sandbox(allow_egress: bool, import_tar: Optional[str]) -> str: + """Create a detached persistent sandbox with a bind-mounted share dir.""" + sandbox_id = f"sb-{uuid.uuid4().hex[:12]}" + share = _share_dir(sandbox_id) + os.makedirs(share, exist_ok=True) + + args = ["run", sandbox_id, "--detach"] + if allow_egress: + args.append("--allow-egress") + if import_tar: + args += ["--import-tar", import_tar] + args += [ + "--mount", + f"type=bind,source={share},destination={SHARE_MOUNT}", + "--write", + "--", + "/bin/sleep", + "infinity", + ] + result = _run_cli(args, timeout=60) + if result.returncode != 0: + shutil.rmtree(share, ignore_errors=True) + raise RuntimeError( + f"sandbox run failed ({result.returncode}): " + f"{result.stderr.decode(errors='replace')[:500]}" + ) + with _sandboxes_lock: + _sandboxes[sandbox_id] = True + return sandbox_id + + +def _delete_sandbox(sandbox_id: str) -> None: + result = _run_cli(["delete", sandbox_id, "--force"], timeout=60) + if result.returncode != 0: + raise RuntimeError( + f"sandbox delete failed ({result.returncode}): " + f"{result.stderr.decode(errors='replace')[:500]}" + ) + with _sandboxes_lock: + _sandboxes.pop(sandbox_id, None) + shutil.rmtree(_share_dir(sandbox_id), ignore_errors=True) + + +def _is_running(sandbox_id: str) -> bool: + with _sandboxes_lock: + if sandbox_id not in _sandboxes: + return False + result = _run_cli(["exec", sandbox_id, "--", "/bin/true"], timeout=30) + return result.returncode == 0 + + +def _exec_in_sandbox( + sandbox_id: str, + argv: List[str], + cwd: Optional[str], + env: Dict[str, str], +) -> "subprocess.Popen[bytes]": + """Launch a command in the sandbox, pipes attached.""" + args = ["exec", sandbox_id] + for key, value in env.items(): + args += ["-e", f"{key}={value}"] + if cwd: + args += ["-w", cwd] + args += ["--", *argv] + return subprocess.Popen( + [SANDBOX_BIN, *args], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + + +def _sandbox_shell(sandbox_id: str, script: str) -> None: + """Run a short shell script inside the sandbox, raising on failure.""" + result = _run_cli( + ["exec", sandbox_id, "--", "/bin/sh", "-c", script], timeout=120 + ) + if result.returncode != 0: + raise RuntimeError( + f"sandbox exec failed ({result.returncode}): " + f"{result.stderr.decode(errors='replace')[:500]}" + ) + + +def _snapshot_to_gcs(sandbox_id: str, gcs_uri: str) -> str: + """Export the sandbox overlay with `sandbox tar` and upload it to GCS.""" + from google.cloud import storage + + match = re.match(r"^gs://([^/]+)/(.+)$", gcs_uri) + if not match: + raise ValueError(f"Invalid gs:// URI: {gcs_uri}") + bucket_name, blob_name = match.groups() + + with tempfile.NamedTemporaryFile(suffix=".tar", delete=False) as tmp: + tar_path = tmp.name + try: + result = _run_cli(["tar", sandbox_id, "--file", tar_path], timeout=600) + if result.returncode != 0: + raise RuntimeError( + f"sandbox tar failed ({result.returncode}): " + f"{result.stderr.decode(errors='replace')[:500]}" + ) + client = storage.Client() + client.bucket(bucket_name).blob(blob_name).upload_from_filename( + tar_path + ) + finally: + os.unlink(tar_path) + return gcs_uri + + +def _download_import_tar(gcs_uri: str) -> str: + """Fetch a snapshot tarball from GCS to a local temp path.""" + from google.cloud import storage + + match = re.match(r"^gs://([^/]+)/(.+)$", gcs_uri) + if not match: + raise ValueError(f"Invalid gs:// URI: {gcs_uri}") + bucket_name, blob_name = match.groups() + + with tempfile.NamedTemporaryFile(suffix=".tar", delete=False) as tmp: + tar_path = tmp.name + client = storage.Client() + client.bucket(bucket_name).blob(blob_name).download_to_filename(tar_path) + return tar_path + + +def _safe_sandbox_path(path: str) -> str: + """Normalize a sandbox-relative path, rejecting traversal.""" + stripped = urllib.parse.unquote(path).lstrip("/") + normalized = posixpath.normpath(stripped) + if normalized != stripped or normalized.split("/", 1)[0] == "..": + raise ValueError(f"Unsafe path: {path}") + return "/" + normalized + + +class BridgeHandler(BaseHTTPRequestHandler): + """Request handler implementing the bridge protocol.""" + + protocol_version = "HTTP/1.1" + + def _send_json(self, status: int, payload: Dict[str, Any]) -> None: + body = json.dumps(payload).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def _send_error_json(self, status: int, message: str) -> None: + self._send_json(status, {"error": message}) + + def _read_body(self) -> bytes: + length = int(self.headers.get("Content-Length") or 0) + if length > MAX_FILE_BYTES: + raise ValueError("Request body exceeds 32 MiB limit") + return self.rfile.read(length) + + def _read_json(self) -> Dict[str, Any]: + body = self._read_body() + return json.loads(body) if body else {} + + def _route(self) -> Optional[Tuple[str, str, str]]: + """Split the path into (sandbox_id, action, rest).""" + parts = self.path.split("?", 1)[0].split("/") + # /v1/sandbox[/[/[/rest...]]] + if len(parts) < 3 or parts[1] != "v1" or parts[2] != "sandbox": + return None + sandbox_id = parts[3] if len(parts) > 3 else "" + action = parts[4] if len(parts) > 4 else "" + rest = "/".join(parts[5:]) + return sandbox_id, action, rest + + def _known_sandbox(self, sandbox_id: str) -> bool: + if not _SANDBOX_ID_RE.match(sandbox_id): + return False + with _sandboxes_lock: + return sandbox_id in _sandboxes + + def do_POST(self) -> None: # noqa: N802 + route = self._route() + if route is None: + self._send_error_json(404, "Not found") + return + sandbox_id, action, _ = route + try: + if not sandbox_id: + self._handle_create() + elif not self._known_sandbox(sandbox_id): + self._send_error_json(404, "Unknown sandbox") + elif action == "exec": + self._handle_exec(sandbox_id) + elif action == "snapshot": + self._handle_snapshot(sandbox_id) + else: + self._send_error_json(404, "Not found") + except Exception as e: + self._send_error_json(500, str(e)) + + def do_GET(self) -> None: # noqa: N802 + route = self._route() + if route is None: + self._send_error_json(404, "Not found") + return + sandbox_id, action, rest = route + try: + if not self._known_sandbox(sandbox_id): + self._send_error_json(404, "Unknown sandbox") + elif action == "running": + self._send_json(200, {"running": _is_running(sandbox_id)}) + elif action == "file" and rest: + self._handle_get_file(sandbox_id, rest) + else: + self._send_error_json(404, "Not found") + except Exception as e: + self._send_error_json(500, str(e)) + + def do_PUT(self) -> None: # noqa: N802 + route = self._route() + if route is None: + self._send_error_json(404, "Not found") + return + sandbox_id, action, rest = route + try: + if not self._known_sandbox(sandbox_id): + self._send_error_json(404, "Unknown sandbox") + elif action == "file" and rest: + self._handle_put_file(sandbox_id, rest) + else: + self._send_error_json(404, "Not found") + except ValueError as e: + self._send_error_json(413, str(e)) + except Exception as e: + self._send_error_json(500, str(e)) + + def do_DELETE(self) -> None: # noqa: N802 + route = self._route() + if route is None: + self._send_error_json(404, "Not found") + return + sandbox_id, action, _ = route + try: + if not sandbox_id or action: + self._send_error_json(404, "Not found") + elif not self._known_sandbox(sandbox_id): + self._send_error_json(404, "Unknown sandbox") + else: + _delete_sandbox(sandbox_id) + self._send_json(200, {"deleted": sandbox_id}) + except Exception as e: + self._send_error_json(500, str(e)) + + def _handle_create(self) -> None: + payload = self._read_json() + import_tar_uri = payload.get("import_tar_uri") + import_tar = ( + _download_import_tar(import_tar_uri) if import_tar_uri else None + ) + try: + sandbox_id = _create_sandbox( + allow_egress=bool(payload.get("allow_egress")), + import_tar=import_tar, + ) + finally: + if import_tar: + os.unlink(import_tar) + self._send_json(200, {"id": sandbox_id}) + + def _handle_exec(self, sandbox_id: str) -> None: + payload = self._read_json() + argv = payload.get("argv") + if not isinstance(argv, list) or not argv: + self._send_error_json(400, "Missing argv") + return + timeout_ms = int(payload.get("timeout_ms", DEFAULT_EXEC_TIMEOUT_MS)) + proc = _exec_in_sandbox( + sandbox_id, + [str(a) for a in argv], + payload.get("cwd"), + dict(payload.get("env") or {}), + ) + + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Cache-Control", "no-cache") + self.send_header("Transfer-Encoding", "chunked") + self.end_headers() + + write_lock = threading.Lock() + + def _emit(kind: str, data: str) -> None: + frame = f"event: {kind}\ndata: {data}\n\n".encode() + with write_lock: + chunk = f"{len(frame):x}\r\n".encode() + frame + b"\r\n" + self.wfile.write(chunk) + self.wfile.flush() + + def _pump(stream: Any, kind: str) -> None: + while True: + chunk = stream.read(16384) + if not chunk: + return + _emit(kind, base64.b64encode(chunk).decode("ascii")) + + threads = [ + threading.Thread( + target=_pump, args=(proc.stdout, "stdout"), daemon=True + ), + threading.Thread( + target=_pump, args=(proc.stderr, "stderr"), daemon=True + ), + ] + for t in threads: + t.start() + + timed_out = False + try: + proc.wait(timeout=timeout_ms / 1000) + except subprocess.TimeoutExpired: + timed_out = True + proc.kill() + proc.wait() + for t in threads: + t.join() + + try: + if timed_out: + _emit( + "error", + json.dumps( + {"error": f"exec timed out after {timeout_ms}ms"} + ), + ) + else: + _emit("exit", json.dumps({"exit_code": proc.returncode})) + with write_lock: + self.wfile.write(b"0\r\n\r\n") + self.wfile.flush() + except BrokenPipeError: + # Client hung up (kill()); the sandbox command was already + # reaped above, so there is nothing to clean up. + pass + + def _handle_put_file(self, sandbox_id: str, rest: str) -> None: + dest = _safe_sandbox_path(rest) + data = self._read_body() + staging_name = f"put-{uuid.uuid4().hex}" + staging_host = os.path.join(_share_dir(sandbox_id), staging_name) + with open(staging_host, "wb") as f: + f.write(data) + try: + _sandbox_shell( + sandbox_id, + f"mkdir -p {shlex.quote(posixpath.dirname(dest) or '/')} && " + f"cp {shlex.quote(f'{SHARE_MOUNT}/{staging_name}')} " + f"{shlex.quote(dest)}", + ) + finally: + os.unlink(staging_host) + self._send_json(200, {"written": dest}) + + def _handle_get_file(self, sandbox_id: str, rest: str) -> None: + src = _safe_sandbox_path(rest) + staging_name = f"get-{uuid.uuid4().hex}" + staging_host = os.path.join(_share_dir(sandbox_id), staging_name) + _sandbox_shell( + sandbox_id, + f"cp {shlex.quote(src)} " + f"{shlex.quote(f'{SHARE_MOUNT}/{staging_name}')}", + ) + try: + size = os.path.getsize(staging_host) + if size > MAX_FILE_BYTES: + self._send_error_json( + 413, f"File is {size} bytes; limit is 32 MiB" + ) + return + self.send_response(200) + self.send_header("Content-Type", "application/octet-stream") + self.send_header("Content-Length", str(size)) + self.end_headers() + with open(staging_host, "rb") as f: + shutil.copyfileobj(f, self.wfile) + finally: + os.unlink(staging_host) + + def _handle_snapshot(self, sandbox_id: str) -> None: + payload = self._read_json() + gcs_uri = payload.get("gcs_uri") + if not gcs_uri: + self._send_error_json(400, "Missing gcs_uri") + return + uri = _snapshot_to_gcs(sandbox_id, gcs_uri) + self._send_json(200, {"uri": uri}) + + def log_message(self, format: str, *args: Any) -> None: + # Route access logs to stdout for Cloud Logging. + print(f"{self.address_string()} - {format % args}") + + +def main() -> None: + port = int(os.environ.get("PORT", "8080")) + os.makedirs(SHARE_ROOT, exist_ok=True) + server = ThreadingHTTPServer(("0.0.0.0", port), BridgeHandler) + print(f"ZenML sandbox bridge listening on :{port}") + server.serve_forever() + + +if __name__ == "__main__": + main() diff --git a/examples/cloudrun_sandbox_bridge/requirements.txt b/examples/cloudrun_sandbox_bridge/requirements.txt new file mode 100644 index 00000000000..ec2a08f31d6 --- /dev/null +++ b/examples/cloudrun_sandbox_bridge/requirements.txt @@ -0,0 +1 @@ +google-cloud-storage>=2.9.0 diff --git a/src/zenml/integrations/gcp/__init__.py b/src/zenml/integrations/gcp/__init__.py index 11bb39f94c9..d89a8d9d021 100644 --- a/src/zenml/integrations/gcp/__init__.py +++ b/src/zenml/integrations/gcp/__init__.py @@ -34,6 +34,7 @@ GCP_VERTEX_ORCHESTRATOR_FLAVOR = "vertex" GCP_VERTEX_STEP_OPERATOR_FLAVOR = "vertex" GCP_DEPLOYER_FLAVOR = "gcp" +GCP_CLOUDRUN_SANDBOX_FLAVOR = "cloudrun" # Service connector constants GCP_CONNECTOR_TYPE = "gcp" @@ -80,6 +81,7 @@ def flavors(cls) -> List[Type[Flavor]]: List of stack component flavors for this integration. """ from zenml.integrations.gcp.flavors import ( + CloudRunSandboxFlavor, GCPArtifactStoreFlavor, GCPDeployerFlavor, GCPImageBuilderFlavor, @@ -89,6 +91,7 @@ def flavors(cls) -> List[Type[Flavor]]: ) return [ + CloudRunSandboxFlavor, GCPArtifactStoreFlavor, GCPDeployerFlavor, GCPImageBuilderFlavor, diff --git a/src/zenml/integrations/gcp/flavors/__init__.py b/src/zenml/integrations/gcp/flavors/__init__.py index a920aa14242..c50f88e9751 100644 --- a/src/zenml/integrations/gcp/flavors/__init__.py +++ b/src/zenml/integrations/gcp/flavors/__init__.py @@ -13,6 +13,11 @@ # permissions and limitations under the License. """GCP integration flavors.""" +from zenml.integrations.gcp.flavors.cloudrun_sandbox_flavor import ( + CloudRunSandboxConfig, + CloudRunSandboxFlavor, + CloudRunSandboxSettings, +) from zenml.integrations.gcp.flavors.gcp_artifact_store_flavor import ( GCPArtifactStoreConfig, GCPArtifactStoreFlavor, @@ -39,6 +44,9 @@ ) __all__ = [ + "CloudRunSandboxFlavor", + "CloudRunSandboxConfig", + "CloudRunSandboxSettings", "GCPArtifactStoreFlavor", "GCPArtifactStoreConfig", "GCPDeployerFlavor", diff --git a/src/zenml/integrations/gcp/flavors/cloudrun_sandbox_flavor.py b/src/zenml/integrations/gcp/flavors/cloudrun_sandbox_flavor.py new file mode 100644 index 00000000000..9f35c97a775 --- /dev/null +++ b/src/zenml/integrations/gcp/flavors/cloudrun_sandbox_flavor.py @@ -0,0 +1,249 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# 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: +# +# https://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. +"""Google Cloud Run sandbox flavor.""" + +from typing import TYPE_CHECKING, Optional, Type +from urllib.parse import urlparse + +from pydantic import Field, field_validator + +from zenml.integrations.gcp import ( + GCP_CLOUDRUN_SANDBOX_FLAVOR, + GCP_CONNECTOR_TYPE, + GCP_RESOURCE_TYPE, +) +from zenml.integrations.gcp.google_credentials_mixin import ( + GoogleCredentialsConfigMixin, +) +from zenml.models import ServiceConnectorRequirements +from zenml.sandboxes import ( + BaseSandboxConfig, + BaseSandboxFlavor, + BaseSandboxSettings, +) + +if TYPE_CHECKING: + from zenml.integrations.gcp.sandboxes import CloudRunSandbox + + +DEFAULT_BRIDGE_TIMEOUT_MS = 120_000 + + +class CloudRunSandboxSettings(BaseSandboxSettings): + """Per-step settings for the Cloud Run sandbox.""" + + timeout_ms: int = Field( + default=DEFAULT_BRIDGE_TIMEOUT_MS, + ge=1, + description="Per-exec timeout in milliseconds, passed to the bridge " + "as timeout_ms on POST /v1/sandbox/:id/exec. Must stay below the " + "Cloud Run service request timeout of the bridge (max 60 minutes). " + "Example: 60000 for a one-minute cap", + ) + cwd: Optional[str] = Field( + default=None, + description="Default working directory for commands executed inside " + "the sandbox. Must be an absolute path in the sandbox filesystem. " + "Example: '/tmp/workspace'", + ) + allow_egress: bool = Field( + default=False, + description="Controls outbound network access for sandboxes created " + "with these settings. Cloud Run sandboxes have no egress by default; " + "setting this to True passes --allow-egress at sandbox creation. " + "Keep disabled when running untrusted code that has no reason to " + "reach the network", + ) + + +class CloudRunSandboxConfig( + BaseSandboxConfig, GoogleCredentialsConfigMixin, CloudRunSandboxSettings +): + """Configuration for the Cloud Run sandbox component.""" + + service_url: str = Field( + description="URL of the Cloud Run service deployed with " + "--sandbox-launcher that runs the ZenML sandbox bridge. Must be an " + "https:// URL. Example: " + "'https://zenml-sandbox-bridge-abc123-ew.a.run.app'", + ) + audience: Optional[str] = Field( + default=None, + description="Audience claim for the Google-signed ID tokens sent to " + "the bridge. Defaults to service_url, which is what Cloud Run IAM " + "expects; only set this when the bridge sits behind a load balancer " + "with a custom audience. Example: 'https://bridge.example.com'", + ) + allow_unauthenticated: bool = Field( + default=False, + description="Skips ID-token authentication entirely. Only for " + "bridges deployed with --allow-unauthenticated, which exposes the " + "sandbox launcher to the public internet and is strongly " + "discouraged outside local experiments", + ) + snapshot_uri_prefix: Optional[str] = Field( + default=None, + description="Cloud Storage URI prefix under which sandbox snapshots " + "are stored, enabling create_snapshot()/restore(). The bridge " + "service account needs read/write access to this location. " + "Example: 'gs://my-bucket/zenml-sandbox-snapshots'", + ) + + @field_validator("service_url") + @classmethod + def _validate_service_url_scheme(cls, value: str) -> str: + """Require https for the bridge URL. + + Args: + value: The configured service URL. + + Returns: + The validated URL, unchanged. + + Raises: + ValueError: If the URL scheme is not https (or http to + localhost/127.0.0.1). + """ + parsed = urlparse(value) + if parsed.scheme == "https": + return value + if parsed.scheme == "http" and parsed.hostname in ( + "localhost", + "127.0.0.1", + ): + return value + raise ValueError( + f"Invalid service_url '{value}': the Cloud Run sandbox bridge " + "must be reached over https (an identity token is sent as a " + "bearer token on every request). Plain http is only allowed " + "for localhost/127.0.0.1 when testing a local bridge." + ) + + @field_validator("snapshot_uri_prefix") + @classmethod + def _validate_snapshot_uri_prefix( + cls, value: Optional[str] + ) -> Optional[str]: + """Require a gs:// URI for the snapshot location. + + Args: + value: The configured snapshot URI prefix. + + Returns: + The validated URI prefix with any trailing slash removed. + + Raises: + ValueError: If the URI does not use the gs:// scheme. + """ + if value is None: + return None + if not value.startswith("gs://"): + raise ValueError( + f"Invalid snapshot_uri_prefix '{value}': must be a Cloud " + "Storage URI starting with 'gs://'." + ) + return value.rstrip("/") + + @property + def is_remote(self) -> bool: + """Cloud Run sandboxes run on Google's infrastructure. + + Returns: + ``True``: the ZenML client is not the host. + """ + return True + + +class CloudRunSandboxFlavor(BaseSandboxFlavor): + """Google Cloud Run sandbox flavor.""" + + @property + def name(self) -> str: + """Flavor name. + + Returns: + ``"cloudrun"``. + """ + return GCP_CLOUDRUN_SANDBOX_FLAVOR + + @property + def display_name(self) -> str: + """Human-readable flavor name. + + Returns: + The display name. + """ + return "Cloud Run" + + @property + def service_connector_requirements( + self, + ) -> Optional[ServiceConnectorRequirements]: + """Service connector requirements for this flavor. + + Returns: + Requirements matching the GCP service connector's generic + resource type. + """ + return ServiceConnectorRequirements( + connector_type=GCP_CONNECTOR_TYPE, + resource_type=GCP_RESOURCE_TYPE, + ) + + @property + def docs_url(self) -> Optional[str]: + """URL to user-facing docs for this flavor. + + Returns: + The flavor docs URL. + """ + return self.generate_default_docs_url() + + @property + def sdk_docs_url(self) -> Optional[str]: + """URL to SDK docs for this flavor. + + Returns: + The flavor SDK docs URL. + """ + return self.generate_default_sdk_docs_url() + + @property + def logo_url(self) -> str: + """Dashboard logo URL. + + Returns: + The flavor logo URL. + """ + return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/sandbox/cloudrun.png" + + @property + def config_class(self) -> Type[CloudRunSandboxConfig]: + """Config class. + + Returns: + ``CloudRunSandboxConfig``. + """ + return CloudRunSandboxConfig + + @property + def implementation_class(self) -> Type["CloudRunSandbox"]: + """Implementation class. + + Returns: + ``CloudRunSandbox``. + """ + from zenml.integrations.gcp.sandboxes import CloudRunSandbox + + return CloudRunSandbox diff --git a/src/zenml/integrations/gcp/sandboxes/__init__.py b/src/zenml/integrations/gcp/sandboxes/__init__.py new file mode 100644 index 00000000000..cea2575074d --- /dev/null +++ b/src/zenml/integrations/gcp/sandboxes/__init__.py @@ -0,0 +1,26 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# 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: +# +# https://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. +"""Cloud Run sandbox for the GCP integration.""" + +from zenml.integrations.gcp.sandboxes.cloudrun_sandbox import ( + CloudRunSandbox, + CloudRunSandboxProcess, + CloudRunSandboxSession, +) + +__all__ = [ + "CloudRunSandbox", + "CloudRunSandboxProcess", + "CloudRunSandboxSession", +] diff --git a/src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py b/src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py new file mode 100644 index 00000000000..73d56142402 --- /dev/null +++ b/src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py @@ -0,0 +1,1206 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# 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: +# +# https://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. +"""Cloud Run sandbox implementation backed by the sandbox bridge HTTP API. + +Cloud Run sandboxes have no standalone REST API: the ``sandbox`` CLI only +exists inside a Cloud Run service deployed with ``--sandbox-launcher``. This +flavor therefore talks to a small "bridge" HTTP service (see +``examples/cloudrun_sandbox_bridge/``) deployed on such a service, using +Google-signed ID tokens for authentication — the standard Cloud Run IAM +invoker flow. +""" + +import base64 +import datetime +import json +import logging +import os +import posixpath +import shlex +import threading +import time +import urllib.parse +from collections import deque +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, + Any, + Deque, + Dict, + FrozenSet, + Iterator, + List, + Optional, + Tuple, + Type, + Union, + cast, +) + +import httpx + +from zenml.config.base_settings import BaseSettings +from zenml.integrations.gcp.flavors.cloudrun_sandbox_flavor import ( + DEFAULT_BRIDGE_TIMEOUT_MS, + CloudRunSandboxConfig, + CloudRunSandboxSettings, +) +from zenml.integrations.gcp.google_credentials_mixin import ( + GoogleCredentialsMixin, +) +from zenml.logger import get_logger +from zenml.sandboxes import ( + BaseSandbox, + BaseSandboxSettings, + SandboxExecError, + SandboxProcess, + SandboxSession, + SandboxSnapshot, +) + +if TYPE_CHECKING: + pass + +logger = get_logger(__name__) + +# Cloud Run caps non-streaming HTTP/1 request bodies at 32 MiB; the bridge +# inherits that limit for file transfer. Larger payloads should go through +# GCS and be pulled from inside the sandbox instead. +_BRIDGE_FILE_MAX_BYTES = 32 * 1024 * 1024 + +# Transient statuses worth retrying; 500s are bridge bugs and surface +# immediately. +_RETRYABLE_STATUSES: FrozenSet[int] = frozenset({429, 502, 503, 504}) +# Only idempotent methods are retried: a POST that failed mid-flight may +# already have executed on the bridge (created a sandbox, launched a +# command), so retrying it risks double execution. +_RETRYABLE_METHODS: FrozenSet[str] = frozenset({"GET", "PUT", "DELETE"}) +_MAX_RETRIES = 3 + +# Per-stream line cap for the SSE pump buffers. A caller that only +# wait()s on a noisy command would otherwise grow these deques without +# bound; past the cap the oldest lines are dropped (logged once). +_STREAM_BUFFER_MAX_LINES = 100_000 + +# Standard Google OAuth2 token endpoint, used when deriving ID-token +# credentials from plain service-account credentials. +_GOOGLE_OAUTH2_TOKEN_ENDPOINT = "https://oauth2.googleapis.com/token" + +# Refresh ID tokens this many seconds before they expire so an exec that +# starts right at the boundary doesn't fail with a 401. +_TOKEN_REFRESH_SLACK_SECONDS = 60 + + +@dataclass(frozen=True) +class _BridgeEvent: + """One decoded SSE event from the bridge exec stream.""" + + # "stdout"/"stderr" events carry decoded text, "exit" carries the + # exit code; error frames raise during parsing and never get here. + kind: str + data: Union[str, int] + + +def _sanitize_remote_path(remote_path: str) -> str: + """Reject paths with `..` segments before they reach the bridge URL. + + Args: + remote_path: A path inside the sandbox filesystem. + + Returns: + The path stripped of any leading slash, ready for URL composition. + + Raises: + ValueError: If the normalized path escapes the sandbox root. + """ + stripped = remote_path.lstrip("/") + normalized = posixpath.normpath(stripped) + if normalized != stripped or normalized.split("/", 1)[0] == "..": + raise ValueError( + f"Remote path '{remote_path}' resolves outside the sandbox " + "filesystem." + ) + return stripped + + +def _parse_sse_stream(lines: Iterator[str]) -> Iterator[_BridgeEvent]: + """Parse the bridge's text/event-stream into typed events. + + Args: + lines: Decoded SSE lines (no trailing newlines). + + Yields: + One ``_BridgeEvent`` per dispatch. + """ + event: Optional[str] = None + data_chunks: List[str] = [] + + def _dispatch() -> Iterator[_BridgeEvent]: + nonlocal event, data_chunks + if event is None and not data_chunks: + return + raw_data = "\n".join(data_chunks) + kind = event or "message" + event = None + data_chunks = [] + + if kind == "stdout" or kind == "stderr": + try: + decoded = base64.b64decode(raw_data).decode( + "utf-8", errors="replace" + ) + except Exception as e: + raise SandboxExecError( + f"Malformed base64 on bridge '{kind}' frame: {e}" + ) from e + yield _BridgeEvent(kind=kind, data=decoded) + elif kind == "exit": + try: + payload = json.loads(raw_data) if raw_data else {} + except json.JSONDecodeError as e: + raise SandboxExecError( + f"Malformed JSON on bridge 'exit' frame: {e}" + ) from e + if "exit_code" not in payload: + # Missing exit_code is a protocol bug — surface it instead + # of silently coercing to 0 (which would mask a failure). + raise SandboxExecError( + f"Bridge 'exit' frame missing 'exit_code': {payload!r}" + ) + yield _BridgeEvent(kind="exit", data=int(payload["exit_code"])) + elif kind == "error": + try: + payload = json.loads(raw_data) if raw_data else {} + except json.JSONDecodeError: + payload = {"error": raw_data} + raise SandboxExecError( + f"Bridge exec error: {payload.get('error', payload)}" + ) + # Unknown event kinds are ignored per SSE spec. + + for line in lines: + if line == "": + yield from _dispatch() + continue + + if line.startswith(":"): + # SSE comment / keep-alive. + continue + if line.startswith("event:"): + event = line[len("event:") :].strip() + elif line.startswith("data:"): + chunk = line[len("data:") :] + if chunk.startswith(" "): + chunk = chunk[1:] + data_chunks.append(chunk) + # Other field types ("id:", "retry:") are accepted but ignored. + + # Flush a trailing buffered event if the stream ended without a final + # blank line. Without this, a bridge that closes the connection + # right after writing `event: exit\ndata: ...\n` would silently lose + # the exit frame and the caller would see exit_code=0. + yield from _dispatch() + + +def _drain_sse_response(resp: "httpx.Response") -> Iterator[_BridgeEvent]: + """Yield events from an already-open SSE response, closing it at the end. + + Args: + resp: An open streaming ``httpx.Response``. + + Yields: + Each decoded ``_BridgeEvent`` from the response body. + """ + try: + yield from _parse_sse_stream(resp.iter_lines()) + finally: + try: + resp.close() + except Exception: + logger.debug("Closing bridge SSE response failed", exc_info=True) + + +class _AdcIdTokenCredentials: + """ID-token credentials minted via ``google.oauth2.id_token``. + + Fallback for credential types with no native ID-token counterpart + (e.g. plain Application Default Credentials on a workstation or the + metadata server on GCE/Cloud Run). Mimics the small slice of the + ``google.auth.credentials.Credentials`` surface the bridge client + uses: ``valid``, ``token`` and ``refresh()``. + """ + + def __init__(self, audience: str) -> None: + """Initialize the credentials. + + Args: + audience: The audience claim for minted ID tokens. + """ + self._audience = audience + self.token: Optional[str] = None + self.expiry: Optional[datetime.datetime] = None + + @property + def valid(self) -> bool: + """Whether the cached token is present and not about to expire. + + Returns: + ``True`` if the token can still be used. + """ + if self.token is None or self.expiry is None: + return False + remaining = self.expiry - datetime.datetime.now(datetime.timezone.utc) + return remaining.total_seconds() > _TOKEN_REFRESH_SLACK_SECONDS + + def refresh(self, request: Any) -> None: + """Fetch a fresh ID token via the ADC-based helper. + + Args: + request: A ``google.auth.transport.Request`` instance. + + Raises: + RuntimeError: If the ambient credentials cannot mint ID tokens. + """ + from google.auth import exceptions as google_auth_exceptions + from google.auth import jwt as google_jwt + from google.oauth2 import id_token as google_id_token + + try: + token = google_id_token.fetch_id_token(request, self._audience) + except google_auth_exceptions.GoogleAuthError as e: + raise RuntimeError( + "Could not mint a Google ID token for the Cloud Run " + f"sandbox bridge (audience '{self._audience}'): {e}. ID " + "tokens require service-account credentials: link a GCP " + "service connector, set service_account_path on the " + "sandbox component, or run with service-account-based " + "Application Default Credentials. User credentials from " + "`gcloud auth application-default login` cannot mint ID " + "tokens." + ) from e + claims = google_jwt.decode(token, verify=False) + self.token = token + self.expiry = datetime.datetime.fromtimestamp( + int(claims["exp"]), tz=datetime.timezone.utc + ) + + +class _CloudRunBridgeClient: + """HTTP wrapper around the bridge's /v1/sandbox/* surface.""" + + def __init__( + self, + service_url: str, + credentials: Optional[Any], + *, + transport: Optional["httpx.BaseTransport"] = None, + ) -> None: + """Initialize the client. + + Args: + service_url: Base URL of the deployed bridge service. + credentials: Google ID-token credentials used to authorize + requests (``None`` for unauthenticated bridges). Must + expose ``valid``, ``token`` and ``refresh(request)``. + transport: Optional httpx transport override, used by tests. + """ + self._credentials = credentials + self._credentials_lock = threading.Lock() + self._client = httpx.Client( + base_url=service_url.rstrip("/"), + timeout=httpx.Timeout(30.0, read=None), + transport=transport, + ) + + def close(self) -> None: + """Release the underlying HTTP client.""" + try: + self._client.close() + except Exception: + logger.debug("Closing bridge HTTP client failed", exc_info=True) + + def _auth_headers(self) -> Dict[str, str]: + """Build the Authorization header from the ID-token credentials. + + Returns: + Headers to merge into the outgoing request. + """ + if self._credentials is None: + return {} + with self._credentials_lock: + if not self._credentials.valid: + from google.auth.transport.requests import Request + + self._credentials.refresh(Request()) + token = self._credentials.token + return {"Authorization": f"Bearer {token}"} + + def _request( + self, + method: str, + path: str, + *, + json_body: Optional[Dict[str, Any]] = None, + content: Optional[bytes] = None, + headers: Optional[Dict[str, str]] = None, + stream: bool = False, + ok_statuses: Tuple[int, ...] = (), + ) -> "httpx.Response": + """Issue a request with the bridge's retry policy. + + Args: + method: HTTP verb. + path: Path under the bridge base URL. + json_body: JSON body for write requests. + content: Raw bytes body (mutually exclusive with json_body). + headers: Extra headers. + stream: When True, return the response without reading it. + ok_statuses: Error statuses that carry meaning for the caller + (e.g. a 404 that answers a liveness probe); returned + as-is instead of raising or retrying. + + Returns: + The httpx Response. + + Raises: + SandboxExecError: On a non-retryable error, or after retries + are exhausted. + """ + retryable_method = method in _RETRYABLE_METHODS + attempt = 0 + while True: + merged_headers = dict(headers or {}) + merged_headers.update(self._auth_headers()) + try: + if stream: + req = self._client.build_request( + method, + path, + json=json_body, + content=content, + headers=merged_headers, + ) + resp = self._client.send(req, stream=True) + else: + resp = self._client.request( + method, + path, + json=json_body, + content=content, + headers=merged_headers, + ) + except httpx.HTTPError as e: + if not retryable_method or attempt >= _MAX_RETRIES: + raise SandboxExecError( + f"Bridge request {method} {path} failed: {e}" + ) from e + time.sleep(0.25 * (2**attempt)) + attempt += 1 + continue + + if resp.status_code in ok_statuses or resp.status_code < 400: + return resp + + retryable = ( + retryable_method and resp.status_code in _RETRYABLE_STATUSES + ) + body_preview = "" + if not stream: + try: + body_preview = resp.text[:500] + except Exception: + body_preview = "" + else: + try: + resp.close() + except Exception: + pass + + if not retryable or attempt >= _MAX_RETRIES: + raise SandboxExecError( + f"Bridge {method} {path} returned " + f"{resp.status_code}: {body_preview}" + ) + time.sleep(0.25 * (2**attempt)) + attempt += 1 + + def create_sandbox( + self, + *, + allow_egress: bool = False, + import_tar_uri: Optional[str] = None, + ) -> str: + """Create a fresh sandbox on the bridge instance. + + Args: + allow_egress: Whether the sandbox gets outbound network access. + import_tar_uri: Optional ``gs://`` URI of a filesystem snapshot + tarball to initialize the sandbox from. + + Returns: + The bridge-assigned sandbox id. + + Raises: + SandboxExecError: If the bridge response is missing an id. + """ + body: Dict[str, Any] = {"allow_egress": allow_egress} + if import_tar_uri is not None: + body["import_tar_uri"] = import_tar_uri + resp = self._request("POST", "/v1/sandbox", json_body=body) + payload = resp.json() + sandbox_id = payload.get("id") + if not sandbox_id: + raise SandboxExecError( + f"Bridge create-sandbox response missing 'id': {payload!r}" + ) + return cast(str, sandbox_id) + + def delete_sandbox(self, sandbox_id: str) -> None: + """Delete a sandbox. + + Args: + sandbox_id: The sandbox to delete. + """ + self._request( + "DELETE", f"/v1/sandbox/{sandbox_id}", ok_statuses=(404, 410) + ) + + def is_running(self, sandbox_id: str) -> bool: + """Check whether the given sandbox is still alive. + + Args: + sandbox_id: The sandbox to check. + + Returns: + True if the bridge reports the sandbox as running, False if it + reports it as stopped or unknown (404/410). + """ + resp = self._request( + "GET", + f"/v1/sandbox/{sandbox_id}/running", + ok_statuses=(404, 410), + ) + if resp.status_code in (404, 410): + return False + return bool(resp.json().get("running", False)) + + def snapshot_sandbox(self, sandbox_id: str, gcs_uri: str) -> str: + """Export the sandbox filesystem overlay to Cloud Storage. + + Args: + sandbox_id: The sandbox to snapshot. + gcs_uri: Destination ``gs://`` URI for the tarball. + + Returns: + The URI of the stored snapshot tarball, as confirmed by the + bridge. + + Raises: + SandboxExecError: If the bridge response is missing the URI. + """ + resp = self._request( + "POST", + f"/v1/sandbox/{sandbox_id}/snapshot", + json_body={"gcs_uri": gcs_uri}, + ) + uri = resp.json().get("uri") + if not uri: + raise SandboxExecError("Bridge snapshot response missing 'uri'") + return cast(str, uri) + + def exec_stream( + self, + sandbox_id: str, + argv: List[str], + *, + cwd: Optional[str] = None, + env: Optional[Dict[str, str]] = None, + timeout_ms: int = DEFAULT_BRIDGE_TIMEOUT_MS, + ) -> Tuple["httpx.Response", Iterator[_BridgeEvent]]: + """Run a command and stream decoded SSE events. + + Args: + sandbox_id: The sandbox to exec into. + argv: Command argv list. + cwd: Optional working directory. + env: Environment variables for this command (passed to the + ``sandbox`` CLI via ``-e``; sandboxes inherit nothing from + the host). + timeout_ms: Per-exec timeout in milliseconds. + + Returns: + The open streaming response (closing it aborts the stream and + unblocks a reader mid-iteration) and an iterator of + ``_BridgeEvent`` instances in dispatch order. + """ + body: Dict[str, Any] = {"argv": argv, "timeout_ms": timeout_ms} + if cwd is not None: + body["cwd"] = cwd + if env: + body["env"] = env + + # POST eagerly so 4xx/auth errors raise at the call site; only SSE + # decoding is deferred. + resp = self._request( + "POST", + f"/v1/sandbox/{sandbox_id}/exec", + json_body=body, + stream=True, + ) + return resp, _drain_sse_response(resp) + + def put_file(self, sandbox_id: str, remote_path: str, data: bytes) -> None: + """Upload raw bytes to a sandbox path. + + Args: + sandbox_id: The sandbox to write into. + remote_path: Path inside the sandbox filesystem; `..` rejected. + data: Raw file bytes; max 32 MiB. + + Raises: + ValueError: If the file exceeds the bridge body limit. + """ + if len(data) > _BRIDGE_FILE_MAX_BYTES: + raise ValueError( + f"File of {len(data)} bytes exceeds the bridge limit of " + f"{_BRIDGE_FILE_MAX_BYTES} bytes (32 MiB)." + ) + safe_path = _sanitize_remote_path(remote_path) + encoded = urllib.parse.quote(safe_path, safe="/") + self._request( + "PUT", + f"/v1/sandbox/{sandbox_id}/file/{encoded}", + content=data, + headers={"Content-Type": "application/octet-stream"}, + ) + + def get_file(self, sandbox_id: str, remote_path: str) -> bytes: + """Download raw bytes from a sandbox path. + + Args: + sandbox_id: The sandbox to read from. + remote_path: Path inside the sandbox filesystem. + + Returns: + The raw file contents (bridge caps at 32 MiB). + """ + safe_path = _sanitize_remote_path(remote_path) + encoded = urllib.parse.quote(safe_path, safe="/") + resp = self._request("GET", f"/v1/sandbox/{sandbox_id}/file/{encoded}") + return resp.content + + +class CloudRunSandboxProcess(SandboxProcess): + """Demuxes a single SSE event iterator into stdout / stderr / exit.""" + + def __init__( + self, + event_iter: Iterator[_BridgeEvent], + *, + session: "CloudRunSandboxSession", + started_at: float, + response: Optional["httpx.Response"] = None, + ) -> None: + """Initialize the process wrapper. + + Args: + event_iter: Iterator yielding bridge SSE events for this exec. + session: Owning session. + started_at: Wall-clock launch time. + response: The underlying streaming SSE response; ``kill()`` + closes it to unblock the pump thread mid-read. + """ + super().__init__(session=session, started_at=started_at) + self._event_iter = event_iter + self._response = response + self._killed = threading.Event() + + # The bridge multiplexes stdout/stderr on one SSE stream. We pump + # it once on a background thread and demux into two buffers so + # `stdout()` and `stderr()` can be iterated independently from the + # caller's thread. + self._lock = threading.Lock() + self._stdout_buf: Deque[str] = deque() + self._stderr_buf: Deque[str] = deque() + self._buffer_truncated = False + self._data_available = threading.Condition(self._lock) + self._done = threading.Event() + self._exit_code: Optional[int] = None + self._pump_error: Optional[BaseException] = None + self._stdout_remainder = "" + self._stderr_remainder = "" + + self._pump = threading.Thread( + target=self._pump_events, + name="cloudrun-bridge-pump", + daemon=True, + ) + self._pump.start() + + def _push_text(self, target: Deque[str], remainder: str, text: str) -> str: + """Line-buffer arbitrary text into the target deque. + + Args: + target: Buffer to append lines to. + remainder: Partial line carried over from the previous chunk. + text: New text chunk. + + Returns: + The trailing partial line (no newline yet) to carry over. + """ + buf = remainder + text + while "\n" in buf: + line, buf = buf.split("\n", 1) + if len(target) >= _STREAM_BUFFER_MAX_LINES: + target.popleft() + if not self._buffer_truncated: + self._buffer_truncated = True + logger.warning( + "Cloud Run exec output exceeded the %d-line " + "buffer; oldest undrained lines are being " + "dropped. Drain process.stdout()/stderr() (or " + "use collect()) to keep full output.", + _STREAM_BUFFER_MAX_LINES, + ) + target.append(line + "\n") + return buf + + def _flush_remainders(self) -> None: + """Push trailing partial lines (no terminating newline) on stream end.""" + if self._stdout_remainder: + self._stdout_buf.append(self._stdout_remainder) + self._stdout_remainder = "" + if self._stderr_remainder: + self._stderr_buf.append(self._stderr_remainder) + self._stderr_remainder = "" + + def _pump_events(self) -> None: + """Background pump: drain SSE iterator into stdout/stderr buffers.""" + try: + for ev in self._event_iter: + with self._data_available: + # Only ``exit`` frames carry an int payload (the exit + # code); stdout/stderr frames carry decoded text. + if isinstance(ev.data, int): + self._exit_code = ev.data + elif ev.kind == "stdout": + self._stdout_remainder = self._push_text( + self._stdout_buf, + self._stdout_remainder, + ev.data, + ) + elif ev.kind == "stderr": + self._stderr_remainder = self._push_text( + self._stderr_buf, + self._stderr_remainder, + ev.data, + ) + self._data_available.notify_all() + except BaseException as e: # noqa: BLE001 + # kill() closes the response under our feet; the resulting + # httpx error is the expected shutdown path, not a failure. + if not self._killed.is_set(): + self._pump_error = e + finally: + with self._data_available: + self._flush_remainders() + self._done.set() + self._data_available.notify_all() + + def _iter_buffer(self, buf: Deque[str]) -> Iterator[str]: + """Block-pop lines from a buffer until the pump signals done. + + Args: + buf: The buffer to drain. + + Yields: + One line per call. + """ + while True: + with self._data_available: + while not buf and not self._done.is_set(): + self._data_available.wait() + if buf: + line = buf.popleft() + else: + return + yield line + + def stdout(self) -> Iterator[str]: + """Yields stdout lines, routed through the session log source. + + Returns: + Line iterator wrapped via ``session._wrap_stream`` so each line + also lands in the per-session ``sandbox:`` log. + """ + return self._session._wrap_stream( + self._iter_buffer(self._stdout_buf), log_level=logging.INFO + ) + + def stderr(self) -> Iterator[str]: + """Yields stderr lines, routed through the session log source. + + Returns: + Line iterator wrapped via ``session._wrap_stream`` so each line + also lands in the per-session ``sandbox:`` log. + """ + return self._session._wrap_stream( + self._iter_buffer(self._stderr_buf), log_level=logging.ERROR + ) + + def wait(self, timeout: Optional[float] = None) -> int: + """Block until the process exits, or ``timeout`` seconds pass. + + Args: + timeout: Optional wall-clock cap. ``None`` waits indefinitely. + + Returns: + The exit code captured from the bridge. + + Raises: + TimeoutError: If ``timeout`` elapsed before the bridge sent an + ``exit`` frame. + SandboxExecError: If the pump captured an error, the stream was + killed via ``kill()``/``close()``, or it ended without an + exit frame. + """ + completed = self._done.wait(timeout) + if not completed: + raise TimeoutError( + f"Cloud Run exec did not complete within {timeout}s. " + "Call process.kill() or session.destroy() to release the " + "bridge stream." + ) + pump_err = self._pump_error + if pump_err is not None: + if isinstance(pump_err, SandboxExecError): + raise SandboxExecError(str(pump_err)) from pump_err + raise SandboxExecError( + f"Bridge SSE pump failed: {pump_err}" + ) from pump_err + if self._exit_code is None: + if self._killed.is_set(): + raise SandboxExecError( + "exec stream was killed via kill() or session close(); " + "the command may keep running in the sandbox until " + "timeout_ms" + ) + # Stream ended cleanly but no exit frame arrived (truncated + # SSE, bridge stalled out, Cloud Run request timeout hit). + # Surface this as a failure rather than returning a bogus 0. + raise SandboxExecError( + "Cloud Run exec finished without an exit frame; the bridge " + "stream may have been truncated by the Cloud Run request " + "timeout or stalled." + ) + return self._exit_code + + def kill(self) -> None: + """Stop reading the SSE stream and unblock consumers.""" + # Set the flag before closing so the pump treats the resulting + # httpx error as a clean shutdown. + self._killed.set() + if self._response is not None: + try: + self._response.close() + except Exception: + logger.debug( + "Closing bridge SSE response during kill() failed", + exc_info=True, + ) + with self._data_available: + self._done.set() + self._data_available.notify_all() + + @property + def exit_code(self) -> Optional[int]: + """Exit code, or ``None`` while still running. + + Returns: + The captured exit code or ``None``. + """ + return self._exit_code + + +class CloudRunSandboxSession(SandboxSession): + """Cloud Run sandbox session over the bridge HTTP API.""" + + def __init__( + self, + sandbox_id: str, + *, + client: _CloudRunBridgeClient, + parent: "CloudRunSandbox", + session_env: Optional[Dict[str, str]] = None, + default_cwd: Optional[str] = None, + default_timeout_ms: int = DEFAULT_BRIDGE_TIMEOUT_MS, + snapshot_uri_prefix: Optional[str] = None, + destroy_on_exit: bool = False, + ) -> None: + """Initialize the session. + + Args: + sandbox_id: The bridge-assigned sandbox id (used as session id). + client: Bridge HTTP client. + parent: Owning sandbox component. + session_env: Env vars merged into every exec. Cloud Run + sandboxes inherit nothing from the host, and the bridge + keeps no session state, so the client resends these on + each command. + default_cwd: Default working directory for execs. + default_timeout_ms: Default per-exec timeout in milliseconds. + snapshot_uri_prefix: ``gs://`` prefix for snapshot tarballs; + ``None`` disables snapshot support. + destroy_on_exit: Whether the context manager destroys the + sandbox on exit instead of just closing the handle. + """ + # Subclass state must be set before super().__init__ so the + # dashboard hook (invoked during base __init__) has what it needs. + self._client = client + self._sandbox_id = sandbox_id + self._session_env = dict(session_env or {}) + self._default_cwd = default_cwd + self._default_timeout_ms = default_timeout_ms + self._snapshot_uri_prefix = snapshot_uri_prefix + self._processes: List[CloudRunSandboxProcess] = [] + super().__init__( + id=sandbox_id, parent=parent, destroy_on_exit=destroy_on_exit + ) + + def _get_dashboard_url(self) -> Optional[str]: + """Bridge sandboxes have no per-sandbox dashboard URL. + + Returns: + ``None``. Sandbox lifecycle events land in the bridge + service's Cloud Logging, but there is no stable per-sandbox + deep link. + """ + return None + + def _exec( + self, + command: Union[str, List[str]], + *, + cwd: Optional[str] = None, + env: Optional[Dict[str, str]] = None, + ) -> SandboxProcess: + """Launch a command in the sandbox. + + Args: + command: A list (argv) or string (shell-split via + ``shlex.split``). + cwd: Working directory inside the sandbox. + env: Per-exec environment variables, merged over the + session-level ``sandbox_environment``. The ``sandbox`` CLI + applies them via ``-e``, so values never appear on the + command line. + + Returns: + A ``CloudRunSandboxProcess``. + """ + argv: List[str] = ( + list(command) + if isinstance(command, list) + else shlex.split(command) + ) + self._log_command(argv) + + merged_env = {**self._session_env, **(env or {})} + effective_cwd = cwd if cwd is not None else self._default_cwd + started_at = time.time() + # exec_stream surfaces every failure mode as SandboxExecError, so + # launch errors propagate to the caller as-is. + response, event_iter = self._client.exec_stream( + self._sandbox_id, + argv, + cwd=effective_cwd, + env=merged_env or None, + timeout_ms=self._default_timeout_ms, + ) + process = CloudRunSandboxProcess( + event_iter, + session=self, + started_at=started_at, + response=response, + ) + self._processes.append(process) + return process + + def _create_snapshot(self) -> SandboxSnapshot: + """Export the sandbox filesystem overlay to Cloud Storage. + + Returns: + A ``SandboxSnapshot`` whose ``ref`` is the ``gs://`` URI of the + tarball produced by ``sandbox tar`` on the bridge. Only + filesystem state is captured — running processes and + environment variables are not part of the snapshot. + + Raises: + NotImplementedError: If the component has no + ``snapshot_uri_prefix`` configured. + """ + if not self._snapshot_uri_prefix: + raise NotImplementedError( + "Snapshots require the 'snapshot_uri_prefix' config " + "attribute on the Cloud Run sandbox component (a gs:// " + "location the bridge service account can write to)." + ) + gcs_uri = ( + f"{self._snapshot_uri_prefix}/{self._sandbox_id}-" + f"{int(time.time())}.tar" + ) + stored_uri = self._client.snapshot_sandbox(self._sandbox_id, gcs_uri) + return SandboxSnapshot( + sandbox_id=self._parent.id, + ref=stored_uri, + metadata={"source_sandbox": self._sandbox_id}, + ) + + def _upload_file(self, local_path: str, remote_path: str) -> None: + """Upload a local file into the sandbox. + + Args: + local_path: Source path on the caller's machine. + remote_path: Destination path inside the sandbox filesystem. + + Raises: + ValueError: If the file exceeds the bridge's body limit. + """ + # Check the size before reading: loading an oversized file into + # memory just to have put_file reject it invites an OOM. + size = os.path.getsize(local_path) + if size > _BRIDGE_FILE_MAX_BYTES: + raise ValueError( + f"File '{local_path}' is {size} bytes, exceeding the " + f"bridge's {_BRIDGE_FILE_MAX_BYTES}-byte upload limit." + ) + with open(local_path, "rb") as f: + data = f.read() + self._client.put_file(self._sandbox_id, remote_path, data) + + def _download_file(self, remote_path: str, local_path: str) -> None: + """Download a file from the sandbox. + + Args: + remote_path: Source path inside the sandbox filesystem. + local_path: Destination path on the caller's machine. + """ + data = self._client.get_file(self._sandbox_id, remote_path) + with open(local_path, "wb") as f: + f.write(data) + + def _close(self) -> None: + """Release local exec streams; the sandbox keeps running.""" + for process in self._processes: + if process.exit_code is None: + process.kill() + + def _destroy(self) -> None: + """Terminate the sandbox on the bridge instance. + + Raises: + RuntimeError: If the bridge fails to delete the sandbox. + """ + try: + self._client.delete_sandbox(self._sandbox_id) + except Exception as e: + raise RuntimeError( + f"Failed to delete Cloud Run sandbox '{self._sandbox_id}': " + f"{e}. It keeps consuming the bridge instance's resources " + "until deleted or the instance is recycled; retry " + "destroy() to terminate it." + ) from e + + +class CloudRunSandbox(BaseSandbox, GoogleCredentialsMixin): + """Sandbox flavor backed by Cloud Run sandboxes via a bridge service.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + """Initialize the Cloud Run sandbox component. + + Args: + *args: Forwarded to ``StackComponent``. + **kwargs: Forwarded to ``StackComponent``. + """ + super().__init__(*args, **kwargs) + self._bridge_client: Optional[_CloudRunBridgeClient] = None + self._bridge_client_lock = threading.Lock() + + @property + def config(self) -> CloudRunSandboxConfig: + """Typed config. + + Returns: + The Cloud-Run-specific config. + """ + return cast(CloudRunSandboxConfig, self._config) + + @property + def settings_class(self) -> Optional[Type["BaseSettings"]]: + """Settings class. + + Returns: + ``CloudRunSandboxSettings``. + """ + return CloudRunSandboxSettings + + def _build_id_token_credentials(self) -> Optional[Any]: + """Derive ID-token credentials for the bridge from the GCP auth. + + Cloud Run IAM authenticates callers with Google-signed ID tokens + (not OAuth2 access tokens), so the component's credentials — from + a linked GCP service connector, a service-account key file, or + Application Default Credentials — must be converted to an + ID-token-minting counterpart. + + Returns: + Credentials exposing ``valid``/``token``/``refresh``, or + ``None`` when ``allow_unauthenticated`` is set. + """ + if self.config.allow_unauthenticated: + return None + + from google.auth import impersonated_credentials + from google.oauth2 import service_account + + audience = self.config.audience or self.config.service_url + credentials, _ = self._get_authentication() + + if isinstance(credentials, impersonated_credentials.Credentials): + return impersonated_credentials.IDTokenCredentials( + credentials, + target_audience=audience, + include_email=True, + ) + if isinstance(credentials, service_account.Credentials): + return service_account.IDTokenCredentials( + signer=credentials.signer, + service_account_email=credentials.service_account_email, + token_uri=_GOOGLE_OAUTH2_TOKEN_ENDPOINT, + target_audience=audience, + ) + # Plain ADC (workstation or metadata server): fall back to the + # generic fetch_id_token helper, which raises a clear error for + # credential types that cannot mint ID tokens (e.g. user creds). + return _AdcIdTokenCredentials(audience) + + def _get_bridge_client(self) -> _CloudRunBridgeClient: + """Return the lazily-built bridge HTTP client. + + Returns: + A cached ``_CloudRunBridgeClient`` bound to this component. + """ + with self._bridge_client_lock: + if self._bridge_client is None: + self._bridge_client = _CloudRunBridgeClient( + self.config.service_url, + self._build_id_token_credentials(), + ) + return self._bridge_client + + def _build_session( + self, + sandbox_id: str, + settings: CloudRunSandboxSettings, + destroy_on_exit: bool = False, + ) -> CloudRunSandboxSession: + """Construct a session handle for an existing sandbox. + + Args: + sandbox_id: The bridge-assigned sandbox id. + settings: Resolved settings for the session. + destroy_on_exit: Whether the session context manager destroys + the sandbox on exit. + + Returns: + A ``CloudRunSandboxSession``. + """ + return CloudRunSandboxSession( + sandbox_id, + client=self._get_bridge_client(), + parent=self, + session_env=self._resolve_session_environment(settings), + default_cwd=settings.cwd, + default_timeout_ms=settings.timeout_ms, + snapshot_uri_prefix=self.config.snapshot_uri_prefix, + destroy_on_exit=destroy_on_exit, + ) + + def create_session( + self, + settings: Optional[BaseSandboxSettings] = None, + destroy_on_exit: bool = False, + ) -> SandboxSession: + """Boot a fresh Cloud Run sandbox via the bridge. + + Args: + settings: Optional per-step overrides. + destroy_on_exit: Whether to destroy the sandbox session when + the session context manager exits. + + Returns: + A ``CloudRunSandboxSession`` bound to the new sandbox. + """ + eff = cast(CloudRunSandboxSettings, self.resolve_settings(settings)) + sandbox_id = self._get_bridge_client().create_sandbox( + allow_egress=eff.allow_egress + ) + return self._build_session( + sandbox_id, eff, destroy_on_exit=destroy_on_exit + ) + + def attach(self, session_id: str) -> SandboxSession: + """Reattach to a still-running Cloud Run sandbox by id. + + Persistent sandboxes only survive as long as the bridge instance + that owns them, so attaching fails when Cloud Run has recycled the + instance in the meantime. + + Args: + session_id: The bridge sandbox id to attach to. + + Returns: + A ``CloudRunSandboxSession``. + + Raises: + RuntimeError: If the bridge reports the sandbox as not running. + """ + if not self._get_bridge_client().is_running(session_id): + raise RuntimeError( + f"Cloud Run sandbox '{session_id}' is not running. The " + "bridge instance that owned it may have been recycled; " + "create a new session (optionally from a snapshot)." + ) + eff = cast(CloudRunSandboxSettings, self.resolve_settings(None)) + return self._build_session(session_id, eff) + + def restore(self, snapshot: SandboxSnapshot) -> SandboxSession: + """Boot a new sandbox from a stored filesystem snapshot. + + Args: + snapshot: A ``SandboxSnapshot`` whose ``ref`` is a ``gs://`` + tarball URI captured via ``create_snapshot()``. + + Returns: + A new ``CloudRunSandboxSession`` initialized from the tarball. + No in-memory state from the original session is preserved, and + env vars are re-applied from the resolved settings. + """ + self._validate_snapshot(snapshot) + eff = cast(CloudRunSandboxSettings, self.resolve_settings(None)) + sandbox_id = self._get_bridge_client().create_sandbox( + allow_egress=eff.allow_egress, + import_tar_uri=snapshot.ref, + ) + return self._build_session(sandbox_id, eff) diff --git a/tests/unit/integrations/gcp/sandboxes/__init__.py b/tests/unit/integrations/gcp/sandboxes/__init__.py new file mode 100644 index 00000000000..9e94df41aaa --- /dev/null +++ b/tests/unit/integrations/gcp/sandboxes/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# 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: +# +# https://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/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py b/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py new file mode 100644 index 00000000000..3a64bfe97dc --- /dev/null +++ b/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py @@ -0,0 +1,463 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# 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: +# +# https://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. +"""Unit tests for the Cloud Run sandbox flavor. + +The bridge HTTP API is mocked via ``httpx.MockTransport`` so these tests +exercise the wiring (SSE parsing, demux, lifecycle, ID-token auth headers) +without a real bridge service or Google credentials. +""" + +import base64 +import json +import os +import tempfile +from typing import Any, Dict, List, Optional +from unittest.mock import MagicMock +from uuid import uuid4 + +import httpx +import pytest + +from zenml.enums import StackComponentType +from zenml.integrations.gcp.flavors import ( + CloudRunSandboxConfig, + CloudRunSandboxFlavor, + CloudRunSandboxSettings, +) +from zenml.integrations.gcp.sandboxes.cloudrun_sandbox import ( + _BRIDGE_FILE_MAX_BYTES, + CloudRunSandboxProcess, + CloudRunSandboxSession, + _BridgeEvent, + _CloudRunBridgeClient, + _parse_sse_stream, +) +from zenml.sandboxes import ( + BaseSandbox, + SandboxExecError, + SandboxSnapshot, +) + +SERVICE_URL = "https://zenml-sandbox-bridge-abc123-ew.a.run.app" + + +class _StaticCredentials: + """Test double for Google ID-token credentials.""" + + def __init__(self, token: str = "test-id-token") -> None: + self.token = token + self.refresh_count = 0 + + @property + def valid(self) -> bool: + return self.refresh_count > 0 + + def refresh(self, request: Any) -> None: + self.refresh_count += 1 + + +def _sse_event(kind: str, data: str) -> str: + return f"event: {kind}\ndata: {data}\n\n" + + +def _b64(text: str) -> str: + return base64.b64encode(text.encode("utf-8")).decode("ascii") + + +def _make_client( + handler: Any, credentials: Optional[_StaticCredentials] = None +) -> _CloudRunBridgeClient: + return _CloudRunBridgeClient( + SERVICE_URL, + credentials, + transport=httpx.MockTransport(handler), + ) + + +def _make_session( + client: _CloudRunBridgeClient, + *, + session_env: Optional[Dict[str, str]] = None, + snapshot_uri_prefix: Optional[str] = None, +) -> CloudRunSandboxSession: + parent = MagicMock(spec=BaseSandbox) + parent.flavor = "cloudrun" + parent.id = uuid4() + return CloudRunSandboxSession( + "sb-0123456789ab", + client=client, + parent=parent, + session_env=session_env, + snapshot_uri_prefix=snapshot_uri_prefix, + ) + + +def _passthrough_wrap_stream(session: CloudRunSandboxSession) -> None: + """Bypass step-log routing so tests can iterate raw process output.""" + session._wrap_stream = lambda lines, **_: lines # type: ignore[method-assign] + + +class TestFlavorMetadata: + def test_name_and_type(self) -> None: + flavor = CloudRunSandboxFlavor() + assert flavor.name == "cloudrun" + assert flavor.type == StackComponentType.SANDBOX + assert flavor.config_class is CloudRunSandboxConfig + + def test_service_connector_requirements(self) -> None: + requirements = CloudRunSandboxFlavor().service_connector_requirements + assert requirements is not None + assert requirements.connector_type == "gcp" + assert requirements.resource_type == "gcp-generic" + + +class TestConfigValidation: + def test_https_url_accepted(self) -> None: + config = CloudRunSandboxConfig(service_url=SERVICE_URL) + assert config.service_url == SERVICE_URL + assert config.is_remote is True + assert config.is_local is False + + def test_http_localhost_accepted(self) -> None: + config = CloudRunSandboxConfig(service_url="http://localhost:8080") + assert config.service_url == "http://localhost:8080" + + def test_plain_http_rejected(self) -> None: + with pytest.raises(ValueError, match="https"): + CloudRunSandboxConfig(service_url="http://bridge.example.com") + + def test_snapshot_prefix_requires_gs_scheme(self) -> None: + with pytest.raises(ValueError, match="gs://"): + CloudRunSandboxConfig( + service_url=SERVICE_URL, + snapshot_uri_prefix="s3://bucket/prefix", + ) + + def test_snapshot_prefix_trailing_slash_stripped(self) -> None: + config = CloudRunSandboxConfig( + service_url=SERVICE_URL, + snapshot_uri_prefix="gs://bucket/prefix/", + ) + assert config.snapshot_uri_prefix == "gs://bucket/prefix" + + def test_settings_defaults(self) -> None: + settings = CloudRunSandboxSettings() + assert settings.timeout_ms == 120_000 + assert settings.cwd is None + assert settings.allow_egress is False + assert settings.sandbox_environment == {} + + +class TestSseParsing: + def test_interleaved_stdout_stderr_and_exit(self) -> None: + stream = ( + _sse_event("stdout", _b64("out line\n")) + + _sse_event("stderr", _b64("err line\n")) + + _sse_event("exit", json.dumps({"exit_code": 3})) + ) + events = list(_parse_sse_stream(iter(stream.split("\n")))) + assert events == [ + _BridgeEvent(kind="stdout", data="out line\n"), + _BridgeEvent(kind="stderr", data="err line\n"), + _BridgeEvent(kind="exit", data=3), + ] + + def test_error_frame_raises(self) -> None: + stream = _sse_event("error", json.dumps({"error": "boom"})) + with pytest.raises(SandboxExecError, match="boom"): + list(_parse_sse_stream(iter(stream.split("\n")))) + + def test_exit_frame_missing_code_raises(self) -> None: + stream = _sse_event("exit", "{}") + with pytest.raises(SandboxExecError, match="exit_code"): + list(_parse_sse_stream(iter(stream.split("\n")))) + + def test_trailing_event_without_blank_line_flushed(self) -> None: + raw = 'event: exit\ndata: {"exit_code": 0}' + events = list(_parse_sse_stream(iter(raw.split("\n")))) + assert events == [_BridgeEvent(kind="exit", data=0)] + + def test_comments_and_unknown_events_ignored(self) -> None: + stream = ( + ": keep-alive\n\n" + + _sse_event("heartbeat", "x") + + _sse_event("exit", json.dumps({"exit_code": 0})) + ) + events = list(_parse_sse_stream(iter(stream.split("\n")))) + assert events == [_BridgeEvent(kind="exit", data=0)] + + +class TestBridgeClient: + def test_create_sandbox_payload_and_auth_header(self) -> None: + seen: Dict[str, Any] = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["path"] = request.url.path + seen["auth"] = request.headers.get("Authorization") + seen["body"] = json.loads(request.content) + return httpx.Response(200, json={"id": "sb-0123456789ab"}) + + credentials = _StaticCredentials() + client = _make_client(handler, credentials) + sandbox_id = client.create_sandbox(allow_egress=True) + + assert sandbox_id == "sb-0123456789ab" + assert seen["path"] == "/v1/sandbox" + assert seen["auth"] == "Bearer test-id-token" + assert seen["body"] == {"allow_egress": True} + assert credentials.refresh_count == 1 + + def test_create_sandbox_with_import_tar(self) -> None: + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert body["import_tar_uri"] == "gs://bucket/snap.tar" + return httpx.Response(200, json={"id": "sb-0123456789ab"}) + + client = _make_client(handler) + client.create_sandbox(import_tar_uri="gs://bucket/snap.tar") + + def test_create_sandbox_missing_id_raises(self) -> None: + client = _make_client(lambda _: httpx.Response(200, json={})) + with pytest.raises(SandboxExecError, match="missing 'id'"): + client.create_sandbox() + + def test_no_auth_header_when_unauthenticated(self) -> None: + def handler(request: httpx.Request) -> httpx.Response: + assert "Authorization" not in request.headers + return httpx.Response(200, json={"id": "sb-0123456789ab"}) + + client = _make_client(handler, credentials=None) + client.create_sandbox() + + def test_delete_tolerates_404(self) -> None: + client = _make_client(lambda _: httpx.Response(404)) + client.delete_sandbox("sb-0123456789ab") + + def test_is_running(self) -> None: + client = _make_client( + lambda _: httpx.Response(200, json={"running": True}) + ) + assert client.is_running("sb-0123456789ab") is True + + client = _make_client(lambda _: httpx.Response(404)) + assert client.is_running("sb-0123456789ab") is False + + def test_snapshot_sandbox(self) -> None: + def handler(request: httpx.Request) -> httpx.Response: + assert request.url.path == "/v1/sandbox/sb-0123456789ab/snapshot" + body = json.loads(request.content) + return httpx.Response(200, json={"uri": body["gcs_uri"]}) + + client = _make_client(handler) + uri = client.snapshot_sandbox( + "sb-0123456789ab", "gs://bucket/snap.tar" + ) + assert uri == "gs://bucket/snap.tar" + + def test_get_retries_transient_errors(self) -> None: + calls: List[int] = [] + + def handler(request: httpx.Request) -> httpx.Response: + calls.append(1) + if len(calls) < 3: + return httpx.Response(503) + return httpx.Response(200, json={"running": True}) + + client = _make_client(handler) + assert client.is_running("sb-0123456789ab") is True + assert len(calls) == 3 + + def test_post_never_retried(self) -> None: + calls: List[int] = [] + + def handler(request: httpx.Request) -> httpx.Response: + calls.append(1) + return httpx.Response(503) + + client = _make_client(handler) + with pytest.raises(SandboxExecError, match="503"): + client.create_sandbox() + assert len(calls) == 1 + + def test_put_file_rejects_oversized_payload(self) -> None: + client = _make_client(lambda _: httpx.Response(200)) + with pytest.raises(ValueError, match="32 MiB"): + client.put_file( + "sb-0123456789ab", + "big.bin", + b"\0" * (_BRIDGE_FILE_MAX_BYTES + 1), + ) + + def test_file_paths_reject_traversal(self) -> None: + client = _make_client(lambda _: httpx.Response(200)) + with pytest.raises(ValueError, match="outside"): + client.get_file("sb-0123456789ab", "../etc/passwd") + + def test_exec_stream_env_and_body(self) -> None: + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert body["argv"] == ["echo", "hi"] + assert body["env"] == {"FOO": "bar"} + assert body["cwd"] == "/tmp" + assert body["timeout_ms"] == 5000 + return httpx.Response( + 200, + content=_sse_event("exit", json.dumps({"exit_code": 0})), + headers={"Content-Type": "text/event-stream"}, + ) + + client = _make_client(handler) + _, events = client.exec_stream( + "sb-0123456789ab", + ["echo", "hi"], + cwd="/tmp", + env={"FOO": "bar"}, + timeout_ms=5000, + ) + assert list(events) == [_BridgeEvent(kind="exit", data=0)] + + +class TestProcess: + def _make_process( + self, events: List[_BridgeEvent] + ) -> CloudRunSandboxProcess: + session = MagicMock() + session._wrap_stream.side_effect = lambda lines, **_: lines + return CloudRunSandboxProcess( + iter(events), session=session, started_at=0.0 + ) + + def test_demux_and_wait(self) -> None: + process = self._make_process( + [ + _BridgeEvent(kind="stdout", data="a\nb\n"), + _BridgeEvent(kind="stderr", data="warn\n"), + _BridgeEvent(kind="exit", data=7), + ] + ) + assert list(process.stdout()) == ["a\n", "b\n"] + assert list(process.stderr()) == ["warn\n"] + assert process.wait(timeout=5) == 7 + assert process.exit_code == 7 + + def test_partial_line_flushed_on_stream_end(self) -> None: + process = self._make_process( + [ + _BridgeEvent(kind="stdout", data="no newline"), + _BridgeEvent(kind="exit", data=0), + ] + ) + assert list(process.stdout()) == ["no newline"] + + def test_wait_surfaces_missing_exit_frame(self) -> None: + process = self._make_process([_BridgeEvent(kind="stdout", data="x\n")]) + with pytest.raises(SandboxExecError, match="without an exit frame"): + process.wait(timeout=5) + + def test_kill_unblocks_wait(self) -> None: + process = self._make_process([_BridgeEvent(kind="stdout", data="x\n")]) + process.kill() + with pytest.raises(SandboxExecError, match="killed"): + process.wait(timeout=5) + + +class TestSession: + def test_exec_merges_session_env(self) -> None: + seen: Dict[str, Any] = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["body"] = json.loads(request.content) + return httpx.Response( + 200, + content=_sse_event("exit", json.dumps({"exit_code": 0})), + headers={"Content-Type": "text/event-stream"}, + ) + + session = _make_session( + _make_client(handler), + session_env={"BASE": "1", "OVERRIDE": "session"}, + ) + _passthrough_wrap_stream(session) + process = session.exec( + "echo hi", env={"OVERRIDE": "exec", "EXTRA": "2"} + ) + assert process.wait(timeout=5) == 0 + assert seen["body"]["env"] == { + "BASE": "1", + "OVERRIDE": "exec", + "EXTRA": "2", + } + assert seen["body"]["argv"] == ["echo", "hi"] + + def test_upload_and_download_roundtrip(self) -> None: + stored: Dict[str, bytes] = {} + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "PUT": + stored[request.url.path] = request.content + return httpx.Response(200, json={"written": "x"}) + return httpx.Response(200, content=stored[request.url.path]) + + session = _make_session(_make_client(handler)) + with tempfile.TemporaryDirectory() as tmp: + src = os.path.join(tmp, "in.txt") + with open(src, "w") as f: + f.write("payload") + session.upload_file(src, "/data/in.txt") + + dst = os.path.join(tmp, "out.txt") + session.download_file("/data/in.txt", dst) + with open(dst) as f: + assert f.read() == "payload" + + def test_snapshot_requires_prefix(self) -> None: + session = _make_session(_make_client(lambda _: httpx.Response(200))) + with pytest.raises(NotImplementedError, match="snapshot_uri_prefix"): + session.create_snapshot() + + def test_snapshot_returns_gcs_ref(self) -> None: + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert body["gcs_uri"].startswith( + "gs://bucket/prefix/sb-0123456789ab-" + ) + return httpx.Response(200, json={"uri": body["gcs_uri"]}) + + session = _make_session( + _make_client(handler), + snapshot_uri_prefix="gs://bucket/prefix", + ) + snapshot = session.create_snapshot() + assert isinstance(snapshot, SandboxSnapshot) + assert snapshot.ref.startswith("gs://bucket/prefix/") + assert snapshot.sandbox_id == session._parent.id + + def test_destroy_deletes_sandbox(self) -> None: + deleted: List[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + deleted.append(request.url.path) + return httpx.Response(200, json={}) + + session = _make_session(_make_client(handler)) + session.destroy() + assert deleted == ["/v1/sandbox/sb-0123456789ab"] + assert session.closed + + def test_closed_session_rejects_exec(self) -> None: + session = _make_session(_make_client(lambda _: httpx.Response(200))) + session.close() + with pytest.raises(Exception, match="closed"): + session.exec("echo hi") From b4386eb8f93c028867e8df72e1edc4f8a555511c Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Thu, 9 Jul 2026 11:33:02 -0700 Subject: [PATCH 2/8] Fail loudly on connector credentials that cannot mint ID tokens A linked GCP service connector using the user-account, oauth2-token, or external-account auth methods returns credentials with no ID-token counterpart. Previously these fell through to the ambient-ADC fallback, which would silently authenticate to the bridge as a different identity than the connector the user configured. Raise an actionable error instead, and keep the ADC fallback only for the no-connector path. Co-Authored-By: Claude Fable 5 --- .../gcp/sandboxes/cloudrun_sandbox.py | 19 ++++++ .../gcp/sandboxes/test_cloudrun_sandbox.py | 65 ++++++++++++++++++- 2 files changed, 83 insertions(+), 1 deletion(-) diff --git a/src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py b/src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py index 73d56142402..7a5026b47a2 100644 --- a/src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py +++ b/src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py @@ -1067,6 +1067,11 @@ def _build_id_token_credentials(self) -> Optional[Any]: Returns: Credentials exposing ``valid``/``token``/``refresh``, or ``None`` when ``allow_unauthenticated`` is set. + + Raises: + RuntimeError: If a linked service connector uses an auth method + whose credentials cannot mint ID tokens (user account, + OAuth2 token, external account). """ if self.config.allow_unauthenticated: return None @@ -1090,6 +1095,20 @@ def _build_id_token_credentials(self) -> Optional[Any]: token_uri=_GOOGLE_OAUTH2_TOKEN_ENDPOINT, target_audience=audience, ) + if self.get_connector() is not None: + # The connector produced credentials we can't convert (user + # account, OAuth2 token, external account). Falling back to the + # ambient ADC below would silently authenticate to the bridge + # as a *different identity* than the connector the user + # configured — fail loudly instead. + raise RuntimeError( + "The linked GCP service connector uses an auth method " + f"({type(credentials).__name__}) that cannot mint the " + "Google ID tokens Cloud Run IAM requires. Reconfigure the " + "connector with the 'service-account' or 'impersonation' " + "auth method, or unlink it to use service_account_path / " + "service-account-based Application Default Credentials." + ) # Plain ADC (workstation or metadata server): fall back to the # generic fetch_id_token helper, which raises a clear error for # credential types that cannot mint ID tokens (e.g. user creds). diff --git a/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py b/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py index 3a64bfe97dc..72494b369af 100644 --- a/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py +++ b/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py @@ -22,8 +22,9 @@ import json import os import tempfile +from datetime import datetime from typing import Any, Dict, List, Optional -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch from uuid import uuid4 import httpx @@ -37,6 +38,7 @@ ) from zenml.integrations.gcp.sandboxes.cloudrun_sandbox import ( _BRIDGE_FILE_MAX_BYTES, + CloudRunSandbox, CloudRunSandboxProcess, CloudRunSandboxSession, _BridgeEvent, @@ -461,3 +463,64 @@ def test_closed_session_rejects_exec(self) -> None: session.close() with pytest.raises(Exception, match="closed"): session.exec("echo hi") + + +def _make_component(**config_kwargs: Any) -> CloudRunSandbox: + """Builds a CloudRunSandbox without going through Stack/Client.""" + return CloudRunSandbox( + name="test-cloudrun", + id=uuid4(), + config=CloudRunSandboxConfig(service_url=SERVICE_URL, **config_kwargs), + flavor="cloudrun", + type=StackComponentType.SANDBOX, + user=None, + created=datetime.now(), + updated=datetime.now(), + environment={}, + secrets=[], + ) + + +class TestIdTokenCredentials: + def test_unauthenticated_config_skips_credentials(self) -> None: + component = _make_component(allow_unauthenticated=True) + assert component._build_id_token_credentials() is None + + def test_service_account_credentials_converted(self) -> None: + from google.oauth2 import service_account + + credentials = MagicMock(spec=service_account.Credentials) + credentials.signer = MagicMock() + credentials.service_account_email = "sa@project.iam" + + component = _make_component() + with ( + patch.object( + component, + "_get_authentication", + return_value=(credentials, "project"), + ), + patch.object( + service_account, "IDTokenCredentials" + ) as id_token_cls, + ): + component._build_id_token_credentials() + + assert id_token_cls.call_args.kwargs["target_audience"] == SERVICE_URL + + def test_unconvertible_connector_credentials_fail_loudly(self) -> None: + """Connector creds that can't mint ID tokens must not silently + fall back to the ambient ADC identity.""" + from google.oauth2.credentials import Credentials as UserCredentials + + component = _make_component() + with ( + patch.object( + component, + "_get_authentication", + return_value=(MagicMock(spec=UserCredentials), "project"), + ), + patch.object(component, "get_connector", return_value=MagicMock()), + ): + with pytest.raises(RuntimeError, match="service-account"): + component._build_id_token_credentials() From 9f42d871fc4ddb1572fad6cd06dc6fa3926e85f1 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Thu, 9 Jul 2026 11:59:45 -0700 Subject: [PATCH 3/8] Apply simplify-review cleanups to the Cloud Run sandbox Reuse: retry backoff now uses exponential_backoff_delays (adds jitter, avoiding retry lockstep across concurrent steps), utc_now, and the existing GCP_PATH_PREFIX constant. Efficiency: token refresh no longer holds its lock on the fast path or rebuilds the google transport per refresh; finished processes are pruned from long-lived sessions; newline-free exec output is flushed past 1 MiB instead of growing the remainder unboundedly; plain bridge calls get a real read timeout (only SSE streams disable it); the bridge example reuses one GCS client, streams uploads to disk, and size-checks downloads inside the sandbox before copying. Simplification: single backoff site in the retry loop, snapshot no longer round-trips an echoed URI, exit frames dispatch on event kind, dead TYPE_CHECKING/dashboard-URL/regex code removed, duplicate 32 MiB check collapsed into the upload path, and a Protocol now types the ID-token credential contract. Co-Authored-By: Claude Fable 5 --- examples/cloudrun_sandbox_bridge/main.py | 89 ++++--- .../gcp/flavors/cloudrun_sandbox_flavor.py | 7 +- .../gcp/sandboxes/cloudrun_sandbox.py | 240 ++++++++++-------- .../gcp/sandboxes/test_cloudrun_sandbox.py | 34 +-- 4 files changed, 214 insertions(+), 156 deletions(-) diff --git a/examples/cloudrun_sandbox_bridge/main.py b/examples/cloudrun_sandbox_bridge/main.py index a38d729c925..4eaf1d7353e 100644 --- a/examples/cloudrun_sandbox_bridge/main.py +++ b/examples/cloudrun_sandbox_bridge/main.py @@ -33,7 +33,7 @@ import urllib.parse import uuid from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Set, Tuple SANDBOX_BIN = os.environ.get("SANDBOX_BIN", "/usr/local/gcp/bin/sandbox") SHARE_ROOT = os.environ.get("SHARE_ROOT", "/tmp/zenml-share") @@ -43,14 +43,35 @@ MAX_FILE_BYTES = 32 * 1024 * 1024 DEFAULT_EXEC_TIMEOUT_MS = 120_000 -_SANDBOX_ID_RE = re.compile(r"^sb-[0-9a-f]{12}$") - # Sandboxes created by this instance. Persistent sandboxes are # instance-local, so an id missing here is gone (or belongs to a recycled # instance) and reads as 404. -_sandboxes: Dict[str, bool] = {} +_sandboxes: Set[str] = set() _sandboxes_lock = threading.Lock() +_gcs_client: Optional[Any] = None +_gcs_client_lock = threading.Lock() + + +def _get_gcs_client() -> Any: + """Return a lazily-built, shared Cloud Storage client.""" + global _gcs_client + with _gcs_client_lock: + if _gcs_client is None: + from google.cloud import storage + + _gcs_client = storage.Client() + return _gcs_client + + +def _parse_gcs_uri(gcs_uri: str) -> Tuple[str, str]: + """Split a gs:// URI into bucket and blob names.""" + match = re.match(r"^gs://([^/]+)/(.+)$", gcs_uri) + if not match: + raise ValueError(f"Invalid gs:// URI: {gcs_uri}") + bucket_name, blob_name = match.groups() + return bucket_name, blob_name + def _run_cli( args: List[str], timeout: Optional[float] = None @@ -92,7 +113,7 @@ def _create_sandbox(allow_egress: bool, import_tar: Optional[str]) -> str: f"{result.stderr.decode(errors='replace')[:500]}" ) with _sandboxes_lock: - _sandboxes[sandbox_id] = True + _sandboxes.add(sandbox_id) return sandbox_id @@ -104,7 +125,7 @@ def _delete_sandbox(sandbox_id: str) -> None: f"{result.stderr.decode(errors='replace')[:500]}" ) with _sandboxes_lock: - _sandboxes.pop(sandbox_id, None) + _sandboxes.discard(sandbox_id) shutil.rmtree(_share_dir(sandbox_id), ignore_errors=True) @@ -148,14 +169,9 @@ def _sandbox_shell(sandbox_id: str, script: str) -> None: ) -def _snapshot_to_gcs(sandbox_id: str, gcs_uri: str) -> str: +def _snapshot_to_gcs(sandbox_id: str, gcs_uri: str) -> None: """Export the sandbox overlay with `sandbox tar` and upload it to GCS.""" - from google.cloud import storage - - match = re.match(r"^gs://([^/]+)/(.+)$", gcs_uri) - if not match: - raise ValueError(f"Invalid gs:// URI: {gcs_uri}") - bucket_name, blob_name = match.groups() + bucket_name, blob_name = _parse_gcs_uri(gcs_uri) with tempfile.NamedTemporaryFile(suffix=".tar", delete=False) as tmp: tar_path = tmp.name @@ -166,27 +182,21 @@ def _snapshot_to_gcs(sandbox_id: str, gcs_uri: str) -> str: f"sandbox tar failed ({result.returncode}): " f"{result.stderr.decode(errors='replace')[:500]}" ) - client = storage.Client() + client = _get_gcs_client() client.bucket(bucket_name).blob(blob_name).upload_from_filename( tar_path ) finally: os.unlink(tar_path) - return gcs_uri def _download_import_tar(gcs_uri: str) -> str: """Fetch a snapshot tarball from GCS to a local temp path.""" - from google.cloud import storage - - match = re.match(r"^gs://([^/]+)/(.+)$", gcs_uri) - if not match: - raise ValueError(f"Invalid gs:// URI: {gcs_uri}") - bucket_name, blob_name = match.groups() + bucket_name, blob_name = _parse_gcs_uri(gcs_uri) with tempfile.NamedTemporaryFile(suffix=".tar", delete=False) as tmp: tar_path = tmp.name - client = storage.Client() + client = _get_gcs_client() client.bucket(bucket_name).blob(blob_name).download_to_filename(tar_path) return tar_path @@ -238,8 +248,8 @@ def _route(self) -> Optional[Tuple[str, str, str]]: return sandbox_id, action, rest def _known_sandbox(self, sandbox_id: str) -> bool: - if not _SANDBOX_ID_RE.match(sandbox_id): - return False + # Membership fully decides: the registry only ever contains ids + # this process minted itself. with _sandboxes_lock: return sandbox_id in _sandboxes @@ -409,11 +419,21 @@ def _pump(stream: Any, kind: str) -> None: def _handle_put_file(self, sandbox_id: str, rest: str) -> None: dest = _safe_sandbox_path(rest) - data = self._read_body() + length = int(self.headers.get("Content-Length") or 0) + if length > MAX_FILE_BYTES: + raise ValueError("Request body exceeds 32 MiB limit") staging_name = f"put-{uuid.uuid4().hex}" staging_host = os.path.join(_share_dir(sandbox_id), staging_name) + # Stream straight to the staging file so concurrent uploads don't + # each hold a full 32 MiB body in memory. with open(staging_host, "wb") as f: - f.write(data) + remaining = length + while remaining > 0: + chunk = self.rfile.read(min(remaining, 1024 * 1024)) + if not chunk: + break + f.write(chunk) + remaining -= len(chunk) try: _sandbox_shell( sandbox_id, @@ -429,18 +449,19 @@ def _handle_get_file(self, sandbox_id: str, rest: str) -> None: src = _safe_sandbox_path(rest) staging_name = f"get-{uuid.uuid4().hex}" staging_host = os.path.join(_share_dir(sandbox_id), staging_name) + # Enforce the size limit inside the sandbox BEFORE copying, so a + # multi-GB file is rejected without paying the copy. + quoted_src = shlex.quote(src) _sandbox_shell( sandbox_id, - f"cp {shlex.quote(src)} " + f"size=$(wc -c < {quoted_src}) && " + f'[ "$size" -le {MAX_FILE_BYTES} ] || ' + f"{{ echo 'file exceeds 32 MiB limit' >&2; exit 92; }} && " + f"cp {quoted_src} " f"{shlex.quote(f'{SHARE_MOUNT}/{staging_name}')}", ) try: size = os.path.getsize(staging_host) - if size > MAX_FILE_BYTES: - self._send_error_json( - 413, f"File is {size} bytes; limit is 32 MiB" - ) - return self.send_response(200) self.send_header("Content-Type", "application/octet-stream") self.send_header("Content-Length", str(size)) @@ -456,8 +477,8 @@ def _handle_snapshot(self, sandbox_id: str) -> None: if not gcs_uri: self._send_error_json(400, "Missing gcs_uri") return - uri = _snapshot_to_gcs(sandbox_id, gcs_uri) - self._send_json(200, {"uri": uri}) + _snapshot_to_gcs(sandbox_id, gcs_uri) + self._send_json(200, {}) def log_message(self, format: str, *args: Any) -> None: # Route access logs to stdout for Cloud Logging. diff --git a/src/zenml/integrations/gcp/flavors/cloudrun_sandbox_flavor.py b/src/zenml/integrations/gcp/flavors/cloudrun_sandbox_flavor.py index 9f35c97a775..ec2f8947c37 100644 --- a/src/zenml/integrations/gcp/flavors/cloudrun_sandbox_flavor.py +++ b/src/zenml/integrations/gcp/flavors/cloudrun_sandbox_flavor.py @@ -23,6 +23,9 @@ GCP_CONNECTOR_TYPE, GCP_RESOURCE_TYPE, ) +from zenml.integrations.gcp.flavors.gcp_artifact_store_flavor import ( + GCP_PATH_PREFIX, +) from zenml.integrations.gcp.google_credentials_mixin import ( GoogleCredentialsConfigMixin, ) @@ -148,10 +151,10 @@ def _validate_snapshot_uri_prefix( """ if value is None: return None - if not value.startswith("gs://"): + if not value.startswith(GCP_PATH_PREFIX): raise ValueError( f"Invalid snapshot_uri_prefix '{value}': must be a Cloud " - "Storage URI starting with 'gs://'." + f"Storage URI starting with '{GCP_PATH_PREFIX}'." ) return value.rstrip("/") diff --git a/src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py b/src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py index 7a5026b47a2..e8ba5639a06 100644 --- a/src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py +++ b/src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py @@ -34,7 +34,6 @@ from collections import deque from dataclasses import dataclass from typing import ( - TYPE_CHECKING, Any, Deque, Dict, @@ -42,6 +41,7 @@ Iterator, List, Optional, + Protocol, Tuple, Type, Union, @@ -68,9 +68,7 @@ SandboxSession, SandboxSnapshot, ) - -if TYPE_CHECKING: - pass +from zenml.utils.time_utils import exponential_backoff_delays, utc_now logger = get_logger(__name__) @@ -92,6 +90,10 @@ # wait()s on a noisy command would otherwise grow these deques without # bound; past the cap the oldest lines are dropped (logged once). _STREAM_BUFFER_MAX_LINES = 100_000 +# Newline-free output (progress bars, minified blobs) never hits the line +# cap; past this many buffered bytes the partial line is flushed as a +# synthetic line so memory and per-chunk copy costs stay bounded. +_STREAM_REMAINDER_MAX_BYTES = 1_048_576 # Standard Google OAuth2 token endpoint, used when deriving ID-token # credentials from plain service-account credentials. @@ -102,6 +104,30 @@ _TOKEN_REFRESH_SLACK_SECONDS = 60 +class _IdTokenCredentials(Protocol): + """Contract the bridge client needs from ID-token credentials. + + Satisfied structurally by ``google.oauth2.service_account + .IDTokenCredentials``, ``google.auth.impersonated_credentials + .IDTokenCredentials`` and ``_AdcIdTokenCredentials``. + """ + + token: Optional[str] + + @property + def valid(self) -> bool: + """Whether the cached token can still be used.""" + ... + + def refresh(self, request: Any) -> None: + """Fetch a fresh token. + + Args: + request: A ``google.auth.transport.Request`` instance. + """ + ... + + @dataclass(frozen=True) class _BridgeEvent: """One decoded SSE event from the bridge exec stream.""" @@ -260,7 +286,7 @@ def valid(self) -> bool: """ if self.token is None or self.expiry is None: return False - remaining = self.expiry - datetime.datetime.now(datetime.timezone.utc) + remaining = self.expiry - utc_now(tz_aware=True) return remaining.total_seconds() > _TOKEN_REFRESH_SLACK_SECONDS def refresh(self, request: Any) -> None: @@ -302,7 +328,7 @@ class _CloudRunBridgeClient: def __init__( self, service_url: str, - credentials: Optional[Any], + credentials: Optional[_IdTokenCredentials], *, transport: Optional["httpx.BaseTransport"] = None, ) -> None: @@ -311,15 +337,20 @@ def __init__( Args: service_url: Base URL of the deployed bridge service. credentials: Google ID-token credentials used to authorize - requests (``None`` for unauthenticated bridges). Must - expose ``valid``, ``token`` and ``refresh(request)``. + requests (``None`` for unauthenticated bridges). transport: Optional httpx transport override, used by tests. """ self._credentials = credentials self._credentials_lock = threading.Lock() + # Reused across token refreshes: each google transport Request + # wraps its own requests.Session, so rebuilding it per refresh + # would pay a fresh TLS handshake to the token endpoint. + self._auth_transport: Optional[Any] = None + # SSE exec streams disable the read timeout per-request; plain + # calls keep the default so a hung bridge can't block forever. self._client = httpx.Client( base_url=service_url.rstrip("/"), - timeout=httpx.Timeout(30.0, read=None), + timeout=httpx.Timeout(30.0), transport=transport, ) @@ -336,15 +367,20 @@ def _auth_headers(self) -> Dict[str, str]: Returns: Headers to merge into the outgoing request. """ - if self._credentials is None: + credentials = self._credentials + if credentials is None: return {} - with self._credentials_lock: - if not self._credentials.valid: - from google.auth.transport.requests import Request - - self._credentials.refresh(Request()) - token = self._credentials.token - return {"Authorization": f"Bearer {token}"} + # Optimistic fast path: only serialize the refresh itself, so + # concurrent requests holding a still-valid token don't contend. + if not credentials.valid: + with self._credentials_lock: + if not credentials.valid: + if self._auth_transport is None: + from google.auth.transport.requests import Request + + self._auth_transport = Request() + credentials.refresh(self._auth_transport) + return {"Authorization": f"Bearer {credentials.token}"} def _request( self, @@ -378,10 +414,14 @@ def _request( are exhausted. """ retryable_method = method in _RETRYABLE_METHODS - attempt = 0 + delays = exponential_backoff_delays( + attempts=_MAX_RETRIES, initial_delay=0.25, max_delay=8.0 + ) while True: merged_headers = dict(headers or {}) merged_headers.update(self._auth_headers()) + cause: Optional[BaseException] = None + retryable = retryable_method try: if stream: req = self._client.build_request( @@ -390,6 +430,7 @@ def _request( json=json_body, content=content, headers=merged_headers, + timeout=httpx.Timeout(30.0, read=None), ) resp = self._client.send(req, stream=True) else: @@ -401,39 +442,36 @@ def _request( headers=merged_headers, ) except httpx.HTTPError as e: - if not retryable_method or attempt >= _MAX_RETRIES: - raise SandboxExecError( - f"Bridge request {method} {path} failed: {e}" - ) from e - time.sleep(0.25 * (2**attempt)) - attempt += 1 - continue - - if resp.status_code in ok_statuses or resp.status_code < 400: - return resp - - retryable = ( - retryable_method and resp.status_code in _RETRYABLE_STATUSES - ) - body_preview = "" - if not stream: - try: - body_preview = resp.text[:500] - except Exception: - body_preview = "" + cause = e + error_message = f"Bridge request {method} {path} failed: {e}" else: - try: - resp.close() - except Exception: - pass + if resp.status_code in ok_statuses or resp.status_code < 400: + return resp - if not retryable or attempt >= _MAX_RETRIES: - raise SandboxExecError( + retryable = ( + retryable_method + and resp.status_code in _RETRYABLE_STATUSES + ) + body_preview = "" + if not stream: + try: + body_preview = resp.text[:500] + except Exception: + body_preview = "" + else: + try: + resp.close() + except Exception: + pass + error_message = ( f"Bridge {method} {path} returned " f"{resp.status_code}: {body_preview}" ) - time.sleep(0.25 * (2**attempt)) - attempt += 1 + + delay = next(delays, None) if retryable else None + if delay is None: + raise SandboxExecError(error_message) from cause + time.sleep(delay) def create_sandbox( self, @@ -495,29 +533,21 @@ def is_running(self, sandbox_id: str) -> bool: return False return bool(resp.json().get("running", False)) - def snapshot_sandbox(self, sandbox_id: str, gcs_uri: str) -> str: + def snapshot_sandbox(self, sandbox_id: str, gcs_uri: str) -> None: """Export the sandbox filesystem overlay to Cloud Storage. + The bridge writes the tarball to exactly the URI it is given, so + a successful response needs no payload. + Args: sandbox_id: The sandbox to snapshot. gcs_uri: Destination ``gs://`` URI for the tarball. - - Returns: - The URI of the stored snapshot tarball, as confirmed by the - bridge. - - Raises: - SandboxExecError: If the bridge response is missing the URI. """ - resp = self._request( + self._request( "POST", f"/v1/sandbox/{sandbox_id}/snapshot", json_body={"gcs_uri": gcs_uri}, ) - uri = resp.json().get("uri") - if not uri: - raise SandboxExecError("Bridge snapshot response missing 'uri'") - return cast(str, uri) def exec_stream( self, @@ -563,19 +593,14 @@ def exec_stream( def put_file(self, sandbox_id: str, remote_path: str, data: bytes) -> None: """Upload raw bytes to a sandbox path. + Size enforcement lives in ``_upload_file`` (the only caller), + which checks before reading the payload into memory. + Args: sandbox_id: The sandbox to write into. remote_path: Path inside the sandbox filesystem; `..` rejected. data: Raw file bytes; max 32 MiB. - - Raises: - ValueError: If the file exceeds the bridge body limit. """ - if len(data) > _BRIDGE_FILE_MAX_BYTES: - raise ValueError( - f"File of {len(data)} bytes exceeds the bridge limit of " - f"{_BRIDGE_FILE_MAX_BYTES} bytes (32 MiB)." - ) safe_path = _sanitize_remote_path(remote_path) encoded = urllib.parse.quote(safe_path, safe="/") self._request( @@ -630,11 +655,10 @@ def __init__( # it once on a background thread and demux into two buffers so # `stdout()` and `stderr()` can be iterated independently from the # caller's thread. - self._lock = threading.Lock() self._stdout_buf: Deque[str] = deque() self._stderr_buf: Deque[str] = deque() self._buffer_truncated = False - self._data_available = threading.Condition(self._lock) + self._data_available = threading.Condition() self._done = threading.Event() self._exit_code: Optional[int] = None self._pump_error: Optional[BaseException] = None @@ -674,6 +698,12 @@ def _push_text(self, target: Deque[str], remainder: str, text: str) -> str: _STREAM_BUFFER_MAX_LINES, ) target.append(line + "\n") + # Newline-free output would otherwise grow the remainder (and the + # per-chunk `remainder + text` copy) without bound; flush it as a + # synthetic line past the cap. + if len(buf) > _STREAM_REMAINDER_MAX_BYTES: + target.append(buf) + return "" return buf def _flush_remainders(self) -> None: @@ -690,21 +720,19 @@ def _pump_events(self) -> None: try: for ev in self._event_iter: with self._data_available: - # Only ``exit`` frames carry an int payload (the exit - # code); stdout/stderr frames carry decoded text. - if isinstance(ev.data, int): - self._exit_code = ev.data + if ev.kind == "exit": + self._exit_code = int(ev.data) elif ev.kind == "stdout": self._stdout_remainder = self._push_text( self._stdout_buf, self._stdout_remainder, - ev.data, + cast(str, ev.data), ) elif ev.kind == "stderr": self._stderr_remainder = self._push_text( self._stderr_buf, self._stderr_remainder, - ev.data, + cast(str, ev.data), ) self._data_available.notify_all() except BaseException as e: # noqa: BLE001 @@ -784,11 +812,12 @@ def wait(self, timeout: Optional[float] = None) -> int: ) pump_err = self._pump_error if pump_err is not None: - if isinstance(pump_err, SandboxExecError): - raise SandboxExecError(str(pump_err)) from pump_err - raise SandboxExecError( - f"Bridge SSE pump failed: {pump_err}" - ) from pump_err + message = ( + str(pump_err) + if isinstance(pump_err, SandboxExecError) + else f"Bridge SSE pump failed: {pump_err}" + ) + raise SandboxExecError(message) from pump_err if self._exit_code is None: if self._killed.is_set(): raise SandboxExecError( @@ -865,8 +894,8 @@ def __init__( destroy_on_exit: Whether the context manager destroys the sandbox on exit instead of just closing the handle. """ - # Subclass state must be set before super().__init__ so the - # dashboard hook (invoked during base __init__) has what it needs. + # Subclass state must be set before super().__init__, which + # publishes session metadata during construction. self._client = client self._sandbox_id = sandbox_id self._session_env = dict(session_env or {}) @@ -878,16 +907,6 @@ def __init__( id=sandbox_id, parent=parent, destroy_on_exit=destroy_on_exit ) - def _get_dashboard_url(self) -> Optional[str]: - """Bridge sandboxes have no per-sandbox dashboard URL. - - Returns: - ``None``. Sandbox lifecycle events land in the bridge - service's Cloud Logging, but there is no stable per-sandbox - deep link. - """ - return None - def _exec( self, command: Union[str, List[str]], @@ -916,6 +935,11 @@ def _exec( ) self._log_command(argv) + # Prune exited processes: a long-lived session running thousands + # of execs would otherwise retain every finished process and its + # output buffers until close(). + self._processes = [p for p in self._processes if p.exit_code is None] + merged_env = {**self._session_env, **(env or {})} effective_cwd = cwd if cwd is not None else self._default_cwd started_at = time.time() @@ -960,10 +984,10 @@ def _create_snapshot(self) -> SandboxSnapshot: f"{self._snapshot_uri_prefix}/{self._sandbox_id}-" f"{int(time.time())}.tar" ) - stored_uri = self._client.snapshot_sandbox(self._sandbox_id, gcs_uri) + self._client.snapshot_sandbox(self._sandbox_id, gcs_uri) return SandboxSnapshot( sandbox_id=self._parent.id, - ref=stored_uri, + ref=gcs_uri, metadata={"source_sandbox": self._sandbox_id}, ) @@ -1055,7 +1079,9 @@ def settings_class(self) -> Optional[Type["BaseSettings"]]: """ return CloudRunSandboxSettings - def _build_id_token_credentials(self) -> Optional[Any]: + def _build_id_token_credentials( + self, + ) -> Optional[_IdTokenCredentials]: """Derive ID-token credentials for the bridge from the GCP auth. Cloud Run IAM authenticates callers with Google-signed ID tokens @@ -1083,17 +1109,23 @@ def _build_id_token_credentials(self) -> Optional[Any]: credentials, _ = self._get_authentication() if isinstance(credentials, impersonated_credentials.Credentials): - return impersonated_credentials.IDTokenCredentials( - credentials, - target_audience=audience, - include_email=True, + return cast( + _IdTokenCredentials, + impersonated_credentials.IDTokenCredentials( + credentials, + target_audience=audience, + include_email=True, + ), ) if isinstance(credentials, service_account.Credentials): - return service_account.IDTokenCredentials( - signer=credentials.signer, - service_account_email=credentials.service_account_email, - token_uri=_GOOGLE_OAUTH2_TOKEN_ENDPOINT, - target_audience=audience, + return cast( + _IdTokenCredentials, + service_account.IDTokenCredentials( + signer=credentials.signer, + service_account_email=credentials.service_account_email, + token_uri=_GOOGLE_OAUTH2_TOKEN_ENDPOINT, + target_audience=audience, + ), ) if self.get_connector() is not None: # The connector produced credentials we can't convert (user diff --git a/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py b/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py index 72494b369af..8edc53f47c9 100644 --- a/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py +++ b/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py @@ -256,16 +256,17 @@ def test_is_running(self) -> None: assert client.is_running("sb-0123456789ab") is False def test_snapshot_sandbox(self) -> None: + seen: Dict[str, Any] = {} + def handler(request: httpx.Request) -> httpx.Response: - assert request.url.path == "/v1/sandbox/sb-0123456789ab/snapshot" - body = json.loads(request.content) - return httpx.Response(200, json={"uri": body["gcs_uri"]}) + seen["path"] = request.url.path + seen["body"] = json.loads(request.content) + return httpx.Response(200, json={}) client = _make_client(handler) - uri = client.snapshot_sandbox( - "sb-0123456789ab", "gs://bucket/snap.tar" - ) - assert uri == "gs://bucket/snap.tar" + client.snapshot_sandbox("sb-0123456789ab", "gs://bucket/snap.tar") + assert seen["path"] == "/v1/sandbox/sb-0123456789ab/snapshot" + assert seen["body"] == {"gcs_uri": "gs://bucket/snap.tar"} def test_get_retries_transient_errors(self) -> None: calls: List[int] = [] @@ -292,14 +293,15 @@ def handler(request: httpx.Request) -> httpx.Response: client.create_sandbox() assert len(calls) == 1 - def test_put_file_rejects_oversized_payload(self) -> None: - client = _make_client(lambda _: httpx.Response(200)) - with pytest.raises(ValueError, match="32 MiB"): - client.put_file( - "sb-0123456789ab", - "big.bin", - b"\0" * (_BRIDGE_FILE_MAX_BYTES + 1), - ) + def test_upload_rejects_oversized_file(self) -> None: + session = _make_session(_make_client(lambda _: httpx.Response(200))) + with tempfile.TemporaryDirectory() as tmp: + big = os.path.join(tmp, "big.bin") + # Sparse file: sets the size without writing 32 MiB of data. + with open(big, "wb") as f: + f.truncate(_BRIDGE_FILE_MAX_BYTES + 1) + with pytest.raises(ValueError, match="upload limit"): + session.upload_file(big, "/data/big.bin") def test_file_paths_reject_traversal(self) -> None: client = _make_client(lambda _: httpx.Response(200)) @@ -434,7 +436,7 @@ def handler(request: httpx.Request) -> httpx.Response: assert body["gcs_uri"].startswith( "gs://bucket/prefix/sb-0123456789ab-" ) - return httpx.Response(200, json={"uri": body["gcs_uri"]}) + return httpx.Response(200, json={}) session = _make_session( _make_client(handler), From 8416040d33cc68cb021f605cedc98749275bc0e1 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Thu, 9 Jul 2026 12:03:06 -0700 Subject: [PATCH 4/8] Drop the working spec from the PR The durable content lives in the docs page, module docstrings, and the bridge README. Co-Authored-By: Claude Fable 5 --- CLOUDRUN_SANDBOX_SPEC.md | 270 --------------------------------------- 1 file changed, 270 deletions(-) delete mode 100644 CLOUDRUN_SANDBOX_SPEC.md diff --git a/CLOUDRUN_SANDBOX_SPEC.md b/CLOUDRUN_SANDBOX_SPEC.md deleted file mode 100644 index a6ca288ddc9..00000000000 --- a/CLOUDRUN_SANDBOX_SPEC.md +++ /dev/null @@ -1,270 +0,0 @@ -# Spec: Google Cloud Run Sandbox integration - -Add a `cloudrun` sandbox flavor to the existing GCP integration, backed by -[Cloud Run sandboxes](https://docs.cloud.google.com/run/docs/code-execution) -(public preview, July 2026). - -## 1. Background - -### 1.1 What Cloud Run sandboxes are - -Cloud Run sandboxes provide fast (<500 ms), strongly isolated environments for -executing untrusted code — LLM/agent-generated code being the headline use -case. Unlike Vercel/Cloudflare/Modal sandboxes there is **no standalone REST -API**: sandboxes are created *inside* a Cloud Run service instance that was -deployed with the sandbox launcher enabled: - -```bash -gcloud beta run deploy SERVICE --image IMAGE_URL --sandbox-launcher -``` - -Inside such an instance a CLI at `/usr/local/gcp/bin/sandbox` is the control -surface: - -| Command | Purpose | -|---|---| -| `sandbox do -- CMD` | ephemeral sandbox: create, run, destroy | -| `sandbox run [CMD] --detach` | create a persistent sandbox | -| `sandbox exec CMD [-e K=V] [-w DIR]` | run a command in a persistent sandbox | -| `sandbox delete [--force]` | destroy a sandbox | -| `sandbox tar --file F` | export the writable overlay as a tarball | -| `sandbox fork ` | clone a running sandbox | - -Key flags: `--env`, `--mount type=bind,source=…,destination=…`, `--write`, -`--allow-egress` (no egress by default), `--import-tar`/`--export-tar`, -`--rootfs`, `--workdir`. Sandboxes share the host instance's CPU/memory and -see the instance's filesystem read-only (plus a writable tmpfs overlay). - -### 1.2 What ZenML's sandbox abstraction expects - -`src/zenml/sandboxes/` (on `develop`) defines the `SANDBOX` stack component -type: - -- `BaseSandbox.create_session(settings, destroy_on_exit=False) -> SandboxSession` - (required), `attach(session_id)` and `restore(snapshot)` (optional). -- `SandboxSession` hooks: `_exec` (required), `_close` (required), - `_upload_file` / `_download_file` / `_create_snapshot` / `_destroy` / - `_get_dashboard_url` (optional). -- `SandboxProcess`: `stdout()` / `stderr()` line iterators, `wait(timeout)`, - `kill()`, `exit_code`; `collect()` is provided by the base class. -- `SandboxSnapshot(sandbox_id: UUID, ref: str, metadata: dict)`. - -Reference implementations: `local`, `modal` (only flavor with snapshots), -`kubernetes`, and `cloudflare` (branch `feature/sandbox-cloudflare`, PR #4907) -which established the **bridge pattern**: a small HTTP service deployed on the -provider mediates between the ZenML client and the provider-local sandbox -primitive. Cloud Run's launcher architecture forces the same pattern. - -## 2. Design - -### 2.1 Architecture - -``` -ZenML client (step code) Google Cloud -┌──────────────────────────┐ HTTPS + ID ┌──────────────────────────────┐ -│ CloudRunSandbox │ token (IAM) │ Cloud Run service │ -│ └ CloudRunSandboxSession├────────────────►│ "sandbox bridge" │ -│ └ CloudRunSandbox- │ SSE streams │ (--sandbox-launcher) │ -│ Process │◄────────────────┤ ┌────────┐ ┌────────┐ │ -└──────────────────────────┘ │ │sandbox1│ │sandbox2│ ... │ - │ └────────┘ └────────┘ │ - └───────────┬──────────────────┘ - │ snapshots (tar) - ▼ - GCS bucket (optional) -``` - -The user deploys the provided **bridge service** (a small synchronous Python -HTTP server that shells out to the `sandbox` CLI) to Cloud Run with -`--sandbox-launcher`. The ZenML `cloudrun` sandbox component talks to it over -HTTPS, authenticating with Google-signed **ID tokens** (standard Cloud Run -IAM invoker auth) minted from the component's Google credentials. - -The bridge speaks the **same wire protocol v1 as the Cloudflare bridge** -(`POST /v1/sandbox`, `POST /v1/sandbox/:id/exec` streaming SSE frames -`stdout`/`stderr` (base64) + `exit`/`error`, `PUT|GET /v1/sandbox/:id/file/*`, -`GET /v1/sandbox/:id/running`, `DELETE /v1/sandbox/:id`), with two -Cloud-Run-specific extensions: - -1. **Per-exec env**: exec body accepts `"env": {K: V}` (the `sandbox` CLI - supports `-e` natively), removing the Cloudflare limitation where per-exec - env raises `NotImplementedError`. Session-level `sandbox_environment` is - merged into every exec by the client — no bridge-side session objects. -2. **Snapshots**: `POST /v1/sandbox/:id/snapshot {"gcs_uri": "gs://…"}` runs - `sandbox tar` and uploads the tarball to GCS; `POST /v1/sandbox` accepts - `"import_tar_uri": "gs://…"` to boot from a snapshot. This gives Cloud Run - parity with Modal (`create_snapshot`/`restore`), which Cloudflare lacks. - -### 2.2 Deployment constraints (documented, not enforced) - -- Persistent sandboxes are **instance-local**. The bridge must be deployed - with `--max-instances 1` (or session affinity for advanced setups) so - `exec` calls for a sandbox land on the instance that owns it. -- `--no-cpu-throttling` and `--min-instances 1` keep detached sandboxes alive - between requests; otherwise Cloud Run may throttle/scale-to-zero and kill - them. -- Cloud Run request timeout caps a single SSE exec stream at 60 min; the - per-exec `timeout_ms` setting must stay below the service timeout. -- Sandbox CPU/memory come out of the bridge instance's allocation — size the - service accordingly. - -### 2.3 Placement: extend the `gcp` integration - -The GCP integration already ships six flavors and depends on -`google-cloud-run` and `google-cloud-storage`; the flavor rides on the -existing `GoogleCredentialsConfigMixin`/`GoogleCredentialsMixin` (service -connector → SA key file → ADC) and the `gcp` service connector -(`resource_type="gcp-generic"`). Flavor name: **`cloudrun`** (precedent: -Vertex components use the service name `vertex`, not `gcp`). - -### 2.4 New classes - -`src/zenml/integrations/gcp/flavors/cloudrun_sandbox_flavor.py` (no google -imports — flavor files must import without the integration installed): - -```python -class CloudRunSandboxSettings(BaseSandboxSettings): - timeout_ms: int = 120_000 # per-exec cap, passed to the bridge - cwd: Optional[str] = None # default workdir inside the sandbox - allow_egress: bool = False # sandbox outbound network (CR default: off) - -class CloudRunSandboxConfig( - BaseSandboxConfig, GoogleCredentialsConfigMixin, CloudRunSandboxSettings -): - service_url: str # https URL of the bridge service - audience: Optional[str] = None # ID-token audience; default service_url - allow_unauthenticated: bool = False # skip ID tokens (public bridge; discouraged) - snapshot_uri_prefix: Optional[str] = None # gs://bucket/prefix enabling snapshots - # + project / service_account_path from GoogleCredentialsConfigMixin - is_remote -> True - -class CloudRunSandboxFlavor(BaseSandboxFlavor): - name = "cloudrun" # GCP_CLOUDRUN_SANDBOX_FLAVOR - service_connector_requirements = ServiceConnectorRequirements( - connector_type=GCP_CONNECTOR_TYPE, resource_type=GCP_RESOURCE_TYPE - ) -``` - -`service_url` gets the same https-or-localhost validator as the Cloudflare -`worker_url`. - -`src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py` (google imports -allowed here): - -- `_CloudRunBridgeClient` — httpx wrapper (httpx is a core dependency), - adapted from `_CloudflareBridgeClient`: same retry policy (idempotent - methods only; 429/502/503/504), same SSE parsing, plus `env` on exec, - `import_tar_uri` on create, and a `snapshot(sandbox_id, gcs_uri)` call. - Auth is a pluggable header callback instead of a static bearer token. -- `CloudRunSandboxProcess` — SSE demux into stdout/stderr line buffers on a - pump thread (same design as `CloudflareSandboxProcess`). -- `CloudRunSandboxSession` — implements `_exec` (with per-exec env support), - `_upload_file`, `_download_file`, `_close`, `_destroy`, `_create_snapshot` - (only when `snapshot_uri_prefix` is configured; `ref` = the GCS tarball URI). -- `CloudRunSandbox(BaseSandbox, GoogleCredentialsMixin)` — implements - `create_session` (threading `destroy_on_exit`, per develop's signature), - `attach`, and `restore` (creates a sandbox with `import_tar_uri=snapshot.ref` - after `_validate_snapshot`). - -**ID-token minting** (`_get_id_token_credentials`): start from -`self._get_authentication()` (mixin), then: - -1. `service_account.Credentials` → `service_account.IDTokenCredentials` - (same signer, `target_audience=audience`). -2. `impersonated_credentials.Credentials` → - `impersonated_credentials.IDTokenCredentials(..., include_email=True)`. -3. Otherwise (plain ADC) → `google.oauth2.id_token.fetch_id_token`. -4. Failure → actionable error naming the three supported paths. - -Tokens are cached and refreshed via `google.auth.transport.requests.Request` -before expiry; the bearer header is injected per request. - -### 2.5 Bridge service (deployable example) - -Following the Cloudflare precedent the bridge is not part of the `zenml` -package. It ships as `examples/cloudrun_sandbox_bridge/`: - -- `main.py` — synchronous stdlib/`ThreadingHTTPServer` implementation of the - wire protocol, shelling out to `/usr/local/gcp/bin/sandbox`: - - create → `sandbox run --detach [--allow-egress] [--import-tar …] - --mount type=bind,source=/tmp/zenml-share/,destination=/mnt/zenml - --write -- sleep infinity` (the bind mount backs file transfer) - - exec → `sandbox exec -e K=V … -w CWD -- argv…`, stdout/stderr piped - and re-emitted as base64 SSE frames, exit code in the `exit` frame - - file PUT/GET → write/read `/tmp/zenml-share//…` on the host plus a - `cp` exec inside the sandbox - - snapshot → `sandbox tar --file /tmp/….tar` then GCS upload - (`google-cloud-storage`) -- `Dockerfile` + `README.md` with the exact deploy command: - - ```bash - gcloud beta run deploy zenml-sandbox-bridge --source . \ - --sandbox-launcher --no-allow-unauthenticated \ - --max-instances 1 --min-instances 1 --no-cpu-throttling - ``` - -The client treats the bridge URL as opaque, so users can substitute their own -implementation of the protocol. - -### 2.6 Registration & metadata - -- `src/zenml/integrations/gcp/__init__.py`: add - `GCP_CLOUDRUN_SANDBOX_FLAVOR = "cloudrun"`; register - `CloudRunSandboxFlavor` in `flavors()`. -- `src/zenml/integrations/gcp/flavors/__init__.py`: export config + flavor. -- Logo: `https://public-flavor-logos.s3.eu-central-1.amazonaws.com/sandbox/cloudrun.png` - (asset upload tracked outside this repo, same as other new flavors). - -## 3. Testing - -`tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py`, mirroring -the Cloudflare suite (`httpx.MockTransport`, no real bridge): - -- flavor metadata / config validation (https enforcement, snapshot prefix - validation) -- bridge client: create/delete/running/exec/file/snapshot request shapes, - retry policy, error surfacing -- SSE parsing: interleaved stdout/stderr, exit frames, error frames, - truncated streams -- process demux: concurrent stdout/stderr draining, `wait` timeout, `kill` -- session lifecycle: per-exec env merging, `destroy_on_exit`, closed-handle - errors, upload/download size guards -- sandbox: `create_session`/`attach`/`restore` wiring, snapshot gating on - `snapshot_uri_prefix`, ID-token path selection (mocked google.auth) - -Live end-to-end testing against a real deployed bridge is manual, per the -repo's policy for external-service integrations. - -## 4. Documentation - -- `docs/book/component-guide/sandboxes/cloudrun.md` — deploy-the-bridge - walkthrough, IAM setup (`roles/run.invoker` for the caller; - GCS access for the bridge SA when snapshots are enabled), registration: - - ```bash - zenml sandbox register my_cloudrun_sandbox --flavor cloudrun \ - --service_url=https://zenml-sandbox-bridge-….run.app - zenml stack update my_stack --sandbox my_cloudrun_sandbox - ``` - -- Add the page to `docs/book/component-guide/toc.md` under Sandboxes. - -## 5. Out of scope (future work) - -- `sandbox fork` → no abstraction hook today (a `fork()` on `SandboxSession` - would benefit Modal too). -- Auto-deploying the bridge from the client via `google-cloud-run` (the - Deployer already owns "deploy a Cloud Run service"; a convenience CLI could - reuse it). -- Multi-instance bridges with session affinity routing. -- `aexec` async execution. - -## 6. Risks - -- **Preview API**: the `sandbox` CLI is pre-GA; flags may change. The CLI - surface is isolated in the bridge (one file), not in the ZenML package. -- **Instance lifetime**: Cloud Run may recycle the bridge instance, killing - persistent sandboxes. `attach()` reports this cleanly via - `GET /running`; docs recommend snapshot checkpointing for long-lived state. -- **Protocol drift** with the Cloudflare bridge: mitigated by versioned - paths (`/v1/…`) and per-flavor unit suites. From ef1d2e7f7eb7ccf8f827bd05b2baef044c112526 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 10 Jul 2026 09:01:51 -0700 Subject: [PATCH 5/8] Fix bandit findings in bridge example and tests The /tmp share-root default and 0.0.0.0 bind in the bridge are both deliberate on Cloud Run (per-instance in-memory /tmp; containers must listen on all interfaces behind IAM), so mark them nosec with the reasoning. Test fixtures use /workspace instead of /tmp. Co-Authored-By: Claude Fable 5 --- examples/cloudrun_sandbox_bridge/main.py | 8 ++++++-- .../integrations/gcp/sandboxes/test_cloudrun_sandbox.py | 4 ++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/examples/cloudrun_sandbox_bridge/main.py b/examples/cloudrun_sandbox_bridge/main.py index 4eaf1d7353e..7fedd883b7b 100644 --- a/examples/cloudrun_sandbox_bridge/main.py +++ b/examples/cloudrun_sandbox_bridge/main.py @@ -36,7 +36,9 @@ from typing import Any, Dict, List, Optional, Set, Tuple SANDBOX_BIN = os.environ.get("SANDBOX_BIN", "/usr/local/gcp/bin/sandbox") -SHARE_ROOT = os.environ.get("SHARE_ROOT", "/tmp/zenml-share") +# /tmp on Cloud Run gen2 is a per-instance in-memory filesystem, not a +# shared multi-user tmp, so a fixed path is safe here. +SHARE_ROOT = os.environ.get("SHARE_ROOT", "/tmp/zenml-share") # nosec B108 # Mount point of the per-sandbox shared directory inside each sandbox, # used to move file payloads across the isolation boundary. SHARE_MOUNT = "/mnt/zenml" @@ -488,7 +490,9 @@ def log_message(self, format: str, *args: Any) -> None: def main() -> None: port = int(os.environ.get("PORT", "8080")) os.makedirs(SHARE_ROOT, exist_ok=True) - server = ThreadingHTTPServer(("0.0.0.0", port), BridgeHandler) + # Cloud Run containers must listen on all interfaces; IAM/ingress + # controls sit in front of the service. + server = ThreadingHTTPServer(("0.0.0.0", port), BridgeHandler) # nosec B104 print(f"ZenML sandbox bridge listening on :{port}") server.serve_forever() diff --git a/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py b/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py index 8edc53f47c9..45fd0c7653e 100644 --- a/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py +++ b/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py @@ -313,7 +313,7 @@ def handler(request: httpx.Request) -> httpx.Response: body = json.loads(request.content) assert body["argv"] == ["echo", "hi"] assert body["env"] == {"FOO": "bar"} - assert body["cwd"] == "/tmp" + assert body["cwd"] == "/workspace" assert body["timeout_ms"] == 5000 return httpx.Response( 200, @@ -325,7 +325,7 @@ def handler(request: httpx.Request) -> httpx.Response: _, events = client.exec_stream( "sb-0123456789ab", ["echo", "hi"], - cwd="/tmp", + cwd="/workspace", env={"FOO": "bar"}, timeout_ms=5000, ) From 5f64fcc76dfb8a9b3efa63e422e34fa10c94415d Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 10 Jul 2026 16:25:46 -0700 Subject: [PATCH 6/8] Address security review findings on the Cloud Run sandbox Critical/high: - Bridge file staging now opens host-side descriptors with O_NOFOLLOW (and O_EXCL on upload) and validates a regular file via fstat, so untrusted sandbox code can't redirect a transfer to a host path by planting a symlink at the staged name. - process.kill() now terminates the remote command: exec carries a client-generated exec id the bridge registers, and kill() calls a new DELETE /v1/sandbox//exec/ endpoint instead of only closing the local SSE stream. - SSE buffers are bounded by total bytes, not line count, and the newline-free flush path also respects the cap, so a flood of long or newline-free output can't OOM the caller. - Sandbox ids are client-generated and bridge creation is idempotent, so a create whose response is lost is cleaned up (best-effort delete) instead of orphaning a sleep-infinity process. - Fail loud when an explicit service_account_path (not only a connector) yields credentials that can't mint ID tokens, instead of silently falling back to ambient ADC. Medium: - Rebuild the bridge client (and its ID-token credentials) once the linked service connector has expired. - Snapshot object names include a uuid so two snapshots within one second don't overwrite each other. - Bridge uses the documented --env/--workdir CLI flags; README flags the preview CLI surface for pre-release verification. Co-Authored-By: Claude Fable 5 --- examples/cloudrun_sandbox_bridge/README.md | 17 +- examples/cloudrun_sandbox_bridge/main.py | 126 ++++++-- .../gcp/sandboxes/cloudrun_sandbox.py | 271 ++++++++++++++---- .../gcp/sandboxes/test_cloudrun_sandbox.py | 131 ++++++++- 4 files changed, 450 insertions(+), 95 deletions(-) diff --git a/examples/cloudrun_sandbox_bridge/README.md b/examples/cloudrun_sandbox_bridge/README.md index e4693e8109f..fd4504624c2 100644 --- a/examples/cloudrun_sandbox_bridge/README.md +++ b/examples/cloudrun_sandbox_bridge/README.md @@ -12,13 +12,26 @@ The bridge translates the ZenML sandbox bridge protocol (v1) into | Endpoint | CLI | |---|---| -| `POST /v1/sandbox` | `sandbox run --detach … -- sleep infinity` | -| `POST /v1/sandbox//exec` | `sandbox exec [-e K=V] [-w DIR] -- argv…` (streamed as SSE) | +| `POST /v1/sandbox` (caller supplies the id) | `sandbox run --detach … -- sleep infinity` | +| `POST /v1/sandbox//exec` | `sandbox exec [--env K=V] [--workdir DIR] -- argv…` (streamed as SSE) | +| `DELETE /v1/sandbox//exec/` | terminates the running command | | `PUT/GET /v1/sandbox//file/` | staged through a bind-mounted share directory | | `POST /v1/sandbox//snapshot` | `sandbox tar ` + upload to Cloud Storage | | `GET /v1/sandbox//running` | `sandbox exec /bin/true` | | `DELETE /v1/sandbox/` | `sandbox delete --force` | +The client generates sandbox ids, so a create whose response is lost can +still be cleaned up. Exec commands carry a client-generated exec id so +`process.kill()` can terminate them. File staging opens host-side +descriptors with `O_NOFOLLOW`, so untrusted sandbox code can't redirect a +transfer to a host path via a symlink. + +> **Verify before release:** the exec CLI flags (`--env`, `--workdir`) and +> the mount/`--import-tar` syntax are taken from the public +> [code-execution docs](https://docs.cloud.google.com/run/docs/code-execution); +> confirm them with `sandbox exec -h` on a deployed instance, since the CLI +> is still in preview. + Authentication is delegated to Cloud Run IAM — the bridge itself contains no auth code. Deploy with `--no-allow-unauthenticated` and Cloud Run verifies the caller's Google ID token before the request reaches the container. diff --git a/examples/cloudrun_sandbox_bridge/main.py b/examples/cloudrun_sandbox_bridge/main.py index 7fedd883b7b..5635b22dc50 100644 --- a/examples/cloudrun_sandbox_bridge/main.py +++ b/examples/cloudrun_sandbox_bridge/main.py @@ -27,13 +27,14 @@ import re import shlex import shutil +import stat import subprocess import tempfile import threading import urllib.parse import uuid from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from typing import Any, Dict, List, Optional, Set, Tuple +from typing import Any, BinaryIO, Dict, List, Optional, Set, Tuple SANDBOX_BIN = os.environ.get("SANDBOX_BIN", "/usr/local/gcp/bin/sandbox") # /tmp on Cloud Run gen2 is a per-instance in-memory filesystem, not a @@ -45,16 +46,56 @@ MAX_FILE_BYTES = 32 * 1024 * 1024 DEFAULT_EXEC_TIMEOUT_MS = 120_000 +# Sandbox ids are caller-generated; validate before they reach the shell, +# a directory name, or a `sandbox run ` argument. +_SANDBOX_ID_RE = re.compile(r"^sb-[0-9a-f]{12}$") + # Sandboxes created by this instance. Persistent sandboxes are # instance-local, so an id missing here is gone (or belongs to a recycled # instance) and reads as 404. _sandboxes: Set[str] = set() _sandboxes_lock = threading.Lock() +# Running exec commands keyed by caller-supplied exec id, so a kill +# request can terminate the underlying process. +_execs: Dict[str, "subprocess.Popen[bytes]"] = {} +_execs_lock = threading.Lock() + _gcs_client: Optional[Any] = None _gcs_client_lock = threading.Lock() +def _open_new_file(path: str) -> BinaryIO: + """Create a staging file for writing, refusing symlinks and clobbering. + + The staging directory is a bind mount the untrusted sandbox can write + to, so O_NOFOLLOW|O_EXCL stops it from pre-planting a symlink that + would redirect the write to a host path. + """ + fd = os.open( + path, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600 + ) + return os.fdopen(fd, "wb") + + +def _open_regular_file(path: str) -> Tuple[BinaryIO, int]: + """Open a staging file for reading, refusing symlinks; verify regular. + + Opening with O_NOFOLLOW and validating via fstat on the resulting + descriptor closes the symlink-swap race: the sandbox cannot make the + bridge read a host file by replacing the staged path with a symlink. + """ + fd = os.open(path, os.O_RDONLY | os.O_NOFOLLOW) + try: + st = os.fstat(fd) + if not stat.S_ISREG(st.st_mode): + raise ValueError("staged path is not a regular file") + return os.fdopen(fd, "rb"), st.st_size + except BaseException: + os.close(fd) + raise + + def _get_gcs_client() -> Any: """Return a lazily-built, shared Cloud Storage client.""" global _gcs_client @@ -88,9 +129,17 @@ def _share_dir(sandbox_id: str) -> str: return os.path.join(SHARE_ROOT, sandbox_id) -def _create_sandbox(allow_egress: bool, import_tar: Optional[str]) -> str: - """Create a detached persistent sandbox with a bind-mounted share dir.""" - sandbox_id = f"sb-{uuid.uuid4().hex[:12]}" +def _create_sandbox( + sandbox_id: str, allow_egress: bool, import_tar: Optional[str] +) -> str: + """Create a detached persistent sandbox with a bind-mounted share dir. + + Idempotent: a caller retrying after a lost response gets the existing + sandbox back instead of a duplicate. + """ + with _sandboxes_lock: + if sandbox_id in _sandboxes: + return sandbox_id share = _share_dir(sandbox_id) os.makedirs(share, exist_ok=True) @@ -148,9 +197,9 @@ def _exec_in_sandbox( """Launch a command in the sandbox, pipes attached.""" args = ["exec", sandbox_id] for key, value in env.items(): - args += ["-e", f"{key}={value}"] + args += ["--env", f"{key}={value}"] if cwd: - args += ["-w", cwd] + args += ["--workdir", cwd] args += ["--", *argv] return subprocess.Popen( [SANDBOX_BIN, *args], @@ -159,6 +208,23 @@ def _exec_in_sandbox( ) +def _kill_exec(exec_id: str) -> bool: + """Terminate a registered exec by id. + + Args: + exec_id: The caller-supplied exec id. + + Returns: + True if a matching running command was found and killed. + """ + with _execs_lock: + proc = _execs.get(exec_id) + if proc is None: + return False + proc.kill() + return True + + def _sandbox_shell(sandbox_id: str, script: str) -> None: """Run a short shell script inside the sandbox, raising on failure.""" result = _run_cli( @@ -316,9 +382,16 @@ def do_DELETE(self) -> None: # noqa: N802 if route is None: self._send_error_json(404, "Not found") return - sandbox_id, action, _ = route + sandbox_id, action, rest = route try: - if not sandbox_id or action: + if action == "exec" and rest: + if not self._known_sandbox(sandbox_id): + self._send_error_json(404, "Unknown sandbox") + elif _kill_exec(rest): + self._send_json(200, {"killed": rest}) + else: + self._send_error_json(404, "Unknown exec") + elif not sandbox_id or action: self._send_error_json(404, "Not found") elif not self._known_sandbox(sandbox_id): self._send_error_json(404, "Unknown sandbox") @@ -330,12 +403,19 @@ def do_DELETE(self) -> None: # noqa: N802 def _handle_create(self) -> None: payload = self._read_json() + sandbox_id = payload.get("id") + if not isinstance(sandbox_id, str) or not _SANDBOX_ID_RE.match( + sandbox_id + ): + self._send_error_json(400, "Missing or invalid sandbox id") + return import_tar_uri = payload.get("import_tar_uri") import_tar = ( _download_import_tar(import_tar_uri) if import_tar_uri else None ) try: - sandbox_id = _create_sandbox( + _create_sandbox( + sandbox_id, allow_egress=bool(payload.get("allow_egress")), import_tar=import_tar, ) @@ -351,12 +431,16 @@ def _handle_exec(self, sandbox_id: str) -> None: self._send_error_json(400, "Missing argv") return timeout_ms = int(payload.get("timeout_ms", DEFAULT_EXEC_TIMEOUT_MS)) + exec_id = payload.get("exec_id") proc = _exec_in_sandbox( sandbox_id, [str(a) for a in argv], payload.get("cwd"), dict(payload.get("env") or {}), ) + if exec_id: + with _execs_lock: + _execs[exec_id] = proc self.send_response(200) self.send_header("Content-Type", "text/event-stream") @@ -398,6 +482,10 @@ def _pump(stream: Any, kind: str) -> None: timed_out = True proc.kill() proc.wait() + finally: + if exec_id: + with _execs_lock: + _execs.pop(exec_id, None) for t in threads: t.join() @@ -427,8 +515,9 @@ def _handle_put_file(self, sandbox_id: str, rest: str) -> None: staging_name = f"put-{uuid.uuid4().hex}" staging_host = os.path.join(_share_dir(sandbox_id), staging_name) # Stream straight to the staging file so concurrent uploads don't - # each hold a full 32 MiB body in memory. - with open(staging_host, "wb") as f: + # each hold a full 32 MiB body in memory. _open_new_file refuses a + # symlink the sandbox may have pre-planted at this path. + with _open_new_file(staging_host) as f: remaining = length while remaining > 0: chunk = self.rfile.read(min(remaining, 1024 * 1024)) @@ -463,12 +552,15 @@ def _handle_get_file(self, sandbox_id: str, rest: str) -> None: f"{shlex.quote(f'{SHARE_MOUNT}/{staging_name}')}", ) try: - size = os.path.getsize(staging_host) - self.send_response(200) - self.send_header("Content-Type", "application/octet-stream") - self.send_header("Content-Length", str(size)) - self.end_headers() - with open(staging_host, "rb") as f: + # Open with O_NOFOLLOW and size from fstat on the fd: the + # sandbox cannot redirect this read to a host file by swapping + # the staged path for a symlink after the cp. + f, size = _open_regular_file(staging_host) + with f: + self.send_response(200) + self.send_header("Content-Type", "application/octet-stream") + self.send_header("Content-Length", str(size)) + self.end_headers() shutil.copyfileobj(f, self.wfile) finally: os.unlink(staging_host) diff --git a/src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py b/src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py index e8ba5639a06..a8eeb764a42 100644 --- a/src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py +++ b/src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py @@ -31,6 +31,7 @@ import threading import time import urllib.parse +import uuid from collections import deque from dataclasses import dataclass from typing import ( @@ -86,13 +87,14 @@ _RETRYABLE_METHODS: FrozenSet[str] = frozenset({"GET", "PUT", "DELETE"}) _MAX_RETRIES = 3 -# Per-stream line cap for the SSE pump buffers. A caller that only +# Per-stream byte cap for the SSE pump buffers. A caller that only # wait()s on a noisy command would otherwise grow these deques without -# bound; past the cap the oldest lines are dropped (logged once). -_STREAM_BUFFER_MAX_LINES = 100_000 -# Newline-free output (progress bars, minified blobs) never hits the line -# cap; past this many buffered bytes the partial line is flushed as a -# synthetic line so memory and per-chunk copy costs stay bounded. +# bound; past the cap the oldest lines are dropped (logged once). The cap +# is on bytes, not line count, because the command is untrusted and can +# emit arbitrarily long (or newline-free) output to exhaust caller memory. +_STREAM_BUFFER_MAX_BYTES = 64 * 1024 * 1024 +# A single newline-free run is flushed as a synthetic line past this size +# so the carried-over remainder (and its per-chunk copy) stays bounded. _STREAM_REMAINDER_MAX_BYTES = 1_048_576 # Standard Google OAuth2 token endpoint, used when deriving ID-token @@ -475,34 +477,30 @@ def _request( def create_sandbox( self, + sandbox_id: str, *, allow_egress: bool = False, import_tar_uri: Optional[str] = None, - ) -> str: - """Create a fresh sandbox on the bridge instance. + ) -> None: + """Create a sandbox on the bridge instance under a caller-chosen id. + + The caller generates the id and the bridge creation is idempotent, + so a create whose response is lost cannot orphan the sandbox: the + caller still knows the id and can delete it. Args: + sandbox_id: The caller-generated sandbox id. allow_egress: Whether the sandbox gets outbound network access. import_tar_uri: Optional ``gs://`` URI of a filesystem snapshot tarball to initialize the sandbox from. - - Returns: - The bridge-assigned sandbox id. - - Raises: - SandboxExecError: If the bridge response is missing an id. """ - body: Dict[str, Any] = {"allow_egress": allow_egress} + body: Dict[str, Any] = { + "id": sandbox_id, + "allow_egress": allow_egress, + } if import_tar_uri is not None: body["import_tar_uri"] = import_tar_uri - resp = self._request("POST", "/v1/sandbox", json_body=body) - payload = resp.json() - sandbox_id = payload.get("id") - if not sandbox_id: - raise SandboxExecError( - f"Bridge create-sandbox response missing 'id': {payload!r}" - ) - return cast(str, sandbox_id) + self._request("POST", "/v1/sandbox", json_body=body) def delete_sandbox(self, sandbox_id: str) -> None: """Delete a sandbox. @@ -557,6 +555,7 @@ def exec_stream( cwd: Optional[str] = None, env: Optional[Dict[str, str]] = None, timeout_ms: int = DEFAULT_BRIDGE_TIMEOUT_MS, + exec_id: Optional[str] = None, ) -> Tuple["httpx.Response", Iterator[_BridgeEvent]]: """Run a command and stream decoded SSE events. @@ -565,9 +564,11 @@ def exec_stream( argv: Command argv list. cwd: Optional working directory. env: Environment variables for this command (passed to the - ``sandbox`` CLI via ``-e``; sandboxes inherit nothing from - the host). + ``sandbox`` CLI via ``--env``; sandboxes inherit nothing + from the host). timeout_ms: Per-exec timeout in milliseconds. + exec_id: Caller-generated id under which the bridge registers + the running command, so ``kill_exec`` can terminate it. Returns: The open streaming response (closing it aborts the stream and @@ -579,6 +580,8 @@ def exec_stream( body["cwd"] = cwd if env: body["env"] = env + if exec_id is not None: + body["exec_id"] = exec_id # POST eagerly so 4xx/auth errors raise at the call site; only SSE # decoding is deferred. @@ -590,6 +593,19 @@ def exec_stream( ) return resp, _drain_sse_response(resp) + def kill_exec(self, sandbox_id: str, exec_id: str) -> None: + """Ask the bridge to terminate a running command. + + Args: + sandbox_id: The sandbox the command runs in. + exec_id: The id passed to ``exec_stream`` for the command. + """ + self._request( + "DELETE", + f"/v1/sandbox/{sandbox_id}/exec/{exec_id}", + ok_statuses=(404, 410), + ) + def put_file(self, sandbox_id: str, remote_path: str, data: bytes) -> None: """Upload raw bytes to a sandbox path. @@ -636,6 +652,9 @@ def __init__( session: "CloudRunSandboxSession", started_at: float, response: Optional["httpx.Response"] = None, + client: Optional["_CloudRunBridgeClient"] = None, + sandbox_id: Optional[str] = None, + exec_id: Optional[str] = None, ) -> None: """Initialize the process wrapper. @@ -645,10 +664,17 @@ def __init__( started_at: Wall-clock launch time. response: The underlying streaming SSE response; ``kill()`` closes it to unblock the pump thread mid-read. + client: Bridge client used by ``kill()`` to terminate the + remote command. + sandbox_id: The sandbox the command runs in. + exec_id: The bridge exec id used by ``kill()``. """ super().__init__(session=session, started_at=started_at) self._event_iter = event_iter self._response = response + self._client = client + self._sandbox_id = sandbox_id + self._exec_id = exec_id self._killed = threading.Event() # The bridge multiplexes stdout/stderr on one SSE stream. We pump @@ -657,6 +683,7 @@ def __init__( # caller's thread. self._stdout_buf: Deque[str] = deque() self._stderr_buf: Deque[str] = deque() + self._buffered_bytes: Dict[str, int] = {"stdout": 0, "stderr": 0} self._buffer_truncated = False self._data_available = threading.Condition() self._done = threading.Event() @@ -672,10 +699,41 @@ def __init__( ) self._pump.start() - def _push_text(self, target: Deque[str], remainder: str, text: str) -> str: + def _enqueue(self, stream: str, target: Deque[str], chunk: str) -> None: + """Append a chunk, dropping oldest data once the byte cap is hit. + + Callers hold ``self._data_available``. + + Args: + stream: ``"stdout"`` or ``"stderr"`` — selects the byte counter. + target: Buffer to append to. + chunk: The text to append. + """ + target.append(chunk) + self._buffered_bytes[stream] += len(chunk) + while ( + self._buffered_bytes[stream] > _STREAM_BUFFER_MAX_BYTES + and len(target) > 1 + ): + dropped = target.popleft() + self._buffered_bytes[stream] -= len(dropped) + if not self._buffer_truncated: + self._buffer_truncated = True + logger.warning( + "Cloud Run exec output exceeded the %d-byte per-stream " + "buffer; oldest undrained output is being dropped. " + "Drain process.stdout()/stderr() (or use collect()) to " + "keep full output.", + _STREAM_BUFFER_MAX_BYTES, + ) + + def _push_text( + self, stream: str, target: Deque[str], remainder: str, text: str + ) -> str: """Line-buffer arbitrary text into the target deque. Args: + stream: ``"stdout"`` or ``"stderr"``. target: Buffer to append lines to. remainder: Partial line carried over from the previous chunk. text: New text chunk. @@ -686,33 +744,22 @@ def _push_text(self, target: Deque[str], remainder: str, text: str) -> str: buf = remainder + text while "\n" in buf: line, buf = buf.split("\n", 1) - if len(target) >= _STREAM_BUFFER_MAX_LINES: - target.popleft() - if not self._buffer_truncated: - self._buffer_truncated = True - logger.warning( - "Cloud Run exec output exceeded the %d-line " - "buffer; oldest undrained lines are being " - "dropped. Drain process.stdout()/stderr() (or " - "use collect()) to keep full output.", - _STREAM_BUFFER_MAX_LINES, - ) - target.append(line + "\n") + self._enqueue(stream, target, line + "\n") # Newline-free output would otherwise grow the remainder (and the # per-chunk `remainder + text` copy) without bound; flush it as a # synthetic line past the cap. if len(buf) > _STREAM_REMAINDER_MAX_BYTES: - target.append(buf) + self._enqueue(stream, target, buf) return "" return buf def _flush_remainders(self) -> None: """Push trailing partial lines (no terminating newline) on stream end.""" if self._stdout_remainder: - self._stdout_buf.append(self._stdout_remainder) + self._enqueue("stdout", self._stdout_buf, self._stdout_remainder) self._stdout_remainder = "" if self._stderr_remainder: - self._stderr_buf.append(self._stderr_remainder) + self._enqueue("stderr", self._stderr_buf, self._stderr_remainder) self._stderr_remainder = "" def _pump_events(self) -> None: @@ -724,12 +771,14 @@ def _pump_events(self) -> None: self._exit_code = int(ev.data) elif ev.kind == "stdout": self._stdout_remainder = self._push_text( + "stdout", self._stdout_buf, self._stdout_remainder, cast(str, ev.data), ) elif ev.kind == "stderr": self._stderr_remainder = self._push_text( + "stderr", self._stderr_buf, self._stderr_remainder, cast(str, ev.data), @@ -746,10 +795,11 @@ def _pump_events(self) -> None: self._done.set() self._data_available.notify_all() - def _iter_buffer(self, buf: Deque[str]) -> Iterator[str]: + def _iter_buffer(self, stream: str, buf: Deque[str]) -> Iterator[str]: """Block-pop lines from a buffer until the pump signals done. Args: + stream: ``"stdout"`` or ``"stderr"`` — selects the byte counter. buf: The buffer to drain. Yields: @@ -761,6 +811,7 @@ def _iter_buffer(self, buf: Deque[str]) -> Iterator[str]: self._data_available.wait() if buf: line = buf.popleft() + self._buffered_bytes[stream] -= len(line) else: return yield line @@ -773,7 +824,8 @@ def stdout(self) -> Iterator[str]: also lands in the per-session ``sandbox:`` log. """ return self._session._wrap_stream( - self._iter_buffer(self._stdout_buf), log_level=logging.INFO + self._iter_buffer("stdout", self._stdout_buf), + log_level=logging.INFO, ) def stderr(self) -> Iterator[str]: @@ -784,7 +836,8 @@ def stderr(self) -> Iterator[str]: also lands in the per-session ``sandbox:`` log. """ return self._session._wrap_stream( - self._iter_buffer(self._stderr_buf), log_level=logging.ERROR + self._iter_buffer("stderr", self._stderr_buf), + log_level=logging.ERROR, ) def wait(self, timeout: Optional[float] = None) -> int: @@ -822,8 +875,8 @@ def wait(self, timeout: Optional[float] = None) -> int: if self._killed.is_set(): raise SandboxExecError( "exec stream was killed via kill() or session close(); " - "the command may keep running in the sandbox until " - "timeout_ms" + "the bridge was asked to terminate the command, but its " + "exit status was not observed" ) # Stream ended cleanly but no exit frame arrived (truncated # SSE, bridge stalled out, Cloud Run request timeout hit). @@ -836,10 +889,25 @@ def wait(self, timeout: Optional[float] = None) -> int: return self._exit_code def kill(self) -> None: - """Stop reading the SSE stream and unblock consumers.""" + """Terminate the remote command and stop reading its stream.""" # Set the flag before closing so the pump treats the resulting # httpx error as a clean shutdown. self._killed.set() + # Terminate the command on the bridge first; closing the local + # response alone would leave it running until timeout_ms. + if ( + self._client is not None + and self._sandbox_id is not None + and self._exec_id is not None + ): + try: + self._client.kill_exec(self._sandbox_id, self._exec_id) + except Exception: + logger.debug( + "Bridge kill request for exec %s failed", + self._exec_id, + exc_info=True, + ) if self._response is not None: try: self._response.close() @@ -942,6 +1010,7 @@ def _exec( merged_env = {**self._session_env, **(env or {})} effective_cwd = cwd if cwd is not None else self._default_cwd + exec_id = uuid.uuid4().hex started_at = time.time() # exec_stream surfaces every failure mode as SandboxExecError, so # launch errors propagate to the caller as-is. @@ -951,12 +1020,16 @@ def _exec( cwd=effective_cwd, env=merged_env or None, timeout_ms=self._default_timeout_ms, + exec_id=exec_id, ) process = CloudRunSandboxProcess( event_iter, session=self, started_at=started_at, response=response, + client=self._client, + sandbox_id=self._sandbox_id, + exec_id=exec_id, ) self._processes.append(process) return process @@ -980,9 +1053,12 @@ def _create_snapshot(self) -> SandboxSnapshot: "attribute on the Cloud Run sandbox component (a gs:// " "location the bridge service account can write to)." ) + # Include a uuid so two snapshots of one sandbox within the same + # second get distinct object names instead of overwriting each + # other (which would silently change what the first restores). gcs_uri = ( f"{self._snapshot_uri_prefix}/{self._sandbox_id}-" - f"{int(time.time())}.tar" + f"{int(time.time())}-{uuid.uuid4().hex[:8]}.tar" ) self._client.snapshot_sandbox(self._sandbox_id, gcs_uri) return SandboxSnapshot( @@ -1127,18 +1203,21 @@ def _build_id_token_credentials( target_audience=audience, ), ) - if self.get_connector() is not None: - # The connector produced credentials we can't convert (user - # account, OAuth2 token, external account). Falling back to the - # ambient ADC below would silently authenticate to the bridge - # as a *different identity* than the connector the user - # configured — fail loudly instead. + if ( + self.get_connector() is not None + or self.config.service_account_path + ): + # An explicitly-chosen credential source (connector or + # service_account_path) produced credentials we can't convert + # (user account, OAuth2 token, external account). Falling back + # to the ambient-ADC helper below would ignore that choice and + # silently authenticate to the bridge as a *different identity* + # (e.g. the metadata server) — fail loudly instead. raise RuntimeError( - "The linked GCP service connector uses an auth method " - f"({type(credentials).__name__}) that cannot mint the " - "Google ID tokens Cloud Run IAM requires. Reconfigure the " - "connector with the 'service-account' or 'impersonation' " - "auth method, or unlink it to use service_account_path / " + f"The configured GCP credentials ({type(credentials).__name__}) " + "cannot mint the Google ID tokens Cloud Run IAM requires. " + "Use a service-account or impersonation service connector, a " + "service-account JSON key at service_account_path, or " "service-account-based Application Default Credentials." ) # Plain ADC (workstation or metadata server): fall back to the @@ -1149,10 +1228,21 @@ def _build_id_token_credentials( def _get_bridge_client(self) -> _CloudRunBridgeClient: """Return the lazily-built bridge HTTP client. + Rebuilds the client (and its baked-in ID-token credentials) once + the linked service connector has expired, so a component does not + keep signing tokens with credentials the connector no longer + vouches for. + Returns: A cached ``_CloudRunBridgeClient`` bound to this component. """ with self._bridge_client_lock: + if ( + self._bridge_client is not None + and self.connector_has_expired() + ): + self._bridge_client.close() + self._bridge_client = None if self._bridge_client is None: self._bridge_client = _CloudRunBridgeClient( self.config.service_url, @@ -1204,8 +1294,10 @@ def create_session( A ``CloudRunSandboxSession`` bound to the new sandbox. """ eff = cast(CloudRunSandboxSettings, self.resolve_settings(settings)) - sandbox_id = self._get_bridge_client().create_sandbox( - allow_egress=eff.allow_egress + sandbox_id = self._new_sandbox_id() + client = self._get_bridge_client() + self._create_sandbox_or_cleanup( + client, sandbox_id, allow_egress=eff.allow_egress ) return self._build_session( sandbox_id, eff, destroy_on_exit=destroy_on_exit @@ -1250,8 +1342,63 @@ def restore(self, snapshot: SandboxSnapshot) -> SandboxSession: """ self._validate_snapshot(snapshot) eff = cast(CloudRunSandboxSettings, self.resolve_settings(None)) - sandbox_id = self._get_bridge_client().create_sandbox( + sandbox_id = self._new_sandbox_id() + client = self._get_bridge_client() + self._create_sandbox_or_cleanup( + client, + sandbox_id, allow_egress=eff.allow_egress, import_tar_uri=snapshot.ref, ) return self._build_session(sandbox_id, eff) + + @staticmethod + def _new_sandbox_id() -> str: + """Generate a fresh sandbox id. + + Returns: + An id of the form ``sb-<12 hex chars>``. + """ + return f"sb-{uuid.uuid4().hex[:12]}" + + @staticmethod + def _create_sandbox_or_cleanup( + client: "_CloudRunBridgeClient", + sandbox_id: str, + *, + allow_egress: bool = False, + import_tar_uri: Optional[str] = None, + ) -> None: + """Create a sandbox, deleting it if the create fails mid-flight. + + Because the id is caller-generated, a create whose response is lost + (the bridge may have created the sandbox before the connection + dropped) can still be cleaned up here instead of orphaning a + ``sleep infinity`` process on the bridge instance. + + Args: + client: Bridge client. + sandbox_id: The caller-generated sandbox id. + allow_egress: Whether the sandbox gets outbound network access. + import_tar_uri: Optional ``gs://`` snapshot tarball to boot from. + + Raises: + Exception: Re-raises whatever the create call raised, after a + best-effort delete of the possibly-created sandbox. + """ + try: + client.create_sandbox( + sandbox_id, + allow_egress=allow_egress, + import_tar_uri=import_tar_uri, + ) + except Exception: + try: + client.delete_sandbox(sandbox_id) + except Exception: + logger.warning( + "Failed to clean up sandbox %s after a failed create; " + "it may linger until the bridge instance is recycled.", + sandbox_id, + ) + raise diff --git a/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py b/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py index 45fd0c7653e..b1d819d5db8 100644 --- a/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py +++ b/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py @@ -212,27 +212,27 @@ def handler(request: httpx.Request) -> httpx.Response: credentials = _StaticCredentials() client = _make_client(handler, credentials) - sandbox_id = client.create_sandbox(allow_egress=True) + client.create_sandbox("sb-0123456789ab", allow_egress=True) - assert sandbox_id == "sb-0123456789ab" assert seen["path"] == "/v1/sandbox" assert seen["auth"] == "Bearer test-id-token" - assert seen["body"] == {"allow_egress": True} + assert seen["body"] == { + "id": "sb-0123456789ab", + "allow_egress": True, + } assert credentials.refresh_count == 1 def test_create_sandbox_with_import_tar(self) -> None: def handler(request: httpx.Request) -> httpx.Response: body = json.loads(request.content) + assert body["id"] == "sb-0123456789ab" assert body["import_tar_uri"] == "gs://bucket/snap.tar" return httpx.Response(200, json={"id": "sb-0123456789ab"}) client = _make_client(handler) - client.create_sandbox(import_tar_uri="gs://bucket/snap.tar") - - def test_create_sandbox_missing_id_raises(self) -> None: - client = _make_client(lambda _: httpx.Response(200, json={})) - with pytest.raises(SandboxExecError, match="missing 'id'"): - client.create_sandbox() + client.create_sandbox( + "sb-0123456789ab", import_tar_uri="gs://bucket/snap.tar" + ) def test_no_auth_header_when_unauthenticated(self) -> None: def handler(request: httpx.Request) -> httpx.Response: @@ -240,7 +240,11 @@ def handler(request: httpx.Request) -> httpx.Response: return httpx.Response(200, json={"id": "sb-0123456789ab"}) client = _make_client(handler, credentials=None) - client.create_sandbox() + client.create_sandbox("sb-0123456789ab") + + def test_kill_exec_tolerates_404(self) -> None: + client = _make_client(lambda _: httpx.Response(404)) + client.kill_exec("sb-0123456789ab", "exec-123") def test_delete_tolerates_404(self) -> None: client = _make_client(lambda _: httpx.Response(404)) @@ -290,7 +294,7 @@ def handler(request: httpx.Request) -> httpx.Response: client = _make_client(handler) with pytest.raises(SandboxExecError, match="503"): - client.create_sandbox() + client.create_sandbox("sb-0123456789ab") assert len(calls) == 1 def test_upload_rejects_oversized_file(self) -> None: @@ -334,12 +338,21 @@ def handler(request: httpx.Request) -> httpx.Response: class TestProcess: def _make_process( - self, events: List[_BridgeEvent] + self, + events: List[_BridgeEvent], + *, + client: Any = None, + exec_id: Optional[str] = None, ) -> CloudRunSandboxProcess: session = MagicMock() session._wrap_stream.side_effect = lambda lines, **_: lines return CloudRunSandboxProcess( - iter(events), session=session, started_at=0.0 + iter(events), + session=session, + started_at=0.0, + client=client, + sandbox_id="sb-0123456789ab", + exec_id=exec_id, ) def test_demux_and_wait(self) -> None: @@ -375,6 +388,34 @@ def test_kill_unblocks_wait(self) -> None: with pytest.raises(SandboxExecError, match="killed"): process.wait(timeout=5) + def test_kill_terminates_remote_command(self) -> None: + client = MagicMock() + process = self._make_process( + [_BridgeEvent(kind="stdout", data="x\n")], + client=client, + exec_id="exec-123", + ) + process.kill() + client.kill_exec.assert_called_once_with("sb-0123456789ab", "exec-123") + + def test_buffer_byte_cap_drops_oldest(self) -> None: + # Two ~1 MiB newline-free chunks each flush as a synthetic line; + # with a smaller cap the oldest is dropped rather than retained. + import zenml.integrations.gcp.sandboxes.cloudrun_sandbox as mod + + chunk = "a" * (mod._STREAM_REMAINDER_MAX_BYTES + 1) + with patch.object(mod, "_STREAM_BUFFER_MAX_BYTES", len(chunk) + 10): + process = self._make_process( + [ + _BridgeEvent(kind="stdout", data=chunk), + _BridgeEvent(kind="stdout", data=chunk), + _BridgeEvent(kind="exit", data=0), + ] + ) + process.wait(timeout=5) + # Only the most recent chunk survives the byte cap. + assert len(list(process.stdout())) == 1 + class TestSession: def test_exec_merges_session_env(self) -> None: @@ -524,5 +565,67 @@ def test_unconvertible_connector_credentials_fail_loudly(self) -> None: ), patch.object(component, "get_connector", return_value=MagicMock()), ): - with pytest.raises(RuntimeError, match="service-account"): + with pytest.raises(RuntimeError, match="cannot mint"): + component._build_id_token_credentials() + + def test_unconvertible_service_account_path_fails_loudly(self) -> None: + """An explicit service_account_path whose creds can't mint ID + tokens must fail rather than fall back to ambient ADC.""" + from google.oauth2.credentials import Credentials as UserCredentials + + component = _make_component( + service_account_path="/keys/service-account.json" + ) + with ( + patch.object( + component, + "_get_authentication", + return_value=(MagicMock(spec=UserCredentials), "project"), + ), + patch.object(component, "get_connector", return_value=None), + ): + with pytest.raises(RuntimeError, match="cannot mint"): component._build_id_token_credentials() + + +class TestComponentLifecycle: + def test_create_session_cleans_up_on_failed_create(self) -> None: + """A create that fails mid-flight must delete the sandbox by its + caller-generated id so it can't orphan.""" + client = MagicMock() + client.create_sandbox.side_effect = SandboxExecError("boom") + component = _make_component() + with patch.object( + component, "_get_bridge_client", return_value=client + ): + with pytest.raises(SandboxExecError, match="boom"): + component.create_session() + created_id = client.create_sandbox.call_args.args[0] + client.delete_sandbox.assert_called_once_with(created_id) + + def test_bridge_client_rebuilt_after_connector_expiry(self) -> None: + component = _make_component() + with ( + patch.object( + component, + "_build_id_token_credentials", + return_value=None, + ), + patch.object( + component, "connector_has_expired", return_value=False + ), + ): + first = component._get_bridge_client() + assert component._get_bridge_client() is first + + with ( + patch.object( + component, + "_build_id_token_credentials", + return_value=None, + ), + patch.object( + component, "connector_has_expired", return_value=True + ), + ): + assert component._get_bridge_client() is not first From 1d4f48060933ace5c68aa975462229dbc6e99cb9 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 17 Jul 2026 11:41:33 +0200 Subject: [PATCH 7/8] Resolve relative sandbox file paths against the session cwd MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The base SandboxSession contract documents, for all flavors, that a relative `remote_path` in `upload_file`/`download_file` resolves against the session's working directory — the same directory `exec()` uses as its default cwd. The Cloud Run flavor diverged: `_exec` honored `default_cwd`/`settings.cwd`, but file transfer anchored relative paths at the sandbox root, so `exec("cat foo.txt")` and `upload_file(src, "foo.txt")` disagreed about where "foo.txt" lived whenever a cwd was configured. Resolve relative remote paths against the session's default cwd before the existing sanitization runs, mirroring exec. Absolute paths are unchanged and the `..`-escape guard still refuses paths that climb above the cwd. Also tighten the per-exec env docstring (values are hidden from the sandboxed process's command line but visible in the trusted bridge host's process list) and document the missing per-task image support in the docs Limitations. Co-Authored-By: Claude Fable 5 --- .../component-guide/sandboxes/cloudrun.md | 1 + .../gcp/sandboxes/cloudrun_sandbox.py | 45 +++++++++-- .../gcp/sandboxes/test_cloudrun_sandbox.py | 76 ++++++++++++++++++- 3 files changed, 114 insertions(+), 8 deletions(-) diff --git a/docs/book/component-guide/sandboxes/cloudrun.md b/docs/book/component-guide/sandboxes/cloudrun.md index 85d2bca8e28..2d6f45ff448 100644 --- a/docs/book/component-guide/sandboxes/cloudrun.md +++ b/docs/book/component-guide/sandboxes/cloudrun.md @@ -108,6 +108,7 @@ Config-only fields on the component: `service_url`, `audience`, `allow_unauthent - Persistent sandboxes are **instance-local**: if Cloud Run recycles the bridge instance, running sandboxes are lost (`attach()` reports this cleanly). Snapshot anything you need to keep. - File transfer through the bridge is capped at 32 MiB per request; move larger data through Cloud Storage. - A single exec is bounded by the Cloud Run request timeout (max 60 minutes). +- No per-task image pinning: unlike the Modal, Kubernetes, and Docker sandbox flavors, this flavor does not implement `image_settings()`, so a per-task image requested by a caller is ignored — the sandbox always runs the rootfs provided by the bridge's launcher. - Cloud Run sandboxes are in public preview; the underlying CLI may change. For the full config surface, see the [SDK docs](https://sdkdocs.zenml.io). diff --git a/src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py b/src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py index a8eeb764a42..3061e5d285d 100644 --- a/src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py +++ b/src/zenml/integrations/gcp/sandboxes/cloudrun_sandbox.py @@ -989,9 +989,11 @@ def _exec( ``shlex.split``). cwd: Working directory inside the sandbox. env: Per-exec environment variables, merged over the - session-level ``sandbox_environment``. The ``sandbox`` CLI - applies them via ``-e``, so values never appear on the - command line. + session-level ``sandbox_environment``. The bridge passes + them to the ``sandbox`` CLI as ``--env`` flags, so they do + not appear on the sandboxed process's own command line + (they are still visible in the trusted bridge host's + process list). Returns: A ``CloudRunSandboxProcess``. @@ -1067,12 +1069,40 @@ def _create_snapshot(self) -> SandboxSnapshot: metadata={"source_sandbox": self._sandbox_id}, ) + def _resolve_remote_path(self, remote_path: str) -> str: + """Resolve a relative remote path against the session's default cwd. + + The base contract documented on ``SandboxSession.upload_file`` / + ``download_file`` specifies that a relative ``remote_path`` + resolves against the session's working directory — the same + directory ``exec()`` uses as its default cwd. Without this, + ``exec`` and file transfer disagree about where a relative path + lives: ``exec`` honors ``default_cwd`` while the bridge anchors + relative file paths at the sandbox root. + + Absolute paths and the no-configured-cwd case are returned + unchanged; the client's ``_sanitize_remote_path`` still rejects + any result that escapes the sandbox root (e.g. a relative + ``../`` that climbs above ``default_cwd``). + + Args: + remote_path: The caller-supplied destination/source path. + + Returns: + The joined absolute path when a default cwd is configured and + the input is relative; otherwise the path unchanged. + """ + if self._default_cwd and not posixpath.isabs(remote_path): + return posixpath.join(self._default_cwd, remote_path) + return remote_path + def _upload_file(self, local_path: str, remote_path: str) -> None: """Upload a local file into the sandbox. Args: local_path: Source path on the caller's machine. remote_path: Destination path inside the sandbox filesystem. + Relative paths resolve against the session's default cwd. Raises: ValueError: If the file exceeds the bridge's body limit. @@ -1087,16 +1117,21 @@ def _upload_file(self, local_path: str, remote_path: str) -> None: ) with open(local_path, "rb") as f: data = f.read() - self._client.put_file(self._sandbox_id, remote_path, data) + self._client.put_file( + self._sandbox_id, self._resolve_remote_path(remote_path), data + ) def _download_file(self, remote_path: str, local_path: str) -> None: """Download a file from the sandbox. Args: remote_path: Source path inside the sandbox filesystem. + Relative paths resolve against the session's default cwd. local_path: Destination path on the caller's machine. """ - data = self._client.get_file(self._sandbox_id, remote_path) + data = self._client.get_file( + self._sandbox_id, self._resolve_remote_path(remote_path) + ) with open(local_path, "wb") as f: f.write(data) diff --git a/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py b/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py index b1d819d5db8..eabfff0aad8 100644 --- a/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py +++ b/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py @@ -92,6 +92,7 @@ def _make_session( *, session_env: Optional[Dict[str, str]] = None, snapshot_uri_prefix: Optional[str] = None, + default_cwd: Optional[str] = None, ) -> CloudRunSandboxSession: parent = MagicMock(spec=BaseSandbox) parent.flavor = "cloudrun" @@ -102,6 +103,7 @@ def _make_session( parent=parent, session_env=session_env, snapshot_uri_prefix=snapshot_uri_prefix, + default_cwd=default_cwd, ) @@ -466,6 +468,71 @@ def handler(request: httpx.Request) -> httpx.Response: with open(dst) as f: assert f.read() == "payload" + def test_relative_paths_resolve_against_default_cwd(self) -> None: + # A relative remote_path must land under the session's default cwd + # (the same directory exec() uses), not at the sandbox root, so + # exec and file transfer agree on where a relative path lives. + stored: Dict[str, bytes] = {} + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "PUT": + stored[request.url.path] = request.content + return httpx.Response(200, json={"written": "x"}) + return httpx.Response(200, content=stored[request.url.path]) + + session = _make_session( + _make_client(handler), default_cwd="/tmp/workspace" + ) + with tempfile.TemporaryDirectory() as tmp: + src = os.path.join(tmp, "in.txt") + with open(src, "w") as f: + f.write("payload") + session.upload_file(src, "foo.txt") + + assert ( + "/v1/sandbox/sb-0123456789ab/file/tmp/workspace/foo.txt" + in (stored) + ) + + dst = os.path.join(tmp, "out.txt") + session.download_file("foo.txt", dst) + with open(dst) as f: + assert f.read() == "payload" + + def test_absolute_paths_ignore_default_cwd(self) -> None: + stored: Dict[str, bytes] = {} + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "PUT": + stored[request.url.path] = request.content + return httpx.Response(200, json={"written": "x"}) + return httpx.Response(200, content=stored[request.url.path]) + + session = _make_session( + _make_client(handler), default_cwd="/tmp/workspace" + ) + with tempfile.TemporaryDirectory() as tmp: + src = os.path.join(tmp, "in.txt") + with open(src, "w") as f: + f.write("payload") + session.upload_file(src, "/data/in.txt") + + # Leading slash is stripped for URL composition, but the path is + # anchored at the sandbox root — not joined onto default_cwd. + assert "/v1/sandbox/sb-0123456789ab/file/data/in.txt" in stored + + def test_relative_escape_above_default_cwd_refused(self) -> None: + session = _make_session( + _make_client(lambda _: httpx.Response(200)), + default_cwd="/tmp/workspace", + ) + with tempfile.TemporaryDirectory() as tmp: + src = os.path.join(tmp, "in.txt") + with open(src, "w") as f: + f.write("payload") + with pytest.raises(ValueError, match="outside"): + session.upload_file(src, "../../etc/passwd") + def test_snapshot_requires_prefix(self) -> None: session = _make_session(_make_client(lambda _: httpx.Response(200))) with pytest.raises(NotImplementedError, match="snapshot_uri_prefix"): @@ -553,7 +620,8 @@ def test_service_account_credentials_converted(self) -> None: def test_unconvertible_connector_credentials_fail_loudly(self) -> None: """Connector creds that can't mint ID tokens must not silently - fall back to the ambient ADC identity.""" + fall back to the ambient ADC identity. + """ from google.oauth2.credentials import Credentials as UserCredentials component = _make_component() @@ -570,7 +638,8 @@ def test_unconvertible_connector_credentials_fail_loudly(self) -> None: def test_unconvertible_service_account_path_fails_loudly(self) -> None: """An explicit service_account_path whose creds can't mint ID - tokens must fail rather than fall back to ambient ADC.""" + tokens must fail rather than fall back to ambient ADC. + """ from google.oauth2.credentials import Credentials as UserCredentials component = _make_component( @@ -591,7 +660,8 @@ def test_unconvertible_service_account_path_fails_loudly(self) -> None: class TestComponentLifecycle: def test_create_session_cleans_up_on_failed_create(self) -> None: """A create that fails mid-flight must delete the sandbox by its - caller-generated id so it can't orphan.""" + caller-generated id so it can't orphan. + """ client = MagicMock() client.create_sandbox.side_effect = SandboxExecError("boom") component = _make_component() From 4b98fd2f0a33bd4aec99fc49dbaf92eb19cb2239 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 17 Jul 2026 12:12:37 +0200 Subject: [PATCH 8/8] Use a non-tmp example cwd in Cloud Run path tests Bandit's B108 flags hardcoded /tmp paths; the literals here are example remote sandbox paths, not real temp-file usage, so switch the example to /srv/workspace instead of suppressing the check. Co-Authored-By: Claude Fable 5 --- .../integrations/gcp/sandboxes/test_cloudrun_sandbox.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py b/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py index eabfff0aad8..d06399fe2bf 100644 --- a/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py +++ b/tests/unit/integrations/gcp/sandboxes/test_cloudrun_sandbox.py @@ -481,7 +481,7 @@ def handler(request: httpx.Request) -> httpx.Response: return httpx.Response(200, content=stored[request.url.path]) session = _make_session( - _make_client(handler), default_cwd="/tmp/workspace" + _make_client(handler), default_cwd="/srv/workspace" ) with tempfile.TemporaryDirectory() as tmp: src = os.path.join(tmp, "in.txt") @@ -490,7 +490,7 @@ def handler(request: httpx.Request) -> httpx.Response: session.upload_file(src, "foo.txt") assert ( - "/v1/sandbox/sb-0123456789ab/file/tmp/workspace/foo.txt" + "/v1/sandbox/sb-0123456789ab/file/srv/workspace/foo.txt" in (stored) ) @@ -509,7 +509,7 @@ def handler(request: httpx.Request) -> httpx.Response: return httpx.Response(200, content=stored[request.url.path]) session = _make_session( - _make_client(handler), default_cwd="/tmp/workspace" + _make_client(handler), default_cwd="/srv/workspace" ) with tempfile.TemporaryDirectory() as tmp: src = os.path.join(tmp, "in.txt") @@ -524,7 +524,7 @@ def handler(request: httpx.Request) -> httpx.Response: def test_relative_escape_above_default_cwd_refused(self) -> None: session = _make_session( _make_client(lambda _: httpx.Response(200)), - default_cwd="/tmp/workspace", + default_cwd="/srv/workspace", ) with tempfile.TemporaryDirectory() as tmp: src = os.path.join(tmp, "in.txt")