diff --git a/clients/python/src/ate_env/__init__.py b/clients/python/src/ate_env/__init__.py index 1e3f6436..d8538ffc 100644 --- a/clients/python/src/ate_env/__init__.py +++ b/clients/python/src/ate_env/__init__.py @@ -33,9 +33,9 @@ async def main(): """ # Layer 1: Core Client & Env -from .client import DEFAULT_ATESPACE, Client - # Layer 2: High-Level Fleet Management & Abstractions +from .async_fleet import AsyncSandboxFleet +from .client import DEFAULT_ATESPACE, Client from .config import FleetConfig from .env import Env, Process @@ -92,6 +92,7 @@ async def main(): # High-level Fleet API "FleetConfig", "SandboxFleet", + "AsyncSandboxFleet", "SandboxHandle", # Low-level Errors "EnvError", @@ -131,4 +132,3 @@ async def main(): "PlanEntry", "FleetPlan", ] - diff --git a/clients/python/src/ate_env/async_fleet.py b/clients/python/src/ate_env/async_fleet.py new file mode 100644 index 00000000..0ba95cc6 --- /dev/null +++ b/clients/python/src/ate_env/async_fleet.py @@ -0,0 +1,276 @@ +# 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. + +"""`AsyncSandboxFleet` — an awaitable, event-loop-native fleet. + +Enables RL post-training frameworks (TorchRL, VeRL, SkyRL) to overlap LLM token +generation with background sandbox pre-warming and non-blocking batch acquisition. +""" + +from __future__ import annotations + +import asyncio +import inspect +import logging +from collections.abc import Callable +from typing import Any + +from .config import FleetConfig +from .fleet import SandboxFleet +from .handle import SandboxHandle +from .types import FleetPlan, Task + +logger = logging.getLogger("sandbox_sdk.async_fleet") + + +class AsyncSandboxFleet: + """Awaitable, natively asynchronous wrapper over `SandboxFleet`. + + Enables asynchronous acquisition, batch claiming, and pipelining + concurrency without blocking the asyncio event loop. + """ + + def __init__( + self, + config: FleetConfig | None = None, + driver: Any | None = None, + *, + sync_fleet: SandboxFleet | None = None, + ): + self._fleet = sync_fleet or SandboxFleet(config, driver) + + def close(self, *, wait: bool = True) -> None: + """Close fleet resources.""" + if hasattr(self.backend, "close"): + self.backend.close() + + def __del__(self) -> None: + try: + self.close(wait=False) + except Exception: + pass + + # --- Synchronous passthroughs --- + @property + def config(self) -> FleetConfig: + return self._fleet.config + + @property + def backend(self) -> Any: + return self._fleet.backend + + @property + def tasks(self) -> list[Task]: + return self._fleet.tasks + + def load_tasks(self, tasks: list[Any]) -> None: + self._fleet.load_tasks(tasks) + + def image_counts(self) -> dict[str, int]: + return self._fleet.image_counts() + + # --- Awaitable Lifecycle Methods --- + async def preflight(self) -> None: + if hasattr(self.backend, "preflight_async"): + await self.backend.preflight_async() + else: + await asyncio.to_thread(self._fleet.preflight) + + async def plan(self) -> FleetPlan: + return await asyncio.to_thread(self._fleet.plan) + + async def setup(self, tasks: list[Any] | None = None) -> AsyncSandboxFleet: + if tasks is not None: + self._fleet.load_tasks(tasks) + await self.preflight() + plan = await self.plan() + for entry in plan.entries: + if entry.replicas > 0: + await self._warm_pool_for_template( + entry.template_id, + entry.replicas, + wait=True, + rootfs_image=getattr(entry, "rootfs_image", None), + ) + return self + + async def _ensure_template_for_image_async(self, image: str) -> tuple[str, str | None]: + cached = self._fleet._image_to_template.get(image) + if cached is not None: + return cached + return await asyncio.to_thread(self._fleet._ensure_template_for_image, image) + + async def _warm_pool_for_template( + self, + template_id: str, + replicas: int, + wait: bool = True, + rootfs_image: str | None = None, + ) -> None: + if hasattr(self.backend, "warm_pool_async"): + await self.backend.warm_pool_async( + template_id, replicas, wait=wait, rootfs_image=rootfs_image + ) + else: + await asyncio.to_thread( + self.backend.warm_pool, + template_id, + replicas, + wait=wait, + rootfs_image=rootfs_image, + ) + + async def warm_images( + self, images: list[str], replicas: int | None = None, wait: bool = True + ) -> None: + rep = replicas if replicas is not None else self.config.max_warmpool_replicas + for img in images: + template_id, rootfs_image = await self._ensure_template_for_image_async(img) + await self._warm_pool_for_template(template_id, rep, wait=wait, rootfs_image=rootfs_image) + + async def unwarm_image(self, image: str) -> None: + entry = self._fleet._image_to_template.get(image) + if entry: + template_id, rootfs_image = entry + if hasattr(self.backend, "unwarm_pool_async"): + await self.backend.unwarm_pool_async(template_id, rootfs_image=rootfs_image) + else: + await asyncio.to_thread( + self.backend.unwarm_pool, template_id, rootfs_image=rootfs_image + ) + + async def acquire(self, task: Task | str, timeout_s: float | None = None) -> SandboxHandle: + """Asynchronously acquire a single sandbox.""" + 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 + rootfs_image = None + else: + template_id, rootfs_image = await self._ensure_template_for_image_async(task_obj.image) + to_s = timeout_s or self.config.acquire_timeout_s + + if hasattr(self.backend, "acquire_async"): + raw_inst = await self.backend.acquire_async( + template_id, self._fleet.run_id, timeout_s=to_s, rootfs_image=rootfs_image + ) + else: + raw_inst = await asyncio.to_thread( + self.backend.acquire, + template_id, + self._fleet.run_id, + timeout_s=to_s, + rootfs_image=rootfs_image, + ) + + try: + runtime, data_plane = self._fleet._build_runtime(raw_inst) + except Exception: + try: + if hasattr(self.backend, "release_async"): + await self.backend.release_async(raw_inst.instance_id, recycle=False) + else: + await asyncio.to_thread(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._fleet.run_id, + backend=self.backend, + runtime=runtime, + data_plane=data_plane, + default_timeout_s=self.config.step_timeout_s, + ) + + async def acquire_batch( + self, tasks: list[Task | str], timeout_s: float | None = None + ) -> list[SandboxHandle]: + """Asynchronously acquire a batch of sandboxes in parallel.""" + return list( + await asyncio.gather(*(self.acquire(t, timeout_s=timeout_s) for t in tasks)) + ) + + async def release(self, handle: SandboxHandle) -> None: + await handle.release_async() + + async def teardown(self) -> None: + if hasattr(self.backend, "reap_async"): + await self.backend.reap_async(self._fleet.run_id) + else: + await asyncio.to_thread(self.backend.reap, self._fleet.run_id) + + # --- Async Context Manager --- + async def __aenter__(self) -> AsyncSandboxFleet: + await self.setup() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: + try: + await self.teardown() + finally: + self.close() + + # --- Parallel Processing --- + async def _call_process_fn( + self, process_fn: Callable[..., Any], task: Task, handle: SandboxHandle + ) -> Any: + if inspect.iscoroutinefunction(process_fn) or inspect.iscoroutinefunction( + getattr(process_fn, "__call__", None) + ): + return await process_fn(task, handle) + result = await asyncio.to_thread(process_fn, task, handle) + if inspect.isawaitable(result): + return await result + return result + + async def run( + self, + process_fn: Callable[[Task, SandboxHandle], Any], + concurrency: int | None = None, + ) -> list[Any]: + """Execute tasks asynchronously with bounded concurrency.""" + c = concurrency or min(self.config.max_concurrent, len(self.tasks) or 1) + sem = asyncio.Semaphore(max(1, c)) + results: list[Any] = [None] * len(self.tasks) + + async def _worker(idx: int, task: Task) -> None: + async with sem: + handle = await self.acquire(task) + try: + res = await self._call_process_fn(process_fn, task, handle) + results[idx] = res + handle.release() + except Exception as e: + logger.error("Async execution failed for task %s: %s", task.id, e) + handle.release() + results[idx] = e + + await asyncio.gather(*(_worker(i, t) for i, t in enumerate(self.tasks))) + return results diff --git a/clients/python/src/ate_env/backend/base.py b/clients/python/src/ate_env/backend/base.py index 2f9f5eff..626632f3 100644 --- a/clients/python/src/ate_env/backend/base.py +++ b/clients/python/src/ate_env/backend/base.py @@ -15,6 +15,7 @@ from __future__ import annotations from abc import ABC, abstractmethod +from typing import Any from ..types import EnvironmentSpec, RawSandboxInstance @@ -36,15 +37,28 @@ 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: + def warm_pool( + self, + template_id: str, + replicas: int, + wait: bool = True, + run_id: str | None = None, + rootfs_image: str | None = None, + ) -> None: """Provision warm replicas (Ready pods in K8s or Paused Actors in Substrate).""" @abstractmethod - def unwarm_pool(self, template_id: str) -> None: + def unwarm_pool(self, template_id: str, rootfs_image: str | None = None) -> 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: + def acquire( + self, + template_id: str, + run_id: str, + timeout_s: float = 180.0, + rootfs_image: str | None = None, + ) -> RawSandboxInstance: """Instantly claim a sandbox instance from warm pool or golden.""" @abstractmethod @@ -54,3 +68,9 @@ def release(self, instance_id: str, recycle: bool = False) -> None: @abstractmethod def reap(self, run_id: str) -> int: """Force delete all resources associated with a run_id.""" + + def build_runtime( + self, instance: RawSandboxInstance, data_plane: str = "ate_env" + ) -> tuple[Any, Any | None]: + """Construct runtime guest hook and data plane coordinates for this sandbox.""" + raise NotImplementedError diff --git a/clients/python/src/ate_env/backend/mock.py b/clients/python/src/ate_env/backend/mock.py index 23e06d9e..b7ebfc0c 100644 --- a/clients/python/src/ate_env/backend/mock.py +++ b/clients/python/src/ate_env/backend/mock.py @@ -31,7 +31,7 @@ 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.warm_pools: dict[str | tuple[str, str | None], int] = collections.defaultdict(int) self.instances: dict[str, RawSandboxInstance] = {} self.reaped_runs: set[str] = set() @@ -44,23 +44,44 @@ def ensure_template(self, env: EnvironmentSpec) -> str: self.templates[tid] = env return tid - def warm_pool(self, template_id: str, replicas: int, wait: bool = True) -> None: + def warm_pool( + self, + template_id: str, + replicas: int, + wait: bool = True, + run_id: str | None = None, + rootfs_image: str | None = None, + ) -> None: self.warm_pools[template_id] = replicas + if rootfs_image is not None: + self.warm_pools[(template_id, rootfs_image)] = replicas - def unwarm_pool(self, template_id: str) -> None: + def unwarm_pool(self, template_id: str, rootfs_image: str | None = None) -> None: self.warm_pools[template_id] = 0 + if rootfs_image is not None: + self.warm_pools.pop((template_id, rootfs_image), None) - def acquire(self, template_id: str, run_id: str, timeout_s: float = 180.0) -> RawSandboxInstance: + def acquire( + self, + template_id: str, + run_id: str, + timeout_s: float = 180.0, + rootfs_image: str | None = None, + ) -> RawSandboxInstance: if self.fail_acquire: raise SandboxStartError("Mock acquire failed: capacity exhausted") inst_id = f"mock-sb-{uuid.uuid4().hex[:8]}" + metadata = {"created_at": time.time(), "template": template_id} + if rootfs_image is not None: + metadata["rootfs_image"] = rootfs_image 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} + metadata=metadata, + rootfs_image=rootfs_image, ) self.instances[inst_id] = inst return inst diff --git a/clients/python/src/ate_env/backend/substrate.py b/clients/python/src/ate_env/backend/substrate.py index 52ff54c9..e6d7065b 100644 --- a/clients/python/src/ate_env/backend/substrate.py +++ b/clients/python/src/ate_env/backend/substrate.py @@ -38,6 +38,36 @@ logger = logging.getLogger("ate_env.backend.substrate") +class _WarmPoolDict(dict): + """Warm paused pool dictionary keyed by (template, image). + + Allows string keys for backwards compatibility, treating 'template' as (template, None). + """ + + def _normalize(self, key: Any) -> Any: + if isinstance(key, str): + return (key, None) + return key + + def __getitem__(self, key: Any) -> Any: + return super().__getitem__(self._normalize(key)) + + def __setitem__(self, key: Any, value: Any) -> None: + super().__setitem__(self._normalize(key), value) + + def __contains__(self, key: Any) -> bool: + return super().__contains__(self._normalize(key)) + + def get(self, key: Any, default: Any = None) -> Any: + return super().get(self._normalize(key), default) + + def pop(self, key: Any, *args: Any) -> Any: + return super().pop(self._normalize(key), *args) + + def setdefault(self, key: Any, default: Any = None) -> Any: + return super().setdefault(self._normalize(key), default) + + class SubstrateBackendDriver(BackendDriver): """Agent Substrate Backend Driver. @@ -89,7 +119,7 @@ def __init__( self._lock = threading.Lock() self._owned_actors: dict[str, RawSandboxInstance] = {} - self._warm_paused_pool: dict[str, list[str]] = {} + self._warm_paused_pool: _WarmPoolDict = _WarmPoolDict() self._background_tasks: set[asyncio.Task[Any]] = set() def _get_loop_runner(self) -> LoopRunner: @@ -170,7 +200,9 @@ def ensure_template(self, env: EnvironmentSpec | str) -> str: return env.template_key() return "default-template" - def _make_instance(self, actor_id: str, template_id: str, run_id: str) -> RawSandboxInstance: + def _make_instance( + self, actor_id: str, template_id: str, run_id: str, rootfs_image: str | None = None + ) -> RawSandboxInstance: planes = { "ate_env": DataPlaneEndpoint( protocol="grpc", @@ -183,6 +215,8 @@ def _make_instance(self, actor_id: str, template_id: str, run_id: str) -> RawSan "template": template_id, "atespace": self.atespace, } + if rootfs_image: + metadata["rootfs_image"] = rootfs_image if self.labels: metadata["labels"] = dict(self.labels) return RawSandboxInstance( @@ -192,37 +226,53 @@ def _make_instance(self, actor_id: str, template_id: str, run_id: str) -> RawSan run_id=run_id, metadata=metadata, data_planes=planes, + rootfs_image=rootfs_image, ) async def acquire_async( - self, template_id: str, run_id: str, timeout_s: float = 180.0 + self, + template_id: str, + run_id: str, + timeout_s: float = 180.0, + rootfs_image: str | None = None, ) -> 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 + self._acquire_internal(template_id, run_id, rootfs_image=rootfs_image), 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: + async def _acquire_internal( + self, + template_id: str, + run_id: str, + rootfs_image: str | None = None, + ) -> 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. """ + pool_key = (template_id, rootfs_image) while True: actor_id: str | None = None with self._lock: - warm_list = self._warm_paused_pool.get(template_id) + warm_list = self._warm_paused_pool.get(pool_key) 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) + logger.info( + "Claimed warm paused actor %s for template %s (rootfs_image=%s)", + actor_id, + template_id, + rootfs_image, + ) # Verify warm actor still exists on the server try: await asyncio.wait_for( @@ -237,16 +287,17 @@ async def _acquire_internal(self, template_id: str, run_id: str) -> RawSandboxIn except Exception as e: # Put back on unexpected error and raise with self._lock: - self._warm_paused_pool.setdefault(template_id, []).append(actor_id) + self._warm_paused_pool.setdefault(pool_key, []).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)...", + "Creating environment %s (template=%s, atespace=%s, rootfs_image=%s)...", actor_id, template_id, self.atespace, + rootfs_image, ) last_err: Exception | None = None for attempt in range(4): @@ -255,6 +306,7 @@ async def _acquire_internal(self, template_id: str, run_id: str) -> RawSandboxIn actor_id, template_name=template_id, atespace=self.atespace, + rootfs_image=rootfs_image, ) last_err = None break @@ -273,23 +325,40 @@ async def _acquire_internal(self, template_id: str, run_id: str) -> RawSandboxIn 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) + instance = self._make_instance(actor_id, template_id, run_id, rootfs_image=rootfs_image) 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 + self, + template_id: str, + run_id: str, + timeout_s: float = 180.0, + rootfs_image: str | None = None, ) -> RawSandboxInstance: - return self._sync(self.acquire_async(template_id, run_id, timeout_s=timeout_s)) + return self._sync( + self.acquire_async(template_id, run_id, timeout_s=timeout_s, rootfs_image=rootfs_image) + ) async def warm_pool_async( - self, template_id: str, replicas: int, wait: bool = True, run_id: str | None = None + self, + template_id: str, + replicas: int, + wait: bool = True, + run_id: str | None = None, + rootfs_image: str | None = 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) + logger.info( + "Sizing warm pool for template %s (rootfs_image=%s) to %d paused actors...", + template_id, + rootfs_image, + replicas, + ) + pool_key = (template_id, rootfs_image) with self._lock: - current = list(self._warm_paused_pool.get(template_id, [])) + current = list(self._warm_paused_pool.get(pool_key, [])) needed = replicas - len(current) if needed <= 0: return @@ -305,6 +374,7 @@ async def _warm_one() -> str | None: actor_id, template_name=template_id, atespace=self.atespace, + rootfs_image=rootfs_image, ) if wait: await self.client.suspend(actor_id, atespace=self.atespace) @@ -333,24 +403,46 @@ def _done(t: asyncio.Task[Any]) -> None: if isinstance(res, SandboxStartError): raise res logger.warning("Error provisioning warm actor: %s", res) - elif res is not None: + elif isinstance(res, str): aid = res with self._lock: - self._warm_paused_pool.setdefault(template_id, []).append(aid) + self._warm_paused_pool.setdefault(pool_key, []).append(aid) self._owned_actors[aid] = self._make_instance( - aid, template_id, run_id=effective_run_id + aid, template_id, run_id=effective_run_id, rootfs_image=rootfs_image ) def warm_pool( - self, template_id: str, replicas: int, wait: bool = True, run_id: str | None = None + self, + template_id: str, + replicas: int, + wait: bool = True, + run_id: str | None = None, + rootfs_image: str | None = None, ) -> None: - self._sync(self.warm_pool_async(template_id, replicas, wait=wait, run_id=run_id)) + self._sync( + self.warm_pool_async( + template_id, + replicas, + wait=wait, + run_id=run_id, + rootfs_image=rootfs_image, + ) + ) - async def unwarm_pool_async(self, template_id: str) -> None: + async def unwarm_pool_async(self, template_id: str, rootfs_image: str | None = None) -> None: """Drain and delete paused actors in the warm pool.""" - logger.info("Unwarming pool for template %s...", template_id) + logger.info("Unwarming pool for template %s (rootfs_image=%s)...", template_id, rootfs_image) with self._lock: - to_delete = self._warm_paused_pool.pop(template_id, []) + if rootfs_image is not None: + to_delete = self._warm_paused_pool.pop((template_id, rootfs_image), []) + else: + keys_to_pop = [ + k for k in list(self._warm_paused_pool.keys()) + if (isinstance(k, tuple) and k[0] == template_id) + ] + to_delete = [] + for k in keys_to_pop: + to_delete.extend(self._warm_paused_pool.pop(k, [])) sem = asyncio.Semaphore(16) @@ -365,8 +457,8 @@ async def _delete_one(actor_id: str) -> 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)) + def unwarm_pool(self, template_id: str, rootfs_image: str | None = None) -> None: + self._sync(self.unwarm_pool_async(template_id, rootfs_image=rootfs_image)) async def release_async(self, instance_id: str, recycle: bool = False) -> None: """Release actor: deletes environment, or suspends if recycling. @@ -380,8 +472,9 @@ async def release_async(self, instance_id: str, recycle: bool = False) -> None: if recycle and instance: try: await self.client.suspend(instance_id, atespace=self.atespace) + pool_key = (instance.template_id, instance.rootfs_image) with self._lock: - self._warm_paused_pool.setdefault(instance.template_id, []).append(instance_id) + self._warm_paused_pool.setdefault(pool_key, []).append(instance_id) self._owned_actors[instance_id] = instance logger.info("Recycled actor %s to paused warm pool", instance_id) return diff --git a/clients/python/src/ate_env/fleet.py b/clients/python/src/ate_env/fleet.py index 0fdca4fa..7e374df7 100644 --- a/clients/python/src/ate_env/fleet.py +++ b/clients/python/src/ate_env/fleet.py @@ -50,7 +50,7 @@ def __init__(self, config: FleetConfig | None = None, driver: BackendDriver | No self.config = config or FleetConfig() self.run_id = uuid.uuid4().hex[:12] self._tasks: list[Task] = [] - self._image_to_template: dict[str, str] = {} + self._image_to_template: dict[str, tuple[str, str | None]] = {} self._plan: FleetPlan | None = None self._lock = threading.Lock() @@ -140,7 +140,7 @@ def plan(self) -> FleetPlan: entries: list[PlanEntry] = [] counts = self.image_counts() for img, task_count in counts.items(): - template_id = self._ensure_template_for_image(img) + template_id, rootfs_image = 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( @@ -148,6 +148,7 @@ def plan(self) -> FleetPlan: template_id=template_id, replicas=replicas, tasks=task_count, + rootfs_image=rootfs_image, )) self._plan = FleetPlan(entries) return self._plan @@ -160,20 +161,26 @@ def setup(self, tasks: list[Any] | None = None) -> None: plan = self.plan() for entry in plan.entries: if entry.replicas > 0: - self.backend.warm_pool(entry.template_id, entry.replicas, wait=True) + self.backend.warm_pool( + entry.template_id, + entry.replicas, + wait=True, + rootfs_image=getattr(entry, "rootfs_image", None), + ) def warm_images(self, images: list[str], replicas: int | None = 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) + template_id, rootfs_image = self._ensure_template_for_image(img) + self.backend.warm_pool(template_id, rep, wait=wait, rootfs_image=rootfs_image) 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) + entry = self._image_to_template.get(image) + if entry: + template_id, rootfs_image = entry + self.backend.unwarm_pool(template_id, rootfs_image=rootfs_image) def acquire( self, @@ -206,11 +213,14 @@ def _acquire_one( if task_obj.template: template_id = task_obj.template + rootfs_image = None else: - template_id = self._ensure_template_for_image(task_obj.image) + template_id, rootfs_image = 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) + raw_inst = self.backend.acquire( + template_id, self.run_id, timeout_s=to_s, rootfs_image=rootfs_image + ) try: runtime, data_plane = self._build_runtime(raw_inst) @@ -278,17 +288,15 @@ def run(self, process_fn: Callable[[Task, SandboxHandle], Any], 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: + def _ensure_template_for_image(self, image: str) -> tuple[str, str | None]: with self._lock: if image in self._image_to_template: return self._image_to_template[image] + res: tuple[str, str | None] 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 + res = (self.config.image_templates[image], None) + else: + base = self.config.template or "default-template" + res = (base, image) + self._image_to_template[image] = res + return res diff --git a/clients/python/src/ate_env/handle.py b/clients/python/src/ate_env/handle.py index 4b336813..8b6167d1 100644 --- a/clients/python/src/ate_env/handle.py +++ b/clients/python/src/ate_env/handle.py @@ -123,6 +123,19 @@ def release(self) -> None: finally: self.runtime.close() + async def release_async(self) -> None: + """Asynchronously release this sandbox back to the backend or terminate it.""" + try: + if hasattr(self.backend, "release_async"): + await self.backend.release_async(self.sandbox_id, recycle=False) + else: + self.backend.release(self.sandbox_id, recycle=False) + finally: + if hasattr(self.runtime, "close_async"): + await self.runtime.close_async() + else: + self.runtime.close() + def recycle(self) -> None: """Return sandbox to warm pool. @@ -135,6 +148,24 @@ def recycle(self) -> None: finally: self.runtime.close() + async def recycle_async(self) -> None: + """Asynchronously 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: + if hasattr(self.backend, "release_async"): + await self.backend.release_async(self.sandbox_id, recycle=True) + else: + self.backend.release(self.sandbox_id, recycle=True) + finally: + if hasattr(self.runtime, "close_async"): + await self.runtime.close_async() + else: + self.runtime.close() + def __enter__(self) -> SandboxHandle: return self @@ -145,4 +176,4 @@ async def __aenter__(self) -> SandboxHandle: return self async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: - self.release() + await self.release_async() diff --git a/clients/python/src/ate_env/types.py b/clients/python/src/ate_env/types.py index c2ed06e1..eee660d0 100644 --- a/clients/python/src/ate_env/types.py +++ b/clients/python/src/ate_env/types.py @@ -242,13 +242,23 @@ def template_key(self) -> str: @dataclass class Task: - """Single workload/task definition (e.g. one SWE-bench instance).""" + """Single workload/task definition (e.g. one SWE-bench instance). + + Attributes: + id: Unique identifier for this task. + image: OCI container image reference. For Substrate backends, if mapped in + image_templates the mapped ActorTemplate is used; otherwise it is passed + as rootfs_image to derive a template on the fly. + metadata: Optional dictionary of arbitrary task metadata. + template: Optional explicit ActorTemplate name override. + """ id: str image: str metadata: dict[str, Any] = field(default_factory=dict) template: str | None = None + @dataclass class ExecResult: """Outcome of a command that the guest actually ran. @@ -319,6 +329,7 @@ class RawSandboxInstance: 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) + rootfs_image: str | None = None @dataclass @@ -328,6 +339,7 @@ class PlanEntry: template_id: str replicas: int tasks: int + rootfs_image: str | None = None class FleetPlan: diff --git a/clients/python/tests/test_async_fleet.py b/clients/python/tests/test_async_fleet.py new file mode 100644 index 00000000..ca5c6b21 --- /dev/null +++ b/clients/python/tests/test_async_fleet.py @@ -0,0 +1,145 @@ +# 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 asyncio + +import pytest + +from ate_env.async_fleet import AsyncSandboxFleet +from ate_env.config import FleetConfig +from ate_env.types import Task + + +@pytest.mark.asyncio +async def test_async_fleet_setup_and_acquire(): + cfg = FleetConfig(backend="mock", max_concurrent=4) + fleet = AsyncSandboxFleet(cfg) + + tasks = [ + Task(id=f"task-{i}", image="repo-img-1") for i in range(4) + ] + fleet.load_tasks(tasks) + + # Async setup + await fleet.setup() + + # Async acquire single + handle = await fleet.acquire(tasks[0]) + assert handle.sandbox_id.startswith("mock-sb-") + + # Command execution + res = handle.exec("echo hello") + assert "echo hello" in res.stdout + + await fleet.release(handle) + await fleet.teardown() + fleet.close() + + +@pytest.mark.asyncio +async def test_async_fleet_acquire_batch(): + cfg = FleetConfig(backend="mock", max_concurrent=4) + fleet = AsyncSandboxFleet(cfg) + + tasks = [ + Task(id=f"task-batch-{i}", image=f"repo-img-{i % 2}") for i in range(4) + ] + fleet.load_tasks(tasks) + await fleet.setup() + + # Async acquire batch concurrently + handles = await fleet.acquire_batch(tasks) + assert len(handles) == 4 + for h in handles: + assert h.sandbox_id.startswith("mock-sb-") + await fleet.release(h) + + await fleet.teardown() + fleet.close() + + +@pytest.mark.asyncio +async def test_async_context_manager(): + cfg = FleetConfig(backend="mock") + async with AsyncSandboxFleet(cfg) as fleet: + task = Task(id="task-ctx", image="repo-img-1") + handle = await fleet.acquire(task) + assert handle is not None + + # Handle async context manager + async with handle: + res = handle.exec("echo in-handle-ctx") + assert "in-handle-ctx" in res.stdout + + +@pytest.mark.asyncio +async def test_async_run_parallel(): + cfg = FleetConfig(backend="mock", max_concurrent=4) + fleet = AsyncSandboxFleet(cfg) + tasks = [Task(id=f"t-{i}", image="img-1") for i in range(6)] + fleet.load_tasks(tasks) + await fleet.setup() + + async def async_worker(task, handle): + await asyncio.sleep(0.01) + return f"done-{task.id}-{handle.sandbox_id[:8]}" + + results = await fleet.run(async_worker, concurrency=3) + assert len(results) == 6 + for i, r in enumerate(results): + assert r.startswith(f"done-t-{i}") + + await fleet.teardown() + fleet.close() + + +@pytest.mark.asyncio +async def test_handle_exec_and_context_manager_release(): + cfg = FleetConfig(backend="mock") + fleet = AsyncSandboxFleet(cfg) + task = Task(id="test-rel", image="img-1") + + # 1. Async context manager defaults to release (not recycle) + async with await fleet.acquire(task) as handle: + res = await handle.exec_async("echo async test") + assert res.ok + assert res.exit_code == 0 + assert "async test" in res.stdout + sb_id = handle.sandbox_id + assert sb_id in fleet.backend.instances + + # Exiting context manager released (deleted) the instance + assert sb_id not in fleet.backend.instances + + # 2. Sync context manager defaults to release + sync_handle = fleet._fleet.acquire(task) + with sync_handle as h: + res = h.exec("echo sync test") + assert res.ok + assert res.exit_code == 0 + assert "sync test" in res.stdout + sb_id2 = h.sandbox_id + assert sb_id2 in fleet.backend.instances + + assert sb_id2 not in fleet.backend.instances + + +@pytest.mark.asyncio +async def test_async_fleet_step_timeout_applies_to_handle(): + cfg = FleetConfig(backend="mock", step_timeout_s=42.0) + fleet = AsyncSandboxFleet(cfg) + h = await fleet.acquire(Task(id="t1", image="img-1")) + assert h.default_timeout_s == 42.0 + await fleet.release(h) + fleet.close() diff --git a/clients/python/tests/test_exec_semantics.py b/clients/python/tests/test_exec_semantics.py index 4f81afaa..08398ce8 100644 --- a/clients/python/tests/test_exec_semantics.py +++ b/clients/python/tests/test_exec_semantics.py @@ -60,6 +60,7 @@ def test_check_raises_command_execution_error_with_result(handle): assert ei.value.result.exit_code == 2 res = handle.exec("failing-cmd", check=False) assert res.exit_code == 2 + assert res.stdout == "" assert res.stderr == "boom" @@ -100,8 +101,6 @@ def test_command_contract(): 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", diff --git a/clients/python/tests/test_fleet_strategies.py b/clients/python/tests/test_fleet_strategies.py index a317c102..b6715dc0 100644 --- a/clients/python/tests/test_fleet_strategies.py +++ b/clients/python/tests/test_fleet_strategies.py @@ -21,7 +21,7 @@ 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 + return res.stdout @pytest.mark.parametrize("strategy", ["none", "naive", "sliding", "pipelined"]) @@ -45,31 +45,39 @@ def test_fleet_strategies(strategy: str): 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(): +def test_user_guide_section_1_synchronous_batch(): 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() + with SandboxFleet(config) as fleet: + sandboxes = fleet.acquire(tasks) + try: + for task, sb in zip(tasks, sandboxes): + res = sb.exec("python3 -c 'print(\"hello world\")'") + assert res.ok + assert res.exit_code == 0 + finally: + fleet.release(sandboxes) + + +async def test_user_guide_section_2_async_rollouts(): + from ate_env.async_fleet import AsyncSandboxFleet + + config = FleetConfig( + backend="mock", + batch_size=4, + max_warmpool_replicas=4, + ) + tasks = [Task(id=f"t-{i}", image="docker.io/library/python:3.11") for i in range(4)] + fleet = AsyncSandboxFleet(config) + await fleet.setup(tasks) + + sandboxes = await fleet.acquire_batch(tasks) + for sb in sandboxes: + async with sb: + await sb.write_file_async("/testbed/calc.py", "def add(a, b): return a + b\n") + res = await sb.exec_async(["python3", "-c", "import calc; assert calc.add(2, 3) == 5"]) + assert res.ok + await fleet.teardown() diff --git a/clients/python/tests/test_substrate_driver.py b/clients/python/tests/test_substrate_driver.py index a5d5d585..4f10a1e3 100644 --- a/clients/python/tests/test_substrate_driver.py +++ b/clients/python/tests/test_substrate_driver.py @@ -91,8 +91,10 @@ async def test_substrate_synchronous_fleet_multiple_acquires(fake_stack): 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")) + img1 = "docker.io/library/python@sha256:" + "a" * 64 + img2 = "docker.io/library/python@sha256:" + "b" * 64 + sb1 = fleet.acquire(Task(id="task-1", image=img1)) + sb2 = fleet.acquire(Task(id="task-2", image=img2)) assert sb1.sandbox_id != sb2.sandbox_id fleet.release([sb1, sb2]) finally: @@ -112,7 +114,8 @@ async def test_substrate_fleet_setup_and_teardown_cleans_warm_pool(fake_stack): def _run_setup_teardown(): fleet = SandboxFleet(cfg) - fleet.load_tasks([Task(id="t1", image="img"), Task(id="t2", image="img")]) + img = "docker.io/library/python@sha256:" + "c" * 64 + fleet.load_tasks([Task(id="t1", image=img), Task(id="t2", image=img)]) fleet.setup() assert len(fakes.environments.environments) == 2 fleet.teardown() @@ -132,24 +135,93 @@ async def test_substrate_preflight_fails_on_closed_port(): def test_substrate_default_template_resolution(): - # B1: A fleet with no template configured falls back to "default-template" + # B1: A fleet with no template configured falls back to ("default-template", image) cfg = FleetConfig(backend="substrate") fleet = SandboxFleet(cfg) - tmpl = fleet._ensure_template_for_image("docker.io/library/python@sha256:2176bb31f3eb") + tmpl, rootfs = fleet._ensure_template_for_image("docker.io/library/python@sha256:2176bb31f3eb") assert tmpl == "default-template" + assert rootfs == "docker.io/library/python@sha256:2176bb31f3eb" - # M5: image_templates mapping wins over config.template + # M5: image_templates mapping gives (mapped, None); any other gives (config.template, image) 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" + assert fleet_mapped._ensure_template_for_image("docker.io/library/python:3.11") == ("python-template", None) + assert fleet_mapped._ensure_template_for_image("other-img") == ("fallback-template", "other-img") # 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" + + +async def test_substrate_rootfs_image_sent_to_server(fake_stack): + from ate_env.async_fleet import AsyncSandboxFleet + + client, fakes = fake_stack + image = "docker.io/library/python@sha256:" + "d" * 64 + + # 1. Default FleetConfig with unmapped task image sends template.name == "default-template" and rootfs_image == image + cfg = FleetConfig(backend="substrate", endpoint=client.endpoint) + + def _run_sync_default(): + fleet = SandboxFleet(cfg) + try: + task = Task(id="task-rootfs", image=image) + sb = fleet.acquire(task) + assert fakes.environments.last_create_request is not None + assert fakes.environments.last_create_request.template.name == "default-template" + assert fakes.environments.last_create_request.rootfs_image == image + fleet.release(sb) + finally: + fleet.close() + + await asyncio.to_thread(_run_sync_default) + + # 2. AsyncSandboxFleet with unmapped task image also sends rootfs_image + async_fleet = AsyncSandboxFleet(cfg) + try: + task_async = Task(id="task-async", image=image) + sb_async = await async_fleet.acquire(task_async) + assert fakes.environments.last_create_request is not None + assert fakes.environments.last_create_request.template.name == "default-template" + assert fakes.environments.last_create_request.rootfs_image == image + await async_fleet.release(sb_async) + finally: + async_fleet.close() + + # 3. Warm pool passes rootfs_image to client.create and keys _warm_paused_pool by (template, image) + driver = SubstrateBackendDriver(client=client, atespace="ate-env") + try: + await driver.warm_pool_async("custom-base", replicas=1, rootfs_image=image) + assert fakes.environments.last_create_request is not None + assert fakes.environments.last_create_request.template.name == "custom-base" + assert fakes.environments.last_create_request.rootfs_image == image + assert len(driver._warm_paused_pool[("custom-base", image)]) == 1 + finally: + await driver.unwarm_pool_async("custom-base", rootfs_image=image) + await driver.close_async() + + # 4. Mapped image gives (mapped, None) and does not send rootfs_image + cfg_mapped = FleetConfig( + backend="substrate", + endpoint=client.endpoint, + image_templates={image: "mapped-template"}, + ) + + def _run_sync_mapped(): + fleet_mapped = SandboxFleet(cfg_mapped) + try: + sb_mapped = fleet_mapped.acquire(Task(id="task-mapped", image=image)) + assert fakes.environments.last_create_request is not None + assert fakes.environments.last_create_request.template.name == "mapped-template" + assert fakes.environments.last_create_request.rootfs_image == "" + fleet_mapped.release(sb_mapped) + finally: + fleet_mapped.close() + + await asyncio.to_thread(_run_sync_mapped)