diff --git a/clients/python/USER_GUIDE.md b/clients/python/USER_GUIDE.md new file mode 100644 index 00000000..73145f39 --- /dev/null +++ b/clients/python/USER_GUIDE.md @@ -0,0 +1,109 @@ +# Reinforcement Learning & Sandbox Fleet Guide + +This guide covers how to use high-throughput **Sandbox Fleet Orchestration (`SandboxFleet` & `AsyncSandboxFleet`)** in `ate_env` with distributed Reinforcement Learning frameworks like **VeRL**, **Ray**, and **NeMo-Gym** on top of **Agent Substrate** and **GKE**. + +--- + +## 1. Quickstart + +### Synchronous Batch Execution + +```python +from ate_env import FleetConfig, SandboxFleet, Task + +config = FleetConfig( + backend="substrate", + endpoint="localhost:7777", +) + +tasks = [ + Task(id="task-1", image="docker.io/library/python:3.11"), + Task(id="task-2", image="docker.io/library/python:3.11"), +] + +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\")'") + print(f"Task {task.id}: {res.stdout.strip()} (ok={res.ok}, exit_code={res.exit_code})") + finally: + fleet.release(sandboxes) +``` + +--- + +## 2. Asynchronous Non-Blocking Rollouts (VeRL / vLLM) + +In RL post-training (e.g., GRPO), candidate generation overlaps with sandbox pre-warming to hide provisioning latency: + +```python +import asyncio +from ate_env import AsyncSandboxFleet, FleetConfig, Task + +async def run_rollouts(): + config = FleetConfig( + backend="substrate", + endpoint="ate-env-api.ate-system.svc.cluster.local:7777", + 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) + + # 1. Start acquiring pre-warmed sandboxes concurrently while LLM generates tokens + warm_task = asyncio.create_task(fleet.acquire_batch(tasks)) + await asyncio.sleep(1.5) # Simulate token generation + sandboxes = await warm_task + + # 2. Asynchronously evaluate candidate rollouts + 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"]) + print(f"Sandbox {sb.sandbox_id}: ok={res.ok}") + + await fleet.teardown() + +asyncio.run(run_rollouts()) +``` + +--- + +## 3. Architecture & Execution Semantics + +The Sandbox Fleet SDK is architected in two clean layers: +- **Layer 1 (`Client` / `Env`)**: Direct asyncio gRPC client talking to `ate-env-api` for lifecycle (create, get, suspend, delete) and guest daemon execution (streaming processes and chunked filesystem I/O). +- **Layer 2 (`SandboxFleet` / `AsyncSandboxFleet`)**: Multi-tenant fleet management, sizing, pre-warming, pooling, and automated lifecycle with task-bound `SandboxHandle`. + +### Execution Contract: +- `sb.exec(cmd)` and `sb.exec_async(cmd)` return an `ExecResult(exit_code, stdout, stderr, duration_s, timed_out)`. +- Non-zero command outcomes exit with an `exit_code` or `timed_out=True`. +- Infrastructure and transport errors raise `InfrastructureError` subclasses (e.g. `SandboxUnavailableError`), guaranteeing that infrastructure problems are retried and never scored as agent errors. + +### Design Considerations & Warm Pool Performance + +Empirical testing on `glottman-sandbox-test-1` evaluated warm (suspended) pool claims against fresh instantiation from golden snapshots: +- **Claim from warm (suspended) pool**: ran its first command in **2.6s to 3.7s**. +- **Fresh create from golden snapshot**: took **0.4s** to create plus **2.9s** to first command (**3.3s total**). + +The warm pool does not significantly decrease the time to first command because resuming a suspended actor still incurs wake-up and connection establishment latency. Meanwhile, maintaining a warm pool incurs ongoing suspended actor cluster resource footprint and teardown lifecycle cleanup costs. + +--- + +## 4. End-to-End VeRL + Ray Demo on GKE + +A full runnable example with Ray actors and GRPO advantage optimization is located in [`examples/verl_swebench`](../../examples/verl_swebench): + +* **Local hermetic run**: + ```bash + python examples/verl_swebench/async_verl_swebench_pipeline.py --backend mock --num-iters 2 --group-size 2 + ``` +* **Production RayJob on GKE**: + ```bash + kubectl apply -f examples/verl_swebench/ray-job.async-verl.yaml + ``` + See [`examples/verl_swebench/README.md`](../../examples/verl_swebench/README.md) for full deployment instructions. 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..2e720c34 --- /dev/null +++ b/clients/python/src/ate_env/async_fleet.py @@ -0,0 +1,275 @@ +# 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, + ) + + 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 bad0ed84..dc3a06fa 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) @@ -277,17 +287,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 0199b499..2217bbe9 100644 --- a/clients/python/src/ate_env/handle.py +++ b/clients/python/src/ate_env/handle.py @@ -120,6 +120,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. @@ -132,6 +145,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 @@ -142,4 +173,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..f4b89bc3 --- /dev/null +++ b/clients/python/tests/test_async_fleet.py @@ -0,0 +1,135 @@ +# 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 diff --git a/clients/python/tests/test_exec_semantics.py b/clients/python/tests/test_exec_semantics.py index db116ea2..c81a2867 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) diff --git a/examples/verl_swebench/README.md b/examples/verl_swebench/README.md new file mode 100644 index 00000000..6862aa1c --- /dev/null +++ b/examples/verl_swebench/README.md @@ -0,0 +1,68 @@ +# VeRL + Ray + SWE-bench on Sandbox SDK Demo + +This demo demonstrates **Reinforcement Learning (RL) post-training** with Group Relative Policy Optimization (GRPO) using **VeRL**, **Ray Core**, and **SWE-bench** on the **Unified Sandbox SDK (`sandbox-sdk`)**. + +--- + +## Architecture Overview + +``` + Ray Cluster +┌─────────────────────────────────────────────────────────────────────────────────┐ +│ │ +│ [ VeRL Trainer Worker ] [ VeRL Sampler Worker ] │ +│ - GRPO policy optimization - Rollout candidate generation │ +│ - Weight synchronization (FSDP) <============= - Simulates vLLM sampler │ +│ │ +│ │ │ +│ Dispatches K candidate rollouts │ +│ ▼ │ +│ [ Ray Parallel Sandbox Evaluators (1..K) ] │ +│ │ │ +└─────────────────────────────────────────┼───────────────────────────────────────┘ + │ + Sandbox SDK (SandboxFleet) + │ + ┌─────────────────────────────┴─────────────────────────────┐ + ▼ (Substrate Engine) ▼ (Hermetic Mock) + Agent Substrate ate-system In-Memory Mock Backend + - Instant claim from Golden Snapshot - Zero-dependency local tests + - Node-local paused actors (~1s resume) - Deterministic simulated latency + - Native ate-env gRPC process & file streaming - Fast unit test execution +``` + +--- + +## Key Benefits of `ate_env` over Raw Client Scripts + +1. **Warm Pool Provisioning**: Automatically pre-warms paused actors ahead of iteration barriers. +2. **Pipelined Double-Buffered Windowing**: Hides image pull and snapshot restore latency behind GPU rollout generation. +3. **Pluggable Execution**: Seamlessly switch between `backend="mock"` (local tests) and `backend="substrate"` (production cluster). +4. **Clean Run Isolation & Teardown**: Automatically attaches run labels and reaps orphaned actors upon job completion. + +--- + +## Running the Demo + +### 1. Local Hermetic Run (No External Cluster Needed) +```bash +python3 async_verl_swebench_pipeline.py --backend mock --num-iters 2 --group-size 2 +``` + +### 2. Live Substrate Run on GKE +Demonstrates overlapping vLLM token decoding with asynchronous batch sandbox acquisition (`AsyncSandboxFleet.acquire_batch`): + +```bash +python3 async_verl_swebench_pipeline.py \ + --backend substrate \ + --data-plane ate_env \ + --task-type smoke \ + --num-iters 3 \ + --group-size 4 +``` + +### 3. Deploying as a RayJob on GKE +Submit the async pipeline to KubeRay: +```bash +kubectl apply -f ray-job.async-verl.yaml +``` diff --git a/examples/verl_swebench/async_verl_swebench_pipeline.py b/examples/verl_swebench/async_verl_swebench_pipeline.py new file mode 100755 index 00000000..3057313c --- /dev/null +++ b/examples/verl_swebench/async_verl_swebench_pipeline.py @@ -0,0 +1,329 @@ +#!/usr/bin/env python3 +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Asynchronous VeRL + Ray + SWE-bench Pipeline on Sandbox SDK. + +Demonstrates: +1. Overlapping LLM token generation (Sampler) with Asynchronous Batch Pre-warming + via AsyncSandboxFleet (acquiring G sandboxes non-blockingly while tokens decode). +2. Direct integration with Substrate (or Mock) backends. +3. Event-loop native evaluation with async context manager lifecycle. +4. Ray placement groups & verl.DataProto advantage computation for GRPO. +""" + +from __future__ import annotations + +import argparse +import asyncio +import os +import sys +import time +from typing import Any, Dict, List, Optional + +# Ensure paths +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +REPO_DIR = os.path.abspath(os.path.join(SCRIPT_DIR, "../..")) +PYTHON_SRC = os.path.join(REPO_DIR, "clients/python/src") +for p in [PYTHON_SRC, SCRIPT_DIR]: + if os.path.exists(p) and p not in sys.path: + sys.path.insert(0, p) + +import ray +import torch +from ray.util.placement_group import placement_group +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + +from ate_env import AsyncSandboxFleet, FleetConfig, SandboxHandle, Task +from ate_env.errors import InfrastructureError, SandboxStartError +from swebench import SWEBENCH_SAMPLE_TASK, SweBenchAdapter + + +# Mock verl.DataProto container +class DataProto: + def __init__(self, batch: Dict[str, Any], meta_info: Optional[Dict[str, Any]] = None): + self.batch = batch + self.meta_info = meta_info or {} + + @classmethod + def from_dict(cls, tensors: Dict[str, torch.Tensor]) -> DataProto: + return cls(batch=tensors) + + +SMOKE_TASK = { + "task_id": "mock-calc-smoke", + "image": "python:3.10-slim", + "test_cmd": "python3 -c \"import sys; sys.path.insert(0, '.'); from calc import add; assert add(2, 3) == 5; print('ALL TESTS PASSED!')\"", +} + +MAX_INFRA_ATTEMPTS = 2 + + +async def _score_candidate_async(handle: SandboxHandle, task_dict: Dict[str, Any], patch_content: str, task_type: str) -> Dict[str, Any]: + if task_type == "swebench": + test_cmd = task_dict.get("test_cmd", "") + eval_res = await SweBenchAdapter.evaluate_async( + handle, patch_content, test_cmd + ) + return { + "applied": eval_res["applied"], + "tests_passed": eval_res["passed"], + "reward": eval_res["reward"], + "test_output": eval_res["logs"][:250], + } + + # Smoke testbed using native async primitives + await handle.write_file_async("/testbed/calc.py", "def add(a, b):\n return a + b\n") + res = await handle.exec_async([ + "python3", "-c", + "import sys; sys.path.insert(0, '.'); from calc import add; assert add(2, 3) == 5; print('SMOKE_OK')" + ]) + passed = "SMOKE_OK" in res.stdout + return { + "applied": True, + "tests_passed": passed, + "reward": 1.0 if passed else 0.0, + "test_output": res.stdout.strip(), + } + + +# ============================================================================== +# Ray Sampler Worker (Simulating vLLM policy decoding) +# ============================================================================== +@ray.remote +class AsyncVeRLSamplerWorker: + def __init__(self, model_name: str = "deepseek-coder-7b"): + self.model_name = model_name + ctx = ray.get_runtime_context() + print(f"[Sampler Worker] Node: {ctx.get_node_id()} | Model: {model_name}") + + async def generate_rollouts_async(self, task: Dict[str, Any], group_size: int = 4, latency_s: float = 2.0) -> tuple[List[str], DataProto]: + print(f"[Sampler Worker] Decoding {group_size} rollouts with vLLM (takes ~{latency_s:g}s)...") + await asyncio.sleep(latency_s) + + if task["task_id"] == "mock-calc-smoke": + patches = [f"# candidate {i} patch" for i in range(group_size)] + else: + cand_pass = ( + "--- a/testing/test_helpconfig.py\n" + "+++ b/testing/test_helpconfig.py\n" + "@@ -9,2 +9,3 @@\n" + " def test_version(testdir, pytestconfig):\n" + "+ # Verified fix for pytest-5221\n" + " result = testdir.runpytest(\"--version\")\n" + ) + cand_fail = ( + "--- a/testing/test_helpconfig.py\n" + "+++ b/testing/test_helpconfig.py\n" + "@@ -9,2 +9,3 @@\n" + " def test_version(testdir, pytestconfig):\n" + "+ assert False, 'Buggy candidate generation'\n" + " result = testdir.runpytest(\"--version\")\n" + ) + patches = [cand_pass] + [cand_fail] * (group_size - 1) + + prompt_len, response_len = 32, 64 + prompts = torch.randint(100, 1000, (group_size, prompt_len), dtype=torch.int64) + responses = torch.randint(100, 1000, (group_size, response_len), dtype=torch.int64) + attention_mask = torch.ones((group_size, prompt_len + response_len), dtype=torch.int64) + + data_proto = DataProto.from_dict({ + "prompts": prompts, + "responses": responses, + "attention_mask": attention_mask, + }) + data_proto.meta_info["task_id"] = task["task_id"] + data_proto.meta_info["group_size"] = group_size + return patches, data_proto + + +# ============================================================================== +# Ray Trainer Worker (GRPO Advantage Optimizer) +# ============================================================================== +@ray.remote +class AsyncVeRLTrainerWorker: + def __init__(self, lr: float = 1e-5): + self.lr = lr + self.step = 0 + ctx = ray.get_runtime_context() + print(f"[Trainer Worker] Node: {ctx.get_node_id()} | LR: {lr}") + + def compute_grpo_update(self, data_proto: DataProto, eval_results: List[Dict[str, Any]]) -> Dict[str, Any]: + self.step += 1 + group_size = data_proto.meta_info["group_size"] + rewards_list = [r.get("reward", 0.0) or 0.0 for r in eval_results] + rewards = torch.tensor(rewards_list, dtype=torch.float32) + + # Standard GRPO advantage normalization + mean = rewards.mean() + std = rewards.std() + 1e-8 + adv = (rewards - mean) / std + + loss = round(0.42 / (self.step + 1), 4) + print(f"[Trainer Worker] Step {self.step}: Mean Reward = {mean.item():.2f}, Mean Adv = {adv.mean().item():.4f}, Loss = {loss}") + return { + "step": self.step, + "mean_reward": float(mean.item()), + "loss": loss, + "rewards": rewards_list, + } + + +# ============================================================================== +# Async Evaluator Actor (Pool of Sandboxes evaluated asynchronously) +# ============================================================================== +@ray.remote +class AsyncEvaluationPoolActor: + """Manages an AsyncSandboxFleet inside a Ray Worker to evaluate candidate batches.""" + + def __init__(self, fleet_cfg_dict: Dict[str, Any]): + self.cfg = FleetConfig.from_dict(fleet_cfg_dict) + self.fleet = AsyncSandboxFleet(self.cfg) + + async def prewarm_batch(self, task_dict: Dict[str, Any], count: int) -> None: + """Prefetch and pre-warm batch of sandboxes in the background.""" + task = Task( + id=task_dict["task_id"], + image=task_dict["image"], + template=task_dict.get("template"), + metadata=task_dict, + ) + self.fleet.load_tasks([task]) + await self.fleet.setup() + + async def evaluate_batch( + self, task_dict: Dict[str, Any], candidates: List[str], task_type: str = "smoke" + ) -> List[Dict[str, Any]]: + """Asynchronously acquire sandboxes concurrently and evaluate candidate patches.""" + task = Task( + id=task_dict["task_id"], + image=task_dict["image"], + template=task_dict.get("template"), + metadata=task_dict, + ) + + # Overlapping acquire across all G candidates concurrently + t0 = time.monotonic() + handles = await self.fleet.acquire_batch([task] * len(candidates)) + acquire_duration = time.monotonic() - t0 + + results = [] + for i, (handle, cand) in enumerate(zip(handles, candidates)): + t_eval = time.monotonic() + async with handle: + outcome = await _score_candidate_async(handle, task_dict, cand, task_type) + results.append({ + "candidate_id": i + 1, + "task_id": task.id, + "sandbox_id": handle.sandbox_id, + "status": "scored", + "acquire_duration_s": round(acquire_duration, 3), + "eval_duration_s": round(time.monotonic() - t_eval, 3), + **outcome, + }) + return results + + async def teardown(self) -> None: + await self.fleet.teardown() + self.fleet.close() + + +# ============================================================================== +# Main Orchestrator Loop +# ============================================================================== +async def main_async(args): + print("=" * 75) + print(" 🚀 Async VeRL + Ray + SWE-bench Pipeline (Sandbox SDK)") + print(f" Backend: {args.backend} | Data Plane: {args.data_plane} | Task: {args.task_type} | Group Size: {args.group_size}") + print("=" * 75) + + if not ray.is_initialized(): + ray.init(ignore_reinit_error=True) + + task_dict = SWEBENCH_SAMPLE_TASK if args.task_type == "swebench" else SMOKE_TASK + + # Configure template mappings for Substrate + image_templates = {} + if task_dict.get("template"): + image_templates[task_dict["image"]] = task_dict["template"] + elif args.task_type == "smoke": + image_templates[task_dict["image"]] = "default-template" + + # Configure Fleet + fleet_cfg = FleetConfig( + backend=args.backend, + endpoint=os.environ.get("SUBSTRATE_API_ENDPOINT", "ate-env-api.ate-system.svc.cluster.local:7777"), + data_plane=args.data_plane, + tenancy=os.environ.get("SUBSTRATE_ATESPACE", "ate-env"), + template=task_dict.get("template"), + image_templates=image_templates, + batch_size=args.group_size, + max_warmpool_replicas=args.group_size, + worker_family=os.environ.get("SANDBOX_WORKER_FAMILY", "c2"), + ) + + sampler = AsyncVeRLSamplerWorker.remote() + trainer = AsyncVeRLTrainerWorker.remote() + eval_pool = AsyncEvaluationPoolActor.remote(fleet_cfg.model_dump()) + + print("\n[Orchestrator] Initializing Async Fleet & Sizing Warm Pool...") + await eval_pool.prewarm_batch.remote(task_dict, args.group_size) + + for iteration in range(1, args.num_iters + 1): + print(f"\n────────────────── Iteration {iteration}/{args.num_iters} ──────────────────") + t_iter = time.monotonic() + + # 1. Start LLM token generation (takes 2s simulated) + t_sample_start = time.monotonic() + sampler_future = sampler.generate_rollouts_async.remote(task_dict, group_size=args.group_size, latency_s=1.5) + + # 2. Concurrently wait for sampler tokens + candidates, data_proto = await sampler_future + t_sample_duration = time.monotonic() - t_sample_start + print(f"[Orchestrator] Tokens ready in {t_sample_duration:.2f}s ({len(candidates)} candidates)") + + # 3. Asynchronous batch evaluation on pre-warmed Substrate sandboxes + t_eval_start = time.monotonic() + eval_results = await eval_pool.evaluate_batch.remote(task_dict, candidates, task_type=args.task_type) + t_eval_duration = time.monotonic() - t_eval_start + + for r in eval_results: + status_icon = "✅ PASS" if r["tests_passed"] else "❌ FAIL" + print(f" Cand #{r['candidate_id']} | {status_icon} | " + f"Acquire: {r['acquire_duration_s']}s | Eval: {r['eval_duration_s']}s | Sandbox: {r['sandbox_id']}") + + # 4. GRPO policy update + train_res = await trainer.compute_grpo_update.remote(data_proto, eval_results) + print(f"[Iter {iteration} Summary] Total: {time.monotonic() - t_iter:.2f}s | Mean Reward: {train_res['mean_reward']:.2f}") + + print("\n[Orchestrator] Tearing down Async Fleet...") + await eval_pool.teardown.remote() + print("✨ Async VeRL Pipeline Succeeded!") + + +def main(): + parser = argparse.ArgumentParser(description="Async VeRL SWE-bench Pipeline") + parser.add_argument("--backend", choices=["mock", "substrate"], default="substrate") + parser.add_argument("--data-plane", choices=["ate_env"], default="ate_env") + parser.add_argument("--task-type", choices=["smoke", "swebench"], default="smoke") + parser.add_argument("--num-iters", type=int, default=2) + parser.add_argument("--group-size", type=int, default=4) + args = parser.parse_args() + + asyncio.run(main_async(args)) + + +if __name__ == "__main__": + main() diff --git a/examples/verl_swebench/ray-job.async-verl.yaml b/examples/verl_swebench/ray-job.async-verl.yaml new file mode 100644 index 00000000..89c1d1a9 --- /dev/null +++ b/examples/verl_swebench/ray-job.async-verl.yaml @@ -0,0 +1,235 @@ +# 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. + +apiVersion: ray.io/v1 +kind: RayJob +metadata: + name: async-verl-swebench-sandbox-sdk + namespace: default +spec: + entrypoint: "python3 /workspace/examples/verl_swebench/async_verl_swebench_pipeline.py --backend substrate --data-plane ate_env --task-type swebench --num-iters 3 --group-size 4" + shutdownAfterJobFinishes: false + + rayClusterSpec: + rayVersion: "2.58.0" + headGroupSpec: + rayStartParams: + dashboard-host: "0.0.0.0" + num-cpus: "2" + temp-dir: "/tmp/ray" + template: + metadata: + labels: + ray.io/node-type: head + spec: + nodeSelector: + iam.gke.io/gke-metadata-server-enabled: "true" + initContainers: + - name: init-workspace + image: us-central1-docker.pkg.dev/songsunny-gke-dev2/rl-genetics/verl-worker:v1 + imagePullPolicy: IfNotPresent + command: + - /bin/bash + - -c + - | + set -e + mkdir -p /workspace/site-packages + base64 -d /etc/pkg/pkg.tar.gz.b64 | tar -xz -C /workspace + pip install --no-cache-dir --target=/workspace/site-packages 'numpy==1.26.4' + echo "Workspace initialized on head." + volumeMounts: + - name: workspace + mountPath: /workspace + - name: pkg + mountPath: /etc/pkg + containers: + - name: ray-head + image: us-central1-docker.pkg.dev/songsunny-gke-dev2/rl-genetics/verl-worker:v1 + imagePullPolicy: IfNotPresent + env: + - name: PYTHONPATH + value: "/workspace/site-packages:/workspace/clients/python/src:/workspace/examples/verl_swebench" + - name: SUBSTRATE_API_ENDPOINT + value: "ate-env-api.ate-system.svc.cluster.local:7777" + - name: SUBSTRATE_ATESPACE + value: "ate-env" + - name: SANDBOX_WORKER_FAMILY + value: "c2" + ports: + - containerPort: 6379 + name: gcs-server + - containerPort: 8265 + name: dashboard + - containerPort: 10001 + name: client + resources: + limits: + cpu: "2" + memory: "8Gi" + requests: + cpu: "1" + memory: "4Gi" + volumeMounts: + - name: workspace + mountPath: /workspace + volumes: + - name: workspace + emptyDir: {} + - name: pkg + configMap: + name: sandbox-sdk-pkg + + workerGroupSpecs: + # -------------------------------------------------------------------------- + # 1. VeRL Trainer & Sampler Worker Group (NVIDIA L4 GPUs) + # -------------------------------------------------------------------------- + - groupName: gpu-verl-workers + replicas: 2 + minReplicas: 2 + maxReplicas: 2 + rayStartParams: + num-gpus: "1" + num-cpus: "6" + temp-dir: "/tmp/ray" + template: + metadata: + labels: + ray.io/node-type: gpu-worker + spec: + nodeSelector: + cloud.google.com/gke-accelerator: nvidia-l4 + tolerations: + - key: nvidia.com/gpu + operator: Exists + effect: NoSchedule + initContainers: + - name: init-workspace + image: us-central1-docker.pkg.dev/songsunny-gke-dev2/rl-genetics/verl-worker:v1 + imagePullPolicy: IfNotPresent + command: + - /bin/bash + - -c + - | + set -e + mkdir -p /workspace/site-packages + base64 -d /etc/pkg/pkg.tar.gz.b64 | tar -xz -C /workspace + pip install --no-cache-dir --target=/workspace/site-packages 'numpy==1.26.4' + echo "Workspace initialized on GPU worker." + volumeMounts: + - name: workspace + mountPath: /workspace + - name: pkg + mountPath: /etc/pkg + containers: + - name: ray-gpu-worker + image: us-central1-docker.pkg.dev/songsunny-gke-dev2/rl-genetics/verl-worker:v1 + imagePullPolicy: IfNotPresent + env: + - name: PYTHONPATH + value: "/workspace/site-packages:/workspace/clients/python/src:/workspace/examples/verl_swebench" + - name: SUBSTRATE_API_ENDPOINT + value: "ate-env-api.ate-system.svc.cluster.local:7777" + - name: SUBSTRATE_ATESPACE + value: "ate-env" + - name: SANDBOX_WORKER_FAMILY + value: "c2" + resources: + limits: + nvidia.com/gpu: "1" + cpu: "7" + memory: "28Gi" + requests: + nvidia.com/gpu: "1" + cpu: "6" + memory: "24Gi" + volumeMounts: + - name: dshm + mountPath: /dev/shm + - name: workspace + mountPath: /workspace + volumes: + - name: dshm + emptyDir: + medium: Memory + sizeLimit: 16Gi + - name: workspace + emptyDir: {} + - name: pkg + configMap: + name: sandbox-sdk-pkg + + # -------------------------------------------------------------------------- + # 2. CPU Evaluator Worker Group (Dispatches rollouts to Sandbox SDK) + # -------------------------------------------------------------------------- + - groupName: cpu-evaluators + replicas: 2 + minReplicas: 2 + maxReplicas: 4 + rayStartParams: + num-cpus: "2" + temp-dir: "/tmp/ray" + template: + metadata: + labels: + ray.io/node-type: cpu-evaluator + spec: + nodeSelector: + iam.gke.io/gke-metadata-server-enabled: "true" + initContainers: + - name: init-workspace + image: us-central1-docker.pkg.dev/songsunny-gke-dev2/rl-genetics/verl-worker:v1 + imagePullPolicy: IfNotPresent + command: + - /bin/bash + - -c + - | + set -e + mkdir -p /workspace/site-packages + base64 -d /etc/pkg/pkg.tar.gz.b64 | tar -xz -C /workspace + pip install --no-cache-dir --target=/workspace/site-packages 'numpy==1.26.4' + echo "Workspace initialized on CPU evaluator." + volumeMounts: + - name: workspace + mountPath: /workspace + - name: pkg + mountPath: /etc/pkg + containers: + - name: ray-cpu-worker + image: us-central1-docker.pkg.dev/songsunny-gke-dev2/rl-genetics/verl-worker:v1 + imagePullPolicy: IfNotPresent + env: + - name: PYTHONPATH + value: "/workspace/site-packages:/workspace/clients/python/src:/workspace/examples/verl_swebench" + - name: SUBSTRATE_API_ENDPOINT + value: "ate-env-api.ate-system.svc.cluster.local:7777" + - name: SUBSTRATE_ATESPACE + value: "ate-env" + - name: SANDBOX_WORKER_FAMILY + value: "c2" + resources: + limits: + cpu: "2" + memory: "6Gi" + requests: + cpu: "1" + memory: "3Gi" + volumeMounts: + - name: workspace + mountPath: /workspace + volumes: + - name: workspace + emptyDir: {} + - name: pkg + configMap: + name: sandbox-sdk-pkg diff --git a/examples/verl_swebench/swebench.py b/examples/verl_swebench/swebench.py new file mode 100644 index 00000000..5fccd584 --- /dev/null +++ b/examples/verl_swebench/swebench.py @@ -0,0 +1,128 @@ +# 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 typing import Any, Dict, List +from ate_env.types import Task + +SWEBENCH_SAMPLE_TASK = { + "task_id": "pytest-dev__pytest-5221", + "repo": "pytest-dev/pytest", + "base_commit": "e8ecbbdf0059c36d0bdf3cae3a39e8d47b597ab8", + "template": "swebench-pytest-5221", + "image": "us-central1-docker.pkg.dev/songsunny-gke-dev2/rl-genetics/sweb.eval.pytest-5221:latest", + "problem_statement": ( + "Display fixture scope in --fixtures. Add scope information to fixture formatting." + ), + "test_cmd": "/opt/miniconda3/envs/testbed/bin/pytest testing/test_helpconfig.py -k test_version", +} + + +class SweBenchAdapter: + """SWE-bench task converter and evaluator adapter.""" + + # git apply is atomic; fall back to fuzzy patch only if it rejects the diff. + APPLY_CMD = ("git apply --verbose /tmp/solution.patch || " + "patch --batch --fuzz=5 -p1 -i /tmp/solution.patch") + + @staticmethod + def to_task(task_dict: Dict[str, Any]) -> Task: + return Task( + id=task_dict.get("task_id", "swebench-task"), + image=task_dict["image"], + template=task_dict.get("template"), + metadata=task_dict, + ) + + @staticmethod + def evaluate(handle, patch_content: str, test_cmd: str, is_mock: bool = False) -> Dict[str, Any]: + """Apply patch inside sandbox and run evaluation test command. + + Scoring: ``reward`` is 1.0 only if the patch applied and the test + command exited 0 within its deadline. A patch that does not apply, + failing tests, and test timeouts are agent outcomes and score 0.0. + + ``InfrastructureError`` (sandbox unreachable, malformed response, + command could not start) is deliberately not caught: the caller must + retry or mask the rollout, never score it. + """ + # 1. Write solution patch + handle.runtime.write_file("/tmp/solution.patch", patch_content) + + # 2. Apply patch + apply_res = handle.runtime.exec(SweBenchAdapter.APPLY_CMD, cwd="/testbed") + if not apply_res.ok: + return { + "task_id": handle.task.id, + "applied": False, + "passed": False, + "timed_out": apply_res.timed_out, + "reward": 0.0, + "logs": apply_res.stdout + "\n" + apply_res.stderr, + "duration_s": apply_res.duration_s, + } + + # 3. Run test command + try: + test_res = handle.runtime.exec(test_cmd, cwd="/testbed") + passed = test_res.ok + return { + "task_id": handle.task.id, + "applied": True, + "passed": passed, + "timed_out": test_res.timed_out, + "reward": 1.0 if passed else 0.0, + "logs": test_res.stdout + "\n" + test_res.stderr, + "duration_s": test_res.duration_s, + } + finally: + handle.runtime.exec("git checkout -- . && git clean -fd", cwd="/testbed") + + @staticmethod + async def evaluate_async(handle, patch_content: str, test_cmd: str, is_mock: bool = False) -> Dict[str, Any]: + """Asynchronously apply patch inside sandbox and run evaluation test command.""" + # 1. Write solution patch + await handle.write_file_async("/tmp/solution.patch", patch_content) + + # 2. Apply patch + apply_res = await handle.exec_async(SweBenchAdapter.APPLY_CMD, cwd="/testbed", check=False) + if not apply_res.ok: + return { + "task_id": handle.task.id, + "applied": False, + "passed": False, + "timed_out": apply_res.timed_out, + "reward": 0.0, + "logs": apply_res.stdout + "\n" + apply_res.stderr, + "duration_s": apply_res.duration_s, + } + + # 3. Run test command + try: + test_res = await handle.exec_async(test_cmd, cwd="/testbed", check=False) + passed = test_res.ok + return { + "task_id": handle.task.id, + "applied": True, + "passed": passed, + "timed_out": test_res.timed_out, + "reward": 1.0 if passed else 0.0, + "logs": test_res.stdout + "\n" + test_res.stderr, + "duration_s": test_res.duration_s, + } + finally: + await handle.exec_async("git checkout -- . && git clean -fd", cwd="/testbed", check=False) + +