From 36f81caf67fdd310a78411b376e22031878fa501 Mon Sep 17 00:00:00 2001 From: Sunny Song Date: Thu, 8 Oct 2026 22:11:13 +0000 Subject: [PATCH] feat(python): build Substrate driver and runtime directly on Client and Env Build SubstrateBackendDriver and SubstrateEnvClientRuntime directly on top of Client and Env, eliminating synthetic/mock stubs and aligning with ate-env-api: - SubstrateBackendDriver directly uses Client for environment lifecycle (create, suspend, delete, and reap) and warm pool management. - Runtime communicates directly with in-guest ProcessService and FileSystemService over gRPC, passing x-env-id / x-env-atespace metadata. - Introduce LoopRunner running a dedicated background event loop for synchronous facades (Fleet, SubstrateBackendDriver, SubstrateEnvClientRuntime) to safely share gRPC channels across sequential calls without event loop collisions. - Wire acquire_timeout_s, ready_timeout_s, labels, and auth_token via gRPC bearer metadata interceptor to Client. - Reaping cleans both run-tagged environments and warm pool actors, draining _warm_paused_pool. - Add context manager support (with SandboxFleet) and setup(tasks=None) to Fleet, supporting single and batch acquire/release. - Drop obsolete HTTP router and synthetic grpc data planes, removing router_url, grpc_endpoint, host, port, ip_address, and interactive session stubs. - Unify exception hierarchy in errors.py (SandboxStartError and PreflightError subclass InfrastructureError) and map gRPC status codes cleanly. - Use standard @dataclass(frozen=True) with custom SecretStr in config.py, avoiding external pydantic dependency. - Default command cwd="" to honor the container image's WORKDIR. - Address PR 1 review feedback: - B1: default template resolution to "default-template", None for unset, and support image digests. - B2: return ExecResult from handle.exec and exec_async; handle __exit__ releases sandbox. - B3: wrap close/teardown/reap in try/finally to guarantee backend driver closure. - M1: kill process on client-side timeout in SubstrateEnvClientRuntime. - M2: drop warm actor and claim next on NotFoundError; document implicit resume. - M3: match NotFoundError and FailedPreconditionError exception types. - M4: LoopRunner stop guard against caller thread, cancel tasks and close cleanly. - M5: config.image_templates takes precedence over config.template in fleet. - M6: remove AsyncSandboxFleet mentions from PR 1 README. - M7: parallelize unwarm_pool_async and reap_async with Semaphore(16). - Cleaned unused exec file helpers and templates cache. - Add comprehensive unit tests against fake ate-env-api gRPC server for driver, runtime, and fleet strategies. --- README.md | 2 +- clients/python/README.md | 9 +- clients/python/src/ate_env/__init__.py | 65 ++- .../python/src/ate_env/backend/__init__.py | 20 + clients/python/src/ate_env/backend/base.py | 57 +++ clients/python/src/ate_env/backend/mock.py | 84 ++++ .../python/src/ate_env/backend/substrate.py | 451 ++++++++++++++++++ clients/python/src/ate_env/client.py | 53 +- clients/python/src/ate_env/config.py | 173 +++++++ clients/python/src/ate_env/errors.py | 111 +++++ clients/python/src/ate_env/exceptions.py | 62 +++ clients/python/src/ate_env/fleet.py | 296 ++++++++++++ clients/python/src/ate_env/handle.py | 146 ++++++ clients/python/src/ate_env/loop_runner.py | 82 ++++ .../python/src/ate_env/runtime/__init__.py | 23 + clients/python/src/ate_env/runtime/base.py | 95 ++++ clients/python/src/ate_env/runtime/mock.py | 76 +++ .../ate_env/runtime/substrate_env_client.py | 339 +++++++++++++ clients/python/src/ate_env/strategies.py | 179 +++++++ clients/python/src/ate_env/types.py | 151 +++++- clients/python/tests/test_config.py | 57 +++ clients/python/tests/test_data_planes.py | 79 +++ clients/python/tests/test_exec_semantics.py | 122 +++++ clients/python/tests/test_fleet_strategies.py | 74 +++ clients/python/tests/test_fleet_types.py | 45 ++ clients/python/tests/test_substrate_driver.py | 154 ++++++ .../python/tests/test_substrate_env_client.py | 125 +++++ 27 files changed, 3118 insertions(+), 12 deletions(-) create mode 100644 clients/python/src/ate_env/backend/__init__.py create mode 100644 clients/python/src/ate_env/backend/base.py create mode 100644 clients/python/src/ate_env/backend/mock.py create mode 100644 clients/python/src/ate_env/backend/substrate.py create mode 100644 clients/python/src/ate_env/config.py create mode 100644 clients/python/src/ate_env/exceptions.py create mode 100644 clients/python/src/ate_env/fleet.py create mode 100644 clients/python/src/ate_env/handle.py create mode 100644 clients/python/src/ate_env/loop_runner.py create mode 100644 clients/python/src/ate_env/runtime/__init__.py create mode 100644 clients/python/src/ate_env/runtime/base.py create mode 100644 clients/python/src/ate_env/runtime/mock.py create mode 100644 clients/python/src/ate_env/runtime/substrate_env_client.py create mode 100644 clients/python/src/ate_env/strategies.py create mode 100644 clients/python/tests/test_config.py create mode 100644 clients/python/tests/test_data_planes.py create mode 100644 clients/python/tests/test_exec_semantics.py create mode 100644 clients/python/tests/test_fleet_strategies.py create mode 100644 clients/python/tests/test_fleet_types.py create mode 100644 clients/python/tests/test_substrate_driver.py create mode 100644 clients/python/tests/test_substrate_env_client.py diff --git a/README.md b/README.md index b0badfab..d20618d1 100644 --- a/README.md +++ b/README.md @@ -28,7 +28,7 @@ while this project adds the environment-shaped API on top. - **`cmd/ate-env-api`** — The API service that manages environments and proxies remote guest requests. - **`cmd/ate-env-guest`** — The daemon server running inside each actor serving command executions, file read/write, and built-in MCP tools. - **`clients/go`** — The Go client library to manage environments, run commands, and perform file operations. -- **`clients/python`** — The async Python client library ([README](clients/python/README.md)). +- **`clients/python`** — Python client for environment lifecycle management and batch sandbox orchestration ([README](clients/python/README.md)). - **`integrations/nemo-gym`** — A [NeMo Gym](https://github.com/NVIDIA-NeMo/Gym) sandbox provider that runs rollout sandboxes as environments, built on the Python client ([README](integrations/nemo-gym/README.md)). ## Installation diff --git a/clients/python/README.md b/clients/python/README.md index 2b8c5e83..e47be4ec 100644 --- a/clients/python/README.md +++ b/clients/python/README.md @@ -1,11 +1,12 @@ -# ate-env-client — Async Python Client +# ate-env-client — Python Client & Sandbox Fleet SDK > [!WARNING] > This is an alpha API and is likely to change until v1.0 is released. -Async Python client for the [Agent Substrate Environment](../../README.md) -API (`ate-env-api`): environment lifecycle, remote command execution, and -streaming file I/O over gRPC. +Python client library for the [Agent Substrate Environment](../../README.md) +API (`ate-env-api`): +1. **Single Environment Lifecycle & Guest Operations**: `Client` and `Env` for fine-grained gRPC control, process execution, and file streaming. +2. **Fleet & Sandbox Orchestration**: `SandboxFleet` for managing batches of sandboxes with pool sizing, pre-warming, and execution strategies. Requires Python >= 3.10. Everything is `asyncio`-native: methods are coroutines, log/file streams are async iterators, and cancellation works diff --git a/clients/python/src/ate_env/__init__.py b/clients/python/src/ate_env/__init__.py index 08375d9c..0d8e0560 100644 --- a/clients/python/src/ate_env/__init__.py +++ b/clients/python/src/ate_env/__init__.py @@ -32,48 +32,103 @@ async def main(): asyncio.run(main()) """ +# Layer 1: Core Client & Env from .client import DEFAULT_ATESPACE, Client from .env import Env, Process + +# Layer 2: High-Level Fleet Management & Abstractions +from .config import FleetConfig +from .fleet import SandboxFleet +from .handle import SandboxHandle + +# Exceptions from .errors import ( + CommandExecutionError, + CommandStartError, + CommandTimeoutError, EnvError, FailedPreconditionError, + InfrastructureError, InvalidArgumentError, NotFoundError, PermissionDeniedError, + PreflightError, ProcessExitedError, RpcError, + SandboxProtocolError, + SandboxStartError, + SandboxUnavailableError, map_rpc_error, ) +from .exceptions import SandboxError + +# Types from .types import ( + DataPlaneEndpoint, EnvironmentInfo, + EnvironmentSpec, EnvironmentStatus, + ExecResult, + FleetPlan, + PlacementSpec, + PlanEntry, ProcessInfo, ProcessOutput, ProcessState, + RawSandboxInstance, + ResourceLimits, ShellResult, Signal, + Task, Template, ) __all__ = [ + # Low-level API "Client", "DEFAULT_ATESPACE", "Env", + "Process", + # High-level Fleet API + "FleetConfig", + "SandboxFleet", + "SandboxHandle", + # Low-level Errors "EnvError", - "EnvironmentInfo", - "EnvironmentStatus", "FailedPreconditionError", "InvalidArgumentError", "NotFoundError", "PermissionDeniedError", - "Process", "ProcessExitedError", + "RpcError", + "map_rpc_error", + # Fleet / Execution Exceptions + "SandboxError", + "CommandExecutionError", + "CommandStartError", + "CommandTimeoutError", + "InfrastructureError", + "PreflightError", + "SandboxProtocolError", + "SandboxStartError", + "SandboxUnavailableError", + # Types + "EnvironmentInfo", + "EnvironmentStatus", "ProcessInfo", "ProcessOutput", "ProcessState", - "RpcError", "ShellResult", "Signal", "Template", - "map_rpc_error", + "ResourceLimits", + "PlacementSpec", + "EnvironmentSpec", + "Task", + "ExecResult", + "DataPlaneEndpoint", + "RawSandboxInstance", + "PlanEntry", + "FleetPlan", ] + diff --git a/clients/python/src/ate_env/backend/__init__.py b/clients/python/src/ate_env/backend/__init__.py new file mode 100644 index 00000000..0dbebec1 --- /dev/null +++ b/clients/python/src/ate_env/backend/__init__.py @@ -0,0 +1,20 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .base import BackendDriver +from .mock import MockBackendDriver +from .substrate import SubstrateBackendDriver + +__all__ = ["BackendDriver", "MockBackendDriver", "SubstrateBackendDriver"] + diff --git a/clients/python/src/ate_env/backend/base.py b/clients/python/src/ate_env/backend/base.py new file mode 100644 index 00000000..42e667c4 --- /dev/null +++ b/clients/python/src/ate_env/backend/base.py @@ -0,0 +1,57 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Optional + +from ..types import EnvironmentSpec, RawSandboxInstance + + +class BackendDriver(ABC): + """ + Control Plane (Backend Protocol) Interface. + + Governs fleet provisioning, template creation, warm-pool sizing, + instance claims, snapshotting, and run isolation / reaping. + """ + + @abstractmethod + def preflight(self) -> None: + """Validate cluster connectivity, permissions, storage, and worker capacity.""" + + @abstractmethod + def ensure_template(self, env: EnvironmentSpec) -> str: + """Ensure template/golden snapshot exists. Returns unique template ID.""" + + @abstractmethod + def warm_pool(self, template_id: str, replicas: int, wait: bool = True) -> None: + """Provision warm replicas (Ready pods in K8s or Paused Actors in Substrate).""" + + @abstractmethod + def unwarm_pool(self, template_id: str) -> None: + """Scale down warm pool for the template to 0.""" + + @abstractmethod + def acquire(self, template_id: str, run_id: str, timeout_s: float = 180.0) -> RawSandboxInstance: + """Instantly claim a sandbox instance from warm pool or golden.""" + + @abstractmethod + def release(self, instance_id: str, recycle: bool = False) -> None: + """Release or recycle instance.""" + + @abstractmethod + def reap(self, run_id: str) -> int: + """Force delete all resources associated with a run_id.""" diff --git a/clients/python/src/ate_env/backend/mock.py b/clients/python/src/ate_env/backend/mock.py new file mode 100644 index 00000000..978fca03 --- /dev/null +++ b/clients/python/src/ate_env/backend/mock.py @@ -0,0 +1,84 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import collections +import time +import uuid +from typing import Dict, List, Optional, Set + +from ..errors import SandboxStartError +from ..types import EnvironmentSpec, RawSandboxInstance +from .base import BackendDriver + + +class MockBackendDriver(BackendDriver): + """In-memory mock backend driver for hermetic unit testing and PoC simulation.""" + + def __init__(self, fail_preflight: bool = False, fail_acquire: bool = False): + self.fail_preflight = fail_preflight + self.fail_acquire = fail_acquire + self.templates: Dict[str, EnvironmentSpec] = {} + self.warm_pools: Dict[str, int] = collections.defaultdict(int) + self.instances: Dict[str, RawSandboxInstance] = {} + self.reaped_runs: Set[str] = set() + + def preflight(self) -> None: + if self.fail_preflight: + raise RuntimeError("Mock preflight failed: cluster unreachable") + + def ensure_template(self, env: EnvironmentSpec) -> str: + tid = env.template_key() + self.templates[tid] = env + return tid + + def warm_pool(self, template_id: str, replicas: int, wait: bool = True) -> None: + self.warm_pools[template_id] = replicas + + def unwarm_pool(self, template_id: str) -> None: + self.warm_pools[template_id] = 0 + + def acquire(self, template_id: str, run_id: str, timeout_s: float = 180.0) -> RawSandboxInstance: + if self.fail_acquire: + raise SandboxStartError("Mock acquire failed: capacity exhausted") + + inst_id = f"mock-sb-{uuid.uuid4().hex[:8]}" + inst = RawSandboxInstance( + instance_id=inst_id, + endpoint=f"http://127.0.0.1:8000/{inst_id}", + template_id=template_id, + run_id=run_id, + metadata={"created_at": time.time(), "template": template_id} + ) + self.instances[inst_id] = inst + return inst + + def release(self, instance_id: str, recycle: bool = False) -> None: + if instance_id in self.instances: + if not recycle: + del self.instances[instance_id] + + def reap(self, run_id: str) -> int: + self.reaped_runs.add(run_id) + to_del = [iid for iid, inst in self.instances.items() if inst.run_id == run_id] + for iid in to_del: + del self.instances[iid] + return len(to_del) + + def build_runtime( + self, instance: RawSandboxInstance, data_plane: str = "mock" + ) -> tuple[Any, Optional[Any]]: + from ..runtime.mock import MockRuntimeHook + return MockRuntimeHook(), None diff --git a/clients/python/src/ate_env/backend/substrate.py b/clients/python/src/ate_env/backend/substrate.py new file mode 100644 index 00000000..6e2e76a7 --- /dev/null +++ b/clients/python/src/ate_env/backend/substrate.py @@ -0,0 +1,451 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import asyncio +import logging +import os +import threading +import time +import uuid +from typing import Any, Dict, List, Mapping, Optional + +from ..client import DEFAULT_ATESPACE, Client +from ..errors import FailedPreconditionError, InfrastructureError, NotFoundError, PreflightError, SandboxStartError +from ..loop_runner import LoopRunner +from ..types import DataPlaneEndpoint, EnvironmentSpec, RawSandboxInstance +from .base import BackendDriver + +logger = logging.getLogger("ate_env.backend.substrate") + + +class SubstrateBackendDriver(BackendDriver): + """Agent Substrate Backend Driver. + + Manages environment lifecycle directly against the ate-env-api gRPC service + via ate_env.Client: creating environments, suspending for pre-warmed pools, + and deleting on release. + """ + + def __init__( + self, + api_endpoint: Optional[str] = None, + atespace: str = DEFAULT_ATESPACE, + worker_family: str = "c2", + auth_token: Optional[str] = None, + *, + client: Optional[Client] = None, + image_templates: Optional[Mapping[str, str]] = None, + default_template: Optional[str] = None, + labels: Optional[Mapping[str, str]] = None, + acquire_timeout_s: float = 180.0, + ready_timeout_s: float = 180.0, + loop_runner: Optional[LoopRunner] = None, + ): + target = ( + api_endpoint + or os.environ.get("ATE_ENV_API_TARGET") + or os.environ.get("SUBSTRATE_API_ENDPOINT") + or "127.0.0.1:7777" + ) + self.api_endpoint = target + self.atespace = atespace or DEFAULT_ATESPACE + self.worker_family = worker_family + self.auth_token = auth_token + self.image_templates = dict(image_templates or {}) + self.default_template = default_template + self.labels = dict(labels or {}) + self.acquire_timeout_s = acquire_timeout_s + self.ready_timeout_s = ready_timeout_s + + self._loop_runner = loop_runner + self._owns_loop_runner = loop_runner is None + + if client is not None: + self._client = client + self._owns_client = False + else: + self._client = None + self._owns_client = True + + self._lock = threading.Lock() + self._owned_actors: Dict[str, RawSandboxInstance] = {} + self._warm_paused_pool: Dict[str, List[str]] = {} + self._background_tasks: set[asyncio.Task[Any]] = set() + + def _get_loop_runner(self) -> LoopRunner: + with self._lock: + if self._loop_runner is None: + self._loop_runner = LoopRunner(name="ate-substrate-driver-sync") + self._owns_loop_runner = True + return self._loop_runner + + @property + def client(self) -> Client: + if self._client is None: + self._client = Client(self.api_endpoint, auth_token=self.auth_token) + return self._client + + def _sync(self, coro: Any) -> Any: + return self._get_loop_runner().run(coro) + + async def close_async(self) -> None: + if self._owns_client and self._client: + await self._client.close() + self._client = None + if self._owns_loop_runner and self._loop_runner: + self._loop_runner.stop() + self._loop_runner = None + + def close(self) -> None: + if self._loop_runner: + if self._owns_client and self._client: + self._loop_runner.run(self._client.close()) + self._client = None + if self._owns_loop_runner: + self._loop_runner.stop() + self._loop_runner = None + elif self._owns_client and self._client: + self._sync(self.close_async()) + + async def preflight_async(self) -> None: + """Validate ate-env-api connectivity.""" + logger.info( + "Running Substrate preflight checks (Atespace: %s, Endpoint: %s)...", + self.atespace, + self.api_endpoint, + ) + try: + await self.client.get(f"preflight-probe-{uuid.uuid4().hex[:6]}", atespace=self.atespace) + except NotFoundError: + # NotFoundError confirms the gRPC server is reachable and active. + logger.debug("Preflight probe successfully reached ate-env-api") + except Exception as e: + raise PreflightError( + f"Substrate preflight check failed to reach ate-env-api at {self.api_endpoint}: {e}" + ) from e + + def preflight(self) -> None: + self._sync(self.preflight_async()) + + def ensure_template(self, env: EnvironmentSpec | str) -> str: + """Resolve ActorTemplate name from template ID or task image. + + Note: ResourceLimits, runtime_bundle, snapshot_storage_uri, and placement + fields in EnvironmentSpec are reserved for future scheduling extensions. + Current Substrate environments are provisioned from pre-registered ActorTemplates. + """ + if isinstance(env, str): + image = env + if image in self.image_templates: + return self.image_templates[image] + return self.default_template or "default-template" + + image = env.image or "" + mapped = self.image_templates.get(image) + if mapped: + return mapped + if self.default_template: + return self.default_template + if isinstance(env, EnvironmentSpec): + return env.template_key() + return "default-template" + + def _make_instance(self, actor_id: str, template_id: str, run_id: str) -> RawSandboxInstance: + planes = { + "ate_env": DataPlaneEndpoint( + protocol="grpc", + address=self.api_endpoint, + headers={"x-env-id": actor_id, "x-env-atespace": self.atespace}, + ), + } + metadata = { + "created_at": time.time(), + "template": template_id, + "atespace": self.atespace, + } + if self.labels: + metadata["labels"] = dict(self.labels) + return RawSandboxInstance( + instance_id=actor_id, + endpoint=self.api_endpoint, + template_id=template_id, + run_id=run_id, + metadata=metadata, + data_planes=planes, + ) + + async def acquire_async( + self, template_id: str, run_id: str, timeout_s: float = 180.0 + ) -> RawSandboxInstance: + """Acquire a live sandbox instance within timeout_s.""" + to_s = timeout_s or self.acquire_timeout_s + try: + return await asyncio.wait_for( + self._acquire_internal(template_id, run_id), timeout=to_s + ) + except asyncio.TimeoutError as e: + raise SandboxStartError( + f"Timed out acquiring sandbox for template {template_id} after {to_s:g}s" + ) from e + + async def _acquire_internal(self, template_id: str, run_id: str) -> RawSandboxInstance: + """Acquire or claim an environment. + + Note on readiness: Substrate resumes suspended environments implicitly on first + guest traffic, so ready_timeout_s does not measure readiness here. + """ + while True: + actor_id: Optional[str] = None + with self._lock: + warm_list = self._warm_paused_pool.get(template_id) + if warm_list: + actor_id = warm_list.pop() + + if actor_id: + logger.info("Claimed warm paused actor %s for template %s", actor_id, template_id) + # Verify warm actor still exists on the server + try: + await asyncio.wait_for( + self.client.get(actor_id, atespace=self.atespace), + timeout=min(self.ready_timeout_s, 15.0), + ) + except NotFoundError: + logger.warning("Warm actor %s not found on server, dropping from pool", actor_id) + with self._lock: + self._owned_actors.pop(actor_id, None) + continue + except Exception as e: + # Put back on unexpected error and raise + with self._lock: + self._warm_paused_pool.setdefault(template_id, []).append(actor_id) + raise SandboxStartError(f"Error checking warm actor {actor_id}: {e}") from e + break + else: + actor_id = f"sb-{uuid.uuid4().hex[:16]}" + logger.info( + "Creating environment %s (template=%s, atespace=%s)...", + actor_id, + template_id, + self.atespace, + ) + last_err: Optional[Exception] = None + for attempt in range(4): + try: + await self.client.create( + actor_id, + template_name=template_id, + atespace=self.atespace, + ) + last_err = None + break + except (NotFoundError, FailedPreconditionError) as e: + raise SandboxStartError( + f"ActorTemplate '{template_id}' invalid or does not exist on Substrate server: {e}" + ) from e + except Exception as e: + last_err = e + if "already exists" in str(e).lower(): + last_err = None + break + if attempt < 3: + await asyncio.sleep(0.05 * (attempt + 1)) + if last_err is not None: + raise SandboxStartError(f"Failed to create environment {actor_id}: {last_err}") from last_err + break + + instance = self._make_instance(actor_id, template_id, run_id) + with self._lock: + self._owned_actors[actor_id] = instance + return instance + + def acquire( + self, template_id: str, run_id: str, timeout_s: float = 180.0 + ) -> RawSandboxInstance: + return self._sync(self.acquire_async(template_id, run_id, timeout_s=timeout_s)) + + async def warm_pool_async( + self, template_id: str, replicas: int, wait: bool = True, run_id: Optional[str] = None + ) -> None: + """Pre-warm a pool of actors by creating and suspending them.""" + logger.info("Sizing warm pool for template %s to %d paused actors...", template_id, replicas) + with self._lock: + current = list(self._warm_paused_pool.get(template_id, [])) + needed = replicas - len(current) + if needed <= 0: + return + + sem = asyncio.Semaphore(16) + effective_run_id = run_id or "warm-pool" + + async def _warm_one() -> Optional[str]: + actor_id = f"warm-{uuid.uuid4().hex[:10]}" + async with sem: + try: + await self.client.create( + actor_id, + template_name=template_id, + atespace=self.atespace, + ) + if wait: + await self.client.suspend(actor_id, atespace=self.atespace) + else: + task = asyncio.create_task(self.client.suspend(actor_id, atespace=self.atespace)) + self._background_tasks.add(task) + def _done(t: asyncio.Task[Any]) -> None: + self._background_tasks.discard(t) + if not t.cancelled() and t.exception(): + logger.warning("Background suspend for warm actor %s failed: %s", actor_id, t.exception()) + task.add_done_callback(_done) + return actor_id + except (NotFoundError, FailedPreconditionError) as e: + raise SandboxStartError( + f"ActorTemplate '{template_id}' invalid or does not exist on Substrate server: {e}" + ) from e + except Exception as e: + logger.warning("Failed to provision warm actor %s: %s", actor_id, e) + return None + + tasks = [_warm_one() for _ in range(needed)] + results = await asyncio.gather(*tasks, return_exceptions=True) + + for res in results: + if isinstance(res, Exception): + if isinstance(res, SandboxStartError): + raise res + logger.warning("Error provisioning warm actor: %s", res) + elif res is not None: + aid = res + with self._lock: + self._warm_paused_pool.setdefault(template_id, []).append(aid) + self._owned_actors[aid] = self._make_instance( + aid, template_id, run_id=effective_run_id + ) + + def warm_pool( + self, template_id: str, replicas: int, wait: bool = True, run_id: Optional[str] = None + ) -> None: + self._sync(self.warm_pool_async(template_id, replicas, wait=wait, run_id=run_id)) + + async def unwarm_pool_async(self, template_id: str) -> None: + """Drain and delete paused actors in the warm pool.""" + logger.info("Unwarming pool for template %s...", template_id) + with self._lock: + to_delete = self._warm_paused_pool.pop(template_id, []) + + sem = asyncio.Semaphore(16) + + async def _delete_one(actor_id: str) -> None: + async with sem: + try: + await self.client.delete(actor_id, atespace=self.atespace) + except Exception as e: + logger.warning("Failed deleting warm actor %s: %s", actor_id, e) + with self._lock: + self._owned_actors.pop(actor_id, None) + + await asyncio.gather(*[_delete_one(aid) for aid in to_delete], return_exceptions=True) + + def unwarm_pool(self, template_id: str) -> None: + self._sync(self.unwarm_pool_async(template_id)) + + async def release_async(self, instance_id: str, recycle: bool = False) -> None: + """Release actor: deletes environment, or suspends if recycling. + + Note: Recycling suspends the actor and returns it to the warm pool + without resetting guest filesystem or processes. + """ + with self._lock: + instance = self._owned_actors.pop(instance_id, None) + + if recycle and instance: + try: + await self.client.suspend(instance_id, atespace=self.atespace) + with self._lock: + self._warm_paused_pool.setdefault(instance.template_id, []).append(instance_id) + self._owned_actors[instance_id] = instance + logger.info("Recycled actor %s to paused warm pool", instance_id) + return + except Exception as e: + logger.warning("Failed to suspend actor %s for recycle, deleting: %s", instance_id, e) + + try: + await self.client.delete(instance_id, atespace=self.atespace) + logger.info("Deleted environment %s", instance_id) + except Exception as e: + logger.warning("Failed to delete environment %s: %s", instance_id, e) + + def release(self, instance_id: str, recycle: bool = False) -> None: + self._sync(self.release_async(instance_id, recycle=recycle)) + + async def reap_async(self, run_id: str) -> int: + """Force delete all environments tagged with this run_id and all warm pool actors.""" + with self._lock: + warm_actors: List[str] = [] + for tid, actors in list(self._warm_paused_pool.items()): + warm_actors.extend(actors) + self._warm_paused_pool.clear() + + to_reap = [ + aid for aid, inst in self._owned_actors.items() + if inst.run_id == run_id or aid in warm_actors or inst.run_id == "warm-pool" + ] + to_reap = list(dict.fromkeys(to_reap + warm_actors)) + + sem = asyncio.Semaphore(16) + failed_actors: List[str] = [] + reaped_count = 0 + reap_lock = asyncio.Lock() + + async def _reap_one(aid: str) -> None: + nonlocal reaped_count + async with sem: + try: + await self.client.delete(aid, atespace=self.atespace) + async with reap_lock: + reaped_count += 1 + except Exception as e: + logger.warning("Failed deleting actor %s during reap: %s", aid, e) + async with reap_lock: + failed_actors.append(aid) + with self._lock: + self._owned_actors.pop(aid, None) + + await asyncio.gather(*[_reap_one(aid) for aid in to_reap], return_exceptions=True) + + if failed_actors: + raise InfrastructureError( + f"Failed to delete {len(failed_actors)} environments during reap: {failed_actors}" + ) + + logger.info("Reaped %d Substrate environments for run %s", reaped_count, run_id) + return reaped_count + + def reap(self, run_id: str) -> int: + return self._sync(self.reap_async(run_id)) + + def build_runtime( + self, instance: RawSandboxInstance, data_plane: str = "ate_env" + ) -> tuple[Any, Optional[Any]]: + from ..runtime.substrate_env_client import SubstrateEnvClientRuntime + + ep = instance.data_planes.get(data_plane) or instance.data_planes.get("ate_env") + return SubstrateEnvClientRuntime( + env_id=instance.instance_id, + endpoint=ep.address if ep else instance.endpoint, + atespace=self.atespace, + client=self.client, + loop_runner=self._get_loop_runner(), + ), ep diff --git a/clients/python/src/ate_env/client.py b/clients/python/src/ate_env/client.py index 373ed4e1..ed1797a9 100644 --- a/clients/python/src/ate_env/client.py +++ b/clients/python/src/ate_env/client.py @@ -37,6 +37,52 @@ def _normalize_target(endpoint: str) -> str: return target +import collections + + +class _ClientCallDetails( + collections.namedtuple( + "_ClientCallDetails", + ["method", "timeout", "metadata", "credentials", "wait_for_ready"], + ), + grpc.aio.ClientCallDetails, +): + pass + + +class _AuthInterceptor( + grpc.aio.UnaryUnaryClientInterceptor, + grpc.aio.UnaryStreamClientInterceptor, + grpc.aio.StreamUnaryClientInterceptor, + grpc.aio.StreamStreamClientInterceptor, +): + def __init__(self, token: str): + self._token = token + + def _inject(self, details: grpc.aio.ClientCallDetails) -> _ClientCallDetails: + md = list(details.metadata or []) + md.append(("authorization", f"Bearer {self._token}")) + return _ClientCallDetails( + details.method, + details.timeout, + md, + details.credentials, + details.wait_for_ready, + ) + + async def intercept_unary_unary(self, continuation, client_call_details, request): + return await continuation(self._inject(client_call_details), request) + + async def intercept_unary_stream(self, continuation, client_call_details, request): + return await continuation(self._inject(client_call_details), request) + + async def intercept_stream_unary(self, continuation, client_call_details, request_iterator): + return await continuation(self._inject(client_call_details), request_iterator) + + async def intercept_stream_stream(self, continuation, client_call_details, request_iterator): + return await continuation(self._inject(client_call_details), request_iterator) + + class Client: """Manages environments against an ate-env-api endpoint over gRPC. @@ -57,6 +103,7 @@ def __init__( endpoint: str = "localhost:7777", *, channel: grpc.aio.Channel | None = None, + auth_token: str | None = None, ): """Create a client for the ate-env-api service at endpoint. @@ -64,13 +111,17 @@ def __init__( client. Pass channel to supply a caller-owned grpc.aio.Channel instead; the client then never closes it. """ + self.endpoint = endpoint if channel is not None: self._channel = channel self._owns_channel = False else: if not endpoint: raise ValueError("ate_env: endpoint is required when channel is not given") - self._channel = grpc.aio.insecure_channel(_normalize_target(endpoint)) + interceptors = [_AuthInterceptor(auth_token)] if auth_token else None + self._channel = grpc.aio.insecure_channel( + _normalize_target(endpoint), interceptors=interceptors + ) self._owns_channel = True self._environments = env_pb2_grpc.EnvironmentServiceStub(self._channel) self._processes = guest_pb2_grpc.ProcessServiceStub(self._channel) diff --git a/clients/python/src/ate_env/config.py b/clients/python/src/ate_env/config.py new file mode 100644 index 00000000..a5404b8c --- /dev/null +++ b/clients/python/src/ate_env/config.py @@ -0,0 +1,173 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from dataclasses import MISSING, dataclass, field, fields +import re +from typing import Any, Dict, List, Literal, Optional + +BackendType = Literal["substrate", "mock"] +StrategyType = Literal["none", "naive", "sliding", "pipelined"] + +# Atespaces (Substrate) are DNS-1123 labels. +_DNS1123_LABEL = re.compile(r"^[a-z0-9]([-a-z0-9]{0,61}[a-z0-9])?$") + +ALLOWED_BACKENDS = {"substrate", "mock"} +ALLOWED_STRATEGIES = {"none", "naive", "sliding", "pipelined"} +ALLOWED_DATA_PLANES = {"ate_env"} + + +class ValidationError(ValueError): + """Configuration validation error.""" + + +class SecretStr: + """Wraps sensitive strings to prevent accidental leakage in repr or str.""" + + def __init__(self, secret_value: str): + if not isinstance(secret_value, str): + raise TypeError("SecretStr value must be a str") + self._secret_value = secret_value + + def get_secret_value(self) -> str: + return self._secret_value + + def __repr__(self) -> str: + return "SecretStr('**********')" + + def __str__(self) -> str: + return "**********" + + def __eq__(self, other: Any) -> bool: + if isinstance(other, SecretStr): + return self._secret_value == other._secret_value + return False + + def __hash__(self) -> int: + return hash(self._secret_value) + + +@dataclass(frozen=True) +class FleetConfig: + """Configuration for SandboxFleet and associated BackendDriver. + + Unknown fields are rejected so misspelled or unsupported options + fail loudly instead of being silently ignored. + """ + + backend: BackendType = "substrate" + endpoint: str = "127.0.0.1:7777" + data_plane: Literal["ate_env"] = "ate_env" + tenancy: str = "ate-env" + strategy: StrategyType = "sliding" + template: Optional[str] = None + image_templates: Dict[str, str] = field(default_factory=dict) + + # Sizing & Rollout parameters + batch_size: int = 8 + max_concurrent: int = 64 + max_warmpool_replicas: int = 4 + window_size: Optional[int] = None + + # Hardware & placement (reserved for template placement and future scheduler extensions) + worker_family: Optional[str] = "c2" + node_selector: Dict[str, str] = field(default_factory=dict) + tolerations: List[Dict[str, Any]] = field(default_factory=list) + labels: Dict[str, str] = field(default_factory=dict) + + # Timeouts & Breakers + acquire_timeout_s: float = 180.0 + ready_timeout_s: float = 180.0 + step_timeout_s: float = 120.0 + auth_token: Optional[SecretStr] = None + + def __init__(self, **kwargs: Any): + valid_fields = {f.name: f for f in fields(self)} + unknown = set(kwargs.keys()) - set(valid_fields.keys()) + if unknown: + raise ValidationError(f"Unknown fields for FleetConfig: {sorted(unknown)}") + for name, f in valid_fields.items(): + if name in kwargs: + val = kwargs[name] + elif f.default is not MISSING: + val = f.default + elif f.default_factory is not MISSING: + val = f.default_factory() + else: + val = None + object.__setattr__(self, name, val) + self.__post_init__() + + def __post_init__(self) -> None: + if self.backend not in ALLOWED_BACKENDS: + raise ValidationError(f"Invalid backend {self.backend!r}, must be one of {ALLOWED_BACKENDS}") + + if self.data_plane not in ALLOWED_DATA_PLANES: + raise ValidationError(f"Invalid data_plane {self.data_plane!r}, must be one of {ALLOWED_DATA_PLANES}") + + if not isinstance(self.tenancy, str) or not _DNS1123_LABEL.match(self.tenancy): + raise ValidationError(f"Invalid tenancy {self.tenancy!r}: must match DNS-1123 label regex") + + if self.strategy not in ALLOWED_STRATEGIES: + raise ValidationError(f"Invalid strategy {self.strategy!r}, must be one of {ALLOWED_STRATEGIES}") + + object.__setattr__(self, "image_templates", dict(self.image_templates)) + + if not isinstance(self.batch_size, int) or self.batch_size < 1: + raise ValidationError(f"batch_size must be >= 1, got {self.batch_size}") + + if not isinstance(self.max_concurrent, int) or self.max_concurrent < 1: + raise ValidationError(f"max_concurrent must be >= 1, got {self.max_concurrent}") + + if not isinstance(self.max_warmpool_replicas, int) or self.max_warmpool_replicas < 0: + raise ValidationError(f"max_warmpool_replicas must be >= 0, got {self.max_warmpool_replicas}") + + if self.window_size is not None and (not isinstance(self.window_size, int) or self.window_size < 1): + raise ValidationError(f"window_size must be >= 1, got {self.window_size}") + + object.__setattr__(self, "node_selector", dict(self.node_selector)) + object.__setattr__(self, "tolerations", list(self.tolerations)) + object.__setattr__(self, "labels", dict(self.labels)) + + if self.acquire_timeout_s <= 0: + raise ValidationError(f"acquire_timeout_s must be > 0, got {self.acquire_timeout_s}") + object.__setattr__(self, "acquire_timeout_s", float(self.acquire_timeout_s)) + + if self.ready_timeout_s <= 0: + raise ValidationError(f"ready_timeout_s must be > 0, got {self.ready_timeout_s}") + object.__setattr__(self, "ready_timeout_s", float(self.ready_timeout_s)) + + if self.step_timeout_s <= 0: + raise ValidationError(f"step_timeout_s must be > 0, got {self.step_timeout_s}") + object.__setattr__(self, "step_timeout_s", float(self.step_timeout_s)) + + if self.auth_token is not None and not isinstance(self.auth_token, SecretStr): + object.__setattr__(self, "auth_token", SecretStr(str(self.auth_token))) + + def to_dict(self) -> Dict[str, Any]: + """Serialize configuration to dict, keeping auth_token as SecretStr for safe round-tripping.""" + res: Dict[str, Any] = {} + for f in fields(self): + val = getattr(self, f.name) + res[f.name] = val + return res + + model_dump = to_dict + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "FleetConfig": + return cls(**data) + + model_validate = from_dict diff --git a/clients/python/src/ate_env/errors.py b/clients/python/src/ate_env/errors.py index 3b8afeaa..cb5a3084 100644 --- a/clients/python/src/ate_env/errors.py +++ b/clients/python/src/ate_env/errors.py @@ -27,6 +27,14 @@ "ProcessExitedError", "RpcError", "map_rpc_error", + "InfrastructureError", + "SandboxUnavailableError", + "SandboxProtocolError", + "CommandStartError", + "SandboxStartError", + "CommandExecutionError", + "CommandTimeoutError", + "PreflightError", ] @@ -82,3 +90,106 @@ def map_rpc_error(err: BaseException) -> BaseException: return ProcessExitedError(details) return FailedPreconditionError(details) return RpcError(details, code) + + +# --------------------------------------------------------------------------- +# Infrastructure and Execution Errors +# +# Contract: ExecResult is returned only when the guest ran the command and it +# either exited or was killed at timeout_s (timed_out=True). Every other outcome +# raises an InfrastructureError. Transport and HTTP/gRPC statuses are never +# encoded as exit codes. +# --------------------------------------------------------------------------- + + +class InfrastructureError(EnvError): + """The SDK could not obtain a verdict from the sandbox for this call. + + Never score these as agent failures (e.g. reward 0). Retry the rollout on a + fresh sandbox when ``retryable`` is True, otherwise mask it. + + Attributes: + sandbox_id: Sandbox the call targeted, when known. + status: Transport status, e.g. "grpc UNAVAILABLE". + retryable: Whether retrying on a fresh sandbox may succeed. + """ + + default_retryable: bool = True + + def __init__( + self, + message: str, + *, + sandbox_id: str | None = None, + status: str | None = None, + retryable: bool | None = None, + ): + super().__init__(message) + self.sandbox_id = sandbox_id + self.status = status + self.retryable = self.default_retryable if retryable is None else retryable + + def __reduce__(self): + return ( + _rebuild_infra_error, + (type(self), str(self), self.sandbox_id, self.status, self.retryable), + ) + + +def _rebuild_infra_error( + cls: type, + message: str, + sandbox_id: str | None, + status: str | None, + retryable: bool, +) -> InfrastructureError: + return cls(message, sandbox_id=sandbox_id, status=status, retryable=retryable) + + +class SandboxUnavailableError(InfrastructureError): + """Sandbox or data plane unreachable, gone, overloaded, or failed mid-call. + + Examples: connection refused/reset, gRPC UNAVAILABLE/UNKNOWN/INTERNAL/ABORTED, + environment not found during lifecycle operations. + """ + + default_retryable = True + + +class SandboxProtocolError(InfrastructureError): + """Request rejected or response malformed: a contract or configuration bug. + + Examples: gRPC INVALID_ARGUMENT, PERMISSION_DENIED. + """ + + default_retryable = False + + +class CommandStartError(InfrastructureError): + """The guest could not start the command (missing executable or cwd inside guest).""" + + default_retryable = False + + +class SandboxStartError(InfrastructureError): + """Raised when sandbox instantiation, template lookup, or boot fails.""" + + default_retryable = True + + +class PreflightError(InfrastructureError): + """Raised when cluster or API connectivity verification fails.""" + + default_retryable = False + + +class CommandExecutionError(EnvError): + """Raised when a non-zero exit code occurs in strict mode.""" + + def __init__(self, message: str, result: object = None): + super().__init__(message) + self.result = result + + +class CommandTimeoutError(CommandExecutionError, TimeoutError): + """Raised in strict mode (check=True) when a command hit its deadline.""" diff --git a/clients/python/src/ate_env/exceptions.py b/clients/python/src/ate_env/exceptions.py new file mode 100644 index 00000000..a4932a21 --- /dev/null +++ b/clients/python/src/ate_env/exceptions.py @@ -0,0 +1,62 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Deprecated: Use ate_env.errors instead. + +This module re-exports error classes from ate_env.errors for backwards compatibility. +""" + +from __future__ import annotations + +import builtins + +from .errors import ( + CommandExecutionError, + CommandStartError, + CommandTimeoutError, + EnvError, + InfrastructureError, + InvalidArgumentError, + NotFoundError, + PermissionDeniedError, + PreflightError, + ProcessExitedError, + RpcError, + SandboxProtocolError, + SandboxStartError, + SandboxUnavailableError, +) + +# Backwards compatibility aliases +SandboxError = EnvError +TimeoutError = builtins.TimeoutError + +__all__ = [ + "SandboxError", + "TimeoutError", + "PreflightError", + "SandboxStartError", + "CommandExecutionError", + "CommandTimeoutError", + "InfrastructureError", + "SandboxUnavailableError", + "SandboxProtocolError", + "CommandStartError", + "EnvError", + "NotFoundError", + "InvalidArgumentError", + "PermissionDeniedError", + "ProcessExitedError", + "RpcError", +] diff --git a/clients/python/src/ate_env/fleet.py b/clients/python/src/ate_env/fleet.py new file mode 100644 index 00000000..6e7e7bdf --- /dev/null +++ b/clients/python/src/ate_env/fleet.py @@ -0,0 +1,296 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import collections +import logging +import threading +import uuid +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + +from .backend.base import BackendDriver +from .backend.mock import MockBackendDriver +from .backend.substrate import SubstrateBackendDriver +from .config import FleetConfig +from .errors import SandboxStartError +from .handle import SandboxHandle +from .runtime.base import RuntimeGuestHook +from .runtime.mock import MockRuntimeHook +from .strategies import STRATEGIES +from .types import ( + DataPlaneEndpoint, + EnvironmentSpec, + FleetPlan, + PlacementSpec, + PlanEntry, + RawSandboxInstance, + Task, +) + +logger = logging.getLogger("ate_env.fleet") + + +class SandboxFleet: + """Unified Sandbox Fleet Orchestrator. + + Coordinates preflight, planning, warm-pool sizing, lifecycle management, + task acquisition, and execution strategies across any backend driver. + """ + + def __init__(self, config: Optional[FleetConfig] = None, driver: Optional[BackendDriver] = None): + self.config = config or FleetConfig() + self.run_id = uuid.uuid4().hex[:12] + self._tasks: List[Task] = [] + self._image_to_template: Dict[str, str] = {} + self._plan: Optional[FleetPlan] = None + self._lock = threading.Lock() + + # Initialize Backend Driver + if driver is not None: + self.backend = driver + elif self.config.backend == "mock": + self.backend = MockBackendDriver() + elif self.config.backend == "substrate": + token = self.config.auth_token + self.backend = SubstrateBackendDriver( + api_endpoint=self.config.endpoint, + atespace=self.config.tenancy, + worker_family=self.config.worker_family or "c2", + auth_token=token.get_secret_value() if token else None, + image_templates=self.config.image_templates, + default_template=self.config.template if self.config.template is not None else "default-template", + labels=self.config.labels, + acquire_timeout_s=self.config.acquire_timeout_s, + ready_timeout_s=self.config.ready_timeout_s, + ) + else: + raise NotImplementedError(f"Backend '{self.config.backend}' is not yet supported or requires extras") + + def __enter__(self) -> "SandboxFleet": + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + try: + self.close() + except Exception: + if exc_type is None: + raise + + def close(self) -> None: + teardown_err: Optional[Exception] = None + try: + self.teardown() + except Exception as e: + teardown_err = e + finally: + if hasattr(self.backend, "close"): + try: + self.backend.close() + except Exception as be: + if teardown_err is None: + teardown_err = be + if teardown_err is not None: + raise teardown_err + + @property + def tasks(self) -> List[Task]: + return self._tasks + + def load_tasks(self, tasks: List[Any]) -> None: + """Load tasks into the fleet (supports list of Task or dicts).""" + normalized: List[Task] = [] + for t in tasks: + if isinstance(t, Task): + normalized.append(t) + elif isinstance(t, dict): + normalized.append(Task( + id=str(t.get("task_id") or t.get("id")), + image=t["image"], + metadata=t, + template=t.get("template"), + )) + else: + raise ValueError(f"Unsupported task item type: {type(t)}") + self._tasks = normalized + logger.info("Fleet loaded %d tasks (%d unique images)", + len(self._tasks), len(self.image_counts())) + + def image_counts(self) -> Dict[str, int]: + """Return task count per unique image.""" + counts: Dict[str, int] = collections.defaultdict(int) + for t in self._tasks: + counts[t.image] += 1 + return dict(counts) + + def preflight(self) -> None: + """Run backend connectivity, permissions, and capacity checks.""" + self.backend.preflight() + + def plan(self) -> FleetPlan: + """Compute provisioning sizing for all tasks.""" + entries: List[PlanEntry] = [] + counts = self.image_counts() + for img, task_count in counts.items(): + template_id = self._ensure_template_for_image(img) + # Size warm pool to min(task_count, max_warmpool_replicas) + replicas = min(task_count, self.config.max_warmpool_replicas) + entries.append(PlanEntry( + image=img, + template_id=template_id, + replicas=replicas, + tasks=task_count, + )) + self._plan = FleetPlan(entries) + return self._plan + + def setup(self, tasks: Optional[List[Any]] = None) -> None: + """Run preflight, plan, and pre-warm all planned images.""" + if tasks is not None: + self.load_tasks(tasks) + self.preflight() + plan = self.plan() + for entry in plan.entries: + if entry.replicas > 0: + self.backend.warm_pool(entry.template_id, entry.replicas, wait=True) + + def warm_images(self, images: List[str], replicas: Optional[int] = None, wait: bool = True) -> None: + """Warm up pools for specific images (used by windowed strategies).""" + rep = replicas if replicas is not None else self.config.max_warmpool_replicas + for img in images: + template_id = self._ensure_template_for_image(img) + self.backend.warm_pool(template_id, rep, wait=wait) + + def unwarm_image(self, image: str) -> None: + """Drain warm pool for a specific image.""" + template_id = self._image_to_template.get(image) + if template_id: + self.backend.unwarm_pool(template_id) + + def acquire( + self, + task: Task | str | Sequence[Task | str], + timeout_s: Optional[float] = None, + ) -> SandboxHandle | List[SandboxHandle]: + """Acquire one or more live sandboxes bound to tasks. + + If a list or tuple of tasks is given, returns a list of SandboxHandle. + Otherwise returns a single SandboxHandle. + """ + if isinstance(task, (list, tuple)): + return [self._acquire_one(t, timeout_s=timeout_s) for t in task] + return self._acquire_one(task, timeout_s=timeout_s) + + def _acquire_one( + self, + task: Task | str, + timeout_s: Optional[float] = None, + ) -> SandboxHandle: + task_obj: Task + if isinstance(task, str): + found = next((t for t in self._tasks if t.id == task), None) + if found: + task_obj = found + else: + task_obj = Task(id=task, image="default") + else: + task_obj = task + + if task_obj.template: + template_id = task_obj.template + else: + template_id = self._ensure_template_for_image(task_obj.image) + + to_s = timeout_s or self.config.acquire_timeout_s + raw_inst = self.backend.acquire(template_id, self.run_id, timeout_s=to_s) + + try: + runtime, data_plane = self._build_runtime(raw_inst) + except Exception: + try: + self.backend.release(raw_inst.instance_id, recycle=False) + except Exception: + logger.warning( + "Failed to release sandbox %s after runtime setup error", + raw_inst.instance_id, + exc_info=True, + ) + raise + + return SandboxHandle( + sandbox_id=raw_inst.instance_id, + endpoint=data_plane.address if data_plane else raw_inst.endpoint, + task=task_obj, + run_id=self.run_id, + backend=self.backend, + runtime=runtime, + data_plane=data_plane, + ) + + def _build_runtime( + self, raw_inst: RawSandboxInstance + ) -> Tuple[RuntimeGuestHook, Optional[DataPlaneEndpoint]]: + """Build the RuntimeGuestHook from this sandbox's own data-plane coordinates.""" + if hasattr(self.backend, "build_runtime"): + return self.backend.build_runtime(raw_inst, data_plane=self.config.data_plane) + + from .runtime.substrate_env_client import SubstrateEnvClientRuntime + + name = self.config.data_plane + ep = raw_inst.data_planes.get(name) or raw_inst.data_planes.get("ate_env") + client = getattr(self.backend, "client", None) + loop_runner = getattr(self.backend, "_loop_runner", None) + return SubstrateEnvClientRuntime( + env_id=raw_inst.instance_id, + endpoint=ep.address if ep else raw_inst.endpoint, + atespace=self.config.tenancy or "ate-env", + client=client, + loop_runner=loop_runner, + ), ep + + def release(self, handle: SandboxHandle | Sequence[SandboxHandle]) -> None: + """Release sandbox instance(s) back to backend and close runtime(s).""" + if isinstance(handle, (list, tuple)): + for h in handle: + h.release() + else: + handle.release() + + def teardown(self) -> None: + """Teardown all resources provisioned by this fleet's run_id.""" + self.backend.reap(self.run_id) + + def run(self, process_fn: Callable[[Task, SandboxHandle], Any], + concurrency: Optional[int] = None) -> List[Any]: + """Execute all tasks using the configured strategy.""" + strat_fn = STRATEGIES.get(self.config.strategy) + if not strat_fn: + raise ValueError(f"Unknown strategy: {self.config.strategy}") + c = concurrency or min(self.config.max_concurrent, len(self._tasks)) + return strat_fn(self, process_fn, max(1, c)) + + def _ensure_template_for_image(self, image: str) -> str: + with self._lock: + if image in self._image_to_template: + return self._image_to_template[image] + if image in self.config.image_templates: + tmpl = self.config.image_templates[image] + self._image_to_template[image] = tmpl + return tmpl + if self.config.template: + self._image_to_template[image] = self.config.template + return self.config.template + tmpl = "default-template" + self._image_to_template[image] = tmpl + return tmpl diff --git a/clients/python/src/ate_env/handle.py b/clients/python/src/ate_env/handle.py new file mode 100644 index 00000000..1373743a --- /dev/null +++ b/clients/python/src/ate_env/handle.py @@ -0,0 +1,146 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +from .backend.base import BackendDriver +from .errors import CommandExecutionError, CommandTimeoutError +from .runtime.base import RuntimeGuestHook +from .types import DataPlaneEndpoint, ExecResult, Task + + +@dataclass +class SandboxHandle: + """Unified handle to a claimed sandbox bound to a specific task. + + Provides high-level command execution, file operations, + and release/recycling through the fleet. + + ``data_plane`` holds this sandbox's own connection coordinates (address + plus routing headers) for harnesses that talk to the sandbox + directly; it is None for the mock backend. + """ + sandbox_id: str + endpoint: str + task: Task + run_id: str + backend: BackendDriver = field(repr=False) + runtime: RuntimeGuestHook = field(repr=False) + data_plane: Optional[DataPlaneEndpoint] = None + + def exec( + self, + command: str | List[str], + cwd: str = "", + env: Optional[Dict[str, str]] = None, + timeout_s: float = 120.0, + check: bool = True, + ) -> ExecResult: + """Execute command inside the sandbox. + + If check=True, raises CommandExecutionError on non-zero exit code + (CommandTimeoutError if the command hit its deadline). + Infrastructure failures always raise InfrastructureError, regardless + of ``check``. + Returns ExecResult. + """ + res = self.runtime.exec(command, cwd=cwd, env=env, timeout_s=timeout_s) + if check and res.timed_out: + raise CommandTimeoutError( + f"Command '{command}' timed out after {timeout_s:g}s:\n" + f"STDOUT: {res.stdout}\nSTDERR: {res.stderr}", + result=res, + ) + if check and res.exit_code != 0: + raise CommandExecutionError( + f"Command '{command}' failed with exit code {res.exit_code}:\n" + f"STDOUT: {res.stdout}\nSTDERR: {res.stderr}", + result=res, + ) + return res + + async def exec_async( + self, + command: str | List[str], + cwd: str = "", + env: Optional[Dict[str, str]] = None, + timeout_s: float = 120.0, + check: bool = True, + ) -> ExecResult: + """Asynchronously execute command inside the sandbox.""" + res = await self.runtime.exec_async(command, cwd=cwd, env=env, timeout_s=timeout_s) + if check and res.timed_out: + raise CommandTimeoutError( + f"Command '{command}' timed out after {timeout_s:g}s:\n" + f"STDOUT: {res.stdout}\nSTDERR: {res.stderr}", + result=res, + ) + if check and res.exit_code != 0: + raise CommandExecutionError( + f"Command '{command}' failed with exit code {res.exit_code}:\n" + f"STDOUT: {res.stdout}\nSTDERR: {res.stderr}", + result=res, + ) + return res + + async def write_file_async(self, path: str, content: bytes | str) -> None: + """Asynchronously write file content directly into sandbox filesystem.""" + await self.runtime.write_file_async(path, content) + + async def read_file_bytes_async(self, path: str) -> bytes: + """Asynchronously read binary file content from sandbox filesystem.""" + return await self.runtime.read_file_bytes_async(path) + + def initialize( + self, + init_script: Optional[str] = None, + cwd: str = "", + timeout_s: float = 120.0, + ) -> ExecResult: + """Run post-boot initialization script inside the guest container.""" + return self.runtime.initialize(init_script=init_script, cwd=cwd, timeout_s=timeout_s) + + def release(self) -> None: + """Release this sandbox back to the backend or terminate it.""" + try: + self.backend.release(self.sandbox_id, recycle=False) + finally: + self.runtime.close() + + def recycle(self) -> None: + """Return sandbox to warm pool. + + Note: This suspends the actor and places it back in the pool without + resetting filesystem or process state. The next claim will receive + the actor in the state left by the previous workload. + """ + try: + self.backend.release(self.sandbox_id, recycle=True) + finally: + self.runtime.close() + + def __enter__(self) -> "SandboxHandle": + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + self.release() + + async def __aenter__(self) -> "SandboxHandle": + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: + self.release() diff --git a/clients/python/src/ate_env/loop_runner.py b/clients/python/src/ate_env/loop_runner.py new file mode 100644 index 00000000..34120f38 --- /dev/null +++ b/clients/python/src/ate_env/loop_runner.py @@ -0,0 +1,82 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import asyncio +import concurrent.futures +import threading +from typing import Any, Coroutine, TypeVar + +T = TypeVar("T") + + +class LoopRunner: + """A dedicated, permanent background thread running an asyncio event loop. + + Allows synchronous facades (like SubstrateBackendDriver and SubstrateEnvClientRuntime) + to share a single grpc.aio channel safely across multiple sequential or concurrent calls + without attaching channels to short-lived loops that get closed. + """ + + def __init__(self, name: str = "ate-env-sync-runner") -> None: + self._loop = asyncio.new_event_loop() + self._thread = threading.Thread(target=self._run_loop, name=name, daemon=True) + self._thread.start() + + def _run_loop(self) -> None: + asyncio.set_event_loop(self._loop) + self._loop.run_forever() + + @property + def loop(self) -> asyncio.AbstractEventLoop: + return self._loop + + def run(self, coro: Coroutine[Any, Any, T], timeout: float | None = None) -> T: + """Run a coroutine on the background event loop and block for the result.""" + future = asyncio.run_coroutine_threadsafe(coro, self._loop) + return future.result(timeout=timeout) + + def submit(self, coro: Coroutine[Any, Any, T]) -> concurrent.futures.Future[T]: + """Submit a coroutine to the background event loop without blocking.""" + return asyncio.run_coroutine_threadsafe(coro, self._loop) + + def stop(self) -> None: + """Stop the background event loop and join the worker thread.""" + if threading.current_thread() is self._thread: + # Prevent deadlock if called from a task running on this runner's loop + return + if self._loop.is_running(): + async def _shutdown() -> None: + tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()] + for t in tasks: + t.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + await self._loop.shutdown_asyncgens() + + future = asyncio.run_coroutine_threadsafe(_shutdown(), self._loop) + try: + future.result(timeout=2.0) + except Exception: + pass + self._loop.call_soon_threadsafe(self._loop.stop) + self._thread.join(timeout=2.0) + if not self._loop.is_closed(): + self._loop.close() + + def __del__(self) -> None: + try: + self.stop() + except Exception: + pass diff --git a/clients/python/src/ate_env/runtime/__init__.py b/clients/python/src/ate_env/runtime/__init__.py new file mode 100644 index 00000000..1262c548 --- /dev/null +++ b/clients/python/src/ate_env/runtime/__init__.py @@ -0,0 +1,23 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .base import RuntimeGuestHook +from .mock import MockRuntimeHook +from .substrate_env_client import SubstrateEnvClientRuntime + +__all__ = [ + "RuntimeGuestHook", + "MockRuntimeHook", + "SubstrateEnvClientRuntime", +] diff --git a/clients/python/src/ate_env/runtime/base.py b/clients/python/src/ate_env/runtime/base.py new file mode 100644 index 00000000..d89b0384 --- /dev/null +++ b/clients/python/src/ate_env/runtime/base.py @@ -0,0 +1,95 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Dict, Iterator, List, Optional + +from ..types import ExecResult + + +def command_to_argv(command: str | List[str]) -> List[str]: + """Command contract shared by every runtime. + + A ``str`` is a shell command and always runs as ``bash -c ``. A + list is executed directly as argv, without a shell. As a result, a typo + in a shell string exits 127 (an agent outcome that is scored), while a + missing ``argv[0]`` raises ``CommandStartError`` (masked). + """ + if isinstance(command, str): + return ["bash", "-c", command] + argv = list(command) + if not argv: + raise ValueError("command must not be empty") + return argv + + +class RuntimeGuestHook(ABC): + """ + Data Plane (Runtime Protocol) Interface. + + Governs in-guest execution: running commands, streaming processes, + transferring files, and exposing interactive sessions. + """ + + @abstractmethod + def exec(self, command: str | List[str], cwd: str = "", + env: Optional[Dict[str, str]] = None, timeout_s: float = 120.0) -> ExecResult: + """Execute command inside the guest container/actor. + + Returns an ExecResult only if the guest ran the command and it either + exited or was killed at ``timeout_s`` (``timed_out=True``). Raises an + ``InfrastructureError`` subclass for every other outcome; transport or + HTTP/gRPC statuses are never encoded as exit codes. + """ + + @abstractmethod + def stream_process(self, argv: List[str], cwd: str = "", + env: Optional[Dict[str, str]] = None) -> Iterator[str]: + """Stream real-time output chunks from a running guest process.""" + + @abstractmethod + def write_file(self, path: str, content: bytes | str) -> None: + """Write file content directly into the guest filesystem.""" + + @abstractmethod + def read_file_bytes(self, path: str) -> bytes: + """Read binary file content from the guest filesystem.""" + + async def exec_async(self, command: str | List[str], cwd: str = "", + env: Optional[Dict[str, str]] = None, timeout_s: float = 120.0) -> ExecResult: + """Asynchronously execute command inside the guest container/actor.""" + import asyncio + return await asyncio.to_thread(self.exec, command, cwd=cwd, env=env, timeout_s=timeout_s) + + async def write_file_async(self, path: str, content: bytes | str) -> None: + """Asynchronously write file content directly into the guest filesystem.""" + import asyncio + await asyncio.to_thread(self.write_file, path, content) + + async def read_file_bytes_async(self, path: str) -> bytes: + """Asynchronously read binary file content from the guest filesystem.""" + import asyncio + return await asyncio.to_thread(self.read_file_bytes, path) + + def initialize(self, init_script: Optional[str] = None, cwd: str = "", + timeout_s: float = 120.0) -> ExecResult: + """Run post-boot initialization script inside the guest.""" + if not init_script: + return ExecResult(exit_code=0, stdout="", stderr="", duration_s=0.0) + return self.exec(["bash", "-c", init_script], cwd=cwd, timeout_s=timeout_s) + + def close(self) -> None: + """Release client-side resources (connections, channels). Idempotent.""" diff --git a/clients/python/src/ate_env/runtime/mock.py b/clients/python/src/ate_env/runtime/mock.py new file mode 100644 index 00000000..0d936c2c --- /dev/null +++ b/clients/python/src/ate_env/runtime/mock.py @@ -0,0 +1,76 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import io +from typing import Dict, Iterator, List, Optional, Union + +from ..types import ExecResult +from .base import RuntimeGuestHook + + +class MockRuntimeHook(RuntimeGuestHook): + """In-memory mock guest hook implementing file operations and command responses.""" + + def __init__(self): + self.files: Dict[str, bytes] = {} + self.executed_commands: List[str] = [] + self.custom_responses: Dict[str, Union[ExecResult, BaseException]] = {} + + def set_response(self, command_substring: str, + result: Union[ExecResult, BaseException]) -> None: + """Preset a result, or an exception (e.g. SandboxUnavailableError) to raise.""" + self.custom_responses[command_substring] = result + + def exec(self, command: str | List[str], cwd: str = "", + env: Optional[Dict[str, str]] = None, timeout_s: float = 120.0) -> ExecResult: + cmd_str = command if isinstance(command, str) else " ".join(command) + self.executed_commands.append(cmd_str) + + # Check for preset custom responses + for sub, res in self.custom_responses.items(): + if sub in cmd_str: + if isinstance(res, BaseException): + raise res + return res + + # Default standard behaviors for common RL / SWE-bench commands + if "pytest" in cmd_str: + if "test_version" in cmd_str: + return ExecResult(exit_code=0, stdout="=== 1 passed in 0.42s ===\n", stderr="") + return ExecResult(exit_code=0, stdout="=== ALL TESTS PASSED ===\n", stderr="") + + if "git apply" in cmd_str or "patch -p1" in cmd_str: + return ExecResult(exit_code=0, stdout="Applied patch successfully\n", stderr="") + + if "git -C /testbed log" in cmd_str: + return ExecResult(exit_code=0, stdout="READY mock-pod 1234abc Base commit\n", stderr="") + + return ExecResult(exit_code=0, stdout=f"Mock exec ok: {cmd_str}\n", stderr="") + + def stream_process(self, argv: List[str], cwd: str = "", + env: Optional[Dict[str, str]] = None) -> Iterator[str]: + cmd_str = " ".join(argv) + yield f"Process started: {cmd_str}\n" + yield "Process completed.\n" + + def write_file(self, path: str, content: bytes | str) -> None: + raw = content.encode("utf-8") if isinstance(content, str) else content + self.files[path] = raw + + def read_file_bytes(self, path: str) -> bytes: + if path not in self.files: + raise FileNotFoundError(f"Mock file not found: {path}") + return self.files[path] diff --git a/clients/python/src/ate_env/runtime/substrate_env_client.py b/clients/python/src/ate_env/runtime/substrate_env_client.py new file mode 100644 index 00000000..247fc285 --- /dev/null +++ b/clients/python/src/ate_env/runtime/substrate_env_client.py @@ -0,0 +1,339 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import asyncio +import builtins +import logging +import queue +import threading +import time +from typing import Any, AsyncIterator, Dict, Iterator, List, Optional + +import grpc + +from ..client import Client as AteClient +from ..env import Env as AteEnv +from ..errors import ( + CommandStartError, + EnvError, + InfrastructureError, + InvalidArgumentError, + NotFoundError, + RpcError, + SandboxProtocolError, + SandboxUnavailableError, +) +from ..loop_runner import LoopRunner +from ..types import ExecResult, ShellResult +from .base import RuntimeGuestHook + +logger = logging.getLogger("ate_env.runtime.substrate_env_client") + + +class SubstrateEnvClientRuntime(RuntimeGuestHook): + """Substrate guest hook backed directly by the official ate_env.Env handle. + + Leverages ate_env.Client and ate_env.Env for process execution, streaming, + and chunked file I/O over gRPC. + """ + + def __init__( + self, + env_id: str, + *, + endpoint: Optional[str] = None, + atespace: str = "ate-env", + client: Optional[AteClient] = None, + env: Optional[AteEnv] = None, + loop_runner: Optional[LoopRunner] = None, + ): + if not env_id: + raise ValueError("env_id is required for SubstrateEnvClientRuntime") + self.env_id = env_id + self.endpoint = endpoint or "localhost:7777" + self.atespace = atespace + + self._loop_runner = loop_runner + self._owns_loop_runner = loop_runner is None + self._lock = threading.Lock() + + if env is not None: + self._env = env + self._client = env._client + self._owns_client = False + elif client is not None: + self._client = client + self._env = client.env(env_id, atespace=atespace) + self._owns_client = False + else: + self._client = None + self._env = None + self._owns_client = True + self._closed = False + + @property + def env(self) -> AteEnv: + self._check_open() + with self._lock: + if self._env is None: + if self._client is None: + self._client = AteClient(self.endpoint) + self._env = self._client.env(self.env_id, atespace=self.atespace) + return self._env + + def _get_loop_runner(self) -> LoopRunner: + with self._lock: + if self._loop_runner is None: + self._loop_runner = LoopRunner(name="ate-substrate-runtime-sync") + self._owns_loop_runner = True + return self._loop_runner + + async def close_async(self) -> None: + """Close client if owned.""" + if not self._closed: + self._closed = True + if self._owns_client and self._client: + await self._client.close() + self._client = None + if self._owns_loop_runner and self._loop_runner: + self._loop_runner.stop() + self._loop_runner = None + + def close(self) -> None: + if not self._closed: + self._closed = True + if self._loop_runner: + if self._owns_client and self._client: + self._loop_runner.run(self._client.close()) + self._client = None + if self._owns_loop_runner: + self._loop_runner.stop() + self._loop_runner = None + elif self._owns_client and self._client: + self._sync(self._client.close()) + + def _check_open(self) -> None: + if self._closed: + raise RuntimeError(f"runtime for environment {self.env_id} is closed") + + def _sync(self, coro: Any) -> Any: + return self._get_loop_runner().run(coro) + + async def exec_async( + self, + command: str | List[str], + cwd: str = "", + env: Optional[Dict[str, str]] = None, + timeout_s: float = 120.0, + ) -> ExecResult: + """Native async execution honoring command contract: argv list runs directly, string runs in shell.""" + self._check_open() + t0 = time.monotonic() + grace_s = 5.0 + active_env = self.env + + try: + active_proc = None + try: + if isinstance(command, (list, tuple)): + argv = list(command) + if not argv: + raise ValueError("command must not be empty") + + async def _run_argv() -> ShellResult: + nonlocal active_proc + proc = await active_env.start_process( + argv, cwd=cwd, env=env, timeout=timeout_s + ) + active_proc = proc + stdout = bytearray() + stderr = bytearray() + exit_info = None + async for out in proc.output(follow=True): + if out.stdout is not None: + stdout.extend(out.stdout) + elif out.stderr is not None: + stderr.extend(out.stderr) + elif out.exit is not None: + exit_info = out.exit + if exit_info is None: + raise RuntimeError("ate_env: output stream ended before process exited") + return ShellResult( + stdout=stdout.decode("utf-8", errors="replace"), + stderr=stderr.decode("utf-8", errors="replace"), + exit_code=exit_info.exit_code, + ) + + if hasattr(asyncio, "timeout"): + async with asyncio.timeout(timeout_s + grace_s): + shell_res = await _run_argv() + else: + shell_res = await asyncio.wait_for(_run_argv(), timeout=timeout_s + grace_s) + else: + if hasattr(asyncio, "timeout"): + async with asyncio.timeout(timeout_s + grace_s): + shell_res = await active_env.shell( + command, cwd=cwd, env=env, timeout=timeout_s + ) + else: + shell_res = await asyncio.wait_for( + active_env.shell(command, cwd=cwd, env=env, timeout=timeout_s), + timeout=timeout_s + grace_s, + ) + except (builtins.TimeoutError, asyncio.TimeoutError): + if active_proc is not None: + try: + await active_proc.kill() + except Exception: + pass + duration = time.monotonic() - t0 + return ExecResult( + exit_code=None, + stdout="", + stderr="", + duration_s=duration, + timed_out=True, + ) + + duration = time.monotonic() - t0 + # 137 = 128 + 9 (SIGKILL sent by guest when guest-side timeout elapsed). + # Note: checking exit_code == 137 and duration >= timeout_s is a heuristic; + # external kills (like OOM) can also produce exit_code 137. + timed_out = shell_res.exit_code == 137 and duration >= timeout_s + return ExecResult( + exit_code=None if timed_out else shell_res.exit_code, + stdout=shell_res.stdout, + stderr=shell_res.stderr, + duration_s=duration, + timed_out=timed_out, + ) + except InvalidArgumentError as e: + raise SandboxProtocolError(str(e), sandbox_id=self.env_id) from e + except NotFoundError as e: + msg = str(e).lower() + if "environment" in msg or "actor" in msg: + raise SandboxUnavailableError(str(e), sandbox_id=self.env_id) from e + raise CommandStartError(str(e), sandbox_id=self.env_id) from e + except RpcError as e: + status_name = f"grpc {e.code.name}" if e.code else "grpc RpcError" + if e.code in ( + grpc.StatusCode.NOT_FOUND, + grpc.StatusCode.UNAVAILABLE, + grpc.StatusCode.DEADLINE_EXCEEDED, + grpc.StatusCode.INTERNAL, + grpc.StatusCode.ABORTED, + ): + raise SandboxUnavailableError( + str(e), sandbox_id=self.env_id, status=status_name + ) from e + raise SandboxProtocolError(str(e), sandbox_id=self.env_id, status=status_name) from e + except InfrastructureError: + raise + except Exception as e: + raise InfrastructureError(str(e), sandbox_id=self.env_id) from e + + def exec( + self, + command: str | List[str], + cwd: str = "", + env: Optional[Dict[str, str]] = None, + timeout_s: float = 120.0, + ) -> ExecResult: + """Synchronous wrapper for exec_async.""" + return self._sync(self.exec_async(command, cwd=cwd, env=env, timeout_s=timeout_s)) + + async def stream_process_async( + self, + argv: List[str], + cwd: str = "", + env: Optional[Dict[str, str]] = None, + ) -> AsyncIterator[str]: + """Stream real-time output chunks from running process inside guest.""" + self._check_open() + try: + proc = await self.env.start_process(argv, cwd=cwd, env=env) + async for chunk in proc.output(follow=True): + if chunk.stdout is not None: + yield chunk.stdout.decode("utf-8", errors="replace") + elif chunk.stderr is not None: + yield chunk.stderr.decode("utf-8", errors="replace") + except (InvalidArgumentError, NotFoundError, RpcError) as e: + raise SandboxUnavailableError(str(e), sandbox_id=self.env_id) from e + except Exception as e: + raise InfrastructureError(str(e), sandbox_id=self.env_id) from e + + def stream_process( + self, + argv: List[str], + cwd: str = "", + env: Optional[Dict[str, str]] = None, + ) -> Iterator[str]: + """Synchronously stream output chunks in real time.""" + self._check_open() + q: queue.Queue[tuple[Optional[str], Optional[Exception]]] = queue.Queue() + + async def _pump(): + try: + async for item in self.stream_process_async(argv, cwd=cwd, env=env): + q.put((item, None)) + q.put((None, None)) + except Exception as e: + q.put((None, e)) + + fut = self._get_loop_runner().submit(_pump()) + try: + while True: + item, err = q.get() + if err is not None: + raise err + if item is None: + break + yield item + finally: + fut.cancel() + + async def write_file_async(self, path: str, content: bytes | str) -> None: + """Write file content directly into guest filesystem via FileSystemService.""" + self._check_open() + raw = content.encode("utf-8") if isinstance(content, str) else content + try: + await self.env.write_file(path, raw) + except InvalidArgumentError as e: + raise SandboxProtocolError(str(e), sandbox_id=self.env_id) from e + except (NotFoundError, RpcError) as e: + raise SandboxUnavailableError(str(e), sandbox_id=self.env_id) from e + except Exception as e: + raise InfrastructureError(str(e), sandbox_id=self.env_id) from e + + def write_file(self, path: str, content: bytes | str) -> None: + self._sync(self.write_file_async(path, content)) + + async def read_file_bytes_async(self, path: str) -> bytes: + """Read binary file content from guest filesystem via FileSystemService.""" + self._check_open() + try: + return await self.env.read_file_bytes(path) + except NotFoundError: + raise FileNotFoundError(f"File not found in sandbox: {path}") + except InvalidArgumentError as e: + raise SandboxProtocolError(str(e), sandbox_id=self.env_id) from e + except RpcError as e: + raise SandboxUnavailableError(str(e), sandbox_id=self.env_id) from e + except Exception as e: + raise InfrastructureError(str(e), sandbox_id=self.env_id) from e + + def read_file_bytes(self, path: str) -> bytes: + return self._sync(self.read_file_bytes_async(path)) diff --git a/clients/python/src/ate_env/strategies.py b/clients/python/src/ate_env/strategies.py new file mode 100644 index 00000000..6652c220 --- /dev/null +++ b/clients/python/src/ate_env/strategies.py @@ -0,0 +1,179 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import collections +import logging +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Any, Callable, Dict, List, Optional + +from .handle import SandboxHandle +from .types import Task + +logger = logging.getLogger("ate_env.strategies") + + +def process_parallel(fleet: Any, tasks: List[Task], process_fn: Callable[[Task, SandboxHandle], Any], + concurrency: int) -> List[Any]: + """Execute acquire -> process_fn -> release across tasks up to concurrency.""" + results: List[Any] = [None] * len(tasks) + + def _execute_single(task: Task) -> Any: + handle = fleet.acquire(task) + try: + return process_fn(task, handle) + finally: + fleet.release(handle) + + if concurrency <= 1: + for i, t in enumerate(tasks): + try: + results[i] = _execute_single(t) + except Exception as e: + logger.error("Task %s failed: %s", t.id, e) + results[i] = e + return results + + with ThreadPoolExecutor(max_workers=concurrency) as executor: + futures = {executor.submit(_execute_single, t): i for i, t in enumerate(tasks)} + for fut in as_completed(futures): + i = futures[fut] + try: + results[i] = fut.result() + except Exception as e: + logger.error("Task %s failed: %s", tasks[i].id, e) + results[i] = e + + return results + + +def run_none(fleet: Any, process_fn: Callable[[Task, SandboxHandle], Any], + concurrency: int, teardown: bool = True) -> List[Any]: + """No pre-warming: on-demand 1:1 acquisition per task.""" + try: + return process_parallel(fleet, fleet.tasks, process_fn, concurrency) + finally: + if teardown: + fleet.teardown() + + +def run_naive(fleet: Any, process_fn: Callable[[Task, SandboxHandle], Any], + concurrency: int, teardown: bool = True) -> List[Any]: + """Pre-warm every image up front, process all tasks in parallel, tear down.""" + try: + fleet.setup() + return process_parallel(fleet, fleet.tasks, process_fn, concurrency) + finally: + if teardown: + fleet.teardown() + + +def run_sliding(fleet: Any, process_fn: Callable[[Task, SandboxHandle], Any], + concurrency: int, teardown: bool = True) -> List[Any]: + """Keep only a sliding window of image pools warm at a time (bounded footprint).""" + fleet.preflight() + fleet.plan() + window = fleet.config.window_size or fleet.config.batch_size + images = list(fleet.image_counts().keys()) + + by_image: Dict[str, List[tuple[int, Task]]] = collections.defaultdict(list) + for i, t in enumerate(fleet.tasks): + by_image[t.image].append((i, t)) + + results: List[Any] = [None] * len(fleet.tasks) + try: + for start in range(0, len(images), window): + batch = images[start:start + window] + fleet.warm_images(batch, wait=True) + batch_pairs = [(i, t) for img in batch for (i, t) in by_image[img]] + batch_tasks = [t for _i, t in batch_pairs] + + logger.info("Sliding window [%d..%d): %d image(s), %d task(s)", + start, start + len(batch), len(batch), len(batch_tasks)) + batch_results = process_parallel(fleet, batch_tasks, process_fn, concurrency) + for (orig_idx, _t), r in zip(batch_pairs, batch_results): + results[orig_idx] = r + + for img in batch: + fleet.unwarm_image(img) + finally: + if teardown: + fleet.teardown() + + return results + + +def run_pipelined(fleet: Any, process_fn: Callable[[Task, SandboxHandle], Any], + concurrency: int, teardown: bool = True) -> List[Any]: + """ + Double-buffered pipelined sliding window. + + While window N tasks run, prefetch window N+1 in the background to overlap + image pull or snapshot restore with GPU rollout time. + """ + fleet.preflight() + fleet.plan() + window = fleet.config.window_size or fleet.config.batch_size + images = list(fleet.image_counts().keys()) + + by_image: Dict[str, List[tuple[int, Task]]] = collections.defaultdict(list) + for i, t in enumerate(fleet.tasks): + by_image[t.image].append((i, t)) + + results: List[Any] = [None] * len(fleet.tasks) + batches = [images[s:s + window] for s in range(0, len(images), window)] + + prefetch_executor = ThreadPoolExecutor(max_workers=1) + try: + if batches: + fleet.warm_images(batches[0], wait=True) + + for n, batch in enumerate(batches): + # Asynchronously prefetch window N+1 + nxt_future = ( + prefetch_executor.submit(fleet.warm_images, batches[n + 1], wait=True) + if n + 1 < len(batches) else None + ) + + batch_pairs = [(i, t) for img in batch for (i, t) in by_image[img]] + batch_tasks = [t for _i, t in batch_pairs] + + logger.info("Pipelined window [%d/%d]: %d image(s), %d task(s)", + n + 1, len(batches), len(batch), len(batch_tasks)) + batch_results = process_parallel(fleet, batch_tasks, process_fn, concurrency) + for (orig_idx, _t), r in zip(batch_pairs, batch_results): + results[orig_idx] = r + + # Unwarm batch N before awaiting next batch to maintain <= 2 windows bound + for img in batch: + fleet.unwarm_image(img) + + if nxt_future is not None: + nxt_future.result() + finally: + prefetch_executor.shutdown(wait=True) + if teardown: + fleet.teardown() + + return results + + +STRATEGIES: Dict[str, Callable[..., List[Any]]] = { + "none": run_none, + "naive": run_naive, + "sliding": run_sliding, + "pipelined": run_pipelined, +} diff --git a/clients/python/src/ate_env/types.py b/clients/python/src/ate_env/types.py index d7aa381e..1377b19b 100644 --- a/clients/python/src/ate_env/types.py +++ b/clients/python/src/ate_env/types.py @@ -21,8 +21,10 @@ from __future__ import annotations import enum -from dataclasses import dataclass +import hashlib +from dataclasses import dataclass, field from datetime import datetime, timezone +from typing import Any, Dict, List, Literal, Optional from ._gen.ateenv.v1alpha import env_pb2, guest_pb2 @@ -35,9 +37,19 @@ "ProcessInfo", "ProcessOutput", "ShellResult", + "ResourceLimits", + "PlacementSpec", + "EnvironmentSpec", + "Task", + "ExecResult", + "DataPlaneEndpoint", + "RawSandboxInstance", + "PlanEntry", + "FleetPlan", ] + class EnvironmentStatus(enum.IntEnum): """Lifecycle status of an environment (ateenv.v1alpha.EnvironmentStatus).""" @@ -193,3 +205,140 @@ def _process_output_from_pb(pb: guest_pb2.ProcessOutput) -> ProcessOutput: if which == "exit": return ProcessOutput(exit=_process_info_from_pb(pb.exit)) raise ValueError(f"ate_env: unexpected process output {pb!r}") + + +@dataclass(frozen=True) +class ResourceLimits: + """CPU and memory limits for a sandbox environment.""" + cpu: str = "2" + memory: str = "4Gi" + + +@dataclass(frozen=True) +class PlacementSpec: + """Placement constraints and worker hardware selectors.""" + node_selector: Dict[str, str] = field(default_factory=dict) + tolerations: List[Dict[str, Any]] = field(default_factory=list) + worker_family: Optional[str] = None # e.g., "c2", "n2", "c3" + + +@dataclass(frozen=True) +class EnvironmentSpec: + """Immutable environment definition representing image, bundle, and constraints.""" + image: str + runtime_bundle: str = "default-guest" + limits: ResourceLimits = field(default_factory=ResourceLimits) + placement: PlacementSpec = field(default_factory=PlacementSpec) + snapshot_storage_uri: Optional[str] = None + + def template_key(self) -> str: + """Derive canonical template ID including CPU family and bundle digest.""" + raw = f"{self.image}::{self.runtime_bundle}::{self.placement.worker_family or 'ambient'}" + img_hash = hashlib.sha256(raw.encode()).hexdigest()[:12] + # Clean alphanumeric template name + clean_img = self.image.split("/")[-1].split(":")[0].replace(".", "-").replace("_", "-") + return f"tmpl-{clean_img[:20]}-{img_hash}" + + +@dataclass +class Task: + """Single workload/task definition (e.g. one SWE-bench instance).""" + id: str + image: str + metadata: Dict[str, Any] = field(default_factory=dict) + template: Optional[str] = None + + +@dataclass +class ExecResult: + """Outcome of a command that the guest actually ran. + + A runtime returns an ExecResult only when the guest started the command + and either observed it exit (``exit_code`` set) or killed it at its + deadline (``timed_out=True`` and ``exit_code=None``; output may be + partial). Anything else (unreachable sandbox, HTTP/gRPC error, malformed + response) raises ``InfrastructureError`` instead, so transport failures + can never be mistaken for agent failures. + + Exit codes follow standard POSIX/Linux conventions: a command killed by + a signal reports 128 plus the signal number (for example 137 for SIGKILL). + Note: checking ``exit_code == 137 and duration >= timeout_s`` to detect + timeouts is a heuristic and may produce false positives if a process was + terminated by an external signal (such as an out-of-memory killer) near + its deadline. + """ + exit_code: Optional[int] + stdout: str + stderr: str + duration_s: float = 0.0 + timed_out: bool = False + + def __post_init__(self) -> None: + if self.timed_out and self.exit_code is not None: + raise ValueError("ExecResult: exit_code must be None when timed_out=True") + if not self.timed_out and self.exit_code is None: + raise ValueError("ExecResult: exit_code is required unless timed_out=True") + + @property + def ok(self) -> bool: + """True iff the command finished within its deadline with exit code 0.""" + return not self.timed_out and self.exit_code == 0 + + +_SENSITIVE_HEADERS = frozenset({"authorization", "proxy-authorization", "cookie", "x-api-key"}) + + +@dataclass(repr=False) +class DataPlaneEndpoint: + """Per-sandbox data-plane coordinates produced by a BackendDriver. + + Attributes: + protocol: ``"grpc"`` (in-guest gRPC server multiplexed via ate-env-api). + address: Host:port of the ate-env-api endpoint. + headers: gRPC metadata routing headers to this specific sandbox, + e.g. ``{"x-env-id": "", "x-env-atespace": ""}``. + """ + protocol: Literal["grpc"] = "grpc" + address: str = "" + headers: Dict[str, str] = field(default_factory=dict) + + def __repr__(self) -> str: + shown = {k: ("" if k.lower() in _SENSITIVE_HEADERS else v) + for k, v in self.headers.items()} + return (f"DataPlaneEndpoint(protocol={self.protocol!r}, " + f"address={self.address!r}, headers={shown!r})") + + +@dataclass +class RawSandboxInstance: + """Low-level sandbox handle returned by BackendDriver.""" + instance_id: str + endpoint: str + template_id: str + run_id: str + metadata: Dict[str, Any] = field(default_factory=dict) + # Keyed by data-plane name matching FleetConfig.data_plane ("router", "grpc"). + data_planes: Dict[str, DataPlaneEndpoint] = field(default_factory=dict) + + +@dataclass +class PlanEntry: + """One environment's calculated provisioning plan.""" + image: str + template_id: str + replicas: int + tasks: int + + +class FleetPlan: + """The complete provisioning plan computed by SandboxFleet.plan().""" + def __init__(self, entries: List[PlanEntry]): + self.entries = entries + self._by_image = {e.image: e for e in entries} + + def for_image(self, image: str) -> Optional[PlanEntry]: + return self._by_image.get(image) + + @property + def total_replicas(self) -> int: + return sum(e.replicas for e in self.entries) diff --git a/clients/python/tests/test_config.py b/clients/python/tests/test_config.py new file mode 100644 index 00000000..7036f66f --- /dev/null +++ b/clients/python/tests/test_config.py @@ -0,0 +1,57 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from ate_env.config import FleetConfig, ValidationError + + +def test_unknown_fields_are_rejected(): + # Previously silently dropped by pydantic (e.g. the NeMo/TML doc snippets). + with pytest.raises(ValidationError, match="warm_pool_size"): + FleetConfig(warm_pool_size=64) + with pytest.raises(ValidationError, match="warmpool_replicas"): + FleetConfig.from_dict({"backend": "mock", "warmpool_replicas": 4}) + + +@pytest.mark.parametrize("kwargs", [ + {"batch_size": 0}, + {"max_concurrent": 0}, + {"max_warmpool_replicas": -1}, + {"acquire_timeout_s": 0}, + {"backend": "kubernetes"}, # only substrate and mock supported + {"tenancy": "Bad_Name"}, + {"tenancy": "-leading-dash"}, + {"strategy": "rolling"}, + {"data_plane": "invalid_plane"}, + {"grpc_endpoint": "{actor}.svc:50051"}, # unknown field +]) +def test_invalid_values_are_rejected(kwargs): + with pytest.raises(ValidationError): + FleetConfig(**kwargs) + + +def test_template_and_image_templates_configured(): + cfg = FleetConfig(template="my-tmpl", image_templates={"python:3.11": "tmpl-python"}) + assert cfg.template == "my-tmpl" + assert cfg.image_templates["python:3.11"] == "tmpl-python" + + +def test_auth_token_is_secret_and_round_trips(): + cfg = FleetConfig(auth_token="s3cr3t") + assert "s3cr3t" not in repr(cfg) + assert "s3cr3t" not in str(cfg.model_dump()) + # model_dump() -> from_dict() is how configs are shipped to Ray workers. + again = FleetConfig.from_dict(cfg.model_dump()) + assert again.auth_token.get_secret_value() == "s3cr3t" diff --git a/clients/python/tests/test_data_planes.py b/clients/python/tests/test_data_planes.py new file mode 100644 index 00000000..cd5bc282 --- /dev/null +++ b/clients/python/tests/test_data_planes.py @@ -0,0 +1,79 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-sandbox data planes: each handle talks to its environment.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock +import pytest + +from ate_env import DataPlaneEndpoint, FleetConfig, SandboxFleet +from ate_env.backend.substrate import SubstrateBackendDriver +from ate_env.errors import SandboxStartError +from ate_env.runtime.mock import MockRuntimeHook +from ate_env.runtime.substrate_env_client import SubstrateEnvClientRuntime +from ate_env.types import EnvironmentSpec + + +def mock_client(): + c = MagicMock() + c.create = AsyncMock() + c.delete = AsyncMock() + c.suspend = AsyncMock() + return c + + +def acquire_two(d): + tid = d.ensure_template(EnvironmentSpec(image="img:1")) + return d.acquire(tid, "run-1"), d.acquire(tid, "run-1") + + +def test_each_sandbox_gets_its_own_coordinates(): + client = mock_client() + d = SubstrateBackendDriver(client=client, atespace="space-a") + a, b = acquire_two(d) + for inst in (a, b): + target = inst.instance_id + assert "ate_env" in inst.data_planes + assert inst.data_planes["ate_env"].headers["x-env-id"] == target + assert inst.data_planes["ate_env"].headers["x-env-atespace"] == "space-a" + assert a.data_planes["ate_env"].headers != b.data_planes["ate_env"].headers + + +def test_fleet_builds_each_runtime_from_its_own_data_plane(): + client = mock_client() + d = SubstrateBackendDriver(client=client, atespace="space-a") + fleet = SandboxFleet( + FleetConfig(backend="substrate", tenancy="space-a", data_plane="ate_env"), + driver=d, + ) + h1, h2 = fleet.acquire("t1"), fleet.acquire("t2") + try: + assert h1.sandbox_id != h2.sandbox_id + for h in (h1, h2): + assert isinstance(h.runtime, SubstrateEnvClientRuntime) + assert h.runtime.env_id == h.sandbox_id + assert h.runtime.atespace == "space-a" + finally: + h1.release() + h2.release() + + +def test_mock_backend_uses_mock_runtime(): + fleet = SandboxFleet(FleetConfig(backend="mock")) + h = fleet.acquire("t1") + assert isinstance(h.runtime, MockRuntimeHook) + assert h.data_plane is None + h.release() diff --git a/clients/python/tests/test_exec_semantics.py b/clients/python/tests/test_exec_semantics.py new file mode 100644 index 00000000..68d22d9e --- /dev/null +++ b/clients/python/tests/test_exec_semantics.py @@ -0,0 +1,122 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Exec result contract, strict-mode errors, command contract, and file helpers.""" + +import pickle +import subprocess + +import pytest + +from ate_env import ( + CommandExecutionError, + CommandStartError, + CommandTimeoutError, + ExecResult, + FleetConfig, + SandboxFleet, + SandboxProtocolError, + SandboxUnavailableError, +) +from ate_env.exceptions import TimeoutError as SdkTimeoutError +from ate_env.runtime.base import ( + RuntimeGuestHook, + command_to_argv, +) +from ate_env.runtime.mock import MockRuntimeHook + + +def test_exec_result_invariants(): + assert ExecResult(exit_code=0, stdout="", stderr="").ok + assert not ExecResult(exit_code=1, stdout="", stderr="").ok + assert not ExecResult(exit_code=None, stdout="partial", stderr="", timed_out=True).ok + with pytest.raises(ValueError): + ExecResult(exit_code=0, stdout="", stderr="", timed_out=True) + with pytest.raises(ValueError): + ExecResult(exit_code=None, stdout="", stderr="") + + +@pytest.fixture +def handle(): + fleet = SandboxFleet(FleetConfig(backend="mock")) + h = fleet.acquire("t1") + yield h + h.release() + + +def test_check_raises_command_execution_error_with_result(handle): + handle.runtime.set_response("failing-cmd", ExecResult(exit_code=2, stdout="", stderr="boom")) + with pytest.raises(CommandExecutionError) as ei: + handle.exec("failing-cmd") + assert ei.value.result.exit_code == 2 + res = handle.exec("failing-cmd", check=False) + assert res.exit_code == 2 + assert res.stderr == "boom" + + +def test_check_raises_command_timeout_error(handle): + handle.runtime.set_response( + "slow-cmd", ExecResult(exit_code=None, stdout="", stderr="", timed_out=True)) + with pytest.raises(CommandTimeoutError) as ei: + handle.exec("slow-cmd") + # Catchable as either a command failure or an SDK timeout. + assert isinstance(ei.value, CommandExecutionError) + assert isinstance(ei.value, SdkTimeoutError) + assert ei.value.result.timed_out + + +@pytest.mark.parametrize("check", [True, False]) +def test_infrastructure_errors_propagate_regardless_of_check(handle, check): + handle.runtime.set_response( + "dead-cmd", SandboxUnavailableError("gone", sandbox_id="sb-1", status="HTTP 503")) + with pytest.raises(SandboxUnavailableError): + handle.exec("dead-cmd", check=check) + + +def test_retryability_defaults_and_pickling(): + assert SandboxUnavailableError("x").retryable + assert not SandboxProtocolError("x").retryable + assert not CommandStartError("x").retryable + err = SandboxUnavailableError("gone", sandbox_id="sb-1", status="HTTP 503", retryable=False) + clone = pickle.loads(pickle.dumps(err)) # e.g. raised inside a Ray worker + assert type(clone) is SandboxUnavailableError + assert (str(clone), clone.sandbox_id, clone.status, clone.retryable) == ( + "gone", "sb-1", "HTTP 503", False) + + +def test_command_contract(): + # Strings always go through bash, so a typo exits 127 instead of failing to start. + assert command_to_argv("lss") == ["bash", "-c", "lss"] + assert command_to_argv("ls -la | wc -l") == ["bash", "-c", "ls -la | wc -l"] + assert command_to_argv(["ls", "-la"]) == ["ls", "-la"] + with pytest.raises(ValueError): + command_to_argv([]) + + +@pytest.mark.parametrize("name", [ + "plain.txt", + "dir with space/it's \"quoted\" $(echo hi) `x`;.txt", + "-leading-dash.txt", +]) +def test_mock_runtime_file_round_trip(name): + rt = MockRuntimeHook() + target = f"/tmp/nested/{name}" + payload = bytes(range(256)) * 4 + rt.write_file(target, payload) + assert rt.read_file_bytes(target) == payload + + +def test_read_missing_file_raises_file_not_found(): + with pytest.raises(FileNotFoundError): + MockRuntimeHook().read_file_bytes("/missing.txt") diff --git a/clients/python/tests/test_fleet_strategies.py b/clients/python/tests/test_fleet_strategies.py new file mode 100644 index 00000000..8bb7b75a --- /dev/null +++ b/clients/python/tests/test_fleet_strategies.py @@ -0,0 +1,74 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +from ate_env.config import FleetConfig +from ate_env.fleet import SandboxFleet +from ate_env.types import Task + + +def dummy_process_fn(task: Task, handle) -> str: + res = handle.exec(f"echo processing {task.id}") + return res.stdout if hasattr(res, "stdout") else res + + +@pytest.mark.parametrize("strategy", ["none", "naive", "sliding", "pipelined"]) +def test_fleet_strategies(strategy: str): + tasks = [ + Task(id=f"task-{i}", image=f"image-{i % 2}") + for i in range(6) + ] + cfg = FleetConfig( + backend="mock", + strategy=strategy, + batch_size=2, + max_warmpool_replicas=2, + ) + fleet = SandboxFleet(cfg) + fleet.load_tasks(tasks) + + results = fleet.run(dummy_process_fn, concurrency=2) + assert len(results) == 6 + for i, r in enumerate(results): + assert f"task-{i}" in r + + +def test_user_guide_section1_synchronous_fleet_with_context_manager(): + config = FleetConfig(backend="mock") + task = Task(id="task-1", image="docker.io/library/python:3.11") + with SandboxFleet(config) as fleet: + sb = fleet.acquire(task) + res = sb.exec("echo hello") + out = res.stdout if hasattr(res, "stdout") else res + assert "echo hello" in out or "hello" in out + + +def test_user_guide_section2_synchronous_setup_and_batch_release(): + config = FleetConfig(backend="mock") + tasks = [ + Task(id="task-1", image="docker.io/library/python:3.11"), + Task(id="task-2", image="docker.io/library/python:3.11"), + ] + fleet = SandboxFleet(config) + fleet.setup(tasks) + sandboxes = fleet.acquire(tasks) + try: + assert len(sandboxes) == 2 + for t, sb in zip(tasks, sandboxes): + res = sb.exec("echo hello") + out = res.stdout if hasattr(res, "stdout") else res + assert "echo hello" in out or "hello" in out + finally: + fleet.release(sandboxes) + fleet.close() diff --git a/clients/python/tests/test_fleet_types.py b/clients/python/tests/test_fleet_types.py new file mode 100644 index 00000000..836ac906 --- /dev/null +++ b/clients/python/tests/test_fleet_types.py @@ -0,0 +1,45 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +from ate_env.types import EnvironmentSpec, PlacementSpec, ResourceLimits, Task + + +def test_environment_spec_template_key_hashing(): + spec1 = EnvironmentSpec( + image="us-central1-docker.pkg.dev/proj/repo/swe-bench:latest", + runtime_bundle="substrate-env:v1", + placement=PlacementSpec(worker_family="c2") + ) + spec2 = EnvironmentSpec( + image="us-central1-docker.pkg.dev/proj/repo/swe-bench:latest", + runtime_bundle="substrate-env:v1", + placement=PlacementSpec(worker_family="c2") + ) + # Different CPU family must produce a different template key to protect snapshot restores + spec3 = EnvironmentSpec( + image="us-central1-docker.pkg.dev/proj/repo/swe-bench:latest", + runtime_bundle="substrate-env:v1", + placement=PlacementSpec(worker_family="c3") + ) + + assert spec1.template_key() == spec2.template_key() + assert spec1.template_key() != spec3.template_key() + assert "c2" not in spec1.template_key() or spec1.template_key().startswith("tmpl-") + + +def test_task_metadata(): + task = Task(id="pytest-1", image="pytest-img", metadata={"repo": "pytest-dev/pytest"}) + assert task.id == "pytest-1" + assert task.metadata["repo"] == "pytest-dev/pytest" diff --git a/clients/python/tests/test_substrate_driver.py b/clients/python/tests/test_substrate_driver.py new file mode 100644 index 00000000..c4020004 --- /dev/null +++ b/clients/python/tests/test_substrate_driver.py @@ -0,0 +1,154 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import asyncio +import pytest + +from ate_env import FleetConfig, SandboxFleet, Task +from ate_env._gen.ateenv.v1alpha import env_pb2 +from ate_env.backend.substrate import SubstrateBackendDriver +from ate_env.errors import PreflightError +from ate_env.types import EnvironmentSpec, PlacementSpec + + +async def test_substrate_driver_lifecycle_with_client(fake_stack): + client, fakes = fake_stack + driver = SubstrateBackendDriver( + client=client, + atespace="ate-env", + worker_family="c2", + ) + + await driver.preflight_async() + + env = EnvironmentSpec( + image="us-central1-docker.pkg.dev/sample/swe-bench:latest", + placement=PlacementSpec(worker_family="c2"), + ) + tid = driver.ensure_template(env) + assert tid.startswith("tmpl-") + + # 1. Warm 2 paused actors: creates and suspends them in ate-env-api + await driver.warm_pool_async(tid, replicas=2) + assert len(driver._warm_paused_pool[tid]) == 2 + # Verify both environments were created and suspended on the server + for aid in driver._warm_paused_pool[tid]: + stored = fakes.environments.environments[("ate-env", aid)] + assert stored.status == env_pb2.ENVIRONMENT_STATUS_SUSPENDED + assert stored.template.name == tid + + # 2. Acquire claims one from warm pool without creating a new one + inst1 = await driver.acquire_async(tid, run_id="run-1") + assert len(driver._warm_paused_pool[tid]) == 1 + assert inst1.instance_id in fakes.environments.environments[("ate-env", inst1.instance_id)].id + + # 3. Acquire another when warm pool is exhausted creates a fresh one + inst2 = await driver.acquire_async(tid, run_id="run-1") + assert inst2.instance_id != inst1.instance_id + assert ("ate-env", inst2.instance_id) in fakes.environments.environments + + # 4. Release with recycle=True calls suspend and returns to warm pool + await driver.release_async(inst1.instance_id, recycle=True) + assert len(driver._warm_paused_pool[tid]) == 1 + assert fakes.environments.environments[("ate-env", inst1.instance_id)].status == env_pb2.ENVIRONMENT_STATUS_SUSPENDED + + # 5. Release with recycle=False deletes environment permanently + await driver.release_async(inst2.instance_id, recycle=False) + assert ("ate-env", inst2.instance_id) not in fakes.environments.environments + + # 6. Reap deletes all remaining actors tagged with run_id + inst3 = await driver.acquire_async(tid, run_id="run-to-reap") + assert ("ate-env", inst3.instance_id) in fakes.environments.environments + reaped = await driver.reap_async("run-to-reap") + assert reaped == 1 + assert ("ate-env", inst3.instance_id) not in fakes.environments.environments + + # 7. Unwarm pool drains remaining paused actors + await driver.warm_pool_async("tmpl-extra", replicas=1) + assert len(driver._warm_paused_pool["tmpl-extra"]) == 1 + await driver.unwarm_pool_async("tmpl-extra") + assert len(driver._warm_paused_pool.get("tmpl-extra", [])) == 0 + + +async def test_substrate_synchronous_fleet_multiple_acquires(fake_stack): + client, fakes = fake_stack + cfg = FleetConfig(backend="substrate", endpoint=client.endpoint, template="tmpl-test") + + def _run_sync_fleet(): + fleet = SandboxFleet(cfg) + try: + sb1 = fleet.acquire(Task(id="task-1", image="img-1")) + sb2 = fleet.acquire(Task(id="task-2", image="img-2")) + assert sb1.sandbox_id != sb2.sandbox_id + fleet.release([sb1, sb2]) + finally: + fleet.close() + + await asyncio.to_thread(_run_sync_fleet) + + +async def test_substrate_fleet_setup_and_teardown_cleans_warm_pool(fake_stack): + client, fakes = fake_stack + cfg = FleetConfig( + backend="substrate", + endpoint=client.endpoint, + template="tmpl-test", + max_warmpool_replicas=2, + ) + + def _run_setup_teardown(): + fleet = SandboxFleet(cfg) + fleet.load_tasks([Task(id="t1", image="img"), Task(id="t2", image="img")]) + fleet.setup() + assert len(fakes.environments.environments) == 2 + fleet.teardown() + assert len(fakes.environments.environments) == 0 + fleet.close() + + await asyncio.to_thread(_run_setup_teardown) + + +async def test_substrate_preflight_fails_on_closed_port(): + driver = SubstrateBackendDriver(api_endpoint="127.0.0.1:59999") + try: + with pytest.raises(PreflightError): + await driver.preflight_async() + finally: + await driver.close_async() + + +def test_substrate_default_template_resolution(): + # B1: A fleet with no template configured falls back to "default-template" + cfg = FleetConfig(backend="substrate") + fleet = SandboxFleet(cfg) + tmpl = fleet._ensure_template_for_image("docker.io/library/python@sha256:2176bb31f3eb") + assert tmpl == "default-template" + + # M5: image_templates mapping wins over config.template + cfg_mapped = FleetConfig( + backend="substrate", + template="fallback-template", + image_templates={"docker.io/library/python:3.11": "python-template"}, + ) + fleet_mapped = SandboxFleet(cfg_mapped) + assert fleet_mapped._ensure_template_for_image("docker.io/library/python:3.11") == "python-template" + assert fleet_mapped._ensure_template_for_image("other-img") == "fallback-template" + + # SubstrateBackendDriver with default_template="default-template" does not treat it as unset + driver = SubstrateBackendDriver(default_template="default-template") + assert driver.ensure_template("docker.io/library/python@sha256:2176bb31f3eb") == "default-template" + driver_none = SubstrateBackendDriver(default_template=None) + assert driver_none.ensure_template("docker.io/library/python@sha256:2176bb31f3eb") == "default-template" diff --git a/clients/python/tests/test_substrate_env_client.py b/clients/python/tests/test_substrate_env_client.py new file mode 100644 index 00000000..80b98b4d --- /dev/null +++ b/clients/python/tests/test_substrate_env_client.py @@ -0,0 +1,125 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock, patch +import grpc +import pytest + +from ate_env.errors import ( + CommandStartError, + InvalidArgumentError, + NotFoundError, + RpcError, + SandboxProtocolError, + SandboxUnavailableError, +) +from ate_env.runtime.substrate_env_client import SubstrateEnvClientRuntime +from .fakes import FakeProc + + +async def test_substrate_env_client_exec(fake_stack): + client, fakes = fake_stack + await client.create("dev1") + + fakes.processes.next_procs.append( + FakeProc( + chunks=[("stdout", b"hello from guest\n")], + exit_code=0, + ) + ) + + rt = SubstrateEnvClientRuntime(client=client, env_id="dev1") + res = await rt.exec_async("echo hello") + assert res.ok + assert res.stdout == "hello from guest\n" + assert res.exit_code == 0 + assert not res.timed_out + + +async def test_substrate_env_client_timeout(fake_stack): + client, fakes = fake_stack + await client.create("dev1") + + # Script a process killed by signal 9 (SIGKILL), exit_code 137 + fakes.processes.next_procs.append( + FakeProc( + chunks=[("stdout", b"partial output")], + signal=9, + ) + ) + + rt = SubstrateEnvClientRuntime(client=client, env_id="dev1") + # Using timeout_s = 0.0 forces duration >= timeout_s + res = await rt.exec_async("sleep 100", timeout_s=0.0) + assert res.timed_out + assert res.exit_code is None + + +async def test_substrate_env_client_file_io(fake_stack): + client, fakes = fake_stack + await client.create("dev1") + + rt = SubstrateEnvClientRuntime(client=client, env_id="dev1") + await rt.write_file_async("/workspace/test.txt", b"sample content") + data = await rt.read_file_bytes_async("/workspace/test.txt") + assert data == b"sample content" + + +async def test_substrate_env_client_error_mapping(fake_stack): + client, fakes = fake_stack + await client.create("dev1") + rt = SubstrateEnvClientRuntime(client=client, env_id="dev1") + + # 1. InvalidArgumentError -> SandboxProtocolError + with patch.object(rt._env, "shell", side_effect=InvalidArgumentError("invalid argument")): + with pytest.raises(SandboxProtocolError): + await rt.exec_async("bad") + + # 2. NotFound on cwd / executable in guest -> CommandStartError (non-retryable) + with patch.object(rt._env, "shell", side_effect=NotFoundError("cwd /missing does not exist")): + with pytest.raises(CommandStartError) as exc_info: + await rt.exec_async("bad") + assert exc_info.value.retryable is False + + # 3. NotFound on environment lifecycle -> SandboxUnavailableError (retryable) + with patch.object(rt._env, "shell", side_effect=NotFoundError('environment "dev1" not found')): + with pytest.raises(SandboxUnavailableError) as exc_info: + await rt.exec_async("bad") + assert exc_info.value.retryable is True + + # 4. RpcError (UNAVAILABLE) -> SandboxUnavailableError + with patch.object(rt._env, "shell", side_effect=RpcError("connection refused", code=grpc.StatusCode.UNAVAILABLE)): + with pytest.raises(SandboxUnavailableError): + await rt.exec_async("bad") + + +async def test_substrate_env_client_sync_exec(fake_stack): + client, fakes = fake_stack + await client.create("dev1") + fakes.processes.next_procs.append( + FakeProc( + chunks=[("stdout", b"sync hello\n")], + exit_code=0, + ) + ) + rt = SubstrateEnvClientRuntime("dev1", endpoint=client.endpoint) + try: + res = await asyncio.to_thread(rt.exec, "echo hello") + assert res.ok + assert res.stdout == "sync hello\n" + finally: + await asyncio.to_thread(rt.close)