From 5d1fcf789b36312716ac85435796101c7b15a308 Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Wed, 3 Jun 2026 14:20:35 +0000 Subject: [PATCH 01/11] feat(sandbox): ECS Fargate sandbox provider MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Port NEL's ECS Fargate sandboxing into Gym as a third sandbox provider, stacked on the provider framework from #1377. - nemo_gym/sandbox/providers/ecs_fargate/{engine,provider}.py — NEL's engine lifted (task-def registration + SSM caching, RunTask with capacity retries, SSH sidecar + reverse-tunnel outside-endpoint routing, in-container exec server, CodeBuild->ECR image build) behind a thin adapter that keeps per-sandbox state in SandboxHandle.raw. - registry loader + `sandbox-ecs` extra (boto3, lazy-imported). - region-only config via SSM autodiscovery (//ecs-sandbox/config), matching NEL. - mini_swe_agent_2 ecs_fargate config + 14 unit tests (AWS/SSH mocked). Fast-follow P0 (separate PR): route the model endpoint over Teleport as an internal option instead of the globally-exposed SSH reverse tunnel that security flagged; SSH tunneling stays available for community use. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Michal Bien --- .../sandbox/providers/ecs_fargate/__init__.py | 32 + .../sandbox/providers/ecs_fargate/engine.py | 2246 +++++++++++++++++ .../sandbox/providers/ecs_fargate/provider.py | 236 ++ nemo_gym/sandbox/providers/registry.py | 9 +- pyproject.toml | 8 + .../configs/mini_swe_agent_ecs_fargate.yaml | 39 + tests/unit_tests/test_ecs_fargate_provider.py | 270 ++ 7 files changed, 2839 insertions(+), 1 deletion(-) create mode 100644 nemo_gym/sandbox/providers/ecs_fargate/__init__.py create mode 100644 nemo_gym/sandbox/providers/ecs_fargate/engine.py create mode 100644 nemo_gym/sandbox/providers/ecs_fargate/provider.py create mode 100644 responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml create mode 100644 tests/unit_tests/test_ecs_fargate_provider.py diff --git a/nemo_gym/sandbox/providers/ecs_fargate/__init__.py b/nemo_gym/sandbox/providers/ecs_fargate/__init__.py new file mode 100644 index 0000000000..a37bc4d8b2 --- /dev/null +++ b/nemo_gym/sandbox/providers/ecs_fargate/__init__.py @@ -0,0 +1,32 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# 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. + +"""ECS Fargate sandbox provider package.""" + +from nemo_gym.sandbox.providers.ecs_fargate.engine import ( + EcsFargateConfig, + SshSidecarConfig, +) +from nemo_gym.sandbox.providers.ecs_fargate.provider import ( + EcsFargateProvider, + engine_config_from_mapping, +) + + +__all__ = [ + "EcsFargateConfig", + "EcsFargateProvider", + "SshSidecarConfig", + "engine_config_from_mapping", +] diff --git a/nemo_gym/sandbox/providers/ecs_fargate/engine.py b/nemo_gym/sandbox/providers/ecs_fargate/engine.py new file mode 100644 index 0000000000..f59dc68e1e --- /dev/null +++ b/nemo_gym/sandbox/providers/ecs_fargate/engine.py @@ -0,0 +1,2246 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# 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. + +"""ECS Fargate sandbox engine. + +Lifted from nemo-evaluator-next (``nemo_evaluator/sandbox/ecs_fargate.py``). +The orchestration is unchanged; only the three host-protocol types it relied on +(``ExecResult``, ``OutsideEndpoint``, ``SandboxSpec``) and the SSM/port +constants are vendored here so the engine stands alone inside the provider. +The Gym-facing adapter lives in ``provider.py``. +""" + +from __future__ import annotations + +import asyncio +import atexit +import base64 +import hashlib +import io +import json +import logging +import os +import random +import re +import shlex +import socket +import subprocess +import tarfile +import tempfile +import threading +import time +import uuid +import zipfile +from dataclasses import dataclass, field +from dataclasses import replace as _dc_replace +from pathlib import Path +from typing import Any, Callable, Self, TypeVar +from urllib.parse import ParseResult, urlparse + +import aiohttp + + +# Matches the exec-server script's TB_EXEC_PORT fallback and the +# exec_server_port written by the reference Terraform. +DEFAULT_EXEC_SERVER_PORT = 19542 +# SSH sidecar port matching the reference Terraform ssh_tunnel_sshd_port. +DEFAULT_SSHD_PORT = 52222 +DEFAULT_SSM_PROJECT = "harbor" + + +@dataclass +class ExecResult: + """Result of executing a command inside a sandbox.""" + + stdout: str + stderr: str + return_code: int + + +@dataclass +class OutsideEndpoint: + """A host-side URL that must be reachable from inside the sandbox. + + The sandbox rewrites *url* for its network topology and exposes the + resolved address as the environment variable *env_var* inside the + container. + """ + + url: str + env_var: str + + +@dataclass +class VolumeMount: + """EFS mount (ECS Fargate). Host bind mounts are unused on Fargate.""" + + host_path: str = "" + container_path: str = "" + readonly: bool = False + efs_filesystem_id: str | None = None + efs_root_directory: str | None = None + efs_access_point_id: str | None = None + + @property + def is_efs(self) -> bool: + return self.efs_filesystem_id is not None + + +@dataclass +class SandboxSpec: + """Per-problem sandbox requirements.""" + + image: str + workdir: str = "/workspace" + env: dict[str, str] = field(default_factory=dict) + files: dict[str, str] = field(default_factory=dict) + entrypoint: str | None = None + volumes: list[VolumeMount] = field(default_factory=list) + environment_dir: str | None = None + + +logger = logging.getLogger(__name__) +T = TypeVar("T") + + +# ── Lazy AWS SDK import ────────────────────────────────────────────── + + +def _require_aws_sdks(): + try: + import importlib + + boto3 = importlib.import_module("boto3") + botocore_config = importlib.import_module("botocore.config") + botocore_exceptions = importlib.import_module("botocore.exceptions") + except ModuleNotFoundError as e: + raise RuntimeError( + "ECS Fargate sandbox requires boto3/botocore. " + "Install them (`pip install boto3`) or use a different sandbox backend." + ) from e + return boto3, getattr(botocore_config, "Config"), getattr(botocore_exceptions, "ClientError") + + +# ── SSM auto-discovery ─────────────────────────────────────────────── + +_ssm_config_cache: dict[str, dict[str, Any]] = {} + + +def resolve_ecs_config_from_ssm( + region: str, + project: str = DEFAULT_SSM_PROJECT, +) -> dict[str, Any]: + """Read ECS sandbox config from SSM Parameter Store. + + Returns a dict matching the JSON structure written by Terraform + (cluster, subnets, security_groups, roles, SSH ARNs, EFS, etc.). + Results are cached per (region, project) for the process lifetime. + """ + cache_key = f"{region}:{project}" + if cache_key in _ssm_config_cache: + return _ssm_config_cache[cache_key] + + boto3, _, ClientError = _require_aws_sdks() + ssm = boto3.client("ssm", region_name=region) + param_name = f"/{project}/ecs-sandbox/config" + try: + resp = _retry_with_backoff( + lambda: ssm.get_parameter(Name=param_name), + operation_name="ssm.get_parameter", + max_retries=5, + ) + except ClientError as exc: + code = (exc.response.get("Error") or {}).get("Code", "") + if code == "ParameterNotFound": + raise RuntimeError( + f"SSM parameter '{param_name}' not found in {region}. " + f"Run 'terraform apply' in the ecs-sandbox stack for this " + f"region, or specify all ECS fields explicitly in your YAML." + ) from exc + raise + + raw = resp["Parameter"]["Value"] + try: + config = json.loads(raw) + except json.JSONDecodeError as exc: + raise RuntimeError(f"SSM parameter '{param_name}' in {region} contains invalid JSON: {exc}") from exc + + _ssm_config_cache[cache_key] = config + logger.info( + "Resolved ECS config from SSM %s in %s (cluster=%s, %d subnets)", + param_name, + region, + config.get("cluster"), + len(config.get("subnets", [])), + ) + return config + + +# ── Config dataclasses ─────────────────────────────────────────────── + + +def _sanitize_id(value: str, max_len: int = 100) -> str: + cleaned = re.sub(r"[^a-zA-Z0-9-]+", "-", value).strip("-") + return cleaned[:max_len] or "task" + + +@dataclass(frozen=True) +class SshSidecarConfig: + """SSH sidecar container configuration. + + exec_server_port set → exec-server mode (one-way tunnel). + exec_server_port None → agent-server mode (two-way tunnel). + """ + + sshd_port: int = 2222 + ssh_ready_timeout_sec: float = 300.0 + public_key_secret_arn: str = "" + private_key_secret_arn: str = "" + image: str | None = None + exec_server_port: int | None = None + + +@dataclass(frozen=True) +class EcsFargateConfig: + """Configuration for the ECS Fargate sandbox.""" + + region: str | None = None + cluster: str = "" + subnets: list[str] = field(default_factory=list) + security_groups: list[str] = field(default_factory=list) + assign_public_ip: bool = False + task_definition: str | None = None + task_definition_family_prefix: str = "ecs-sandbox" + image_template: str | None = None + container_name: str = "main" + container_port: int | None = None + cpu: str = "4096" + memory: str = "8192" + ephemeral_storage_gib: int | None = None + platform_version: str | None = None + execution_role_arn: str | None = None + task_role_arn: str | None = None + extra_env: dict[str, str] | None = None + log_group: str | None = None + log_stream_prefix: str | None = None + max_task_lifetime_sec: int = 14400 + startup_timeout_sec: float = 300.0 + poll_interval_sec: float = 2.0 + run_task_max_retries: int = 30 + ssh_sidecar: SshSidecarConfig | None = None + s3_bucket: str | None = None + s3_prefix: str | None = None + ecr_repository: str | None = None + environment_dir: str | None = None + codebuild_project: str | None = None + codebuild_service_role: str | None = None + codebuild_compute_type: str = "BUILD_GENERAL1_MEDIUM" + codebuild_build_timeout: int = 60 + dockerhub_secret_arn: str | None = None + build_parallelism: int = 50 + efs_filesystem_id: str | None = None + efs_access_point_id: str | None = None + ssm_project: str = DEFAULT_SSM_PROJECT + + +@dataclass(frozen=True) +class _OutsideEndpointRoute: + endpoint: OutsideEndpoint + source_netloc: str + host: str + target_port: int + remote_port: int + scheme: str + + @classmethod + def for_endpoint(cls, endpoint: OutsideEndpoint, *, remote_port: int | None = None) -> _OutsideEndpointRoute: + parsed = urlparse(endpoint.url) + host = parsed.hostname + if not host: + raise ValueError(f"Cannot resolve hostname from OutsideEndpoint: {endpoint.url}") + target_port = _port_from_url(parsed) + return cls( + endpoint=endpoint, + source_netloc=parsed.netloc, + host=host, + target_port=target_port, + remote_port=remote_port or target_port, + scheme=parsed.scheme or "http", + ) + + def resolved_endpoint_url(self) -> str: + return self._rewrite(urlparse(self.endpoint.url)) + + def resolve_url(self, url: str) -> str: + return self._rewrite(urlparse(url)) + + def _rewrite(self, parsed: ParseResult) -> str: + return parsed._replace( + scheme=self.scheme, + netloc=f"127.0.0.1:{self.remote_port}", + ).geturl() + + +def _port_from_url(parsed: ParseResult) -> int: + return parsed.port or (443 if parsed.scheme == "https" else 80) + + +@dataclass(frozen=True) +class _OutsideEndpointRouting: + endpoints: tuple[OutsideEndpoint, ...] = () + _routes_by_env: dict[str, _OutsideEndpointRoute] = field(default_factory=dict) + _reverse_specs: tuple[str, ...] = () + _agent_tunnel_port: int | None = None + + @classmethod + def empty(cls, endpoints: list[OutsideEndpoint] | None = None) -> _OutsideEndpointRouting: + return cls(endpoints=tuple(endpoints or [])) + + @classmethod + def for_exec_server( + cls, + endpoints: list[OutsideEndpoint], + sidecar: SshSidecarConfig, + ) -> _OutsideEndpointRouting: + reverse_specs: list[str] = [] + routes_by_env: dict[str, _OutsideEndpointRoute] = {} + used_ports = {sidecar.sshd_port} + if sidecar.exec_server_port is not None: + used_ports.add(sidecar.exec_server_port) + + target_port_map: dict[tuple[str, int], int] = {} + for ep in endpoints: + target = _OutsideEndpointRoute.for_endpoint(ep) + key = (target.host, target.target_port) + remote_port = target_port_map.get(key) + if remote_port is None: + remote_port = cls._allocate_reverse_port(target.target_port, used_ports) + target_port_map[key] = remote_port + reverse_specs.append(f"{remote_port}:{target.host}:{target.target_port}") + route = _OutsideEndpointRoute.for_endpoint(ep, remote_port=remote_port) + logger.info( + "Reverse tunnel: container :%d → host %s:%d (%s)", + route.remote_port, + route.host, + route.target_port, + ep.env_var, + ) + routes_by_env[ep.env_var] = route + + return cls(endpoints=tuple(endpoints), _routes_by_env=routes_by_env, _reverse_specs=tuple(reverse_specs)) + + @classmethod + def for_agent_server(cls, endpoints: list[OutsideEndpoint]) -> _OutsideEndpointRouting: + if len(endpoints) > 1: + raise ValueError("Agent-server mode supports only one OutsideEndpoint") + if not endpoints: + raise ValueError("Agent-server mode requires OutsideEndpoint passed to start()") + route = _OutsideEndpointRoute.for_endpoint(endpoints[0]) + return cls( + endpoints=tuple(endpoints), + _routes_by_env={route.endpoint.env_var: route}, + _agent_tunnel_port=route.target_port, + ) + + @property + def reverse_specs(self) -> list[str]: + return list(self._reverse_specs) + + @property + def agent_tunnel_port(self) -> int | None: + return self._agent_tunnel_port + + def agent_tunnel_target(self) -> tuple[str, int]: + if not self.endpoints: + raise ValueError("Agent-server mode requires OutsideEndpoint passed to start()") + route = self._routes_by_env[self.endpoints[0].env_var] + return route.host, route.target_port + + def env_overrides(self) -> dict[str, str]: + return {env_var: route.resolved_endpoint_url() for env_var, route in self._routes_by_env.items()} + + def resolved_endpoint_url(self, env_var: str) -> str | None: + route = self._routes_by_env.get(env_var) + if route is None: + return None + return route.resolved_endpoint_url() + + def resolve_url(self, url: str) -> str: + parsed = urlparse(url) + for route in self._routes_by_env.values(): + if route.source_netloc == parsed.netloc: + return route.resolve_url(url) + if self._agent_tunnel_port is not None: + return parsed._replace(netloc=f"127.0.0.1:{self._agent_tunnel_port}").geturl() + raise RuntimeError("resolve_outside_endpoint() requires SSH reverse tunnel") + + @staticmethod + def _allocate_reverse_port(preferred: int, used_ports: set[int]) -> int: + if 0 < preferred <= 65535 and preferred not in used_ports: + used_ports.add(preferred) + return preferred + for candidate in range(20000, 61000): + if candidate not in used_ports: + used_ports.add(candidate) + return candidate + raise RuntimeError("No available local port for ECS reverse tunnel") + + +# ── Retry utilities ────────────────────────────────────────────────── + +_RETRYABLE_CODES = frozenset( + { + "ThrottlingException", + "TooManyRequestsException", + "ServiceUnavailable", + "RequestLimitExceeded", + } +) +_RETRYABLE_MESSAGES = ( + "capacity is unavailable", + "rate exceeded", + "too many concurrent", + "throttl", + "connect timeout", + "read timeout", + "connection reset", + "endpointconnectionerror", +) + + +def _is_retryable_error(exc: Exception) -> bool: + msg = str(exc).lower() + code = "" + if hasattr(exc, "response"): + code = (exc.response.get("Error") or {}).get("Code", "") # type: ignore[union-attr] + return code in _RETRYABLE_CODES or any(m in msg for m in _RETRYABLE_MESSAGES) + + +def _retry_with_backoff( + func: Callable[[], T], + *, + operation_name: str, + max_retries: int | None = None, + base_delay: float = 1.0, + max_delay: float = 60.0, + jitter: float = 0.5, +) -> T: + attempt = 0 + while True: + try: + return func() + except Exception as exc: + if not _is_retryable_error(exc): + raise + attempt += 1 + if max_retries is not None and attempt > max_retries: + logger.error("%s failed after %d retries: %s", operation_name, attempt - 1, exc) + raise + delay = min(base_delay * (2 ** (attempt - 1)), max_delay) + delay *= 1 + random.uniform(-jitter, jitter) + logger.warning("%s throttled (attempt %d), retrying in %.1fs: %s", operation_name, attempt, delay, exc) + time.sleep(delay) + + +# ── SSH helpers ────────────────────────────────────────────────────── + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def download_secret_to_file(secret_arn: str, region: str | None = None) -> str: + """Fetch a Secrets Manager secret → temp file (mode 0600).""" + key_material = download_secret_to_string(secret_arn, region=region) + fd, path = tempfile.mkstemp(prefix="ecs-ssh-", suffix=".key") + try: + os.write(fd, key_material.encode()) + finally: + os.close(fd) + os.chmod(path, 0o600) + return path + + +def download_secret_to_string(secret_arn: str, region: str | None = None) -> str: + boto3, *_ = _require_aws_sdks() + sm = boto3.client("secretsmanager", region_name=region) + return _retry_with_backoff( + lambda: sm.get_secret_value(SecretId=secret_arn)["SecretString"], + operation_name="secretsmanager.get_secret_value", + max_retries=5, + ) + + +# ── SSH tunnel ─────────────────────────────────────────────────────── + + +class SshTunnel: + """Manages an ``ssh -N`` subprocess with ``-L`` / ``-R`` tunnels.""" + + def __init__( + self, + *, + host: str, + port: int = 2222, + user: str = "root", + key_file: str, + forward_port: int | None = None, + forwards: list[str] | None = None, + reverses: list[str] | None = None, + local_port_override: int | None = None, + ) -> None: + self._host = host + self._port = port + self._user = user + self._key_file = key_file + self._simple_forward_port = forward_port + self._forwards = list(forwards or []) + self._reverses = list(reverses or []) + self._local_port: int | None = local_port_override + self._proc: subprocess.Popen[bytes] | None = None + + @property + def local_port(self) -> int: + if self._local_port is None: + raise RuntimeError("Tunnel not open yet — call open() first") + return self._local_port + + @property + def is_open(self) -> bool: + return self._proc is not None and self._proc.poll() is None + + def open(self, *, max_retries: int = 15, initial_backoff: float = 5.0) -> None: + if self.is_open: + return + use_simple = self._simple_forward_port is not None + last_err = "" + backoff = initial_backoff + for attempt in range(1, max_retries + 1): + if use_simple: + self._local_port = _free_port() + cmd = self._build_ssh_cmd() + logger.info("SSH tunnel attempt %d/%d: %s", attempt, max_retries, " ".join(cmd)) + self._proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE) + time.sleep(3) + if self._proc.poll() is None: + if self._local_port: + try: + self._wait_for_local_port(self._local_port, timeout=15.0) + except Exception as port_exc: + logger.warning("SSH alive but forward port %d not open: %s", self._local_port, port_exc) + self._kill() + last_err = str(port_exc) + time.sleep(min(5.0, attempt * 1.5)) + continue + logger.info("SSH tunnel started (pid=%d, attempt %d/%d)", self._proc.pid, attempt, max_retries) + return + stderr = self._proc.stderr.read().decode(errors="replace") if self._proc.stderr else "" + last_err = stderr.strip() + self._proc = None + if not any( + m in last_err + for m in ( + "Connection refused", + "Connection timed out", + "No route to host", + "Connection reset", + ) + ): + raise RuntimeError(f"SSH tunnel exited immediately (attempt {attempt}): {last_err}") + logger.warning( + "SSH tunnel attempt %d/%d failed: %s — retrying in %.0fs", attempt, max_retries, last_err, backoff + ) + time.sleep(backoff) + backoff = min(30.0, backoff * 1.5) + raise RuntimeError(f"SSH tunnel failed after {max_retries} attempts: {last_err}") + + def close(self) -> None: + self._kill() + + def wait_ready(self, *, health_url: str | None = None, timeout: float = 300.0) -> None: + if health_url: + self._poll_health(health_url, timeout) + elif self._local_port: + self._wait_for_local_port(self._local_port, timeout) + + def check_health(self) -> bool: + return self.is_open + + def __enter__(self) -> SshTunnel: + self.open() + return self + + def __exit__(self, *exc: object) -> None: + self.close() + + def _build_ssh_cmd(self) -> list[str]: + cmd = [ + "ssh", + "-N", + "-o", + "StrictHostKeyChecking=no", + "-o", + "UserKnownHostsFile=/dev/null", + "-o", + "ServerAliveInterval=30", + "-o", + "ServerAliveCountMax=20", + "-o", + "ConnectTimeout=15", + "-o", + "ExitOnForwardFailure=yes", + "-o", + "LogLevel=ERROR", + "-i", + self._key_file, + "-p", + str(self._port), + ] + if self._simple_forward_port is not None: + cmd += ["-L", f"127.0.0.1:{self._local_port}:127.0.0.1:{self._simple_forward_port}"] + for spec in self._forwards: + cmd += ["-L", spec] + for spec in self._reverses: + cmd += ["-R", spec] + cmd.append(f"{self._user}@{self._host}") + return cmd + + def _kill(self) -> None: + if self._proc is None: + return + try: + self._proc.terminate() + try: + self._proc.wait(timeout=5) + except subprocess.TimeoutExpired: + self._proc.kill() + logger.info("SSH tunnel closed (pid=%d)", self._proc.pid) + except ProcessLookupError: + pass + finally: + self._proc = None + + def _wait_for_local_port(self, port: int, timeout: float = 30.0) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if self._proc and self._proc.poll() is not None: + raise RuntimeError("SSH tunnel process exited while waiting for port") + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.settimeout(1.0) + s.connect(("127.0.0.1", port)) + return + except OSError: + time.sleep(0.3) + raise TimeoutError(f"Local port 127.0.0.1:{port} not open after {timeout:.0f}s") + + def _poll_health(self, url: str, timeout: float) -> None: + import urllib.error + import urllib.request + + deadline = time.monotonic() + timeout + attempt = 0 + while time.monotonic() < deadline: + attempt += 1 + if not self.is_open: + raise RuntimeError("SSH tunnel died while waiting for health endpoint") + try: + with urllib.request.urlopen(urllib.request.Request(url, method="GET"), timeout=5) as resp: + if resp.status == 200: + logger.info("Health endpoint ready (attempt %d): %s", attempt, url) + return + except (urllib.error.URLError, OSError, TimeoutError): + pass + time.sleep(min(3.0, 1.0 + attempt * 0.5)) + raise TimeoutError(f"Health endpoint not reachable after {timeout:.0f}s: {url}") + + +# ── SSH sidecar container builder ──────────────────────────────────── + + +def build_ssh_sidecar_container( + sidecar_cfg: SshSidecarConfig, + *, + public_key_value: str, + max_lifetime_sec: int, + log_group: str | None = None, + log_region: str = "us-east-1", + log_stream_prefix: str = "ecs-sandbox", +) -> dict[str, Any]: + port = sidecar_cfg.sshd_port + image = sidecar_cfg.image or "alpine:latest" + sshd_cfg = ( + f"Port {port}\\nPermitRootLogin prohibit-password\\n" + "PasswordAuthentication no\\nAllowTcpForwarding yes\\n" + "PermitListen any\\nGatewayPorts clientspecified\\n" + "X11Forwarding no\\nPrintMotd no\\nLogLevel ERROR\\n" + "ClientAliveInterval 30\\nClientAliveCountMax 20\\n" + "TCPKeepAlive yes\\nUseDNS no\\nMaxSessions 50\\n" + ) + watchdog = "" + if max_lifetime_sec > 0: + watchdog = ( + f"( sleep {max_lifetime_sec}; " + f"echo 'sidecar watchdog: TTL ({max_lifetime_sec}s) reached'; " + "kill 1 2>/dev/null; sleep 3; kill -9 1 2>/dev/null ) & " + ) + sshd_cmd = ( + "set -e; apk add --no-cache openssh-server netcat-openbsd; " + "mkdir -p /root/.ssh; chmod 700 /root/.ssh; " + 'printf "%s\\n" "$SSH_PUBLIC_KEY" > /root/.ssh/authorized_keys; ' + "chmod 600 /root/.ssh/authorized_keys; ssh-keygen -A; " + f"printf '{sshd_cfg}' > /etc/ssh/sshd_config; " + f"{watchdog}exec /usr/sbin/sshd -D -e -p {port}" + ) + container: dict[str, Any] = { + "name": "ssh-tunnel", + "image": image, + "essential": True, + "entryPoint": ["sh", "-c"], + "command": [sshd_cmd], + "environment": [{"name": "SSH_PUBLIC_KEY", "value": public_key_value}], + "healthCheck": { + "command": ["CMD-SHELL", f"nc -z localhost {port} || exit 1"], + "interval": 5, + "timeout": 3, + "retries": 10, + "startPeriod": 30, + }, + } + if log_group: + container["logConfiguration"] = { + "logDriver": "awslogs", + "options": { + "awslogs-group": log_group, + "awslogs-region": log_region, + "awslogs-stream-prefix": f"{log_stream_prefix}-tunnel", + "awslogs-create-group": "true", + }, + } + return container + + +# ── Exec server — embedded script + HTTP client ───────────────────── + +EXEC_SERVER_SCRIPT = r'''#!/usr/bin/env python3 +"""Zero-dependency HTTP exec server for sandbox containers.""" +from __future__ import annotations +import base64, json, os, shutil, subprocess +from http.server import BaseHTTPRequestHandler, HTTPServer, ThreadingHTTPServer +from urllib.parse import parse_qs, urlparse +_BASH = shutil.which("bash") +_PORT = int(os.environ.get("TB_EXEC_PORT", "19542")) +_BIND = os.environ.get("TB_EXEC_BIND", "127.0.0.1") +class _H(BaseHTTPRequestHandler): + def log_message(self, fmt, *a): pass + def do_GET(self): + p = urlparse(self.path) + if p.path == "/health": self._ok({"ok": True}) + elif p.path == "/download": + qs = parse_qs(p.query) + paths = qs.get("path", []) + if not paths: self._err(400, "missing ?path=") + else: self._dl(paths[0]) + else: self._err(404, f"not found: {p.path}") + def do_POST(self): + p = urlparse(self.path) + body = self._body() + if p.path == "/exec": self._exec(body) + elif p.path == "/upload": self._up(body) + else: self._err(404, f"not found: {p.path}") + def _exec(self, b): + cmd = b.get("cmd") + if not cmd: self._err(400, "missing 'cmd'"); return + t = b.get("timeout", 300) + try: + cp = subprocess.run(cmd, shell=True, executable=_BASH, capture_output=True, timeout=t) + self._ok({"stdout": cp.stdout.decode("utf-8", errors="replace"), + "stderr": cp.stderr.decode("utf-8", errors="replace"), + "rc": cp.returncode}) + except subprocess.TimeoutExpired: + self._ok({"stdout":"","stderr":f"timed out after {t}s","rc":124}) + except Exception as e: + self._ok({"stdout":"","stderr":str(e),"rc":-1}) + def _up(self, b): + path, c = b.get("path"), b.get("content") + if not path or c is None: self._err(400, "missing path/content"); return + try: + data = base64.b64decode(c) + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + with open(path, "wb") as f: f.write(data) + m = b.get("mode") + if m: os.chmod(path, int(m, 8)) + self._ok({"ok": True}) + except Exception as e: self._err(500, str(e)) + def _dl(self, path): + if not os.path.isfile(path): self._err(404, f"not found: {path}"); return + try: + with open(path, "rb") as f: data = f.read() + self.send_response(200) + self.send_header("Content-Type", "application/octet-stream") + self.send_header("Content-Length", str(len(data))) + self.end_headers(); self.wfile.write(data) + except Exception as e: self._err(500, str(e)) + def _body(self): + n = int(self.headers.get("Content-Length", 0)) + if n == 0: return {} + try: return json.loads(self.rfile.read(n)) + except Exception: return {} + def _ok(self, obj): + p = json.dumps(obj).encode() + self.send_response(200) + self.send_header("Content-Type","application/json") + self.send_header("Content-Length",str(len(p))) + self.end_headers(); self.wfile.write(p) + def _err(self, code, msg): + p = json.dumps({"error": msg}).encode() + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(p))) + self.end_headers(); self.wfile.write(p) +if __name__ == "__main__": + s = ThreadingHTTPServer((_BIND, _PORT), _H) + print(f"exec_server on {_BIND}:{_PORT}", flush=True) + try: s.serve_forever() + except KeyboardInterrupt: pass + finally: s.server_close() +''' + +_EXEC_SERVER_B64 = base64.b64encode(EXEC_SERVER_SCRIPT.encode()).decode() + +_TRANSIENT_ERRORS = ( + ConnectionResetError, + ConnectionRefusedError, + ConnectionAbortedError, + BrokenPipeError, + TimeoutError, + OSError, +) + + +class ExecClient: + """Async HTTP client for the exec server (through the SSH tunnel). + + Methods are coroutines so each in-flight request occupies an event-loop + slot, not an executor thread — required to scale past asyncio's default + thread-pool cap when many long agent commands run concurrently (FEP-886). + """ + + def __init__(self, *, port: int, connect_timeout: float = 30.0) -> None: + self._base = f"http://127.0.0.1:{port}" + self._timeout = connect_timeout + # Lazy: aiohttp.ClientSession must be created inside a running event + # loop, but ExecClient is constructed from `_do_start` (which runs in + # asyncio.to_thread). Defer creation to the first request. + self._session: aiohttp.ClientSession | None = None + + async def _ensure_session(self) -> aiohttp.ClientSession: + if self._session is None or self._session.closed: + self._session = aiohttp.ClientSession() + return self._session + + async def close(self) -> None: + if self._session is not None and not self._session.closed: + await self._session.close() + self._session = None + + async def exec(self, cmd: str, *, timeout: int = 300) -> ExecResult: + resp = await self._post("/exec", {"cmd": cmd, "timeout": timeout}) + return ExecResult( + stdout=resp.get("stdout", ""), + stderr=resp.get("stderr", ""), + return_code=resp.get("rc", -1), + ) + + async def upload( + self, remote_path: str, data: bytes | Path, *, mode: str | None = None, max_retries: int = 3 + ) -> None: + if isinstance(data, Path): + data = data.read_bytes() + body: dict[str, Any] = {"path": remote_path, "content": base64.b64encode(data).decode()} + if mode is not None: + body["mode"] = mode + payload_mb = len(body["content"]) / (1024 * 1024) + upload_timeout = max(self._timeout, 60.0 + payload_mb * 2.0) + last_err: Exception | None = None + for attempt in range(1, max_retries + 1): + try: + resp = await self._post("/upload", body, timeout_override=upload_timeout) + if not resp.get("ok"): + raise RuntimeError(f"upload to {remote_path} failed: {resp}") + return + except (TimeoutError, OSError, RuntimeError) as exc: + last_err = exc + if attempt < max_retries: + logger.warning("upload %s attempt %d/%d: %s", remote_path, attempt, max_retries, exc) + await asyncio.sleep(2.0 * attempt) + raise RuntimeError(f"upload to {remote_path} failed after {max_retries} attempts: {last_err}") + + async def download(self, remote_path: str, *, max_retries: int = 3) -> bytes: + import urllib.parse + + url = f"{self._base}/download?path={urllib.parse.quote(remote_path)}" + return await self._request( + label=f"download {remote_path}", url=url, method="GET", timeout=self._timeout, max_retries=max_retries + ) + + async def health(self) -> bool: + try: + await self._request(label="health", url=f"{self._base}/health", method="GET", timeout=5, max_retries=1) + return True + except (ConnectionError, OSError, TimeoutError, RuntimeError): + return False + + async def _post( + self, path: str, body: dict[str, Any], *, timeout_override: float | None = None, max_retries: int = 4 + ) -> dict[str, Any]: + url = f"{self._base}{path}" + payload = json.dumps(body).encode() + if timeout_override is not None: + http_timeout = timeout_override + else: + cmd_timeout = body.get("timeout") + http_timeout = ( + max(self._timeout, cmd_timeout + 30) if isinstance(cmd_timeout, (int, float)) else self._timeout + ) + raw = await self._request( + label=f"POST {path}", + url=url, + method="POST", + data=payload, + headers={"Content-Type": "application/json"}, + timeout=http_timeout, + max_retries=max_retries, + ) + return json.loads(raw) + + async def _request( + self, + *, + label: str, + url: str, + method: str, + data: bytes | None = None, + headers: dict[str, str] | None = None, + timeout: float, + max_retries: int, + ) -> bytes: + session = await self._ensure_session() + client_timeout = aiohttp.ClientTimeout(total=timeout) + last_err: Exception | None = None + for attempt in range(1, max_retries + 1): + try: + async with session.request(method, url, data=data, headers=headers, timeout=client_timeout) as resp: + body = await resp.read() + if resp.status >= 400: + raise RuntimeError(f"{label} failed (HTTP {resp.status}): {body.decode(errors='replace')}") + return body + except RuntimeError: + raise + except (aiohttp.ClientError, TimeoutError, OSError) as exc: + last_err = exc + if attempt < max_retries: + wait = min(15.0, 2.0 ** (attempt - 1)) + logger.warning("%s attempt %d/%d: %s — retry in %.1fs", label, attempt, max_retries, exc, wait) + await asyncio.sleep(wait) + continue + raise ConnectionError(f"{label} failed after {max_retries} attempts: {last_err}") from last_err + raise ConnectionError(f"{label} unreachable") + + +# ── Image builder — AWS CodeBuild + ECR caching ───────────────────── + + +class ImageBuilder: + """Build Docker images via CodeBuild → ECR with content-hash caching.""" + + _lock = threading.Lock() + _inflight_builds: dict[str, threading.Event] = {} + _build_semaphore: threading.Semaphore | None = None + _build_semaphore_size: int = 0 + + @staticmethod + def get_ecr_image_tag(environment_dir: str | Path, environment_name: str) -> str: + h = hashlib.sha256() + root = Path(environment_dir) + for p in sorted(root.rglob("*")): + if p.is_file(): + h.update(str(p.relative_to(root)).encode()) + h.update(p.read_bytes()) + return f"{environment_name}__{h.hexdigest()[:8]}" + + @staticmethod + def image_exists_in_ecr(ecr_repository: str, tag: str, region: str | None = None) -> bool: + boto3, _, ClientError = _require_aws_sdks() + ecr_region = ImageBuilder._ecr_region(ecr_repository, fallback=region) + ecr = boto3.client("ecr", region_name=ecr_region) + repo_name = ecr_repository.split("/", 1)[1] if "/" in ecr_repository else ecr_repository + try: + _retry_with_backoff( + lambda: ecr.describe_images(repositoryName=repo_name, imageIds=[{"imageTag": tag}]), + operation_name="ecr.describe_images", + max_retries=5, + ) + return True + except ClientError as exc: + code = exc.response.get("Error", {}).get("Code", "") + if code in ("ImageNotFoundException", "RepositoryNotFoundException"): + return False + raise + + @staticmethod + def _ecr_region(ecr_repository: str, fallback: str | None = None) -> str | None: + """Extract the region from an ECR repo URL like '123.dkr.ecr.us-west-2.amazonaws.com/repo'.""" + parts = ecr_repository.split(".") + if len(parts) >= 4 and parts[1] == "dkr" and parts[2] == "ecr": + return parts[3] + return fallback + + @staticmethod + def list_ecr_tags(ecr_repository: str, region: str | None = None) -> set[str]: + """Return all image tags present in an ECR repository. + + Uses paginated ``list_images`` to fetch every tagged image in + a handful of API calls, rather than one ``describe_images`` + call per tag. + """ + boto3, _, ClientError = _require_aws_sdks() + ecr_region = ImageBuilder._ecr_region(ecr_repository, fallback=region) + ecr = boto3.client("ecr", region_name=ecr_region) + repo_name = ecr_repository.split("/", 1)[1] if "/" in ecr_repository else ecr_repository + + def _fetch_all_tags() -> set[str]: + tags: set[str] = set() + paginator = ecr.get_paginator("list_images") + for page in paginator.paginate( + repositoryName=repo_name, + filter={"tagStatus": "TAGGED"}, + ): + for img_id in page.get("imageIds", []): + if tag := img_id.get("imageTag"): + tags.add(tag) + return tags + + try: + return _retry_with_backoff(_fetch_all_tags, operation_name="ecr.list_images", max_retries=5) + except ClientError as exc: + if exc.response.get("Error", {}).get("Code") == "RepositoryNotFoundException": + return set() + raise + + @staticmethod + def ecr_docker_login(ecr_repository: str, region: str | None = None) -> None: + """Authenticate the local Docker daemon against an ECR registry.""" + ecr_region = ImageBuilder._ecr_region(ecr_repository, fallback=region) + registry = ecr_repository.split("/")[0] + region_flag = f" --region {ecr_region}" if ecr_region else "" + cmd = f"aws ecr get-login-password{region_flag} | docker login --username AWS --password-stdin {registry}" + result = subprocess.run(cmd, shell=True, capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError(f"ECR docker login failed: {result.stderr.strip()}") + logger.info("ECR docker login succeeded for %s", registry) + + @staticmethod + def docker_push_to_ecr(local_image: str, ecr_repository: str, tag: str) -> str: + """Tag a local Docker image and push it to ECR. Returns the ECR URL.""" + ecr_url = f"{ecr_repository}:{tag}" + subprocess.run(["docker", "tag", local_image, ecr_url], check=True, capture_output=True) + result = subprocess.run(["docker", "push", ecr_url], capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError(f"docker push {ecr_url} failed: {result.stderr.strip()}") + logger.info("Pushed %s -> %s", local_image, ecr_url) + return ecr_url + + @classmethod + def ensure_image_built(cls, *, cfg: EcsFargateConfig, environment_name: str, force_build: bool = False) -> str: + ecr_repo = cfg.ecr_repository + env_dir = cfg.environment_dir + if not ecr_repo or not env_dir: + raise ValueError("ecr_repository and environment_dir are required for image building") + tag = cls.get_ecr_image_tag(env_dir, environment_name) + image_url = f"{ecr_repo}:{tag}" + + # Dedup: check if another thread is already building this tag + with cls._lock: + if tag in cls._inflight_builds: + event = cls._inflight_builds[tag] + event.wait() + return image_url + if not force_build and cls.image_exists_in_ecr(ecr_repo, tag, cfg.region): + logger.info("ECR cache hit — skipping build: %s", image_url) + return image_url + + # Register as builder + event = threading.Event() + with cls._lock: + if tag in cls._inflight_builds: + cls._inflight_builds[tag].wait() + return image_url + cls._inflight_builds[tag] = event + with cls._lock: + if cls._build_semaphore is None or cls._build_semaphore_size != cfg.build_parallelism: + cls._build_semaphore = threading.Semaphore(cfg.build_parallelism) + cls._build_semaphore_size = cfg.build_parallelism + try: + cls._build_semaphore.acquire() # type: ignore[union-attr] + try: + if not force_build and cls.image_exists_in_ecr(ecr_repo, tag, cfg.region): + return image_url + cls._build_and_push(cfg=cfg, environment_name=environment_name, tag=tag, image_url=image_url) + finally: + cls._build_semaphore.release() # type: ignore[union-attr] + finally: + event.set() + with cls._lock: + cls._inflight_builds.pop(tag, None) + return image_url + + @staticmethod + def _upload_build_context(cfg: EcsFargateConfig, environment_name: str, nonce: str) -> str: + boto3, *_ = _require_aws_sdks() + env_dir = Path(cfg.environment_dir or ".") + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf: + for item in env_dir.rglob("*"): + if item.is_file(): + zf.write(item, arcname=str(item.relative_to(env_dir))) + buf.seek(0) + s3 = boto3.client("s3", region_name=cfg.region) + s3_prefix = cfg.s3_prefix or "ecs-sandbox" + s3_key = f"{s3_prefix}/codebuild/{environment_name}-{nonce}.zip" + body = buf.read() + _retry_with_backoff( + lambda: s3.put_object(Bucket=cfg.s3_bucket, Key=s3_key, Body=body), + operation_name="s3.put_object(build_context)", + max_retries=5, + ) + return s3_key + + @staticmethod + def _resolve_codebuild_project(cfg: EcsFargateConfig, cb: Any, nonce: str) -> str: + _, _, ClientError = _require_aws_sdks() + if cfg.codebuild_project: + return cfg.codebuild_project + if not cfg.codebuild_service_role: + raise RuntimeError("codebuild_project or codebuild_service_role is required") + project_name = f"ecs-sandbox-build-{nonce}" + try: + _retry_with_backoff( + lambda: cb.create_project( + name=project_name, + source={"type": "NO_SOURCE", "buildspec": "version: 0.2"}, + artifacts={"type": "NO_ARTIFACTS"}, + environment={ + "type": "LINUX_CONTAINER", + "image": "aws/codebuild/amazonlinux-x86_64-standard:5.0", + "computeType": cfg.codebuild_compute_type, + "privilegedMode": True, + }, + serviceRole=cfg.codebuild_service_role, + timeoutInMinutes=cfg.codebuild_build_timeout, + ), + operation_name="codebuild.create_project", + max_retries=5, + ) + except ClientError as e: + if "already exists" not in str(e).lower(): + raise + return project_name + + @staticmethod + def _generate_buildspec(cfg: EcsFargateConfig, repo_name: str, tag: str, image_url: str) -> str: + ecr_registry = (cfg.ecr_repository or "").split("/")[0] + ecr_region = ImageBuilder._ecr_region(cfg.ecr_repository or "", fallback="$AWS_DEFAULT_REGION") + pre_build_cmds = [ + f"aws ecr get-login-password --region {ecr_region}" + f" | docker login --username AWS --password-stdin {ecr_registry}", + ] + if cfg.dockerhub_secret_arn: + pre_build_cmds.append( + f"DOCKERHUB_CREDS=$(aws secretsmanager get-secret-value" + f" --secret-id {cfg.dockerhub_secret_arn}" + f" --query SecretString --output text --region $AWS_DEFAULT_REGION)" + f' && DH_USER=$(echo "$DOCKERHUB_CREDS" | python3 -c' + """ "import sys,json;print(json.load(sys.stdin)['username'])")""" + f' && if [ -n "$DH_USER" ]; then echo "$DOCKERHUB_CREDS" | python3 -c' + """ "import sys,json;print(json.load(sys.stdin)['password'])" """ + f'| docker login -u "$DH_USER" --password-stdin; fi' + f' || echo "Docker Hub login failed — continuing without auth"' + ) + pre_yaml = "\n".join(f" - {c}" for c in pre_build_cmds) + build_cmd = ( + f"for i in 1 2 3; do docker build -t {repo_name}:{tag} . && break; " + f'echo "build failed ($i/3), retry in 30s"; sleep 30; done' + ) + return ( + "version: 0.2\nphases:\n pre_build:\n commands:\n" + f"{pre_yaml}\n build:\n commands:\n" + f" - {build_cmd}\n - docker tag {repo_name}:{tag} {image_url}\n" + f" post_build:\n commands:\n - docker push {image_url}\n" + ) + + @staticmethod + def _poll_codebuild(cb: Any, build_id: str, image_url: str) -> None: + while True: + time.sleep(10 + random.uniform(0, 5)) + build = _retry_with_backoff( + lambda: cb.batch_get_builds(ids=[build_id])["builds"][0], + operation_name=f"BatchGetBuilds({build_id})", + max_retries=8, + base_delay=2.0, + max_delay=120.0, + ) + status = build["buildStatus"] + if status == "SUCCEEDED": + logger.info("CodeBuild succeeded: %s", build_id) + return + if status in ("FAILED", "FAULT", "STOPPED", "TIMED_OUT"): + phases = build.get("phases", []) + failed = [p for p in phases if p.get("phaseStatus") not in (None, "SUCCEEDED")] + ctx = "; ".join(f"{p['phaseType']}: {p.get('phaseStatus')}" for p in failed) or status + raise RuntimeError(f"CodeBuild failed for {image_url}: {ctx} (build: {build_id})") + logger.debug("CodeBuild %s — phase=%s status=%s", build_id, build.get("currentPhase"), status) + + @classmethod + def _build_and_push(cls, *, cfg: EcsFargateConfig, environment_name: str, tag: str, image_url: str) -> None: + boto3, *_ = _require_aws_sdks() + ecr_repo = cfg.ecr_repository or "" + repo_name = ecr_repo.split("/", 1)[1] if "/" in ecr_repo else ecr_repo + nonce = uuid.uuid4().hex[:8] + logger.info("Building image via CodeBuild: %s", image_url) + s3_key = cls._upload_build_context(cfg, environment_name, nonce) + cb = boto3.client("codebuild", region_name=cfg.region) + project_name = cls._resolve_codebuild_project(cfg, cb, nonce) + buildspec = cls._generate_buildspec(cfg, repo_name, tag, image_url) + resp = _retry_with_backoff( + lambda: cb.start_build( + projectName=project_name, + sourceTypeOverride="S3", + sourceLocationOverride=f"{cfg.s3_bucket}/{s3_key}", + buildspecOverride=buildspec, + timeoutInMinutesOverride=cfg.codebuild_build_timeout, + privilegedModeOverride=True, + environmentTypeOverride="LINUX_CONTAINER", + imageOverride="aws/codebuild/amazonlinux-x86_64-standard:5.0", + computeTypeOverride=cfg.codebuild_compute_type, + ), + operation_name="codebuild.start_build", + max_retries=5, + ) + build_id = resp["build"]["id"] + logger.info("CodeBuild started: %s", build_id) + cls._poll_codebuild(cb, build_id, image_url) + + @classmethod + def run_buildspec_via_codebuild( + cls, + *, + cfg: EcsFargateConfig, + buildspec: str, + job_label: str = "harness-build", + timeout_minutes: int | None = None, + ) -> None: + """Run an arbitrary buildspec via CodeBuild (privileged mode for DinD). + + Unlike :meth:`_build_and_push` which uploads a Dockerfile context to S3, + this method uses ``NO_SOURCE`` — the buildspec is fully self-contained + (e.g. it installs packages and runs a harness that builds Docker images + internally). + """ + boto3, *_ = _require_aws_sdks() + nonce = uuid.uuid4().hex[:8] + cb = boto3.client("codebuild", region_name=cfg.region) + project_name = cls._resolve_codebuild_project(cfg, cb, nonce) + timeout = timeout_minutes or cfg.codebuild_build_timeout + + logger.info("Starting CodeBuild harness build: %s (timeout=%dm)", job_label, timeout) + resp = _retry_with_backoff( + lambda: cb.start_build( + projectName=project_name, + sourceTypeOverride="NO_SOURCE", + buildspecOverride=buildspec, + timeoutInMinutesOverride=timeout, + privilegedModeOverride=True, + environmentTypeOverride="LINUX_CONTAINER", + imageOverride="aws/codebuild/amazonlinux-x86_64-standard:5.0", + computeTypeOverride=cfg.codebuild_compute_type, + ), + operation_name="codebuild.start_build(harness)", + max_retries=5, + ) + build_id = resp["build"]["id"] + logger.info("CodeBuild harness build started: %s (build=%s)", job_label, build_id) + cls._poll_codebuild(cb, build_id, job_label) + + +# ── Core sandbox ───────────────────────────────────────────────────── + +_active_sandboxes: dict[int, Any] = {} +_cleanup_lock = threading.RLock() +_atexit_registered = False +_PROCESS_NONCE = f"{int(time.time())}-{uuid.uuid4().hex[:8]}" +_exec_server_url_cache: dict[str, str] = {} + +_task_def_cache: dict[str, str] = {} +_task_def_cache_lock = threading.Lock() +_task_def_inflight: dict[str, threading.Event] = {} + +# Env vars whose values vary per sandbox invocation (workspace session id, etc). +# Routed via RunTask containerOverrides so the underlying task definition stays +# content-stable and the FEP-866 hash cache hits across invocations of the same +# task. Per-invocation env keys discovered dynamically from OutsideEndpoint +# routing (e.g. MODEL_BASE_URL with session-scoped URLs) are merged in at call +# time; this set holds the keys that are NOT visible to OutsideEndpoint routing. +_PER_INVOCATION_ENV_KEYS: frozenset[str] = frozenset({"_NEL_EFS_SESSION"}) + + +def _compute_task_def_hash(payload: dict[str, Any]) -> str: + # Strip logConfiguration from every container definition before hashing. + # Log config (group, stream-prefix, region) is a visibility annotation — it has + # no effect on what the sandbox does. Two task defs that differ only in + # log_stream_prefix or log_group are functionally identical and should share + # a cache entry so cross-run SSM cache hits work across differently-named runs. + def _strip_log_cfg(containers: list) -> list: + return [{k: v for k, v in c.items() if k != "logConfiguration"} for c in containers] + + canonical = { + k: (_strip_log_cfg(v) if k == "containerDefinitions" else v) for k, v in payload.items() if k != "family" + } + blob = json.dumps(canonical, sort_keys=True, default=str).encode() + return hashlib.sha256(blob).hexdigest()[:24] + + +def _emergency_cleanup() -> None: + with _cleanup_lock: + for sb in list(_active_sandboxes.values()): + try: + sb._sync_stop() + except Exception: + logger.debug("Emergency cleanup failed for sandbox %s", id(sb), exc_info=True) + + +class EcsFargateSandbox: + """ECS Fargate sandbox — async :class:`Sandbox` protocol.""" + + def __init__(self, spec: SandboxSpec, *, ecs_config: EcsFargateConfig) -> None: + self._spec = spec + self._cfg = ecs_config + self._task_arn: str | None = None + self._task_def_arn: str | None = None + self._task_ip: str | None = None + self._ssh_key_file: str | None = None + self._ssh_tunnel: SshTunnel | None = None + self._exec_client: ExecClient | None = None + self._started = False + self._stopped = False + self._ecs: Any = None + self._ec2: Any = None + self._ssm: Any = None + self._runtime_container_env: dict[str, str] = {} + self._ssh_tunnel_port: int | None = None + self._agent_forward_port: int | None = None + self._outside_endpoints: list[OutsideEndpoint] = [] + self._outside_endpoint_routing = _OutsideEndpointRouting.empty() + self._run_id = uuid.uuid4().hex[:12] + + # ── Protocol properties ────────────────────────────────────────── + + @property + def spec(self) -> SandboxSpec: + return self._spec + + def resolved_endpoint_url(self, env_var: str) -> str | None: + return self._outside_endpoint_routing.resolved_endpoint_url(env_var) + + @property + def is_running(self) -> bool: + return self._started and not self._stopped + + @property + def container_ip(self) -> str | None: + return self._task_ip + + # ── Extra properties ───────────────────────────────────────────── + + @property + def task_arn(self) -> str | None: + return self._task_arn + + @property + def local_port(self) -> int | None: + if self._ssh_tunnel: + try: + return self._ssh_tunnel.local_port + except RuntimeError: + pass + return None + + @property + def ssh_tunnel(self) -> SshTunnel | None: + return self._ssh_tunnel + + @property + def exec_client(self) -> ExecClient | None: + return self._exec_client + + @property + def model_tunnel_port(self) -> int | None: + return self._ssh_tunnel_port + + # ── Protocol async lifecycle ───────────────────────────────────── + + async def start(self, *, outside_endpoints: list[OutsideEndpoint] | None = None) -> None: + if self._started: + return + self._outside_endpoints = outside_endpoints or [] + self._outside_endpoint_routing = _OutsideEndpointRouting.empty(self._outside_endpoints) + sidecar = self._cfg.ssh_sidecar + if sidecar and sidecar.exec_server_port is None: + _OutsideEndpointRouting.for_agent_server(self._outside_endpoints) + try: + await asyncio.to_thread(self._do_start) + self._started = True + except Exception: + if self._exec_client is not None: + await self._exec_client.close() + await asyncio.to_thread(self._cleanup) + raise + + async def stop(self) -> None: + if self._stopped: + return + self._stopped = True + if self._exec_client is not None: + await self._exec_client.close() + await asyncio.to_thread(self._cleanup) + self._unregister_from_cleanup() + + async def exec( + self, + command: str, + timeout_sec: float = 180, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + user: str | int | None = None, + ) -> ExecResult: + self._require_exec_client() + shell_cmd = command + if env: + exports = " ".join(f"{k}={v}" for k, v in env.items()) + shell_cmd = f"export {exports} && {shell_cmd}" + if cwd: + shell_cmd = f"cd {cwd} && {shell_cmd}" + if user is not None: + if isinstance(user, int): + shell_cmd = f'su -s /bin/bash "$(getent passwd {user} | cut -d: -f1)" -c {shlex.quote(shell_cmd)}' + else: + shell_cmd = f"su -s /bin/bash {shlex.quote(str(user))} -c {shlex.quote(shell_cmd)}" + try: + return await self._exec_client.exec(shell_cmd, timeout=int(timeout_sec)) # type: ignore[union-attr] + except ConnectionError: + if self._ssh_tunnel and not self._ssh_tunnel.is_open: + logger.warning("SSH tunnel dead — attempting reconnect before re-raising") + try: + await self.reconnect_tunnel() + sidecar = self._cfg.ssh_sidecar + if sidecar and sidecar.exec_server_port is not None: + health_url = f"http://127.0.0.1:{self._ssh_tunnel.local_port}/health" # type: ignore[union-attr] + self._ssh_tunnel.wait_ready(health_url=health_url, timeout=60.0) # type: ignore[union-attr] + old_client = self._exec_client + self._exec_client = ExecClient(port=self._ssh_tunnel.local_port) # type: ignore[union-attr] + if old_client is not None: + await old_client.close() + return await self._exec_client.exec(shell_cmd, timeout=int(timeout_sec)) + except Exception as reconnect_err: + logger.warning("Tunnel reconnect failed: %s", reconnect_err) + raise + + async def upload(self, local_path: Path, remote_path: str) -> None: + self._require_exec_client() + local = Path(local_path) + if local.is_dir(): + for child in local.rglob("*"): + if child.is_file(): + await self.upload(child, f"{remote_path}/{child.relative_to(local)}") + return + if local.stat().st_size > 512 * 1024 and self._cfg.s3_bucket: + await self._upload_via_s3([local], os.path.dirname(remote_path) or "/tmp") + else: + await self._exec_client.upload(remote_path, local) # type: ignore[union-attr] + + async def download(self, remote_path: str, local_path: Path) -> None: + self._require_exec_client() + data = await self._exec_client.download(remote_path) # type: ignore[union-attr] + dest = Path(local_path) + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_bytes(data) + + def resolve_outside_endpoint(self, url: str) -> str: + return self._outside_endpoint_routing.resolve_url(url) + + async def __aenter__(self) -> Self: + await self.start() + return self + + async def __aexit__(self, *exc: object) -> None: + await self.stop() + + # ── Extra public methods ───────────────────────────────────────── + + async def reconnect_tunnel(self) -> None: + if self._stopped or not self._started: + raise RuntimeError("Cannot reconnect tunnel on a stopped/unstarted sandbox") + sidecar = self._cfg.ssh_sidecar + if sidecar is None: + return + if self._ssh_tunnel: + self._ssh_tunnel.close() + self._ssh_tunnel = None + await asyncio.to_thread(self._open_tunnel, sidecar) + + # ── Sync start (runs via asyncio.to_thread) ────────────────────── + + def _do_start(self) -> None: + cfg = self._cfg + sidecar = cfg.ssh_sidecar + if sidecar is None: + raise ValueError("ssh_sidecar must be configured") + self._init_aws_clients() + + built_image: str | None = None + env_dir = cfg.environment_dir or self._spec.environment_dir + if cfg.ecr_repository and env_dir: + per_task_cfg = _dc_replace(cfg, environment_dir=env_dir) + built_image = ImageBuilder.ensure_image_built( + cfg=per_task_cfg, environment_name=_sanitize_id(self._spec.image or "sandbox") + ) + image = self._resolve_image(built_image) + + if not sidecar.private_key_secret_arn or not sidecar.public_key_secret_arn: + raise ValueError("ssh_sidecar private_key_secret_arn and public_key_secret_arn are required") + self._ssh_key_file = download_secret_to_file(sidecar.private_key_secret_arn, cfg.region) + ssh_public_key_value = download_secret_to_string(sidecar.public_key_secret_arn, cfg.region) + + has_exec_server = sidecar.exec_server_port is not None + if not has_exec_server: + self._outside_endpoint_routing = _OutsideEndpointRouting.for_agent_server(self._outside_endpoints) + self._ssh_tunnel_port = self._outside_endpoint_routing.agent_tunnel_port + else: + self._outside_endpoint_routing = _OutsideEndpointRouting.for_exec_server(self._outside_endpoints, sidecar) + + command = self._build_container_command(sidecar) + env = self._build_env_vars() + stable_env, self._runtime_container_env = self._split_env(env) + log_region = cfg.region or os.environ.get("AWS_DEFAULT_REGION", "us-east-1") + sidecar_def = build_ssh_sidecar_container( + sidecar, + public_key_value=ssh_public_key_value, + max_lifetime_sec=cfg.max_task_lifetime_sec, + log_group=cfg.log_group, + log_region=log_region, + log_stream_prefix=cfg.log_stream_prefix or "ecs-sandbox", + ) + self._task_def_arn = self._register_task_definition( + image=image, command=command, env=stable_env, sidecar_def=sidecar_def + ) + self._task_arn = self._run_task(self._task_def_arn) + self._register_for_cleanup() + self._wait_for_running() + self._task_ip = self._get_task_public_ip() + self._wait_for_ssh_ready(self._task_ip, sidecar.sshd_port, sidecar.ssh_ready_timeout_sec) + self._open_tunnel(sidecar) + + if has_exec_server: + health_url = f"http://127.0.0.1:{self._ssh_tunnel.local_port}/health" # type: ignore[union-attr] + self._ssh_tunnel.wait_ready(health_url=health_url, timeout=sidecar.ssh_ready_timeout_sec) # type: ignore[union-attr] + self._exec_client = ExecClient(port=self._ssh_tunnel.local_port) # type: ignore[union-attr] + + # ── Internal helpers ───────────────────────────────────────────── + + def _init_aws_clients(self) -> None: + boto3, Config, _ = _require_aws_sdks() + boto_cfg = Config(connect_timeout=30, read_timeout=60, retries={"max_attempts": 8, "mode": "adaptive"}) + self._ecs = boto3.client("ecs", region_name=self._cfg.region, config=boto_cfg) + self._ec2 = boto3.client("ec2", region_name=self._cfg.region, config=boto_cfg) + self._ssm = boto3.client("ssm", region_name=self._cfg.region, config=boto_cfg) + + def _resolve_image(self, built_image: str | None = None) -> str: + if built_image: + return built_image + cfg = self._cfg + if cfg.image_template: + sanitized = _sanitize_id(self._spec.image or "sandbox") + fmt_keys = { + "task_id": self._spec.image or "", + "task_id_sanitized": sanitized, + **(self._spec.env or {}), + } + try: + return cfg.image_template.format_map(fmt_keys) + except KeyError as exc: + raise ValueError( + f"ecs.image_template placeholder {exc} not found in " + f"available keys: {sorted(fmt_keys)}. " + f"Hint: use sandbox.image_template (resolved via seed " + f"metadata) instead of ecs.image_template for task-specific " + f"placeholders like {{task_id}}." + ) from exc + if cfg.ecr_repository and self._spec.image: + return f"{cfg.ecr_repository}:{_sanitize_id(self._spec.image)}" + if self._spec.image: + return self._spec.image + if not cfg.task_definition: + raise ValueError( + "No image available: set image on SandboxSpec, image_template, " + "ecr_repository + environment_dir, or task_definition" + ) + return "" + + def _upload_exec_server(self) -> str: + cfg = self._cfg + if not cfg.s3_bucket: + raise ValueError("s3_bucket is required for exec server upload") + cache_key = f"{cfg.s3_bucket}/{self._run_id}" + if cache_key in _exec_server_url_cache: + return _exec_server_url_cache[cache_key] + boto3, *_ = _require_aws_sdks() + s3 = boto3.client("s3", region_name=cfg.region) + prefix = cfg.s3_prefix or "ecs-sandbox" + key = f"{prefix}/{self._run_id}-{_PROCESS_NONCE}/_exec_server/exec_server.py" + _retry_with_backoff( + lambda: s3.put_object(Bucket=cfg.s3_bucket, Key=key, Body=EXEC_SERVER_SCRIPT.encode()), + operation_name="s3.put_object(exec_server)", + max_retries=5, + ) + url = s3.generate_presigned_url("get_object", Params={"Bucket": cfg.s3_bucket, "Key": key}, ExpiresIn=21600) + _exec_server_url_cache[cache_key] = url + logger.info("Uploaded exec server → s3://%s/%s", cfg.s3_bucket, key) + return url + + def _build_container_command(self, sidecar: SshSidecarConfig) -> list[str] | None: + if sidecar.exec_server_port is None: + return None + exec_port = sidecar.exec_server_port or 19542 + hostname = re.sub(r"[^A-Za-z0-9._-]", "-", self._spec.image or "sandbox")[:63] + setup = ( + f"hostname {shlex.quote(hostname)} 2>/dev/null || true; " + f"echo '{_EXEC_SERVER_B64}' | base64 -d > /tmp/_exec_server.py; " + "if ! command -v python3 >/dev/null 2>&1; then " + " if command -v apt-get >/dev/null 2>&1; then " + " apt-get update -qq && apt-get install -y -qq --no-install-recommends python3 bash; " + " elif command -v apk >/dev/null 2>&1; then " + " apk add --no-cache python3 bash; " + " elif command -v yum >/dev/null 2>&1; then " + " yum install -y python3 bash; " + " elif command -v dnf >/dev/null 2>&1; then " + " dnf install -y python3 bash; " + " fi; " + "fi; " + "if ! command -v python3 >/dev/null 2>&1; then " + " echo 'FATAL: exec_server bootstrap failed — python3 not available' >&2; " + " exit 1; " + "fi; " + f"TB_EXEC_PORT={exec_port} TB_EXEC_BIND=127.0.0.1 " + "exec python3 /tmp/_exec_server.py" + ) + return ["sh", "-lc", setup] + + def _build_env_vars(self) -> dict[str, str]: + env: dict[str, str] = dict(self._spec.env) + cfg = self._cfg + if cfg.extra_env: + for k, v in cfg.extra_env.items(): + env[k] = self._render_env_value(v) + env.update(self._outside_endpoint_routing.env_overrides()) + return env + + def _split_env(self, env: dict[str, str]) -> tuple[dict[str, str], dict[str, str]]: + runtime_keys = _PER_INVOCATION_ENV_KEYS | set(self._outside_endpoint_routing.env_overrides().keys()) + stable = {k: v for k, v in env.items() if k not in runtime_keys} + runtime = {k: v for k, v in env.items() if k in runtime_keys} + return stable, runtime + + def _render_env_value(self, value: str) -> str: + if self._ssh_tunnel_port is not None: + value = value.replace("{ssh_tunnel_port}", str(self._ssh_tunnel_port)) + if self._task_ip: + value = value.replace("{task_ip}", self._task_ip) + value = value.replace("{image}", self._spec.image or "") + return value + + # ── Task definition registration ───────────────────────────────── + + def _register_task_definition( + self, *, image: str, command: list[str] | None, env: dict[str, str], sidecar_def: dict[str, Any] + ) -> str: + cfg = self._cfg + log_region = cfg.region or os.environ.get("AWS_DEFAULT_REGION", "us-east-1") + log_cfg: dict[str, Any] | None = None + if cfg.log_group: + log_cfg = { + "logDriver": "awslogs", + "options": { + "awslogs-group": cfg.log_group, + "awslogs-region": log_region, + "awslogs-stream-prefix": cfg.log_stream_prefix or "ecs-sandbox", + "awslogs-create-group": "true", + }, + } + + _, _, ClientError = _require_aws_sdks() + base: dict[str, Any] | None = None + if cfg.task_definition: + try: + base = _retry_with_backoff( + lambda: self._ecs.describe_task_definition(taskDefinition=cfg.task_definition)["taskDefinition"], + operation_name="ecs.describe_task_definition", + max_retries=5, + ) + except ClientError as exc: + if exc.response.get("Error", {}).get("Code") == "ClientException": + logger.warning("Base task definition %s not found, registering from scratch", cfg.task_definition) + else: + raise + + if base is not None: + return self._register_from_base( + base=base, image=image, command=command, env=env, sidecar_def=sidecar_def, log_cfg=log_cfg + ) + return self._register_from_scratch( + image=image, command=command, env=env, sidecar_def=sidecar_def, log_cfg=log_cfg + ) + + def _register_from_base( + self, + *, + base: dict, + image: str, + command: list[str] | None, + env: dict[str, str], + sidecar_def: dict, + log_cfg: dict | None, + ) -> str: + cfg = self._cfg + containers = list(base.get("containerDefinitions") or []) + target = next((cd for cd in containers if cd.get("name") == cfg.container_name), None) + if target is None: + raise RuntimeError( + f"Base task-def has no container '{cfg.container_name}'. " + f"Available: {[c.get('name') for c in containers]}" + ) + if image: + target["image"] = image + if command is not None: + target["command"] = command + target.pop("entryPoint", None) + if env: + existing = {e["name"]: e["value"] for e in target.get("environment", [])} + existing.update(env) + target["environment"] = [{"name": k, "value": v} for k, v in sorted(existing.items())] + if log_cfg: + target["logConfiguration"] = log_cfg + target["dependsOn"] = [{"containerName": "ssh-tunnel", "condition": "HEALTHY"}] + containers = [c for c in containers if c.get("name") != "ssh-tunnel"] + containers.append(sidecar_def) + + task_volumes, mount_points = self._build_efs_volumes() + if mount_points: + existing_mounts = target.get("mountPoints") or [] + target["mountPoints"] = existing_mounts + mount_points + + family = self._make_family_name() + payload: dict[str, Any] = { + "family": family, + "networkMode": base.get("networkMode", "awsvpc"), + "requiresCompatibilities": base.get("requiresCompatibilities", ["FARGATE"]), + "cpu": str(max(int(base.get("cpu") or "256"), int(cfg.cpu))), + "memory": str(max(int(base.get("memory") or "512"), int(cfg.memory))), + "containerDefinitions": containers, + "ephemeralStorage": { + "sizeInGiB": max( + (base.get("ephemeralStorage") or {}).get("sizeInGiB", 20), cfg.ephemeral_storage_gib or 20 + ) + }, + } + for k in ("taskRoleArn", "executionRoleArn", "runtimePlatform", "volumes"): + if base.get(k) is not None: + payload[k] = base[k] + if task_volumes: + existing_vols = payload.get("volumes") or [] + payload["volumes"] = existing_vols + task_volumes + if cfg.execution_role_arn: + payload["executionRoleArn"] = cfg.execution_role_arn + if cfg.task_role_arn: + payload["taskRoleArn"] = cfg.task_role_arn + return self._do_register(payload) + + def _register_from_scratch( + self, *, image: str, command: list[str] | None, env: dict[str, str], sidecar_def: dict, log_cfg: dict | None + ) -> str: + cfg = self._cfg + if not cfg.execution_role_arn: + raise RuntimeError("execution_role_arn required when no base task definition provided") + container_def: dict[str, Any] = { + "name": cfg.container_name, + "essential": True, + "dependsOn": [{"containerName": "ssh-tunnel", "condition": "HEALTHY"}], + } + if image: + container_def["image"] = image + if command is not None: + container_def["command"] = command + if cfg.container_port: + container_def["portMappings"] = [{"containerPort": cfg.container_port, "protocol": "tcp"}] + if env: + container_def["environment"] = [{"name": k, "value": v} for k, v in sorted(env.items())] + if log_cfg: + container_def["logConfiguration"] = log_cfg + + task_volumes, mount_points = self._build_efs_volumes() + if mount_points: + container_def["mountPoints"] = mount_points + + payload: dict[str, Any] = { + "family": self._make_family_name(), + "networkMode": "awsvpc", + "requiresCompatibilities": ["FARGATE"], + "cpu": cfg.cpu, + "memory": cfg.memory, + "executionRoleArn": cfg.execution_role_arn, + "containerDefinitions": [container_def, sidecar_def], + } + if task_volumes: + payload["volumes"] = task_volumes + if cfg.task_role_arn: + payload["taskRoleArn"] = cfg.task_role_arn + if cfg.ephemeral_storage_gib: + payload["ephemeralStorage"] = {"sizeInGiB": cfg.ephemeral_storage_gib} + return self._do_register(payload) + + def _build_efs_volumes(self) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + """Build EFS volume definitions and mount points from spec.volumes.""" + task_volumes: list[dict[str, Any]] = [] + mount_points: list[dict[str, Any]] = [] + for i, vol in enumerate(self._spec.volumes): + if not vol.is_efs: + continue + vol_name = f"efs-{i}" + efs_cfg: dict[str, Any] = { + "fileSystemId": vol.efs_filesystem_id, + "transitEncryption": "ENABLED", + } + if vol.efs_access_point_id: + efs_cfg["authorizationConfig"] = { + "accessPointId": vol.efs_access_point_id, + "iam": "ENABLED", + } + elif vol.efs_root_directory: + efs_cfg["rootDirectory"] = vol.efs_root_directory + task_volumes.append({"name": vol_name, "efsVolumeConfiguration": efs_cfg}) + mount_points.append( + { + "sourceVolume": vol_name, + "containerPath": vol.container_path, + "readOnly": vol.readonly, + } + ) + return task_volumes, mount_points + + def _ssm_lookup_task_def(self, h: str) -> str | None: + _, _, ClientError = _require_aws_sdks() + param_name = f"/{self._cfg.ssm_project}/task-defs/{h}" + try: + resp = self._ssm.get_parameter(Name=param_name) + except ClientError as exc: + code = exc.response["Error"]["Code"] + if code == "ParameterNotFound": + return None + logger.warning("SSM GetParameter %s failed (%s); falling through to register", param_name, code) + return None + + arn = resp["Parameter"]["Value"] + try: + desc = self._ecs.describe_task_definition(taskDefinition=arn) + except ClientError as exc: + logger.warning( + "Cached task def %s no longer describable (%s); re-registering", + arn, + exc.response["Error"]["Code"], + ) + return None + if desc["taskDefinition"]["status"] != "ACTIVE": + logger.info("SSM cache entry %s is not ACTIVE; re-registering (hash %s)", arn, h) + return None + logger.info("Reusing task def from SSM cache: %s (hash %s)", arn, h) + return arn + + def _ssm_write_task_def(self, h: str, arn: str) -> None: + _, _, ClientError = _require_aws_sdks() + param_name = f"/{self._cfg.ssm_project}/task-defs/{h}" + try: + self._ssm.put_parameter(Name=param_name, Value=arn, Type="String", Overwrite=True) + logger.info("Wrote SSM task-def cache entry: %s -> %s", param_name, arn) + except ClientError as exc: + code = exc.response["Error"]["Code"] + logger.warning("SSM PutParameter %s failed (%s); cache entry not written", param_name, code) + + def _do_register(self, payload: dict[str, Any]) -> str: + h = _compute_task_def_hash(payload) + + while True: + with _task_def_cache_lock: + if h in _task_def_cache: + arn = _task_def_cache[h] + logger.info("Reusing cached task def %s (hash %s)", arn, h) + return arn + if h in _task_def_inflight: + event = _task_def_inflight[h] + else: + event = threading.Event() + _task_def_inflight[h] = event + break + event.wait() + + try: + arn = self._ssm_lookup_task_def(h) or self._register_task_def_fresh(payload, h) + with _task_def_cache_lock: + _task_def_cache[h] = arn + return arn + finally: + self._release_inflight(h, event) + + def _register_task_def_fresh(self, payload: dict[str, Any], h: str) -> str: + resp = _retry_with_backoff( + lambda: self._ecs.register_task_definition(**payload), + operation_name="register_task_definition", + max_retries=25, + ) + arn = resp["taskDefinition"]["taskDefinitionArn"] + logger.info("Registered task def: %s (hash %s)", arn, h) + self._ssm_write_task_def(h, arn) + return arn + + def _release_inflight(self, h: str, event: threading.Event) -> None: + with _task_def_cache_lock: + _task_def_inflight.pop(h, None) + event.set() + + def _make_family_name(self) -> str: + nonce = uuid.uuid4().hex[:12] + raw = f"{self._cfg.task_definition_family_prefix}-{_sanitize_id(self._spec.image or 'sandbox')}-{nonce}" + family = re.sub(r"[^A-Za-z0-9_-]", "_", raw)[:255] + if not family or not re.match(r"^[A-Za-z0-9]", family): + family = f"ecs_{family}" + return family + + # ── Run task + wait ────────────────────────────────────────────── + + def _run_task(self, task_def_arn: str) -> str: + cfg = self._cfg + run_kwargs: dict[str, Any] = { + "cluster": cfg.cluster, + "taskDefinition": task_def_arn, + "launchType": "FARGATE", + "networkConfiguration": { + "awsvpcConfiguration": { + "subnets": cfg.subnets, + "securityGroups": cfg.security_groups, + "assignPublicIp": "ENABLED" if cfg.assign_public_ip else "DISABLED", + } + }, + } + if self._runtime_container_env: + run_kwargs["overrides"] = { + "containerOverrides": [ + { + "name": cfg.container_name, + "environment": [ + {"name": k, "value": v} for k, v in sorted(self._runtime_container_env.items()) + ], + } + ] + } + has_efs = any(v.is_efs for v in self._spec.volumes) + if cfg.platform_version: + run_kwargs["platformVersion"] = cfg.platform_version + elif has_efs: + run_kwargs["platformVersion"] = "1.4.0" + + last_failures: Any = None + for attempt in range(1, cfg.run_task_max_retries + 1): + try: + resp = _retry_with_backoff( + lambda: self._ecs.run_task(**run_kwargs), operation_name="run_task", max_retries=3 + ) + except Exception as exc: + if not _is_retryable_error(exc) or attempt >= cfg.run_task_max_retries: + raise + delay = min(60.0, 2.0 ** min(6, attempt - 1)) + random.random() * 2 + logger.warning( + "run_task failed (%d/%d): %s — retry in %.1fs", attempt, cfg.run_task_max_retries, exc, delay + ) + time.sleep(delay) + continue + failures = resp.get("failures") or [] + if not failures: + tasks = resp.get("tasks") or [] + if not tasks: + raise RuntimeError("run_task returned no tasks") + task_arn = tasks[0]["taskArn"] + logger.info("Started ECS task: %s", task_arn) + return task_arn + last_failures = failures + reasons = " | ".join(str(f.get("reason", "")) for f in failures) + if not any(m in reasons.lower() for m in _RETRYABLE_MESSAGES) or attempt >= cfg.run_task_max_retries: + raise RuntimeError(f"run_task failures: {failures}") + delay = min(60.0, 2.0 ** min(6, attempt - 1)) + random.random() * 2 + logger.warning( + "run_task capacity issue (%d/%d): %s — retry in %.1fs", + attempt, + cfg.run_task_max_retries, + reasons, + delay, + ) + time.sleep(delay) + raise RuntimeError(f"run_task failed after {cfg.run_task_max_retries} retries: {last_failures}") + + def _wait_for_running(self) -> None: + cfg = self._cfg + start = time.monotonic() + poll = 5.0 + last_status = "" + while True: + elapsed = time.monotonic() - start + if elapsed > cfg.startup_timeout_sec: + raise TimeoutError(f"ECS task not RUNNING after {elapsed:.0f}s (last: {last_status})") + try: + resp = self._ecs.describe_tasks(cluster=cfg.cluster, tasks=[self._task_arn]) + except Exception as exc: + if _is_retryable_error(exc): + time.sleep(poll + random.random() * 3) + continue + raise + tasks = resp.get("tasks") or [] + if not tasks: + raise RuntimeError("ECS task disappeared") + status = tasks[0].get("lastStatus", "UNKNOWN") + if status == "RUNNING": + logger.info("ECS task RUNNING after %.0fs", elapsed) + return + if status == "STOPPED": + raise RuntimeError(f"ECS task stopped: {tasks[0].get('stoppedReason')}") + if status != last_status: + logger.info("ECS task %s (%.0fs)", status, elapsed) + last_status = status + time.sleep(poll + random.random() * 3) + poll = min(15.0, poll + 0.5) + + def _get_task_public_ip(self) -> str: + max_retries = 10 + for attempt in range(1, max_retries + 1): + try: + resp = self._ecs.describe_tasks(cluster=self._cfg.cluster, tasks=[self._task_arn]) + tasks = resp.get("tasks") or [] + if not tasks: + raise RuntimeError("Task not found") + eni_id = None + for att in tasks[0].get("attachments") or []: + if att.get("type") == "ElasticNetworkInterface": + for d in att.get("details") or []: + if d.get("name") == "networkInterfaceId": + eni_id = d["value"] + break + if eni_id: + break + if not eni_id: + for att in tasks[0].get("attachments") or []: + for d in att.get("details") or []: + if d.get("name") == "privateIPv4Address" and d.get("value"): + return d["value"] + raise RuntimeError("No ENI/IP yet") + iface = self._ec2.describe_network_interfaces(NetworkInterfaceIds=[eni_id])["NetworkInterfaces"][0] + pub = (iface.get("Association") or {}).get("PublicIp") + if pub: + logger.info("Container public IP: %s", pub) + return pub + priv = iface.get("PrivateIpAddress") + if priv: + logger.info("Container private IP: %s", priv) + return priv + raise RuntimeError(f"ENI {eni_id} has no IP") + except Exception as exc: + if attempt >= max_retries: + raise + if _is_retryable_error(exc): + time.sleep(min(15.0, 2.0**attempt + random.random())) + else: + logger.warning("get_task_ip attempt %d/%d: %s", attempt, max_retries, exc) + time.sleep(min(15.0, 3.0 + attempt * 2)) + raise RuntimeError("get_task_ip exhausted retries") + + @staticmethod + def _wait_for_ssh_ready(host: str, port: int, timeout: float) -> None: + deadline = time.monotonic() + timeout + logger.info("Waiting for SSH at %s:%d", host, port) + while time.monotonic() < deadline: + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.settimeout(5.0) + s.connect((host, port)) + s.settimeout(5.0) + if b"SSH" in s.recv(256): + logger.info("SSH ready at %s:%d", host, port) + return + except OSError: + pass + time.sleep(2.0) + raise TimeoutError(f"SSH not ready at {host}:{port} after {timeout:.0f}s") + + def _open_tunnel(self, sidecar: SshSidecarConfig) -> None: + assert self._task_ip is not None and self._ssh_key_file is not None + if sidecar.exec_server_port is not None: + self._ssh_tunnel = SshTunnel( + host=self._task_ip, + port=sidecar.sshd_port, + user="root", + key_file=self._ssh_key_file, + forward_port=sidecar.exec_server_port, + reverses=self._outside_endpoint_routing.reverse_specs, + ) + self._ssh_tunnel.open() + else: + remote_host, remote_port = self._outside_endpoint_routing.agent_tunnel_target() + assert self._ssh_tunnel_port is not None + self._agent_forward_port = _free_port() + container_port = self._cfg.container_port + if not container_port: + raise ValueError("container_port is required in agent-server mode") + self._ssh_tunnel = SshTunnel( + host=self._task_ip, + port=sidecar.sshd_port, + user="root", + key_file=self._ssh_key_file, + forwards=[f"{self._agent_forward_port}:localhost:{container_port}"], + reverses=[f"{self._ssh_tunnel_port}:{remote_host}:{remote_port}"], + local_port_override=self._agent_forward_port, + ) + self._ssh_tunnel.open() + + # ── Cleanup ────────────────────────────────────────────────────── + + def _cleanup(self) -> None: + if self._ssh_tunnel: + try: + self._ssh_tunnel.close() + except Exception: + logger.debug("Failed to close SSH tunnel", exc_info=True) + self._ssh_tunnel = None + if self._task_arn and self._ecs: + try: + _retry_with_backoff( + lambda: self._ecs.stop_task( + cluster=self._cfg.cluster, task=self._task_arn, reason="sandbox cleanup" + ), + operation_name="stop_task", + max_retries=10, + ) + logger.info("Stopped ECS task: %s", self._task_arn) + except Exception as exc: + logger.warning("Failed to stop task %s: %s", self._task_arn, exc) + if self._ssh_key_file: + try: + os.remove(self._ssh_key_file) + except Exception: + logger.debug("Failed to remove SSH key file %s", self._ssh_key_file, exc_info=True) + self._ssh_key_file = None + + def _sync_stop(self) -> None: + """Synchronous stop for emergency cleanup (atexit handler).""" + if self._stopped: + return + self._stopped = True + self._cleanup() + + def _require_exec_client(self) -> None: + if self._exec_client is None: + raise RuntimeError( + "exec()/upload()/download() require exec-server mode " + "(ssh_sidecar.exec_server_port). In agent-server mode use sandbox.ssh_tunnel." + ) + + async def _upload_via_s3(self, paths: list[Path], dest_dir: str) -> None: + cfg = self._cfg + if not cfg.s3_bucket: + raise ValueError("s3_bucket is required for S3 staging") + + def _pack() -> bytes: + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tar: + for p in paths: + if p.is_file(): + tar.add(str(p), arcname=p.name) + elif p.is_dir(): + for child in p.rglob("*"): + if child.is_file(): + tar.add(str(child), arcname=str(child.relative_to(p))) + buf.seek(0) + return buf.read() + + body = await asyncio.to_thread(_pack) + boto3, *_ = _require_aws_sdks() + s3 = boto3.client("s3", region_name=cfg.region) + prefix = cfg.s3_prefix or "ecs-sandbox" + nonce = uuid.uuid4().hex[:12] + key = f"{prefix}/{self._run_id}/upload-{nonce}.tar.gz" + await asyncio.to_thread( + _retry_with_backoff, + lambda: s3.put_object(Bucket=cfg.s3_bucket, Key=key, Body=body), + operation_name="s3.put_object(upload)", + max_retries=5, + ) + url = await asyncio.to_thread( + s3.generate_presigned_url, + "get_object", + Params={"Bucket": cfg.s3_bucket, "Key": key}, + ExpiresIn=21600, + ) + dl_cmd = ( + f"mkdir -p {shlex.quote(dest_dir)} && TGZ=/tmp/_upload_$$.tar.gz && " + f"( curl -sf -L --max-time 300 -o $TGZ {shlex.quote(url)} 2>/dev/null || " + f"python3 -c 'import urllib.request as u,sys;u.urlretrieve(sys.argv[1],sys.argv[2])' " + f"{shlex.quote(url)} $TGZ ) && " + f"tar xzf $TGZ -C {shlex.quote(dest_dir)} && rm -f $TGZ && echo ok" + ) + result = await self._exec_client.exec(dl_cmd, timeout=360) # type: ignore[union-attr] + if "ok" not in result.stdout: + raise RuntimeError( + f"S3 upload extraction failed (rc={result.return_code}): {result.stderr or result.stdout}" + ) + + # ── Atexit cleanup ─────────────────────────────────────────────── + + def _register_for_cleanup(self) -> None: + global _atexit_registered + with _cleanup_lock: + _active_sandboxes[id(self)] = self + if not _atexit_registered: + atexit.register(_emergency_cleanup) + _atexit_registered = True + + def _unregister_from_cleanup(self) -> None: + with _cleanup_lock: + _active_sandboxes.pop(id(self), None) diff --git a/nemo_gym/sandbox/providers/ecs_fargate/provider.py b/nemo_gym/sandbox/providers/ecs_fargate/provider.py new file mode 100644 index 0000000000..08af81d396 --- /dev/null +++ b/nemo_gym/sandbox/providers/ecs_fargate/provider.py @@ -0,0 +1,236 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# 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. + +"""ECS Fargate sandbox provider. + +Adapts the lifted :mod:`engine` (one stateful ``EcsFargateSandbox`` per +sandbox) to Gym's stateless ``SandboxProvider`` contract: per-sandbox engine +state lives in ``SandboxHandle.raw``; the provider methods delegate to it. + +The model endpoint is assumed reachable from inside the sandbox directly, or +routed via the SSH reverse tunnel when ``outside_endpoints`` are supplied +through ``spec.provider_options``. See the package README for the planned +Teleport follow-up. +""" + +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path +from typing import Any + +from nemo_gym.sandbox.providers.base import ( + SandboxCreateError, + SandboxExecResult, + SandboxHandle, + SandboxSpec, + SandboxStatus, +) +from nemo_gym.sandbox.providers.ecs_fargate import engine + + +def _outside_endpoints(spec: SandboxSpec) -> list[engine.OutsideEndpoint]: + raw = spec.provider_options.get("outside_endpoints") or [] + endpoints = [] + for item in raw: + if isinstance(item, engine.OutsideEndpoint): + endpoints.append(item) + else: + endpoints.append(engine.OutsideEndpoint(url=item["url"], env_var=item["env_var"])) + return endpoints + + +def _volumes(spec: SandboxSpec) -> list[engine.VolumeMount]: + raw = spec.provider_options.get("volumes") or [] + volumes = [] + for item in raw: + if isinstance(item, engine.VolumeMount): + volumes.append(item) + else: + volumes.append(engine.VolumeMount(**item)) + return volumes + + +def _engine_spec(spec: SandboxSpec) -> engine.SandboxSpec: + if spec.image is None: + raise SandboxCreateError("ECS Fargate sandbox requires SandboxSpec.image") + entrypoint = " ".join(spec.entrypoint) if spec.entrypoint else None + return engine.SandboxSpec( + image=spec.image, + workdir=spec.workdir or "/workspace", + env=dict(spec.env), + files=dict(spec.files), + entrypoint=entrypoint, + volumes=_volumes(spec), + environment_dir=spec.provider_options.get("environment_dir"), + ) + + +class EcsFargateProvider: + """Run sandboxes as AWS ECS Fargate tasks behind an SSH sidecar.""" + + name = "ecs_fargate" + + def __init__(self, **config: Any) -> None: + self._cfg = engine_config_from_mapping(config) + + async def create(self, spec: SandboxSpec) -> SandboxHandle: + cfg = self._cfg + if spec.ready_timeout_s is not None: + cfg = replace(cfg, startup_timeout_sec=float(spec.ready_timeout_s)) + sandbox = engine.EcsFargateSandbox(_engine_spec(spec), ecs_config=cfg) + try: + await sandbox.start(outside_endpoints=_outside_endpoints(spec)) + except Exception as e: # noqa: BLE001 — uniform create failure surface + raise SandboxCreateError(f"ECS Fargate create failed: {e}") from e + return SandboxHandle( + sandbox_id=sandbox.task_arn or "", + provider_name=self.name, + raw=sandbox, + ) + + async def exec( + self, + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | float | None = None, + user: str | int | None = None, + ) -> SandboxExecResult: + sandbox: engine.EcsFargateSandbox = handle.raw + result = await sandbox.exec( + command, + timeout_sec=180 if timeout_s is None else float(timeout_s), + cwd=cwd, + env=env, + user=user, + ) + return SandboxExecResult( + stdout=result.stdout, + stderr=result.stderr, + return_code=result.return_code, + ) + + async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: + sandbox: engine.EcsFargateSandbox = handle.raw + await sandbox.upload(Path(source_path), target_path) + + async def download_file(self, handle: SandboxHandle, source_path: str, target_path: Path) -> None: + sandbox: engine.EcsFargateSandbox = handle.raw + await sandbox.download(source_path, Path(target_path)) + + async def status(self, handle: SandboxHandle) -> SandboxStatus: + sandbox: engine.EcsFargateSandbox = handle.raw + return SandboxStatus.RUNNING if sandbox.is_running else SandboxStatus.STOPPED + + async def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: + sandbox: engine.EcsFargateSandbox = handle.raw + await sandbox.stop() + + async def aclose(self) -> None: + return None + + +def engine_config_from_mapping(config: dict[str, Any]) -> engine.EcsFargateConfig: + """Build an ``EcsFargateConfig`` from provider kwargs. + + When ``region`` is given but ``cluster`` is omitted, infrastructure is + auto-discovered from SSM (written by the reference Terraform); explicit + kwargs always win over SSM. Mirrors NEL's orchestrator resolution. + """ + config = dict(config) + region = config.get("region") + cluster = config.get("cluster") + ssm_project = config.get("ssm_project", engine.DEFAULT_SSM_PROJECT) + + ssm: dict[str, Any] = {} + if region is not None and cluster is None: + ssm = engine.resolve_ecs_config_from_ssm(region, ssm_project) + + def pick(key: str, default: Any = None) -> Any: + val = config.get(key) + if val is not None: + return val + return ssm.get(key, default) + + ssh_sidecar = _sidecar_config(config.get("ssh_sidecar"), ssm.get("ssh_sidecar", {})) + + return engine.EcsFargateConfig( + region=region, + cluster=pick("cluster", ""), + subnets=config.get("subnets") or ssm.get("subnets", []), + security_groups=config.get("security_groups") or ssm.get("security_groups", []), + assign_public_ip=pick("assign_public_ip", True), + task_definition=config.get("task_definition"), + task_definition_family_prefix=config.get("task_definition_family_prefix", "ecs-sandbox"), + image_template=config.get("image_template"), + container_name=config.get("container_name", "main"), + container_port=config.get("container_port"), + cpu=str(config.get("cpu", "4096")), + memory=str(config.get("memory", "8192")), + ephemeral_storage_gib=config.get("ephemeral_storage_gib"), + platform_version=config.get("platform_version"), + execution_role_arn=pick("execution_role_arn"), + task_role_arn=pick("task_role_arn"), + extra_env=config.get("extra_env"), + log_group=pick("log_group"), + log_stream_prefix=config.get("log_stream_prefix"), + max_task_lifetime_sec=config.get("max_task_lifetime_sec") or 14400, + startup_timeout_sec=float(config.get("startup_timeout_sec", 300.0)), + ssh_sidecar=ssh_sidecar, + s3_bucket=pick("s3_bucket"), + s3_prefix=config.get("s3_prefix"), + ecr_repository=pick("ecr_repository"), + environment_dir=config.get("environment_dir"), + codebuild_project=config.get("codebuild_project"), + codebuild_service_role=pick("codebuild_service_role"), + codebuild_compute_type=config.get("codebuild_compute_type") or "BUILD_GENERAL1_MEDIUM", + codebuild_build_timeout=config.get("codebuild_build_timeout") or 60, + dockerhub_secret_arn=pick("dockerhub_secret_arn"), + efs_filesystem_id=pick("efs_filesystem_id"), + efs_access_point_id=pick("efs_access_point_id"), + ssm_project=ssm_project, + ) + + +def _sidecar_config(yaml_sidecar: Any, ssm_ssh: dict[str, Any]) -> engine.SshSidecarConfig | None: + if yaml_sidecar is not None: + sc = dict(yaml_sidecar) if isinstance(yaml_sidecar, dict) else yaml_sidecar + if isinstance(sc, engine.SshSidecarConfig): + return sc + pub = sc.get("public_key_secret_arn") or ssm_ssh.get("public_key_secret_arn", "") + priv = sc.get("private_key_secret_arn") or ssm_ssh.get("private_key_secret_arn", "") + if not pub or not priv: + raise ValueError( + "ssh_sidecar.public_key_secret_arn and ssh_sidecar.private_key_secret_arn " + "are required (set explicitly or auto-discovered from SSM)." + ) + return engine.SshSidecarConfig( + sshd_port=sc.get("sshd_port", engine.DEFAULT_SSHD_PORT), + ssh_ready_timeout_sec=sc.get("ssh_ready_timeout_sec", 300.0), + public_key_secret_arn=pub, + private_key_secret_arn=priv, + image=sc.get("image"), + exec_server_port=sc.get("exec_server_port", engine.DEFAULT_EXEC_SERVER_PORT), + ) + if ssm_ssh.get("public_key_secret_arn") and ssm_ssh.get("private_key_secret_arn"): + return engine.SshSidecarConfig( + sshd_port=ssm_ssh.get("sshd_port", engine.DEFAULT_SSHD_PORT), + public_key_secret_arn=ssm_ssh["public_key_secret_arn"], + private_key_secret_arn=ssm_ssh["private_key_secret_arn"], + exec_server_port=ssm_ssh.get("exec_server_port", engine.DEFAULT_EXEC_SERVER_PORT), + ) + return None diff --git a/nemo_gym/sandbox/providers/registry.py b/nemo_gym/sandbox/providers/registry.py index 12347e2614..2d46a682e7 100644 --- a/nemo_gym/sandbox/providers/registry.py +++ b/nemo_gym/sandbox/providers/registry.py @@ -137,5 +137,12 @@ def _load_apptainer_provider() -> ProviderClass: return ApptainerProvider -_BUILTIN_PROVIDER_LOADERS["apptainer"] = _load_apptainer_provider +def _load_ecs_fargate_provider() -> ProviderClass: + from nemo_gym.sandbox.providers.ecs_fargate import EcsFargateProvider + + return EcsFargateProvider + + _BUILTIN_PROVIDER_LOADERS["opensandbox"] = _load_opensandbox_provider +_BUILTIN_PROVIDER_LOADERS["apptainer"] = _load_apptainer_provider +_BUILTIN_PROVIDER_LOADERS["ecs_fargate"] = _load_ecs_fargate_provider diff --git a/pyproject.toml b/pyproject.toml index 1bcd8bcfa0..1cb3ad4f6d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -250,6 +250,14 @@ sandbox = [ "opensandbox>=0.1.9", ] +sandbox-ecs = [ + # boto3: AWS SDK used by the ECS Fargate sandbox provider for ECS/EC2/ECR/ + # CodeBuild/S3/SSM/Secrets Manager calls. + # Updated: Tue Jun 03, 2026 with boto3>=1.34 + # License: Apache 2.0 https://github.com/boto/boto3/blob/develop/LICENSE + "boto3>=1.34", +] + # We include dev dependencies as an extra since technically each server module is a consumer (which means we cannot use dependency groups, which are intended to be within a project). dev = [ # Pre-commit: Used for pre-commit hooks. diff --git a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml new file mode 100644 index 0000000000..89f8e3876c --- /dev/null +++ b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml @@ -0,0 +1,39 @@ +mini_swe_agent_2: + responses_api_agents: + mini_swe_agent_2: + entrypoint: app.py + domain: coding + description: Software engineering tasks driven by mini-swe-agent harness on ECS Fargate. + value: Improve agentic software engineering capabilities. + model_server: + type: responses_api_models + name: policy_model + concurrency: 8 + env: sandbox + # Region-only config: the provider auto-discovers cluster/subnets/SGs/roles + # and the SSH-sidecar key ARNs from SSM (//ecs-sandbox/config, + # ssm_project defaults to "harbor"), exactly like NEL. + sandbox_provider: + ecs_fargate: + region: ${oc.env:AWS_REGION} + cpu: "2048" + memory: "8192" + ephemeral_storage_gib: 50 + sandbox_spec: + timeout_s: 18000 + ready_timeout_s: 1200 + metadata: + benchmark: swebench-verified + harness: mini-swe-agent + sandbox-api: ecs-fargate + sandbox_environment_kwargs: + cwd: /testbed + conda_env: testbed + activate_conda: true + user: root + delete: true + run_golden: false + step_timeout: 600 + eval_timeout: 1800 + skip_if_exists: false + step_limit: 250 diff --git a/tests/unit_tests/test_ecs_fargate_provider.py b/tests/unit_tests/test_ecs_fargate_provider.py new file mode 100644 index 0000000000..6b78b9b689 --- /dev/null +++ b/tests/unit_tests/test_ecs_fargate_provider.py @@ -0,0 +1,270 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# 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. + +"""ECS Fargate provider tests — all AWS/SSH/network calls mocked.""" + +from __future__ import annotations + +import contextlib +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from nemo_gym.sandbox import AsyncSandbox +from nemo_gym.sandbox.providers import ( + SandboxSpec, + SandboxStatus, + create_provider, + get_provider_class, + list_providers, +) +from nemo_gym.sandbox.providers.ecs_fargate import EcsFargateProvider, engine +from nemo_gym.sandbox.providers.ecs_fargate.provider import engine_config_from_mapping + + +_ENG = "nemo_gym.sandbox.providers.ecs_fargate.engine" + + +def _provider_config(**overrides): + cfg = dict( + region="us-west-2", + cluster="test-cluster", + subnets=["subnet-aaa"], + security_groups=["sg-bbb"], + assign_public_ip=True, + execution_role_arn="arn:aws:iam::1234:role/ecsTaskExec", + task_role_arn="arn:aws:iam::1234:role/ecsTask", + ssh_sidecar={ + "sshd_port": 2222, + "public_key_secret_arn": "arn:aws:secretsmanager:us-east-1:1234:secret:pub", + "private_key_secret_arn": "arn:aws:secretsmanager:us-east-1:1234:secret:priv", + "exec_server_port": 5000, + }, + ) + cfg.update(overrides) + return {"ecs_fargate": cfg} + + +@contextlib.contextmanager +def _mock_engine_start(exec_result=None): + """Patch every AWS/SSH seam so ``EcsFargateSandbox.start`` runs offline. + + Yields the fake exec client so delegation can be asserted. + """ + tunnel = MagicMock() + tunnel.local_port = 19000 + exec_client = MagicMock() + exec_client.exec = AsyncMock(return_value=exec_result or engine.ExecResult("out", "err", 0)) + exec_client.upload = AsyncMock() + exec_client.download = AsyncMock(return_value=b"payload") + exec_client.close = AsyncMock() + + with ( + patch.object(engine.EcsFargateSandbox, "_init_aws_clients"), + patch.object(engine.EcsFargateSandbox, "_resolve_image", return_value="python:3.12"), + patch.object(engine.EcsFargateSandbox, "_register_task_definition", return_value="task-def-arn"), + patch.object(engine.EcsFargateSandbox, "_run_task", return_value="task-arn"), + patch.object(engine.EcsFargateSandbox, "_register_for_cleanup"), + patch.object(engine.EcsFargateSandbox, "_wait_for_running"), + patch.object(engine.EcsFargateSandbox, "_get_task_public_ip", return_value="10.0.0.10"), + patch.object(engine.EcsFargateSandbox, "_wait_for_ssh_ready"), + patch(f"{_ENG}.download_secret_to_file", return_value="/tmp/key"), + patch(f"{_ENG}.download_secret_to_string", return_value="ssh-rsa fake"), + patch(f"{_ENG}.build_ssh_sidecar_container", return_value={"name": "ssh-tunnel"}), + patch(f"{_ENG}._free_port", return_value=19001), + patch(f"{_ENG}.SshTunnel", return_value=tunnel) as tunnel_cls, + patch(f"{_ENG}.ExecClient", return_value=exec_client), + ): + yield exec_client, tunnel_cls + + +# ── Registration ────────────────────────────────────────────────────── + + +def test_provider_registered(): + assert "ecs_fargate" in list_providers() + assert get_provider_class("ecs_fargate") is EcsFargateProvider + assert EcsFargateProvider.name == "ecs_fargate" + + +# ── Config resolution ───────────────────────────────────────────────── + + +def test_config_explicit_values_no_ssm(): + p = create_provider(_provider_config()) + cfg = p._cfg + assert cfg.cluster == "test-cluster" + assert cfg.subnets == ["subnet-aaa"] + assert cfg.assign_public_ip is True + assert cfg.ssh_sidecar.exec_server_port == 5000 + assert cfg.ssh_sidecar.sshd_port == 2222 + + +def test_config_ssm_autodiscovery_merges_and_yaml_wins(): + ssm_blob = { + "cluster": "ssm-cluster", + "subnets": ["subnet-ssm"], + "security_groups": ["sg-ssm"], + "execution_role_arn": "arn:ssm:exec", + "ssh_sidecar": { + "public_key_secret_arn": "arn:ssm:pub", + "private_key_secret_arn": "arn:ssm:priv", + "sshd_port": 52222, + }, + } + with patch(f"{_ENG}.resolve_ecs_config_from_ssm", return_value=ssm_blob) as resolve: + # region set, cluster omitted -> SSM is consulted. + p = create_provider({"ecs_fargate": {"region": "us-west-2", "subnets": ["subnet-override"]}}) + resolve.assert_called_once_with("us-west-2", "harbor") + cfg = p._cfg + assert cfg.cluster == "ssm-cluster" # filled from SSM + assert cfg.subnets == ["subnet-override"] # explicit YAML wins + assert cfg.execution_role_arn == "arn:ssm:exec" + assert cfg.ssh_sidecar.public_key_secret_arn == "arn:ssm:pub" + assert cfg.ssh_sidecar.sshd_port == 52222 + + +def test_config_no_ssm_when_cluster_present(): + with patch(f"{_ENG}.resolve_ecs_config_from_ssm") as resolve: + create_provider(_provider_config()) + resolve.assert_not_called() + + +def test_sidecar_missing_key_arns_raises(): + with pytest.raises(ValueError, match="public_key_secret_arn"): + engine_config_from_mapping({"cluster": "c", "ssh_sidecar": {"sshd_port": 2222}}) + + +# ── create / lifecycle delegation ───────────────────────────────────── + + +async def test_create_returns_handle_with_running_sandbox(): + provider = create_provider(_provider_config()) + spec = SandboxSpec(image="python:3.12") + with _mock_engine_start(): + handle = await provider.create(spec) + assert handle.provider_name == "ecs_fargate" + assert handle.sandbox_id == "task-arn" + assert isinstance(handle.raw, engine.EcsFargateSandbox) + assert await provider.status(handle) == SandboxStatus.RUNNING + await provider.close(handle) + assert await provider.status(handle) == SandboxStatus.STOPPED + + +async def test_exec_maps_engine_result(): + provider = create_provider(_provider_config()) + spec = SandboxSpec(image="python:3.12", workdir="/work") + with _mock_engine_start(exec_result=engine.ExecResult("hello\n", "", 0)) as (exec_client, _): + handle = await provider.create(spec) + result = await provider.exec(handle, "echo hello", timeout_s=42) + assert (result.stdout, result.stderr, result.return_code) == ("hello\n", "", 0) + # engine.exec received the gym timeout as timeout_sec + _, kwargs = exec_client.exec.call_args + assert kwargs["timeout"] == 42 + + +async def test_upload_and_download_delegate(tmp_path): + provider = create_provider(_provider_config()) + spec = SandboxSpec(image="python:3.12") + src = tmp_path / "in.txt" + src.write_text("data") + dest = tmp_path / "out.txt" + with _mock_engine_start() as (exec_client, _): + handle = await provider.create(spec) + await provider.upload_file(handle, src, "/remote/in.txt") + await provider.download_file(handle, "/remote/out.txt", dest) + exec_client.upload.assert_awaited_once() + assert dest.read_bytes() == b"payload" + + +async def test_outside_endpoints_build_reverse_tunnel(): + provider = create_provider(_provider_config()) + spec = SandboxSpec( + image="python:3.12", + provider_options={"outside_endpoints": [{"url": "http://127.0.0.1:4000/v1", "env_var": "MODEL_BASE_URL"}]}, + ) + with _mock_engine_start() as (_, tunnel_cls): + handle = await provider.create(spec) + # exec-server mode opens a forward tunnel to the exec server + a reverse + # tunnel for each outside endpoint. + _, kwargs = tunnel_cls.call_args + assert kwargs["forward_port"] == 5000 + # reverse spec format: "::" + assert any(s.endswith(":127.0.0.1:4000") for s in kwargs["reverses"]) + # the resolved endpoint is injected into the container env + routing = handle.raw._outside_endpoint_routing + assert routing.resolved_endpoint_url("MODEL_BASE_URL").startswith("http://127.0.0.1:") + + +async def test_create_missing_image_raises(): + provider = create_provider(_provider_config()) + from nemo_gym.sandbox.providers import SandboxCreateError + + with pytest.raises(SandboxCreateError, match="requires SandboxSpec.image"): + await provider.create(SandboxSpec(image=None)) + + +# ── Public AsyncSandbox surface ─────────────────────────────────────── + + +async def test_async_sandbox_end_to_end(): + spec = SandboxSpec(image="python:3.12", files={"/app/run.sh": "echo hi"}) + with _mock_engine_start() as (exec_client, _): + sb = AsyncSandbox(_provider_config(), spec) + await sb.start() + # initial files uploaded via the provider during start() + exec_client.upload.assert_awaited() + res = await sb.exec("ls") + assert res.return_code == 0 + assert await sb.status() == SandboxStatus.RUNNING + await sb.stop() + assert await sb.status() == SandboxStatus.STOPPED + + +# ── Pure engine helpers (no AWS) ────────────────────────────────────── + + +def test_task_def_hash_ignores_log_config(): + base = { + "family": "f1", + "containerDefinitions": [ + {"name": "main", "image": "x", "logConfiguration": {"options": {"awslogs-group": "g1"}}} + ], + } + other = { + "family": "f2", + "containerDefinitions": [ + {"name": "main", "image": "x", "logConfiguration": {"options": {"awslogs-group": "g2"}}} + ], + } + assert engine._compute_task_def_hash(base) == engine._compute_task_def_hash(other) + + +def test_generate_buildspec_pushes_to_ecr(): + cfg = engine.EcsFargateConfig( + region="us-west-2", + ecr_repository="123.dkr.ecr.us-west-2.amazonaws.com/sandbox", + ) + spec = engine.ImageBuilder._generate_buildspec(cfg, "sandbox", "tag1", f"{cfg.ecr_repository}:tag1") + assert "docker build -t sandbox:tag1" in spec + assert f"docker push {cfg.ecr_repository}:tag1" in spec + + +def test_get_ecr_image_tag_is_content_addressed(tmp_path): + (tmp_path / "Dockerfile").write_text("FROM scratch") + tag1 = engine.ImageBuilder.get_ecr_image_tag(tmp_path, "env") + tag2 = engine.ImageBuilder.get_ecr_image_tag(tmp_path, "env") + assert tag1 == tag2 and tag1.startswith("env__") + (tmp_path / "Dockerfile").write_text("FROM alpine") + assert engine.ImageBuilder.get_ecr_image_tag(tmp_path, "env") != tag1 From 21d8aa205a78eba806c3bb5f0c893e10da187aa2 Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Mon, 8 Jun 2026 11:51:05 +0200 Subject: [PATCH 02/11] feat(sandbox): auto-mirror ECS images on demand + fix conda activation - ECS Fargate provider mirrors a missing public image into the ECR mirror on demand during create (ImageBuilder.ensure_mirrored, gated by the new auto_mirror config, default on), with the same in-flight dedup as the environment_dir build path. Callers no longer pre-stage images. - Fix MiniSWESandboxEnvironment conda activation: source conda from known install roots instead of `conda info --base`, which fails in the ECS exec server's non-login shell and was silently dropping the golden gold-patch apply (zeroing reward). - Add the sandbox-ecs extra to mini_swe_agent_2 requirements so the agent's per-server venv can run the ECS provider under ng_run. - Docs: ECS Fargate provider README (general, provider-scoped) + a concise SWE-bench-on-ECS run guide in the agent README. - Tests for mirror resolution/dedup/auto-mirror hook and conda command build. Signed-off-by: Michal Bien --- .../sandbox/providers/ecs_fargate/README.md | 79 +++++++++++ .../sandbox/providers/ecs_fargate/engine.py | 134 +++++++++++++++++- .../sandbox/providers/ecs_fargate/provider.py | 1 + .../mini_swe_agent_2/README.md | 51 +++++++ .../mini_swe_agent_2/requirements.txt | 2 +- .../mini_swe_agent_2/sandbox_environment.py | 12 +- .../tests/test_sandbox_environment.py | 36 ++++- tests/unit_tests/test_ecs_fargate_provider.py | 102 +++++++++++++ 8 files changed, 408 insertions(+), 9 deletions(-) create mode 100644 nemo_gym/sandbox/providers/ecs_fargate/README.md diff --git a/nemo_gym/sandbox/providers/ecs_fargate/README.md b/nemo_gym/sandbox/providers/ecs_fargate/README.md new file mode 100644 index 0000000000..5010060028 --- /dev/null +++ b/nemo_gym/sandbox/providers/ecs_fargate/README.md @@ -0,0 +1,79 @@ +# ECS Fargate sandbox provider + +Runs each `nemo_gym.sandbox` sandbox as an AWS ECS Fargate task behind an SSH +sidecar. It implements the provider-neutral `SandboxProvider` contract, so any +sandbox-backed agent (not just mini-swe-agent) can use it by setting +`sandbox_provider.ecs_fargate` in its config. + +## Prerequisites + +- **Infrastructure** provisioned in the target account/region. The reference + Terraform stack publishes its outputs to SSM at + `//ecs-sandbox/config` (`ssm_project` defaults to `harbor`): + cluster, subnets, security groups, task/execution roles, ECR mirror, EFS, and + the SSH-sidecar key ARNs. A missing parameter raises an actionable error. +- **Credentials**: `AWS_ACCESS_KEY_ID`/`AWS_SECRET_ACCESS_KEY` (or an instance + role) plus `AWS_REGION`. +- **Network**: the host must reach each task's SSH sidecar port (`52222`) — run + inside the sandbox VPC/peered network, or allow the host IP on the sidecar + security group. Exec, file transfer, and model reverse-tunnels ride this SSH + connection. + +## Configuration + +Region-only is enough; everything else auto-discovers from SSM (explicit YAML +always wins): + +```yaml +sandbox_provider: + ecs_fargate: + region: ${oc.env:AWS_REGION} + cpu: "2048" + memory: "8192" + ephemeral_storage_gib: 50 +``` + +Common fields: + +| Field | Default | Purpose | +| --- | --- | --- | +| `region` | — | AWS region; enables SSM auto-discovery when `cluster` is omitted | +| `cpu` / `memory` | `"4096"` / `"8192"` | Fargate task size (CPU units / MiB) | +| `ephemeral_storage_gib` | task default | Task ephemeral disk | +| `auto_mirror` | `true` | Pull a missing public image into the ECR mirror on demand (see below) | +| `ssm_project` | `harbor` | SSM namespace for auto-discovery | +| `environment_dir` | — | Build a task image from a Dockerfile dir via CodeBuild instead of using a prebuilt image | + +Per-sandbox `ready_timeout_s`, `env`, `files`, `metadata`, and +`provider_options` (e.g. `platform`, `outside_endpoints`) come from the +`SandboxSpec`. + +## Images and on-demand mirroring + +ECS pulls task images from the account ECR mirror, not their origin registry. A +bare/public image (e.g. `docker.io/swebench/sweb.eval.x86_64.:latest`) +resolves to the mirror tag `:`. Resolution order: + +1. `environment_dir` set → build the image via CodeBuild and use it. +2. Image is already an ECR reference → use verbatim (never re-mirrored). +3. Bare/public name + `auto_mirror=true` → mirror into ECR on demand (CodeBuild + pull → retag → push) during `create`, then launch. Concurrent tasks for the + same image de-duplicate onto one build. + +Set `auto_mirror: false` to require a pre-populated mirror and fail fast on a +miss. The first task for a new image waits on a one-time build (~1–3 min for +typical SWE-bench images); later tasks hit the ECR cache. + +## Lifecycle + +`create` launches the task + SSH sidecar and returns once the exec server is +healthy. `exec`, `upload_file`, `download_file`, `status`, and `close` delegate +to the per-sandbox engine over the SSH tunnel. `outside_endpoints` (via +`spec.provider_options`) open reverse tunnels so an in-sandbox process can reach +a host-side endpoint (e.g. a model server). + +## Security + +The sidecar security group in the reference stack allows `0.0.0.0/0` on `52222` +for convenience. Restrict it to the orchestrator's egress IP (or move to a +private/Teleport path) before non-smoke use. diff --git a/nemo_gym/sandbox/providers/ecs_fargate/engine.py b/nemo_gym/sandbox/providers/ecs_fargate/engine.py index f59dc68e1e..49d19609c1 100644 --- a/nemo_gym/sandbox/providers/ecs_fargate/engine.py +++ b/nemo_gym/sandbox/providers/ecs_fargate/engine.py @@ -195,6 +195,19 @@ def _sanitize_id(value: str, max_len: int = 100) -> str: return cleaned[:max_len] or "task" +_ECR_IMAGE_REF_RE = re.compile(r"^[0-9]{12}\.dkr\.ecr\.[a-z0-9-]+\.amazonaws\.com/", re.IGNORECASE) + + +def _is_ecr_image_ref(image: str) -> bool: + """True for references already pointing at an ECR registry host. + + Such references (e.g. an image resolved against the configured ECR mirror) + must be used as-is; reapplying the ``ecr_repository`` + sanitize rewrite + would corrupt the tag. + """ + return bool(_ECR_IMAGE_REF_RE.match(image)) + + @dataclass(frozen=True) class SshSidecarConfig: """SSH sidecar container configuration. @@ -249,6 +262,12 @@ class EcsFargateConfig: codebuild_build_timeout: int = 60 dockerhub_secret_arn: str | None = None build_parallelism: int = 50 + # When a public/bare image is routed to the ECR mirror and is not yet + # present, pull it into ECR on demand (via CodeBuild) during create. The + # mirror normally runs ahead of time, but this keeps create self-healing so + # callers never have to pre-stage images manually. Set False to require a + # pre-populated mirror and fail fast on a miss. + auto_mirror: bool = True efs_filesystem_id: str | None = None efs_access_point_id: str | None = None ssm_project: str = DEFAULT_SSM_PROJECT @@ -1108,6 +1127,97 @@ def ensure_image_built(cls, *, cfg: EcsFargateConfig, environment_name: str, for cls._inflight_builds.pop(tag, None) return image_url + @classmethod + def ensure_mirrored(cls, *, cfg: EcsFargateConfig, src_image: str, force: bool = False) -> str: + """Ensure ``src_image`` is present in the ECR mirror, pulling it if not. + + Public/bare image references are served from the ECR mirror tag + (``{ecr_repository}:{sanitize(src_image)}``) rather than pulled from + their origin registry at task-launch time. This mirrors the image into + ECR on demand via a self-contained CodeBuild job (privileged DinD that + logs into Docker Hub + ECR, pulls, retags, and pushes), so callers never + have to pre-stage images. Concurrent callers for the same tag dedup on a + shared in-flight event, exactly like :meth:`ensure_image_built`. + """ + ecr_repo = cfg.ecr_repository + if not ecr_repo: + raise ValueError("ecr_repository is required to mirror images") + tag = _sanitize_id(src_image) + image_url = f"{ecr_repo}:{tag}" + + with cls._lock: + if tag in cls._inflight_builds: + cls._inflight_builds[tag].wait() + return image_url + if not force and cls.image_exists_in_ecr(ecr_repo, tag, cfg.region): + logger.info("ECR mirror hit — skipping pull: %s", image_url) + return image_url + + event = threading.Event() + with cls._lock: + if tag in cls._inflight_builds: + cls._inflight_builds[tag].wait() + return image_url + cls._inflight_builds[tag] = event + with cls._lock: + if cls._build_semaphore is None or cls._build_semaphore_size != cfg.build_parallelism: + cls._build_semaphore = threading.Semaphore(cfg.build_parallelism) + cls._build_semaphore_size = cfg.build_parallelism + try: + cls._build_semaphore.acquire() # type: ignore[union-attr] + try: + if not force and cls.image_exists_in_ecr(ecr_repo, tag, cfg.region): + return image_url + logger.info("Mirroring %s -> %s via CodeBuild ...", src_image, image_url) + buildspec = cls._generate_mirror_buildspec(cfg, src_image, image_url) + cls.run_buildspec_via_codebuild( + cfg=cfg, + buildspec=buildspec, + job_label=f"mirror::{tag}", + timeout_minutes=cfg.codebuild_build_timeout, + ) + logger.info("Mirrored OK: %s -> %s", src_image, image_url) + finally: + cls._build_semaphore.release() # type: ignore[union-attr] + finally: + event.set() + with cls._lock: + cls._inflight_builds.pop(tag, None) + return image_url + + @staticmethod + def _generate_mirror_buildspec(cfg: EcsFargateConfig, src_image: str, ecr_url: str) -> str: + ecr_registry = (cfg.ecr_repository or "").split("/")[0] + ecr_region = ImageBuilder._ecr_region(cfg.ecr_repository or "", fallback="$AWS_DEFAULT_REGION") + pre_build_cmds = [ + f"aws ecr get-login-password --region {ecr_region}" + f" | docker login --username AWS --password-stdin {ecr_registry}", + ] + if cfg.dockerhub_secret_arn: + pre_build_cmds.append( + f"DOCKERHUB_CREDS=$(aws secretsmanager get-secret-value" + f" --secret-id {cfg.dockerhub_secret_arn}" + f" --query SecretString --output text --region $AWS_DEFAULT_REGION)" + f' && DH_USER=$(echo "$DOCKERHUB_CREDS" | python3 -c' + """ "import sys,json;print(json.load(sys.stdin)['username'])")""" + f' && if [ -n "$DH_USER" ]; then echo "$DOCKERHUB_CREDS" | python3 -c' + """ "import sys,json;print(json.load(sys.stdin)['password'])" """ + f'| docker login -u "$DH_USER" --password-stdin; fi' + f' || echo "Docker Hub login failed — continuing without auth"' + ) + pull_cmd = ( + f"for i in 1 2 3; do docker pull --platform linux/amd64 {src_image} && break; " + f'echo "pull failed ($i/3), retry in 30s"; sleep 30; done' + ) + pre_yaml = "\n".join(f" - {c}" for c in pre_build_cmds) + return ( + "version: 0.2\nphases:\n pre_build:\n commands:\n" + f"{pre_yaml}\n build:\n commands:\n" + f" - {pull_cmd}\n" + f" - docker tag {src_image} {ecr_url}\n" + f" post_build:\n commands:\n - docker push {ecr_url}\n" + ) + @staticmethod def _upload_build_context(cfg: EcsFargateConfig, environment_name: str, nonce: str) -> str: boto3, *_ = _require_aws_sdks() @@ -1528,6 +1638,17 @@ def _do_start(self) -> None: built_image = ImageBuilder.ensure_image_built( cfg=per_task_cfg, environment_name=_sanitize_id(self._spec.image or "sandbox") ) + elif ( + cfg.auto_mirror + and cfg.ecr_repository + and self._spec.image + and not cfg.image_template + and not _is_ecr_image_ref(self._spec.image) + ): + # The bare/public image is served from the ECR mirror tag; pull it + # into ECR on demand so a missing mirror entry self-heals instead of + # failing the task's image pull. + ImageBuilder.ensure_mirrored(cfg=cfg, src_image=self._spec.image) image = self._resolve_image(built_image) if not sidecar.private_key_secret_arn or not sidecar.public_key_secret_arn: @@ -1599,9 +1720,18 @@ def _resolve_image(self, built_image: str | None = None) -> str: f"metadata) instead of ecs.image_template for task-specific " f"placeholders like {{task_id}}." ) from exc - if cfg.ecr_repository and self._spec.image: - return f"{cfg.ecr_repository}:{_sanitize_id(self._spec.image)}" if self._spec.image: + # A reference that already points at an ECR registry (e.g. the + # configured mirror) is used verbatim — re-prefixing it under + # ecr_repository and re-sanitizing would corrupt the tag and make + # an existing image unresolvable. + if _is_ecr_image_ref(self._spec.image): + return self._spec.image + # Bare / public names are routed to the ECR mirror tag rather than + # pulled directly from their origin registry (avoids Docker Hub + # rate limits when many tasks start concurrently). + if cfg.ecr_repository: + return f"{cfg.ecr_repository}:{_sanitize_id(self._spec.image)}" return self._spec.image if not cfg.task_definition: raise ValueError( diff --git a/nemo_gym/sandbox/providers/ecs_fargate/provider.py b/nemo_gym/sandbox/providers/ecs_fargate/provider.py index 08af81d396..6c4013702f 100644 --- a/nemo_gym/sandbox/providers/ecs_fargate/provider.py +++ b/nemo_gym/sandbox/providers/ecs_fargate/provider.py @@ -199,6 +199,7 @@ def pick(key: str, default: Any = None) -> Any: codebuild_service_role=pick("codebuild_service_role"), codebuild_compute_type=config.get("codebuild_compute_type") or "BUILD_GENERAL1_MEDIUM", codebuild_build_timeout=config.get("codebuild_build_timeout") or 60, + auto_mirror=config.get("auto_mirror", True), dockerhub_secret_arn=pick("dockerhub_secret_arn"), efs_filesystem_id=pick("efs_filesystem_id"), efs_access_point_id=pick("efs_access_point_id"), diff --git a/responses_api_agents/mini_swe_agent_2/README.md b/responses_api_agents/mini_swe_agent_2/README.md index ac285143ac..ac64da39c5 100644 --- a/responses_api_agents/mini_swe_agent_2/README.md +++ b/responses_api_agents/mini_swe_agent_2/README.md @@ -24,6 +24,7 @@ over the older Docker/Singularity mini-SWE integration. - [Run One-Example Smoke](#run-one-example-smoke) - [Expected Outputs](#expected-outputs) - [Repeated Rollouts](#repeated-rollouts) + - [Running SWE-bench on ECS Fargate](#running-swe-bench-on-ecs-fargate) - [Sandbox Environment Adapter](#sandbox-environment-adapter) - [Environment Lifecycle](#environment-lifecycle) - [Contributing](#contributing) @@ -484,6 +485,56 @@ Use the agent's `step_timeout` and `eval_timeout` config values or CLI overrides to bound tool and verification execution. If you launch from a custom Kubernetes wrapper, add any outer per-sample guard there. +## Running SWE-bench on ECS Fargate + +Swap OpenSandbox for ECS Fargate by using +`configs/mini_swe_agent_ecs_fargate.yaml` — the agent loop and SWE-bench +verifier are unchanged. For provider setup (AWS infra/SSM, credentials, the +`:52222` network requirement, and automatic image mirroring via `auto_mirror`) +see the [ECS Fargate provider README](../../nemo_gym/sandbox/providers/ecs_fargate/README.md). +SWE-bench task images are pulled into the ECR mirror on first use, so no manual +image staging is needed. + +Each input row needs the SWE-bench instance fields shown under +[Dataset Information](#dataset-information) plus `subset` (`verified`), `split` +(`test`), and `responses_create_params`. + +**Golden smoke (no model)** — `run_golden=true` applies the gold patch and runs +the verifier in-sandbox, so the model is never called: + +```bash +CONFIG_PATHS="responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml,responses_api_models/vllm_model/configs/vllm_model.yaml" + +AWS_REGION=us-east-1 ng_run "+config_paths=[$CONFIG_PATHS]" \ + ++mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.run_golden=true + +ng_collect_rollouts +agent_name=mini_swe_agent_2 \ + +input_jsonl_fpath=data/swe_verified_smoke.jsonl \ + +output_jsonl_fpath=results/ecs_golden.jsonl \ + +limit=1 +num_repeats=1 +num_samples_in_parallel=1 +``` + +A resolved instance returns reward `1.0` with `tests_status` populated. + +**Real rollout** — drop `run_golden` and point the model server at a live +OpenAI-compatible endpoint (`policy_model_name` is the model id sent upstream): + +```bash +AWS_REGION=us-east-1 ng_run "+config_paths=[$CONFIG_PATHS]" \ + ++policy_base_url=https:///v1 ++policy_api_key= ++policy_model_name= + +ng_collect_rollouts +agent_name=mini_swe_agent_2 \ + +input_jsonl_fpath=data/swe_verified_smoke.jsonl \ + +output_jsonl_fpath=results/ecs_rollout.jsonl \ + +limit=8 +num_repeats=1 +num_samples_in_parallel=8 \ + '+responses_create_params={max_output_tokens: 16384, temperature: 0.6, top_p: 0.95}' +``` + +Reasoning models often only accept the default `temperature` (`1`) and reject a +custom `top_p`; in that case use +`'+responses_create_params={temperature: 1, max_output_tokens: 16384}'` and keep +`max_output_tokens` large enough for reasoning tokens. + ## Sandbox Environment Adapter `MiniSWESandboxEnvironment` adapts mini-swe-agent's synchronous environment diff --git a/responses_api_agents/mini_swe_agent_2/requirements.txt b/responses_api_agents/mini_swe_agent_2/requirements.txt index 2f82dbd4c7..22fb4a3220 100644 --- a/responses_api_agents/mini_swe_agent_2/requirements.txt +++ b/responses_api_agents/mini_swe_agent_2/requirements.txt @@ -1,3 +1,3 @@ --e nemo-gym[dev,sandbox] @ ../../ +-e nemo-gym[dev,sandbox,sandbox-ecs] @ ../../ mini-swe-agent==2.1.0 swebench==4.1.0 diff --git a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py index 704122b9ce..e18501eb32 100644 --- a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py +++ b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py @@ -124,7 +124,17 @@ def _command(self, command: str) -> str: if not self.config.activate_conda or not self.config.conda_env: return command quoted_env = shlex.quote(self.config.conda_env) - return f"source $(conda info --base)/etc/profile.d/conda.sh && conda activate {quoted_env} && {command}" + # Resolve conda from common install roots before activating: the ECS exec shell + # is non-login, so `conda` isn't on PATH and `conda info --base` can't be relied + # on. The grouped loop (not an `&&` chain) keeps a missing root from aborting the + # command; cwd is handled by exec(cwd=...), so we don't `cd` here. + return ( + '{ for __base in /opt/miniconda3 /opt/conda "$HOME/miniconda3" ' + '"$(command -v conda >/dev/null 2>&1 && conda info --base 2>/dev/null)"; do ' + '[ -n "$__base" ] && [ -f "$__base/etc/profile.d/conda.sh" ] && ' + '. "$__base/etc/profile.d/conda.sh" && break; done; } && ' + f"conda activate {quoted_env} && {command}" + ) def execute( self, diff --git a/responses_api_agents/mini_swe_agent_2/tests/test_sandbox_environment.py b/responses_api_agents/mini_swe_agent_2/tests/test_sandbox_environment.py index 0419e15b85..c4ba516003 100644 --- a/responses_api_agents/mini_swe_agent_2/tests/test_sandbox_environment.py +++ b/responses_api_agents/mini_swe_agent_2/tests/test_sandbox_environment.py @@ -57,7 +57,34 @@ def test_check_finished_ignores_nonzero_submit_sentinel() -> None: ) -def test_execute_passes_configured_cwd_without_conda_cd() -> None: +def _env_with(activate_conda: bool, conda_env): + env = MiniSWESandboxEnvironment.__new__(MiniSWESandboxEnvironment) + env.config = MiniSWESandboxEnvironmentConfig.__new__(MiniSWESandboxEnvironmentConfig) + env.config.activate_conda = activate_conda + env.config.conda_env = conda_env + return env + + +def test_command_passthrough_when_conda_disabled() -> None: + env = _env_with(activate_conda=False, conda_env="testbed") + assert env._command("git apply patch.diff") == "git apply patch.diff" + + +def test_command_resolves_conda_without_relying_on_path() -> None: + # Non-login ECS exec shell: conda isn't on PATH, so source conda.sh from known roots + # via a grouped loop (not an `&&` chain that aborts on a missing root). cwd is passed + # to exec(cwd=...), so the command itself must not `cd`. + env = _env_with(activate_conda=True, conda_env="testbed") + wrapped = env._command("git apply patch.diff") + + assert "for __base in" in wrapped + assert "/opt/miniconda3" in wrapped # one of the search roots + assert "conda activate testbed && git apply patch.diff" in wrapped + assert wrapped.count("&&") >= 2 + assert not wrapped.startswith("cd ") # cwd handled by exec(cwd=...), not via cd + + +def test_execute_passes_configured_cwd_to_exec() -> None: class FakeSandbox: def __init__(self) -> None: self.calls: list[dict[str, Any]] = [] @@ -83,8 +110,7 @@ def exec(self, command: str, **kwargs: Any): env.config.activate_conda = True env.config.conda_env = "testbed" env.execute("python -V", cwd="/repo") - assert fake_sandbox.calls[-1]["command"] == ( - "source $(conda info --base)/etc/profile.d/conda.sh && conda activate testbed && python -V" - ) - assert "cd /repo" not in fake_sandbox.calls[-1]["command"] + cmd = fake_sandbox.calls[-1]["command"] + assert "for __base in" in cmd and "conda activate testbed && python -V" in cmd + assert "cd /repo" not in cmd assert fake_sandbox.calls[-1]["cwd"] == "/repo" diff --git a/tests/unit_tests/test_ecs_fargate_provider.py b/tests/unit_tests/test_ecs_fargate_provider.py index 6b78b9b689..fe73607d5e 100644 --- a/tests/unit_tests/test_ecs_fargate_provider.py +++ b/tests/unit_tests/test_ecs_fargate_provider.py @@ -261,6 +261,108 @@ def test_generate_buildspec_pushes_to_ecr(): assert f"docker push {cfg.ecr_repository}:tag1" in spec +def _resolve_image_for(image, **cfg_overrides): + cfg = engine.EcsFargateConfig(region="us-east-1", **cfg_overrides) + sandbox = engine.EcsFargateSandbox(engine.SandboxSpec(image=image), ecs_config=cfg) + return sandbox._resolve_image() + + +def test_resolve_image_routes_bare_name_to_ecr_mirror(): + # Bare/public names are mirrored to the ECR tag, never pulled directly. + ecr = "463701203462.dkr.ecr.us-east-1.amazonaws.com/harbor-us-east-1" + resolved = _resolve_image_for("docker.io/swebench/sweb.eval.x86_64.astropy_1776_astropy-12907:latest", ecr_repository=ecr) + assert resolved == f"{ecr}:{engine._sanitize_id('docker.io/swebench/sweb.eval.x86_64.astropy_1776_astropy-12907:latest')}" + assert "docker.io" not in resolved.split(":", 1)[1] # origin registry not used for the pull + + +def test_resolve_image_passes_through_existing_ecr_ref(): + # A reference already in the ECR mirror is used verbatim (tag preserved, + # including the double underscore the sanitizer would otherwise collapse). + ecr = "463701203462.dkr.ecr.us-east-1.amazonaws.com/harbor-us-east-1" + existing = f"{ecr}:nel-harbor-tasks-swe-bench-astropy-1-1ccf0d50cb33__1ccf0d50" + assert _resolve_image_for(existing, ecr_repository=ecr) == existing + + +def test_resolve_image_template_takes_precedence(): + ecr = "463701203462.dkr.ecr.us-east-1.amazonaws.com/harbor-us-east-1" + resolved = _resolve_image_for( + "anything", ecr_repository=ecr, image_template="{task_id}-built" + ) + assert resolved == "anything-built" + + +def test_is_ecr_image_ref_matches_only_ecr_hosts(): + assert engine._is_ecr_image_ref("463701203462.dkr.ecr.us-east-1.amazonaws.com/repo:tag") + assert not engine._is_ecr_image_ref("docker.io/swebench/sweb.eval:latest") + assert not engine._is_ecr_image_ref("ubuntu:24.04") + + +def test_generate_mirror_buildspec_pulls_tags_and_pushes(): + cfg = engine.EcsFargateConfig( + region="us-east-1", + ecr_repository="123.dkr.ecr.us-east-1.amazonaws.com/mirror", + dockerhub_secret_arn="arn:aws:secretsmanager:us-east-1:123:secret:dh", + ) + src = "docker.io/swebench/sweb.eval.x86_64.astropy_1776_astropy-12907:latest" + ecr_url = f"{cfg.ecr_repository}:{engine._sanitize_id(src)}" + spec = engine.ImageBuilder._generate_mirror_buildspec(cfg, src, ecr_url) + assert f"docker pull --platform linux/amd64 {src}" in spec + assert f"docker tag {src} {ecr_url}" in spec + assert f"docker push {ecr_url}" in spec + assert "get-login-password" in spec # ECR login + assert "secretsmanager get-secret-value" in spec # Docker Hub login + + +def test_ensure_mirrored_skips_when_already_present(): + cfg = engine.EcsFargateConfig(region="us-east-1", ecr_repository="123.dkr.ecr.us-east-1.amazonaws.com/mirror") + with ( + patch.object(engine.ImageBuilder, "image_exists_in_ecr", return_value=True), + patch.object(engine.ImageBuilder, "run_buildspec_via_codebuild") as cb, + ): + url = engine.ImageBuilder.ensure_mirrored(cfg=cfg, src_image="ubuntu:24.04") + cb.assert_not_called() + assert url == f"{cfg.ecr_repository}:{engine._sanitize_id('ubuntu:24.04')}" + + +def test_ensure_mirrored_runs_codebuild_when_missing(): + cfg = engine.EcsFargateConfig(region="us-east-1", ecr_repository="123.dkr.ecr.us-east-1.amazonaws.com/mirror") + with ( + patch.object(engine.ImageBuilder, "image_exists_in_ecr", return_value=False), + patch.object(engine.ImageBuilder, "run_buildspec_via_codebuild") as cb, + ): + url = engine.ImageBuilder.ensure_mirrored(cfg=cfg, src_image="ubuntu:24.04") + cb.assert_called_once() + assert url.endswith(engine._sanitize_id("ubuntu:24.04")) + + +async def test_create_auto_mirrors_missing_public_image(): + ecr = "123.dkr.ecr.us-east-1.amazonaws.com/mirror" + provider = create_provider(_provider_config(ecr_repository=ecr)) + spec = SandboxSpec(image="docker.io/swebench/sweb.eval:latest") + with _mock_engine_start(), patch.object(engine.ImageBuilder, "ensure_mirrored") as m: + await provider.create(spec) + m.assert_called_once() + assert m.call_args.kwargs["src_image"] == "docker.io/swebench/sweb.eval:latest" + + +async def test_create_skips_mirror_for_existing_ecr_ref(): + ecr = "463701203462.dkr.ecr.us-east-1.amazonaws.com/mirror" + provider = create_provider(_provider_config(ecr_repository=ecr)) + spec = SandboxSpec(image=f"{ecr}:already-mirrored") + with _mock_engine_start(), patch.object(engine.ImageBuilder, "ensure_mirrored") as m: + await provider.create(spec) + m.assert_not_called() + + +async def test_create_skips_mirror_when_auto_mirror_disabled(): + ecr = "123.dkr.ecr.us-east-1.amazonaws.com/mirror" + provider = create_provider(_provider_config(ecr_repository=ecr, auto_mirror=False)) + spec = SandboxSpec(image="docker.io/swebench/sweb.eval:latest") + with _mock_engine_start(), patch.object(engine.ImageBuilder, "ensure_mirrored") as m: + await provider.create(spec) + m.assert_not_called() + + def test_get_ecr_image_tag_is_content_addressed(tmp_path): (tmp_path / "Dockerfile").write_text("FROM scratch") tag1 = engine.ImageBuilder.get_ecr_image_tag(tmp_path, "env") From adb9e360f996aa6012d9e9cc361ddd54a08674ca Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Mon, 22 Jun 2026 15:55:38 +0200 Subject: [PATCH 03/11] test(sandbox-ecs): engine.py unit coverage 38%->99% + secrets/format CI fixes Add tests/unit_tests/test_ecs_fargate_engine.py (208 boto3/SSH/HTTP-mocked tests asserting real call args, retries, parsing and error paths), raising ecs_fargate/engine.py to 99% and the repo total to 97.4% (>=96 gate). Tag secretsmanager-ARN test fixtures with allowlist-secret pragmas and apply ruff-format. Fixes the Test (coverage), lint and secrets-detector checks. Signed-off-by: Michal Bien --- tests/unit_tests/test_ecs_fargate_engine.py | 2673 +++++++++++++++++ tests/unit_tests/test_ecs_fargate_provider.py | 25 +- 2 files changed, 2687 insertions(+), 11 deletions(-) create mode 100644 tests/unit_tests/test_ecs_fargate_engine.py diff --git a/tests/unit_tests/test_ecs_fargate_engine.py b/tests/unit_tests/test_ecs_fargate_engine.py new file mode 100644 index 0000000000..79a83d95e4 --- /dev/null +++ b/tests/unit_tests/test_ecs_fargate_engine.py @@ -0,0 +1,2673 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# 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. + +"""Unit tests for the ECS Fargate sandbox *engine* internals. + +Every AWS (boto3), SSH (subprocess), socket, and HTTP (aiohttp/urllib) seam is +mocked — no test touches a real network, process, or AWS account. The tests +drive the REAL method/function under test and assert on its observable +behavior (return values, raised errors, boto3/SSH call arguments, retry and +parsing logic). They complement ``test_ecs_fargate_provider.py`` (which covers +the Gym-facing adapter) by exercising the engine module directly. +""" + +from __future__ import annotations + +import base64 +import contextlib +import io +import json +import threading +import types +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch +from urllib.parse import urlparse + +import aiohttp +import pytest + +from nemo_gym.sandbox.providers.ecs_fargate import engine + + +_ENG = "nemo_gym.sandbox.providers.ecs_fargate.engine" + + +class ClientError(Exception): + """Local stand-in for ``botocore.exceptions.ClientError``. + + CI's base venv has no boto3/botocore (it's the ``sandbox-ecs`` extra), so we + don't import it. engine.py only references ``ClientError`` via + ``_require_aws_sdks()`` — which ``_patch_aws`` patches to return this class — + so engine's ``except ClientError`` catches these instances unchanged. + """ + + def __init__(self, error_response, operation_name="Operation"): + self.response = error_response + self.operation_name = operation_name + super().__init__(f"{operation_name}: {error_response}") + + +# ── Shared fixtures / helpers ───────────────────────────────────────── + + +@pytest.fixture(autouse=True) +def _clear_engine_module_state(request): + """Reset process-global engine caches and stub the AWS SDK accessor. + + CI's base venv has no boto3/botocore (it's the sandbox-ecs extra), so patch + ``engine._require_aws_sdks`` to hand back mocks + the local fake ``ClientError`` + for every test — so engine's ``except ClientError`` catches the fakes the tests + raise without real boto3. Tests needing specific clients override it via + ``_patch_aws`` (nested, wins); the two tests that exercise ``_require_aws_sdks`` + itself opt out by name. + """ + caches = ( + engine._ssm_config_cache, + engine._exec_server_url_cache, + engine._task_def_cache, + engine._task_def_inflight, + engine._active_sandboxes, + engine.ImageBuilder._inflight_builds, + ) + for cache in caches: + cache.clear() + engine.ImageBuilder._build_semaphore = None + engine.ImageBuilder._build_semaphore_size = 0 + if request.node.name in {"test_require_aws_sdks_returns_triple", "test_require_aws_sdks_raises_when_missing"}: + yield + else: + with patch( + f"{_ENG}._require_aws_sdks", + return_value=(MagicMock(name="boto3"), MagicMock(name="Config"), ClientError), + ): + yield + for cache in caches: + cache.clear() + + +def _client_error(code: str, op: str = "Operation", msg: str = "boom") -> ClientError: + return ClientError({"Error": {"Code": code, "Message": msg}}, op) + + +@contextlib.contextmanager +def _patch_aws(**clients): + """Patch ``engine._require_aws_sdks`` so ``boto3.client(name)`` returns mocks. + + Returns the fake ``boto3`` module. ``Config`` is a callable mock and + ``ClientError`` is the *real* class so ``except ClientError`` works. + """ + boto3 = MagicMock(name="boto3") + + def _client(name, **_kw): + if name not in clients: + raise AssertionError(f"unexpected boto3.client({name!r})") + return clients[name] + + boto3.client.side_effect = _client + with patch(f"{_ENG}._require_aws_sdks", return_value=(boto3, MagicMock(name="Config"), ClientError)): + yield boto3 + + +def _exec_sidecar(**overrides) -> engine.SshSidecarConfig: + base = dict( + sshd_port=2222, + public_key_secret_arn="arn:aws:secretsmanager:us-east-1:1234:secret:pub", # pragma: allowlist secret + private_key_secret_arn="arn:aws:secretsmanager:us-east-1:1234:secret:priv", # pragma: allowlist secret + exec_server_port=5000, + ) + base.update(overrides) + return engine.SshSidecarConfig(**base) + + +def _agent_sidecar(**overrides) -> engine.SshSidecarConfig: + base = dict( + sshd_port=2222, + public_key_secret_arn="arn:aws:secretsmanager:us-east-1:1234:secret:pub", # pragma: allowlist secret + private_key_secret_arn="arn:aws:secretsmanager:us-east-1:1234:secret:priv", # pragma: allowlist secret + exec_server_port=None, + ) + base.update(overrides) + return engine.SshSidecarConfig(**base) + + +def _make_sandbox(spec: engine.SandboxSpec | None = None, **cfg_overrides) -> engine.EcsFargateSandbox: + cfg_kwargs = dict( + region="us-west-2", + cluster="test-cluster", + subnets=["subnet-a"], + security_groups=["sg-a"], + ) + cfg_kwargs.update(cfg_overrides) + cfg = engine.EcsFargateConfig(**cfg_kwargs) + spec = spec or engine.SandboxSpec(image="python:3.12") + return engine.EcsFargateSandbox(spec, ecs_config=cfg) + + +def _attach_exec_client(sb: engine.EcsFargateSandbox, result: engine.ExecResult | None = None): + ec = MagicMock(name="exec_client") + ec.exec = AsyncMock(return_value=result or engine.ExecResult("out", "err", 0)) + ec.upload = AsyncMock() + ec.download = AsyncMock(return_value=b"payload") + ec.close = AsyncMock() + sb._exec_client = ec + return ec + + +# ── _require_aws_sdks ───────────────────────────────────────────────── + + +def test_require_aws_sdks_returns_triple(): + pytest.importorskip("boto3") # real boto3/botocore only with the sandbox-ecs extra + from botocore.exceptions import ClientError as RealClientError + + boto3, config_cls, client_error = engine._require_aws_sdks() + assert hasattr(boto3, "client") + assert client_error is RealClientError + assert callable(config_cls) + + +def test_require_aws_sdks_raises_when_missing(): + with patch("importlib.import_module", side_effect=ModuleNotFoundError("no boto3")): + with pytest.raises(RuntimeError, match="requires boto3/botocore"): + engine._require_aws_sdks() + + +# ── resolve_ecs_config_from_ssm ─────────────────────────────────────── + + +def test_resolve_ecs_config_from_ssm_parses_and_caches(): + blob = {"cluster": "c1", "subnets": ["s1", "s2"]} + ssm = MagicMock() + ssm.get_parameter.return_value = {"Parameter": {"Value": json.dumps(blob)}} + with _patch_aws(ssm=ssm): + out = engine.resolve_ecs_config_from_ssm("us-west-2", project="harbor") + # Second call is served from the per-(region, project) cache. + out2 = engine.resolve_ecs_config_from_ssm("us-west-2", project="harbor") + assert out == blob + assert out2 is out + ssm.get_parameter.assert_called_once_with(Name="/harbor/ecs-sandbox/config") + + +def test_resolve_ecs_config_from_ssm_parameter_not_found(): + ssm = MagicMock() + ssm.get_parameter.side_effect = _client_error("ParameterNotFound", "GetParameter") + with _patch_aws(ssm=ssm): + with pytest.raises(RuntimeError, match="not found in us-west-2"): + engine.resolve_ecs_config_from_ssm("us-west-2") + + +def test_resolve_ecs_config_from_ssm_other_client_error_reraised(): + ssm = MagicMock() + ssm.get_parameter.side_effect = _client_error("AccessDenied", "GetParameter") + with _patch_aws(ssm=ssm): + with pytest.raises(ClientError): + engine.resolve_ecs_config_from_ssm("us-west-2") + + +def test_resolve_ecs_config_from_ssm_invalid_json(): + ssm = MagicMock() + ssm.get_parameter.return_value = {"Parameter": {"Value": "{not-json"}} + with _patch_aws(ssm=ssm): + with pytest.raises(RuntimeError, match="invalid JSON"): + engine.resolve_ecs_config_from_ssm("us-east-1") + + +# ── _sanitize_id / _is_ecr_image_ref / _port_from_url ───────────────── + + +def test_sanitize_id_replaces_and_truncates(): + assert engine._sanitize_id("docker.io/foo/bar:latest") == "docker-io-foo-bar-latest" + assert engine._sanitize_id("---") == "task" # empty after strip -> fallback + assert len(engine._sanitize_id("a" * 200, max_len=10)) == 10 + + +def test_is_ecr_image_ref(): + assert engine._is_ecr_image_ref("463701203462.dkr.ecr.us-east-1.amazonaws.com/repo:tag") + assert not engine._is_ecr_image_ref("docker.io/library/ubuntu:24.04") + assert not engine._is_ecr_image_ref("ubuntu:24.04") + + +def test_port_from_url_defaults(): + assert engine._port_from_url(urlparse("http://h/x")) == 80 + assert engine._port_from_url(urlparse("https://h/x")) == 443 + assert engine._port_from_url(urlparse("http://h:8123/x")) == 8123 + + +# ── _OutsideEndpointRoute / _OutsideEndpointRouting ─────────────────── + + +def test_route_for_endpoint_and_rewrite(): + ep = engine.OutsideEndpoint(url="https://api.host:9000/v1", env_var="MODEL_BASE_URL") + route = engine._OutsideEndpointRoute.for_endpoint(ep, remote_port=20001) + assert route.host == "api.host" + assert route.target_port == 9000 + assert route.remote_port == 20001 + assert route.scheme == "https" + assert route.resolved_endpoint_url() == "https://127.0.0.1:20001/v1" + assert route.resolve_url("https://api.host:9000/other") == "https://127.0.0.1:20001/other" + + +def test_route_for_endpoint_rejects_hostless_url(): + ep = engine.OutsideEndpoint(url="/relative/only", env_var="X") + with pytest.raises(ValueError, match="Cannot resolve hostname"): + engine._OutsideEndpointRoute.for_endpoint(ep) + + +def test_routing_for_exec_server_dedupes_and_allocates(): + endpoints = [ + engine.OutsideEndpoint(url="http://10.0.0.1:4000/v1", env_var="A"), + engine.OutsideEndpoint(url="http://10.0.0.1:4000/v2", env_var="B"), + engine.OutsideEndpoint(url="http://10.0.0.2:5000/v1", env_var="C"), + ] + sidecar = _exec_sidecar(sshd_port=2222, exec_server_port=5000) + routing = engine._OutsideEndpointRouting.for_exec_server(endpoints, sidecar) + specs = routing.reverse_specs + # A and B share (host, port) so only two reverse specs are created. + assert len(specs) == 2 + assert any(s.endswith(":10.0.0.1:4000") for s in specs) + # 5000 is reserved by the exec server, so :5000 cannot be reused as a remote port. + assert not any(s.startswith("5000:") for s in specs) + overrides = routing.env_overrides() + assert set(overrides) == {"A", "B", "C"} + # A and B share the same reverse tunnel port (deduped) but keep their paths. + assert urlparse(overrides["A"]).netloc == urlparse(overrides["B"]).netloc + assert overrides["A"].endswith("/v1") and overrides["B"].endswith("/v2") + assert all(v.startswith("http://127.0.0.1:") for v in overrides.values()) + + +def test_routing_for_agent_server_validation_and_target(): + with pytest.raises(ValueError, match="only one OutsideEndpoint"): + engine._OutsideEndpointRouting.for_agent_server( + [engine.OutsideEndpoint("http://h:1/", "A"), engine.OutsideEndpoint("http://h:2/", "B")] + ) + with pytest.raises(ValueError, match="requires OutsideEndpoint"): + engine._OutsideEndpointRouting.for_agent_server([]) + + routing = engine._OutsideEndpointRouting.for_agent_server( + [engine.OutsideEndpoint("http://model.host:7000/v1", "MODEL")] + ) + assert routing.agent_tunnel_port == 7000 + assert routing.agent_tunnel_target() == ("model.host", 7000) + assert routing.resolved_endpoint_url("MODEL") == "http://127.0.0.1:7000/v1" + assert routing.resolved_endpoint_url("MISSING") is None + + +def test_routing_agent_tunnel_target_requires_endpoints(): + routing = engine._OutsideEndpointRouting.empty() + with pytest.raises(ValueError, match="requires OutsideEndpoint"): + routing.agent_tunnel_target() + + +def test_routing_resolve_url_paths(): + # Reverse-tunnel match by netloc. + endpoints = [engine.OutsideEndpoint("http://10.0.0.1:4000/v1", "A")] + routing = engine._OutsideEndpointRouting.for_exec_server(endpoints, _exec_sidecar()) + assert routing.resolve_url("http://10.0.0.1:4000/chat").startswith("http://127.0.0.1:") + + # Agent-tunnel fallback rewrites netloc to the agent port. + agent = engine._OutsideEndpointRouting.for_agent_server([engine.OutsideEndpoint("http://m:7000/v1", "M")]) + assert agent.resolve_url("http://anything:1/path") == "http://127.0.0.1:7000/path" + + # No routes and no agent tunnel -> error. + with pytest.raises(RuntimeError, match="requires SSH reverse tunnel"): + engine._OutsideEndpointRouting.empty().resolve_url("http://x:1/") + + +def test_allocate_reverse_port_prefers_then_scans_then_exhausts(): + used: set[int] = set() + assert engine._OutsideEndpointRouting._allocate_reverse_port(4000, used) == 4000 + # 4000 now reserved -> a second request for it scans for a free port. + second = engine._OutsideEndpointRouting._allocate_reverse_port(4000, used) + assert second != 4000 and 20000 <= second < 61000 + # Invalid preferred (0) with the whole scan range occupied -> RuntimeError. + full = set(range(20000, 61000)) + with pytest.raises(RuntimeError, match="No available local port"): + engine._OutsideEndpointRouting._allocate_reverse_port(0, full) + + +# ── _is_retryable_error / _retry_with_backoff ───────────────────────── + + +def test_is_retryable_error_codes_and_messages(): + assert engine._is_retryable_error(_client_error("ThrottlingException")) + assert engine._is_retryable_error(RuntimeError("Rate exceeded for op")) + assert engine._is_retryable_error(RuntimeError("read timeout while connecting")) + assert not engine._is_retryable_error(RuntimeError("totally fatal")) + assert not engine._is_retryable_error(_client_error("AccessDenied", msg="nope")) + + +def test_retry_with_backoff_success_first_try(): + fn = MagicMock(return_value="ok") + assert engine._retry_with_backoff(fn, operation_name="op") == "ok" + fn.assert_called_once() + + +def test_retry_with_backoff_non_retryable_raises_immediately(): + fn = MagicMock(side_effect=ValueError("fatal")) + with pytest.raises(ValueError): + engine._retry_with_backoff(fn, operation_name="op") + fn.assert_called_once() + + +def test_retry_with_backoff_retries_then_succeeds(): + fn = MagicMock(side_effect=[_client_error("ThrottlingException"), "done"]) + with patch(f"{_ENG}.time.sleep") as sleep: + out = engine._retry_with_backoff(fn, operation_name="op", base_delay=0.01) + assert out == "done" + assert fn.call_count == 2 + sleep.assert_called_once() + + +def test_retry_with_backoff_exhausts_max_retries(): + fn = MagicMock(side_effect=_client_error("ThrottlingException")) + with patch(f"{_ENG}.time.sleep"): + with pytest.raises(ClientError): + engine._retry_with_backoff(fn, operation_name="op", max_retries=2) + # 1 initial + 2 retries. + assert fn.call_count == 3 + + +# ── _free_port / secrets ────────────────────────────────────────────── + + +def test_free_port_returns_bound_port(): + sock = MagicMock() + sock.__enter__.return_value = sock + sock.getsockname.return_value = ("127.0.0.1", 54321) + with patch(f"{_ENG}.socket.socket", return_value=sock): + assert engine._free_port() == 54321 + sock.bind.assert_called_once_with(("127.0.0.1", 0)) + + +def test_download_secret_to_string(): + sm = MagicMock() + sm.get_secret_value.return_value = {"SecretString": "the-key-material"} # pragma: allowlist secret + with _patch_aws(secretsmanager=sm): + out = engine.download_secret_to_string("arn:secret:priv", region="us-east-1") # pragma: allowlist secret + assert out == "the-key-material" + sm.get_secret_value.assert_called_once_with(SecretId="arn:secret:priv") # pragma: allowlist secret + + +def test_download_secret_to_file_writes_0600(): + with patch(f"{_ENG}.download_secret_to_string", return_value="PRIVATE-KEY"): # pragma: allowlist secret + path = engine.download_secret_to_file("arn:secret:priv", region="us-east-1") # pragma: allowlist secret + p = Path(path) + try: + assert p.read_text() == "PRIVATE-KEY" + assert (p.stat().st_mode & 0o777) == 0o600 + finally: + p.unlink() + + +# ── SshTunnel ───────────────────────────────────────────────────────── + + +def _alive_proc(pid: int = 4242): + proc = MagicMock(name="proc") + proc.poll.return_value = None + proc.pid = pid + return proc + + +def _dead_proc(stderr: bytes): + proc = MagicMock(name="proc") + proc.poll.return_value = 1 + proc.stderr.read.return_value = stderr + return proc + + +def test_ssh_tunnel_build_cmd_simple_forward_and_extra(): + t = engine.SshTunnel( + host="1.2.3.4", + port=2222, + key_file="/tmp/k", + forward_port=5000, + forwards=["6000:localhost:6000"], + reverses=["7000:10.0.0.1:7000"], + local_port_override=18000, + ) + cmd = t._build_ssh_cmd() + assert cmd[0] == "ssh" and "-N" in cmd + assert "/tmp/k" in cmd + assert "-L" in cmd and "127.0.0.1:18000:127.0.0.1:5000" in cmd + assert "6000:localhost:6000" in cmd + assert "7000:10.0.0.1:7000" in cmd + assert cmd[-1] == "root@1.2.3.4" + + +def test_ssh_tunnel_local_port_property_before_open(): + t = engine.SshTunnel(host="h", key_file="/k") + with pytest.raises(RuntimeError, match="Tunnel not open yet"): + _ = t.local_port + + +def test_ssh_tunnel_open_noop_when_already_open(): + t = engine.SshTunnel(host="h", key_file="/k", forward_port=5000) + t._proc = _alive_proc() + with patch(f"{_ENG}.subprocess.Popen") as popen: + t.open() + popen.assert_not_called() + + +def test_ssh_tunnel_open_success_simple_forward(): + t = engine.SshTunnel(host="h", key_file="/k", forward_port=5000) + proc = _alive_proc(pid=99) + with ( + patch(f"{_ENG}.subprocess.Popen", return_value=proc) as popen, + patch(f"{_ENG}._free_port", return_value=18055), + patch(f"{_ENG}.time.sleep"), + patch.object(engine.SshTunnel, "_wait_for_local_port") as wait_port, + ): + t.open() + assert t.is_open + assert t.local_port == 18055 + wait_port.assert_called_once() + popen.assert_called_once() + + +def test_ssh_tunnel_open_retries_when_forward_port_not_ready(): + t = engine.SshTunnel(host="h", key_file="/k", forward_port=5000) + with ( + patch(f"{_ENG}.subprocess.Popen", side_effect=[_alive_proc(), _alive_proc()]), + patch(f"{_ENG}._free_port", side_effect=[18001, 18002]), + patch(f"{_ENG}.time.sleep"), + patch.object(engine.SshTunnel, "_kill"), + patch.object(engine.SshTunnel, "_wait_for_local_port", side_effect=[TimeoutError("nope"), None]), + ): + t.open(max_retries=3) + assert t.local_port == 18002 + + +def test_ssh_tunnel_open_retries_on_connection_refused_then_succeeds(): + t = engine.SshTunnel(host="h", key_file="/k", forward_port=5000) + with ( + patch( + f"{_ENG}.subprocess.Popen", side_effect=[_dead_proc(b"ssh: connect: Connection refused"), _alive_proc()] + ), + patch(f"{_ENG}._free_port", side_effect=[18001, 18002]), + patch(f"{_ENG}.time.sleep"), + patch.object(engine.SshTunnel, "_wait_for_local_port"), + ): + t.open(max_retries=3, initial_backoff=0.01) + assert t.is_open + + +def test_ssh_tunnel_open_raises_on_fatal_stderr(): + t = engine.SshTunnel(host="h", key_file="/k", forward_port=5000) + with ( + patch(f"{_ENG}.subprocess.Popen", return_value=_dead_proc(b"Permission denied (publickey)")), + patch(f"{_ENG}._free_port", return_value=18001), + patch(f"{_ENG}.time.sleep"), + ): + with pytest.raises(RuntimeError, match="exited immediately"): + t.open(max_retries=3) + + +def test_ssh_tunnel_open_exhausts_retries(): + t = engine.SshTunnel(host="h", key_file="/k", forward_port=5000) + with ( + patch(f"{_ENG}.subprocess.Popen", return_value=_dead_proc(b"Connection timed out")), + patch(f"{_ENG}._free_port", return_value=18001), + patch(f"{_ENG}.time.sleep"), + ): + with pytest.raises(RuntimeError, match="failed after 2 attempts"): + t.open(max_retries=2, initial_backoff=0.01) + + +def test_ssh_tunnel_kill_terminates_and_force_kills(): + t = engine.SshTunnel(host="h", key_file="/k") + proc = MagicMock() + proc.wait.side_effect = engine.subprocess.TimeoutExpired(cmd="ssh", timeout=5) + t._proc = proc + t.close() + proc.terminate.assert_called_once() + proc.kill.assert_called_once() + assert t._proc is None + + +def test_ssh_tunnel_kill_swallows_process_lookup_error(): + t = engine.SshTunnel(host="h", key_file="/k") + proc = MagicMock() + proc.terminate.side_effect = ProcessLookupError() + t._proc = proc + t.close() # must not raise + assert t._proc is None + + +def test_ssh_tunnel_kill_noop_when_no_proc(): + t = engine.SshTunnel(host="h", key_file="/k") + t.close() # _proc is None -> early return, no error + assert t._proc is None + + +def _connect_sock(recv: bytes | None = None, *, raise_oserror: bool = False): + sock = MagicMock() + sock.__enter__.return_value = sock + if raise_oserror: + sock.connect.side_effect = OSError("refused") + if recv is not None: + sock.recv.return_value = recv + return sock + + +def test_ssh_tunnel_wait_for_local_port_success(): + t = engine.SshTunnel(host="h", key_file="/k") + with patch(f"{_ENG}.socket.socket", return_value=_connect_sock()): + t._wait_for_local_port(18000, timeout=5.0) # connects immediately, returns + + +def test_ssh_tunnel_wait_for_local_port_proc_died(): + t = engine.SshTunnel(host="h", key_file="/k") + t._proc = MagicMock() + t._proc.poll.return_value = 1 + with patch(f"{_ENG}.time.monotonic", side_effect=[0.0, 1.0]): + with pytest.raises(RuntimeError, match="exited while waiting"): + t._wait_for_local_port(18000, timeout=30.0) + + +def test_ssh_tunnel_wait_for_local_port_timeout(): + t = engine.SshTunnel(host="h", key_file="/k") + with ( + patch(f"{_ENG}.socket.socket", return_value=_connect_sock(raise_oserror=True)), + patch(f"{_ENG}.time.monotonic", side_effect=[0.0, 0.1, 100.0]), + patch(f"{_ENG}.time.sleep"), + ): + with pytest.raises(TimeoutError, match="not open after"): + t._wait_for_local_port(18000, timeout=5.0) + + +class _UrlResp: + def __init__(self, status: int) -> None: + self.status = status + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +def test_ssh_tunnel_poll_health_success(): + t = engine.SshTunnel(host="h", key_file="/k") + t._proc = _alive_proc() + with ( + patch("urllib.request.urlopen", return_value=_UrlResp(200)), + patch(f"{_ENG}.time.monotonic", side_effect=[0.0, 0.1]), + patch(f"{_ENG}.time.sleep"), + ): + t._poll_health("http://127.0.0.1:18000/health", timeout=5.0) + + +def test_ssh_tunnel_poll_health_dies_midway(): + t = engine.SshTunnel(host="h", key_file="/k") + t._proc = MagicMock() + t._proc.poll.return_value = 1 # not open + with patch(f"{_ENG}.time.monotonic", side_effect=[0.0, 0.1]): + with pytest.raises(RuntimeError, match="died while waiting"): + t._poll_health("http://127.0.0.1:18000/health", timeout=5.0) + + +def test_ssh_tunnel_poll_health_timeout(): + import urllib.error + + t = engine.SshTunnel(host="h", key_file="/k") + t._proc = _alive_proc() + with ( + patch("urllib.request.urlopen", side_effect=urllib.error.URLError("down")), + patch(f"{_ENG}.time.monotonic", side_effect=[0.0, 0.1, 100.0]), + patch(f"{_ENG}.time.sleep"), + ): + with pytest.raises(TimeoutError, match="not reachable"): + t._poll_health("http://127.0.0.1:18000/health", timeout=5.0) + + +def test_ssh_tunnel_wait_ready_dispatch(): + t = engine.SshTunnel(host="h", key_file="/k") + with patch.object(engine.SshTunnel, "_poll_health") as poll: + t.wait_ready(health_url="http://x/health", timeout=10.0) + poll.assert_called_once() + + t2 = engine.SshTunnel(host="h", key_file="/k", local_port_override=18000) + with patch.object(engine.SshTunnel, "_wait_for_local_port") as wait_port: + t2.wait_ready(timeout=10.0) + wait_port.assert_called_once() + + # Neither health_url nor local_port -> no-op (no error). + engine.SshTunnel(host="h", key_file="/k").wait_ready() + + +def test_ssh_tunnel_check_health_and_context_manager(): + t = engine.SshTunnel(host="h", key_file="/k", forward_port=5000) + assert t.check_health() is False # no proc + proc = _alive_proc() + with ( + patch(f"{_ENG}.subprocess.Popen", return_value=proc), + patch(f"{_ENG}._free_port", return_value=18000), + patch(f"{_ENG}.time.sleep"), + patch.object(engine.SshTunnel, "_wait_for_local_port"), + patch.object(engine.SshTunnel, "_kill") as kill, + ): + with t as ctx: + assert ctx is t + assert t.check_health() is True + kill.assert_called_once() # __exit__ -> close -> _kill + + +# ── build_ssh_sidecar_container ─────────────────────────────────────── + + +def test_build_ssh_sidecar_container_defaults_and_watchdog(): + cfg = _exec_sidecar(sshd_port=52222, image=None) + c = engine.build_ssh_sidecar_container(cfg, public_key_value="ssh-rsa KEY", max_lifetime_sec=3600) + assert c["name"] == "ssh-tunnel" + assert c["image"] == "alpine:latest" + assert c["environment"] == [{"name": "SSH_PUBLIC_KEY", "value": "ssh-rsa KEY"}] + assert "sleep 3600" in c["command"][0] # watchdog present + assert "Port 52222" in c["command"][0] + assert c["healthCheck"]["command"] == ["CMD-SHELL", "nc -z localhost 52222 || exit 1"] + assert "logConfiguration" not in c + + +def test_build_ssh_sidecar_container_no_watchdog_and_logs(): + cfg = _exec_sidecar(sshd_port=2222, image="myimg:1") + c = engine.build_ssh_sidecar_container( + cfg, + public_key_value="K", + max_lifetime_sec=0, + log_group="/aws/ecs/sb", + log_region="us-west-2", + log_stream_prefix="pref", + ) + assert c["image"] == "myimg:1" + assert "sidecar watchdog" not in c["command"][0] + assert c["logConfiguration"]["options"]["awslogs-group"] == "/aws/ecs/sb" + assert c["logConfiguration"]["options"]["awslogs-stream-prefix"] == "pref-tunnel" + + +# ── ExecClient ──────────────────────────────────────────────────────── + + +class _FakeAiohttpCtx: + """Async context manager standing in for ``session.request(...)``.""" + + def __init__(self, result): + self._result = result # _FakeResp or Exception + + async def __aenter__(self): + if isinstance(self._result, Exception): + raise self._result + return self._result + + async def __aexit__(self, *exc): + return False + + +class _FakeResp: + def __init__(self, status: int, body: bytes): + self.status = status + self._body = body + + async def read(self) -> bytes: + return self._body + + +def _fake_session(results): + session = MagicMock(name="session") + session.closed = False + session.request = MagicMock(side_effect=[_FakeAiohttpCtx(r) for r in results]) + return session + + +async def test_exec_client_exec_maps_result(): + client = engine.ExecClient(port=18000) + with patch.object(client, "_post", AsyncMock(return_value={"stdout": "o", "stderr": "e", "rc": 7})) as post: + out = await client.exec("echo hi", timeout=11) + assert (out.stdout, out.stderr, out.return_code) == ("o", "e", 7) + post.assert_awaited_once_with("/exec", {"cmd": "echo hi", "timeout": 11}) + + +async def test_exec_client_upload_bytes_and_mode(): + client = engine.ExecClient(port=18000) + with patch.object(client, "_post", AsyncMock(return_value={"ok": True})) as post: + await client.upload("/remote/f", b"hello", mode="755") + _, body = post.await_args.args + assert body["path"] == "/remote/f" + assert base64.b64decode(body["content"]) == b"hello" + assert body["mode"] == "755" + + +async def test_exec_client_upload_reads_path(tmp_path): + f = tmp_path / "blob.bin" + f.write_bytes(b"frompath") + client = engine.ExecClient(port=18000) + with patch.object(client, "_post", AsyncMock(return_value={"ok": True})) as post: + await client.upload("/remote/blob", f) + _, body = post.await_args.args + assert base64.b64decode(body["content"]) == b"frompath" + + +async def test_exec_client_upload_retries_then_fails(): + client = engine.ExecClient(port=18000) + with ( + patch.object(client, "_post", AsyncMock(side_effect=TimeoutError("slow"))), + patch(f"{_ENG}.asyncio.sleep", AsyncMock()) as sleep, + ): + with pytest.raises(RuntimeError, match="failed after 2 attempts"): + await client.upload("/remote/f", b"x", max_retries=2) + sleep.assert_awaited() # backed off between attempts + + +async def test_exec_client_upload_raises_when_server_reports_not_ok(): + client = engine.ExecClient(port=18000) + with ( + patch.object(client, "_post", AsyncMock(return_value={"ok": False, "error": "disk full"})), + patch(f"{_ENG}.asyncio.sleep", AsyncMock()), + ): + with pytest.raises(RuntimeError, match="failed after 1 attempts"): + await client.upload("/remote/f", b"x", max_retries=1) + + +async def test_exec_client_download_quotes_path(): + client = engine.ExecClient(port=18000) + with patch.object(client, "_request", AsyncMock(return_value=b"DL")) as req: + out = await client.download("/a b/c") + assert out == b"DL" + assert req.await_args.kwargs["url"].endswith("/download?path=/a%20b/c") + + +async def test_exec_client_health_true_and_false(): + client = engine.ExecClient(port=18000) + with patch.object(client, "_request", AsyncMock(return_value=b"{}")): + assert await client.health() is True + with patch.object(client, "_request", AsyncMock(side_effect=ConnectionError("x"))): + assert await client.health() is False + + +async def test_exec_client_post_computes_timeout_and_parses_json(): + client = engine.ExecClient(port=18000, connect_timeout=30.0) + with patch.object(client, "_request", AsyncMock(return_value=b'{"rc": 0}')) as req: + out = await client._post("/exec", {"cmd": "x", "timeout": 100}) + assert out == {"rc": 0} + # cmd timeout (100) + 30 dominates the connect timeout (30). + assert req.await_args.kwargs["timeout"] == 130 + + +async def test_exec_client_post_honors_timeout_override(): + client = engine.ExecClient(port=18000, connect_timeout=30.0) + with patch.object(client, "_request", AsyncMock(return_value=b"{}")) as req: + await client._post("/upload", {"path": "x"}, timeout_override=222.0) + assert req.await_args.kwargs["timeout"] == 222.0 + + +async def test_exec_client_request_success(): + client = engine.ExecClient(port=18000) + session = _fake_session([_FakeResp(200, b"BODY")]) + with patch.object(client, "_ensure_session", AsyncMock(return_value=session)): + out = await client._request(label="t", url="http://x/y", method="GET", timeout=5, max_retries=2) + assert out == b"BODY" + + +async def test_exec_client_request_http_error_not_retried(): + client = engine.ExecClient(port=18000) + session = _fake_session([_FakeResp(500, b"oops")]) + with patch.object(client, "_ensure_session", AsyncMock(return_value=session)): + with pytest.raises(RuntimeError, match="HTTP 500"): + await client._request(label="t", url="http://x/y", method="GET", timeout=5, max_retries=3) + # 4xx/5xx is terminal — exactly one request attempt. + assert session.request.call_count == 1 + + +async def test_exec_client_request_retries_transient_then_succeeds(): + client = engine.ExecClient(port=18000) + session = _fake_session([aiohttp.ClientError("reset"), _FakeResp(200, b"OK")]) + with ( + patch.object(client, "_ensure_session", AsyncMock(return_value=session)), + patch(f"{_ENG}.asyncio.sleep", AsyncMock()) as sleep, + ): + out = await client._request(label="t", url="http://x/y", method="GET", timeout=5, max_retries=3) + assert out == b"OK" + sleep.assert_awaited_once() + + +async def test_exec_client_request_exhausts_to_connection_error(): + client = engine.ExecClient(port=18000) + session = _fake_session([aiohttp.ClientError("a"), aiohttp.ClientError("b")]) + with ( + patch.object(client, "_ensure_session", AsyncMock(return_value=session)), + patch(f"{_ENG}.asyncio.sleep", AsyncMock()), + ): + with pytest.raises(ConnectionError, match="failed after 2 attempts"): + await client._request(label="t", url="http://x/y", method="GET", timeout=5, max_retries=2) + + +async def test_exec_client_request_zero_retries_is_unreachable(): + # Defensive guard: a non-positive retry budget yields a clear ConnectionError + # rather than silently returning None. + client = engine.ExecClient(port=18000) + session = _fake_session([]) + with patch.object(client, "_ensure_session", AsyncMock(return_value=session)): + with pytest.raises(ConnectionError, match="unreachable"): + await client._request(label="t", url="http://x/y", method="GET", timeout=5, max_retries=0) + session.request.assert_not_called() + + +async def test_exec_client_ensure_session_creates_and_reuses(): + client = engine.ExecClient(port=18000) + fake = MagicMock(closed=False) + fake.close = AsyncMock() + with patch(f"{_ENG}.aiohttp.ClientSession", return_value=fake) as ctor: + s1 = await client._ensure_session() + s2 = await client._ensure_session() + assert s1 is fake and s2 is fake + ctor.assert_called_once() # reused, not recreated + await client.close() + fake.close.assert_awaited_once() + assert client._session is None + + +# ── Embedded exec server (_H request handler) ───────────────────────── + + +@pytest.fixture(scope="module") +def exec_server_handler(): + """Compile the embedded exec-server script and return its ``_H`` handler.""" + ns: dict = {} + exec(compile(engine.EXEC_SERVER_SCRIPT, "", "exec"), ns) + return ns["_H"] + + +def _drive_handler(handler_cls, *, path: str, method: str, body: dict | None = None): + """Invoke a handler method with stubbed HTTP plumbing; return (code, payload).""" + h = handler_cls.__new__(handler_cls) + raw = json.dumps(body).encode() if body is not None else b"" + h.path = path + h.rfile = io.BytesIO(raw) + h.wfile = io.BytesIO() + h.headers = {"Content-Length": str(len(raw))} + captured: dict = {} + h.send_response = lambda code: captured.__setitem__("code", code) + h.send_header = lambda *a, **k: None + h.end_headers = lambda: None + if method == "GET": + h.do_GET() + else: + h.do_POST() + out = h.wfile.getvalue() + parsed = json.loads(out) if out and path != "/download" else out + return captured.get("code"), parsed + + +def test_exec_server_health(exec_server_handler): + code, payload = _drive_handler(exec_server_handler, path="/health", method="GET") + assert code == 200 and payload == {"ok": True} + + +def test_exec_server_exec_runs_command(exec_server_handler): + fake = types.SimpleNamespace(stdout=b"hello\n", stderr=b"", returncode=0) + with patch("subprocess.run", return_value=fake): + code, payload = _drive_handler(exec_server_handler, path="/exec", method="POST", body={"cmd": "echo hello"}) + assert code == 200 + assert payload == {"stdout": "hello\n", "stderr": "", "rc": 0} + + +def test_exec_server_exec_missing_cmd(exec_server_handler): + code, payload = _drive_handler(exec_server_handler, path="/exec", method="POST", body={"timeout": 5}) + assert code == 400 + assert "missing 'cmd'" in payload["error"] + + +def test_exec_server_exec_timeout(exec_server_handler): + import subprocess as _sp + + with patch("subprocess.run", side_effect=_sp.TimeoutExpired(cmd="x", timeout=3)): + code, payload = _drive_handler( + exec_server_handler, path="/exec", method="POST", body={"cmd": "sleep 9", "timeout": 3} + ) + assert code == 200 + assert payload["rc"] == 124 and "timed out" in payload["stderr"] + + +def test_exec_server_exec_generic_error(exec_server_handler): + with patch("subprocess.run", side_effect=RuntimeError("boom")): + code, payload = _drive_handler(exec_server_handler, path="/exec", method="POST", body={"cmd": "x"}) + assert code == 200 + assert payload["rc"] == -1 and "boom" in payload["stderr"] + + +def test_exec_server_upload_and_chmod(exec_server_handler, tmp_path): + target = tmp_path / "sub" / "out.txt" + content = base64.b64encode(b"written").decode() + code, payload = _drive_handler( + exec_server_handler, + path="/upload", + method="POST", + body={"path": str(target), "content": content, "mode": "600"}, + ) + assert code == 200 and payload == {"ok": True} + assert target.read_bytes() == b"written" + assert (target.stat().st_mode & 0o777) == 0o600 + + +def test_exec_server_upload_missing_fields(exec_server_handler): + code, payload = _drive_handler(exec_server_handler, path="/upload", method="POST", body={"path": "/x"}) + assert code == 400 and "missing path/content" in payload["error"] + + +def test_exec_server_download_roundtrip(exec_server_handler, tmp_path): + f = tmp_path / "dl.bin" + f.write_bytes(b"DATA-BYTES") + code, raw = _drive_handler(exec_server_handler, path="/download", method="GET", body=None) + # No ?path -> 400 (driver appends none); test explicit path separately below. + assert code == 400 + + h = exec_server_handler.__new__(exec_server_handler) + h.path = f"/download?path={f}" + h.wfile = io.BytesIO() + h.headers = {} + captured: dict = {} + h.send_response = lambda code: captured.__setitem__("code", code) + h.send_header = lambda *a, **k: None + h.end_headers = lambda: None + h.do_GET() + assert captured["code"] == 200 + assert h.wfile.getvalue() == b"DATA-BYTES" + + +def test_exec_server_download_not_found(exec_server_handler, tmp_path): + missing = tmp_path / "nope.bin" + h = exec_server_handler.__new__(exec_server_handler) + h.path = f"/download?path={missing}" + h.wfile = io.BytesIO() + h.headers = {} + captured: dict = {} + h.send_response = lambda code: captured.__setitem__("code", code) + h.send_header = lambda *a, **k: None + h.end_headers = lambda: None + h.do_GET() + assert captured["code"] == 404 + + +def test_exec_server_unknown_routes(exec_server_handler): + code, payload = _drive_handler(exec_server_handler, path="/bogus", method="GET") + assert code == 404 + code, payload = _drive_handler(exec_server_handler, path="/bogus", method="POST", body={}) + assert code == 404 + + +# ── ImageBuilder: ECR queries ───────────────────────────────────────── + + +def _ecr_repo(region: str = "us-west-2") -> str: + return f"123456789012.dkr.ecr.{region}.amazonaws.com/sandbox" + + +def _zip_names(data: bytes) -> set[str]: + import zipfile + + with zipfile.ZipFile(io.BytesIO(data)) as zf: + return set(zf.namelist()) + + +def test_ecr_region_extraction_and_fallback(): + assert engine.ImageBuilder._ecr_region(_ecr_repo("eu-west-1")) == "eu-west-1" + assert engine.ImageBuilder._ecr_region("plainrepo", fallback="us-east-2") == "us-east-2" + + +def test_image_exists_in_ecr_true(): + ecr = MagicMock() + ecr.describe_images.return_value = {"imageDetails": [{}]} + with _patch_aws(ecr=ecr) as boto3: + assert engine.ImageBuilder.image_exists_in_ecr(_ecr_repo(), "tag1") is True + # region parsed from the repo URL, repo_name stripped of the registry host. + boto3.client.assert_called_once_with("ecr", region_name="us-west-2") + ecr.describe_images.assert_called_once_with(repositoryName="sandbox", imageIds=[{"imageTag": "tag1"}]) + + +@pytest.mark.parametrize("code", ["ImageNotFoundException", "RepositoryNotFoundException"]) +def test_image_exists_in_ecr_false_on_missing(code): + ecr = MagicMock() + ecr.describe_images.side_effect = _client_error(code, "DescribeImages") + with _patch_aws(ecr=ecr): + assert engine.ImageBuilder.image_exists_in_ecr(_ecr_repo(), "tag1") is False + + +def test_image_exists_in_ecr_reraises_other(): + ecr = MagicMock() + ecr.describe_images.side_effect = _client_error("AccessDeniedException", "DescribeImages") + with _patch_aws(ecr=ecr): + with pytest.raises(ClientError): + engine.ImageBuilder.image_exists_in_ecr(_ecr_repo(), "tag1") + + +def test_list_ecr_tags_paginates(): + ecr = MagicMock() + paginator = MagicMock() + paginator.paginate.return_value = [ + {"imageIds": [{"imageTag": "a"}, {"imageDigest": "sha256:x"}]}, + {"imageIds": [{"imageTag": "b"}, {"imageTag": "a"}]}, + ] + ecr.get_paginator.return_value = paginator + with _patch_aws(ecr=ecr): + tags = engine.ImageBuilder.list_ecr_tags(_ecr_repo()) + assert tags == {"a", "b"} + paginator.paginate.assert_called_once_with(repositoryName="sandbox", filter={"tagStatus": "TAGGED"}) + + +def test_list_ecr_tags_empty_on_repo_not_found(): + ecr = MagicMock() + paginator = MagicMock() + paginator.paginate.side_effect = _client_error("RepositoryNotFoundException", "ListImages") + ecr.get_paginator.return_value = paginator + with _patch_aws(ecr=ecr): + assert engine.ImageBuilder.list_ecr_tags(_ecr_repo()) == set() + + +def test_list_ecr_tags_reraises_other(): + ecr = MagicMock() + paginator = MagicMock() + paginator.paginate.side_effect = _client_error("AccessDenied", "ListImages") + ecr.get_paginator.return_value = paginator + with _patch_aws(ecr=ecr): + with pytest.raises(ClientError): + engine.ImageBuilder.list_ecr_tags(_ecr_repo()) + + +def test_ecr_docker_login_success_builds_command(): + result = types.SimpleNamespace(returncode=0, stderr="") + with patch(f"{_ENG}.subprocess.run", return_value=result) as run: + engine.ImageBuilder.ecr_docker_login(_ecr_repo("us-west-2")) + cmd = run.call_args.args[0] + assert "aws ecr get-login-password --region us-west-2" in cmd + assert "docker login --username AWS --password-stdin 123456789012.dkr.ecr.us-west-2.amazonaws.com" in cmd + + +def test_ecr_docker_login_no_region_flag_for_plain_repo(): + result = types.SimpleNamespace(returncode=0, stderr="") + with patch(f"{_ENG}.subprocess.run", return_value=result) as run: + engine.ImageBuilder.ecr_docker_login("plainrepo") + assert "--region" not in run.call_args.args[0] + + +def test_ecr_docker_login_failure_raises(): + result = types.SimpleNamespace(returncode=1, stderr="bad creds") + with patch(f"{_ENG}.subprocess.run", return_value=result): + with pytest.raises(RuntimeError, match="ECR docker login failed: bad creds"): + engine.ImageBuilder.ecr_docker_login(_ecr_repo()) + + +def test_docker_push_to_ecr_success(): + tag_res = types.SimpleNamespace(returncode=0, stderr="") + push_res = types.SimpleNamespace(returncode=0, stderr="") + with patch(f"{_ENG}.subprocess.run", side_effect=[tag_res, push_res]) as run: + url = engine.ImageBuilder.docker_push_to_ecr("local:img", _ecr_repo(), "t1") + assert url == f"{_ecr_repo()}:t1" + assert run.call_args_list[0].args[0] == ["docker", "tag", "local:img", f"{_ecr_repo()}:t1"] + assert run.call_args_list[1].args[0] == ["docker", "push", f"{_ecr_repo()}:t1"] + + +def test_docker_push_to_ecr_failure_raises(): + tag_res = types.SimpleNamespace(returncode=0, stderr="") + push_res = types.SimpleNamespace(returncode=1, stderr="denied") + with patch(f"{_ENG}.subprocess.run", side_effect=[tag_res, push_res]): + with pytest.raises(RuntimeError, match="docker push .* failed: denied"): + engine.ImageBuilder.docker_push_to_ecr("local:img", _ecr_repo(), "t1") + + +# ── ImageBuilder: build orchestration ───────────────────────────────── + + +def _build_cfg(**overrides) -> engine.EcsFargateConfig: + base = dict( + region="us-west-2", + ecr_repository=_ecr_repo(), + environment_dir="/env", + s3_bucket="bkt", + codebuild_service_role="arn:aws:iam::123:role/cb", + build_parallelism=4, + ) + base.update(overrides) + return engine.EcsFargateConfig(**base) + + +def test_ensure_image_built_requires_repo_and_dir(): + with pytest.raises(ValueError, match="ecr_repository and environment_dir are required"): + engine.ImageBuilder.ensure_image_built(cfg=engine.EcsFargateConfig(region="us-west-2"), environment_name="env") + + +def test_ensure_image_built_cache_hit_skips_build(): + cfg = _build_cfg() + with ( + patch.object(engine.ImageBuilder, "get_ecr_image_tag", return_value="env__abcd1234"), + patch.object(engine.ImageBuilder, "image_exists_in_ecr", return_value=True), + patch.object(engine.ImageBuilder, "_build_and_push") as build, + ): + url = engine.ImageBuilder.ensure_image_built(cfg=cfg, environment_name="env") + assert url == f"{cfg.ecr_repository}:env__abcd1234" + build.assert_not_called() + + +def test_ensure_image_built_builds_on_miss(): + cfg = _build_cfg() + with ( + patch.object(engine.ImageBuilder, "get_ecr_image_tag", return_value="env__abcd1234"), + patch.object(engine.ImageBuilder, "image_exists_in_ecr", return_value=False), + patch.object(engine.ImageBuilder, "_build_and_push") as build, + ): + url = engine.ImageBuilder.ensure_image_built(cfg=cfg, environment_name="env") + assert url == f"{cfg.ecr_repository}:env__abcd1234" + build.assert_called_once() + assert build.call_args.kwargs["tag"] == "env__abcd1234" + + +def test_ensure_image_built_force_skips_cache_check(): + cfg = _build_cfg() + with ( + patch.object(engine.ImageBuilder, "get_ecr_image_tag", return_value="env__abcd1234"), + patch.object(engine.ImageBuilder, "image_exists_in_ecr") as exists, + patch.object(engine.ImageBuilder, "_build_and_push") as build, + ): + engine.ImageBuilder.ensure_image_built(cfg=cfg, environment_name="env", force_build=True) + exists.assert_not_called() + build.assert_called_once() + + +def test_ensure_image_built_dedupes_on_inflight(): + cfg = _build_cfg() + tag = "env__abcd1234" + done = threading.Event() + done.set() + engine.ImageBuilder._inflight_builds[tag] = done + with ( + patch.object(engine.ImageBuilder, "get_ecr_image_tag", return_value=tag), + patch.object(engine.ImageBuilder, "image_exists_in_ecr") as exists, + patch.object(engine.ImageBuilder, "_build_and_push") as build, + ): + url = engine.ImageBuilder.ensure_image_built(cfg=cfg, environment_name="env") + assert url == f"{cfg.ecr_repository}:{tag}" + exists.assert_not_called() + build.assert_not_called() + + +def test_ensure_image_built_dedupes_on_second_check(): + # A concurrent builder registers the tag between the first cache check and + # the second in-flight check; we must wait on it rather than rebuild. + cfg = _build_cfg() + tag = "env__deadbeef" + + def _register_inflight(*_a, **_k): + ev = threading.Event() + ev.set() + engine.ImageBuilder._inflight_builds[tag] = ev + return False + + with ( + patch.object(engine.ImageBuilder, "get_ecr_image_tag", return_value=tag), + patch.object(engine.ImageBuilder, "image_exists_in_ecr", side_effect=_register_inflight), + patch.object(engine.ImageBuilder, "_build_and_push") as build, + ): + url = engine.ImageBuilder.ensure_image_built(cfg=cfg, environment_name="env") + assert url == f"{cfg.ecr_repository}:{tag}" + build.assert_not_called() + + +def test_ensure_image_built_cache_filled_after_semaphore(): + # The image appears in ECR (built by a peer) by the time we acquire the + # build semaphore -> skip the build. + cfg = _build_cfg() + with ( + patch.object(engine.ImageBuilder, "get_ecr_image_tag", return_value="env__cafe"), + patch.object(engine.ImageBuilder, "image_exists_in_ecr", side_effect=[False, True]), + patch.object(engine.ImageBuilder, "_build_and_push") as build, + ): + url = engine.ImageBuilder.ensure_image_built(cfg=cfg, environment_name="env") + assert url == f"{cfg.ecr_repository}:env__cafe" + build.assert_not_called() + + +def test_ensure_mirrored_requires_repo(): + with pytest.raises(ValueError, match="ecr_repository is required"): + engine.ImageBuilder.ensure_mirrored(cfg=engine.EcsFargateConfig(region="us-west-2"), src_image="ubuntu:24.04") + + +def test_ensure_mirrored_dedupes_on_inflight(): + cfg = _build_cfg() + tag = engine._sanitize_id("ubuntu:24.04") + done = threading.Event() + done.set() + engine.ImageBuilder._inflight_builds[tag] = done + with ( + patch.object(engine.ImageBuilder, "image_exists_in_ecr") as exists, + patch.object(engine.ImageBuilder, "run_buildspec_via_codebuild") as cb, + ): + url = engine.ImageBuilder.ensure_mirrored(cfg=cfg, src_image="ubuntu:24.04") + assert url.endswith(tag) + exists.assert_not_called() + cb.assert_not_called() + + +def test_ensure_mirrored_dedupes_on_second_check(): + cfg = _build_cfg() + tag = engine._sanitize_id("ubuntu:24.04") + + def _register_inflight(*_a, **_k): + ev = threading.Event() + ev.set() + engine.ImageBuilder._inflight_builds[tag] = ev + return False + + with ( + patch.object(engine.ImageBuilder, "image_exists_in_ecr", side_effect=_register_inflight), + patch.object(engine.ImageBuilder, "run_buildspec_via_codebuild") as cb, + ): + url = engine.ImageBuilder.ensure_mirrored(cfg=cfg, src_image="ubuntu:24.04") + assert url.endswith(tag) + cb.assert_not_called() + + +def test_ensure_mirrored_cache_filled_after_semaphore(): + cfg = _build_cfg() + with ( + patch.object(engine.ImageBuilder, "image_exists_in_ecr", side_effect=[False, True]), + patch.object(engine.ImageBuilder, "run_buildspec_via_codebuild") as cb, + ): + engine.ImageBuilder.ensure_mirrored(cfg=cfg, src_image="ubuntu:24.04") + cb.assert_not_called() + + +def test_generate_buildspec_with_dockerhub_login(): + cfg = engine.EcsFargateConfig( + region="us-west-2", + ecr_repository=_ecr_repo(), + dockerhub_secret_arn="arn:aws:secretsmanager:us-west-2:123:secret:dh", # pragma: allowlist secret + ) + spec = engine.ImageBuilder._generate_buildspec(cfg, "sandbox", "t1", f"{cfg.ecr_repository}:t1") + assert "docker build -t sandbox:t1" in spec + assert "secretsmanager get-secret-value" in spec # Docker Hub auth branch + assert f"docker push {cfg.ecr_repository}:t1" in spec + + +def test_upload_build_context_zips_and_uploads(tmp_path): + (tmp_path / "Dockerfile").write_text("FROM scratch") + (tmp_path / "sub").mkdir() + (tmp_path / "sub" / "f.txt").write_text("hi") + cfg = _build_cfg(environment_dir=str(tmp_path), s3_prefix="pfx") + s3 = MagicMock() + with _patch_aws(s3=s3): + key = engine.ImageBuilder._upload_build_context(cfg, "env", "nonce1") + assert key == "pfx/codebuild/env-nonce1.zip" + put = s3.put_object.call_args.kwargs + assert put["Bucket"] == "bkt" and put["Key"] == key + # The uploaded payload is a real zip containing the env files. + names = _zip_names(put["Body"]) + assert "Dockerfile" in names and "sub/f.txt" in names + + +def test_resolve_codebuild_project_uses_explicit(): + cfg = _build_cfg(codebuild_project="my-project") + with _patch_aws(): + # No client calls expected when project is explicit. + assert engine.ImageBuilder._resolve_codebuild_project(cfg, MagicMock(), "nonce") == "my-project" + + +def test_resolve_codebuild_project_requires_role(): + cfg = _build_cfg(codebuild_service_role=None, codebuild_project=None) + with _patch_aws(): + with pytest.raises(RuntimeError, match="codebuild_project or codebuild_service_role"): + engine.ImageBuilder._resolve_codebuild_project(cfg, MagicMock(), "nonce") + + +def test_resolve_codebuild_project_creates_project(): + cfg = _build_cfg() + cb = MagicMock() + with _patch_aws(): + name = engine.ImageBuilder._resolve_codebuild_project(cfg, cb, "abc123") + assert name == "ecs-sandbox-build-abc123" + cb.create_project.assert_called_once() + assert cb.create_project.call_args.kwargs["serviceRole"] == "arn:aws:iam::123:role/cb" + + +def test_resolve_codebuild_project_tolerates_already_exists(): + cfg = _build_cfg() + cb = MagicMock() + cb.create_project.side_effect = _client_error( + "ResourceAlreadyExistsException", "CreateProject", msg="already exists" + ) + with _patch_aws(): + name = engine.ImageBuilder._resolve_codebuild_project(cfg, cb, "abc123") + assert name == "ecs-sandbox-build-abc123" + + +def test_resolve_codebuild_project_reraises_other_error(): + cfg = _build_cfg() + cb = MagicMock() + cb.create_project.side_effect = _client_error("AccessDenied", "CreateProject", msg="forbidden") + with _patch_aws(): + with pytest.raises(ClientError): + engine.ImageBuilder._resolve_codebuild_project(cfg, cb, "abc123") + + +def test_poll_codebuild_success(): + cb = MagicMock() + cb.batch_get_builds.return_value = {"builds": [{"buildStatus": "SUCCEEDED"}]} + with patch(f"{_ENG}.time.sleep"), patch(f"{_ENG}.random.uniform", return_value=0): + engine.ImageBuilder._poll_codebuild(cb, "bid", "img:tag") # returns without error + + +def test_poll_codebuild_in_progress_then_success(): + cb = MagicMock() + cb.batch_get_builds.side_effect = [ + {"builds": [{"buildStatus": "IN_PROGRESS", "currentPhase": "BUILD"}]}, + {"builds": [{"buildStatus": "SUCCEEDED"}]}, + ] + with patch(f"{_ENG}.time.sleep"), patch(f"{_ENG}.random.uniform", return_value=0): + engine.ImageBuilder._poll_codebuild(cb, "bid", "img:tag") + assert cb.batch_get_builds.call_count == 2 + + +def test_poll_codebuild_failure_reports_phases(): + cb = MagicMock() + cb.batch_get_builds.return_value = { + "builds": [ + { + "buildStatus": "FAILED", + "phases": [ + {"phaseType": "BUILD", "phaseStatus": "FAILED"}, + {"phaseType": "DOWNLOAD_SOURCE", "phaseStatus": "SUCCEEDED"}, + ], + } + ] + } + with patch(f"{_ENG}.time.sleep"), patch(f"{_ENG}.random.uniform", return_value=0): + with pytest.raises(RuntimeError, match="CodeBuild failed for img:tag: BUILD: FAILED"): + engine.ImageBuilder._poll_codebuild(cb, "bid", "img:tag") + + +def test_build_and_push_orchestration(): + cfg = _build_cfg() + cb = MagicMock() + cb.start_build.return_value = {"build": {"id": "build-99"}} + with ( + _patch_aws(codebuild=cb), + patch.object(engine.ImageBuilder, "_upload_build_context", return_value="pfx/ctx.zip"), + patch.object(engine.ImageBuilder, "_resolve_codebuild_project", return_value="proj"), + patch.object(engine.ImageBuilder, "_generate_buildspec", return_value="version: 0.2"), + patch.object(engine.ImageBuilder, "_poll_codebuild") as poll, + ): + engine.ImageBuilder._build_and_push( + cfg=cfg, environment_name="env", tag="t1", image_url=f"{cfg.ecr_repository}:t1" + ) + sb = cb.start_build.call_args.kwargs + assert sb["projectName"] == "proj" + assert sb["sourceTypeOverride"] == "S3" + assert sb["sourceLocationOverride"] == "bkt/pfx/ctx.zip" + poll.assert_called_once_with(cb, "build-99", f"{cfg.ecr_repository}:t1") + + +def test_run_buildspec_via_codebuild_no_source(): + cfg = _build_cfg() + cb = MagicMock() + cb.start_build.return_value = {"build": {"id": "build-7"}} + with ( + _patch_aws(codebuild=cb), + patch.object(engine.ImageBuilder, "_resolve_codebuild_project", return_value="proj"), + patch.object(engine.ImageBuilder, "_poll_codebuild") as poll, + ): + engine.ImageBuilder.run_buildspec_via_codebuild( + cfg=cfg, buildspec="version: 0.2", job_label="mirror::x", timeout_minutes=12 + ) + sb = cb.start_build.call_args.kwargs + assert sb["sourceTypeOverride"] == "NO_SOURCE" + assert sb["timeoutInMinutesOverride"] == 12 + poll.assert_called_once_with(cb, "build-7", "mirror::x") + + +# ── EcsFargateSandbox: client init + env/command builders ───────────── + + +def test_init_aws_clients_creates_three_clients(): + sb = _make_sandbox() + with _patch_aws(ecs=MagicMock(), ec2=MagicMock(), ssm=MagicMock()) as boto3: + sb._init_aws_clients() + assert sb._ecs is not None and sb._ec2 is not None and sb._ssm is not None + created = {c.args[0] for c in boto3.client.call_args_list} + assert created == {"ecs", "ec2", "ssm"} + for c in boto3.client.call_args_list: + assert c.kwargs["region_name"] == "us-west-2" + + +def test_build_container_command_none_without_exec_server(): + sb = _make_sandbox() + assert sb._build_container_command(_agent_sidecar()) is None + + +def test_build_container_command_bootstraps_exec_server(): + sb = _make_sandbox(spec=engine.SandboxSpec(image="my/img:1")) + cmd = sb._build_container_command(_exec_sidecar(exec_server_port=5000)) + assert cmd[:2] == ["sh", "-lc"] + setup = cmd[2] + assert "base64 -d > /tmp/_exec_server.py" in setup + assert "TB_EXEC_PORT=5000" in setup + assert "exec python3 /tmp/_exec_server.py" in setup + + +def test_build_env_vars_merges_spec_extra_and_routing(): + spec = engine.SandboxSpec(image="img:1", env={"FROM_SPEC": "1"}) + sb = _make_sandbox(spec=spec, extra_env={"RENDERED": "ip={task_ip} img={image}"}) + sb._task_ip = "9.9.9.9" + sb._outside_endpoint_routing = engine._OutsideEndpointRouting.for_exec_server( + [engine.OutsideEndpoint("http://10.0.0.1:4000/v1", "MODEL_BASE_URL")], _exec_sidecar() + ) + env = sb._build_env_vars() + assert env["FROM_SPEC"] == "1" + assert env["RENDERED"] == "ip=9.9.9.9 img=img:1" + assert env["MODEL_BASE_URL"].startswith("http://127.0.0.1:") + + +def test_split_env_separates_runtime_keys(): + sb = _make_sandbox() + sb._outside_endpoint_routing = engine._OutsideEndpointRouting.for_exec_server( + [engine.OutsideEndpoint("http://10.0.0.1:4000/v1", "MODEL_BASE_URL")], _exec_sidecar() + ) + env = {"STABLE": "a", "_NEL_EFS_SESSION": "sess", "MODEL_BASE_URL": "http://127.0.0.1:4000/v1"} + stable, runtime = sb._split_env(env) + assert stable == {"STABLE": "a"} + assert runtime == {"_NEL_EFS_SESSION": "sess", "MODEL_BASE_URL": "http://127.0.0.1:4000/v1"} + + +def test_render_env_value_substitutes_placeholders(): + sb = _make_sandbox(spec=engine.SandboxSpec(image="img:2")) + sb._ssh_tunnel_port = 7000 + sb._task_ip = "1.2.3.4" + assert sb._render_env_value("p={ssh_tunnel_port} ip={task_ip} i={image}") == "p=7000 ip=1.2.3.4 i=img:2" + + +def test_make_family_name_sanitizes_and_prefixes(): + sb = _make_sandbox(spec=engine.SandboxSpec(image="weird/Name:v1"), task_definition_family_prefix="ecs-sandbox") + fam = sb._make_family_name() + assert fam.startswith("ecs-sandbox-") + assert all(ch.isalnum() or ch in "_-" for ch in fam) + + +def test_make_family_name_prepends_ecs_when_non_alnum_start(): + sb = _make_sandbox(spec=engine.SandboxSpec(image=""), task_definition_family_prefix="") + fam = sb._make_family_name() + assert fam.startswith("ecs_") + + +# ── EcsFargateSandbox: _resolve_image branches ──────────────────────── + + +def test_resolve_image_built_passthrough(): + sb = _make_sandbox() + assert sb._resolve_image("built:img") == "built:img" + + +def test_resolve_image_template_key_error(): + sb = _make_sandbox(spec=engine.SandboxSpec(image="x"), image_template="{nonexistent_key}-img") + with pytest.raises(ValueError, match="placeholder"): + sb._resolve_image() + + +def test_resolve_image_bare_without_ecr_is_verbatim(): + sb = _make_sandbox(spec=engine.SandboxSpec(image="ubuntu:24.04")) + assert sb._resolve_image() == "ubuntu:24.04" + + +def test_resolve_image_empty_with_task_definition(): + sb = _make_sandbox(spec=engine.SandboxSpec(image=""), task_definition="arn:td") + assert sb._resolve_image() == "" + + +def test_resolve_image_no_source_raises(): + sb = _make_sandbox(spec=engine.SandboxSpec(image="")) + with pytest.raises(ValueError, match="No image available"): + sb._resolve_image() + + +# ── EcsFargateSandbox: exec-server upload ───────────────────────────── + + +def test_upload_exec_server_requires_bucket(): + sb = _make_sandbox() + with pytest.raises(ValueError, match="s3_bucket is required"): + sb._upload_exec_server() + + +def test_upload_exec_server_uploads_and_caches(): + sb = _make_sandbox(s3_bucket="bkt", s3_prefix="pfx") + s3 = MagicMock() + s3.generate_presigned_url.return_value = "https://signed.example/exec" + with _patch_aws(s3=s3): + url1 = sb._upload_exec_server() + url2 = sb._upload_exec_server() + assert url1 == "https://signed.example/exec" + assert url2 == url1 + s3.put_object.assert_called_once() # cached on the second call + assert s3.put_object.call_args.kwargs["Bucket"] == "bkt" + + +# ── EcsFargateSandbox: EFS volumes ──────────────────────────────────── + + +def test_build_efs_volumes_access_point_and_root_dir(): + spec = engine.SandboxSpec( + image="img:1", + volumes=[ + engine.VolumeMount( + container_path="/mnt/a", readonly=True, efs_filesystem_id="fs-1", efs_access_point_id="ap-1" + ), + engine.VolumeMount(container_path="/mnt/b", efs_filesystem_id="fs-2", efs_root_directory="/data"), + engine.VolumeMount(host_path="/h", container_path="/mnt/c"), # non-EFS -> skipped + ], + ) + sb = _make_sandbox(spec=spec) + volumes, mounts = sb._build_efs_volumes() + assert len(volumes) == 2 and len(mounts) == 2 + assert volumes[0]["efsVolumeConfiguration"]["authorizationConfig"]["accessPointId"] == "ap-1" + assert volumes[1]["efsVolumeConfiguration"]["rootDirectory"] == "/data" + assert mounts[0] == {"sourceVolume": "efs-0", "containerPath": "/mnt/a", "readOnly": True} + + +# ── EcsFargateSandbox: task definition registration ─────────────────── + + +def test_register_task_definition_from_scratch_when_no_base(): + sb = _make_sandbox(execution_role_arn="arn:exec") + sb._ecs = MagicMock() + with patch.object(engine.EcsFargateSandbox, "_register_from_scratch", return_value="scratch-arn") as scratch: + arn = sb._register_task_definition(image="img", command=["sh"], env={}, sidecar_def={"name": "ssh-tunnel"}) + assert arn == "scratch-arn" + scratch.assert_called_once() + + +def test_register_task_definition_from_base_when_describable(): + sb = _make_sandbox(task_definition="arn:td", log_group="/g") + sb._ecs = MagicMock() + base = {"taskDefinition": {"containerDefinitions": [{"name": "main"}]}} + sb._ecs.describe_task_definition.return_value = base + with patch.object(engine.EcsFargateSandbox, "_register_from_base", return_value="base-arn") as from_base: + arn = sb._register_task_definition( + image="img", command=None, env={"K": "V"}, sidecar_def={"name": "ssh-tunnel"} + ) + assert arn == "base-arn" + assert from_base.call_args.kwargs["base"] == base["taskDefinition"] + # log_group set -> a log configuration is threaded through. + assert from_base.call_args.kwargs["log_cfg"]["options"]["awslogs-group"] == "/g" + + +def test_register_task_definition_missing_base_falls_back_to_scratch(): + sb = _make_sandbox(task_definition="arn:missing") + sb._ecs = MagicMock() + sb._ecs.describe_task_definition.side_effect = _client_error("ClientException", "DescribeTaskDefinition") + with patch.object(engine.EcsFargateSandbox, "_register_from_scratch", return_value="scratch-arn") as scratch: + arn = sb._register_task_definition(image="img", command=None, env={}, sidecar_def={"name": "ssh-tunnel"}) + assert arn == "scratch-arn" + scratch.assert_called_once() + + +def test_register_task_definition_describe_other_error_raises(): + sb = _make_sandbox(task_definition="arn:td") + sb._ecs = MagicMock() + sb._ecs.describe_task_definition.side_effect = _client_error("AccessDenied", "DescribeTaskDefinition") + with pytest.raises(ClientError): + sb._register_task_definition(image="img", command=None, env={}, sidecar_def={"name": "ssh-tunnel"}) + + +def test_register_from_base_builds_payload(): + sb = _make_sandbox( + spec=engine.SandboxSpec(image="img:1"), + execution_role_arn="arn:exec", + task_role_arn="arn:taskrole", + cpu="8192", + memory="16384", + ) + base = { + "containerDefinitions": [ + {"name": "main", "environment": [{"name": "OLD", "value": "1"}]}, + {"name": "ssh-tunnel"}, # stale sidecar — must be dropped + ], + "networkMode": "awsvpc", + "requiresCompatibilities": ["FARGATE"], + "cpu": "256", + "memory": "512", + "executionRoleArn": "arn:base-exec", + } + with patch.object(engine.EcsFargateSandbox, "_do_register", return_value="arn") as do_reg: + arn = sb._register_from_base( + base=base, + image="newimg:2", + command=["sh", "-lc", "x"], + env={"NEW": "2"}, + sidecar_def={"name": "ssh-tunnel", "image": "alpine"}, + log_cfg={"logDriver": "awslogs"}, + ) + assert arn == "arn" + payload = do_reg.call_args.args[0] + names = [c["name"] for c in payload["containerDefinitions"]] + assert names.count("ssh-tunnel") == 1 # old sidecar replaced + main = next(c for c in payload["containerDefinitions"] if c["name"] == "main") + assert main["image"] == "newimg:2" + assert main["command"] == ["sh", "-lc", "x"] + assert main["dependsOn"] == [{"containerName": "ssh-tunnel", "condition": "HEALTHY"}] + assert {"name": "NEW", "value": "2"} in main["environment"] + assert {"name": "OLD", "value": "1"} in main["environment"] + # cpu/memory are the max of base and config. + assert payload["cpu"] == "8192" and payload["memory"] == "16384" + assert payload["executionRoleArn"] == "arn:exec" + assert payload["taskRoleArn"] == "arn:taskrole" + + +def test_register_from_base_missing_container_raises(): + sb = _make_sandbox(container_name="main") + base = {"containerDefinitions": [{"name": "other"}]} + with pytest.raises(RuntimeError, match="no container 'main'"): + sb._register_from_base( + base=base, image="i", command=None, env={}, sidecar_def={"name": "ssh-tunnel"}, log_cfg=None + ) + + +def test_register_from_base_appends_efs_volumes_to_existing(): + spec = engine.SandboxSpec( + image="img:1", volumes=[engine.VolumeMount(container_path="/mnt/efs", efs_filesystem_id="fs-9")] + ) + sb = _make_sandbox(spec=spec) + base = { + "containerDefinitions": [{"name": "main", "mountPoints": [{"sourceVolume": "pre", "containerPath": "/pre"}]}], + "volumes": [{"name": "pre"}], + } + with patch.object(engine.EcsFargateSandbox, "_do_register", return_value="arn") as do_reg: + sb._register_from_base( + base=base, image="i", command=None, env={}, sidecar_def={"name": "ssh-tunnel"}, log_cfg=None + ) + payload = do_reg.call_args.args[0] + main = next(c for c in payload["containerDefinitions"] if c["name"] == "main") + # Existing mounts/volumes are preserved and the EFS ones are appended. + assert {"sourceVolume": "pre", "containerPath": "/pre"} in main["mountPoints"] + assert any(m["containerPath"] == "/mnt/efs" for m in main["mountPoints"]) + vol_names = {v["name"] for v in payload["volumes"]} + assert vol_names == {"pre", "efs-0"} + + +def test_register_from_scratch_requires_execution_role(): + sb = _make_sandbox(execution_role_arn=None) + with pytest.raises(RuntimeError, match="execution_role_arn required"): + sb._register_from_scratch(image="i", command=None, env={}, sidecar_def={"name": "ssh-tunnel"}, log_cfg=None) + + +def test_register_from_scratch_builds_payload(): + spec = engine.SandboxSpec( + image="img:1", + volumes=[engine.VolumeMount(container_path="/mnt", efs_filesystem_id="fs-1")], + ) + sb = _make_sandbox( + spec=spec, + execution_role_arn="arn:exec", + task_role_arn="arn:taskrole", + container_port=8080, + ephemeral_storage_gib=40, + ) + with patch.object(engine.EcsFargateSandbox, "_do_register", return_value="arn") as do_reg: + sb._register_from_scratch( + image="img:1", + command=["sh"], + env={"K": "V"}, + sidecar_def={"name": "ssh-tunnel"}, + log_cfg={"logDriver": "awslogs"}, + ) + payload = do_reg.call_args.args[0] + main = payload["containerDefinitions"][0] + assert main["name"] == "main" + assert main["portMappings"] == [{"containerPort": 8080, "protocol": "tcp"}] + assert main["environment"] == [{"name": "K", "value": "V"}] + assert main["logConfiguration"] == {"logDriver": "awslogs"} + assert main["mountPoints"][0]["containerPath"] == "/mnt" + assert payload["executionRoleArn"] == "arn:exec" + assert payload["taskRoleArn"] == "arn:taskrole" + assert payload["ephemeralStorage"] == {"sizeInGiB": 40} + assert payload["volumes"][0]["name"] == "efs-0" + + +# ── EcsFargateSandbox: SSM task-def cache + _do_register ─────────────── + + +def test_ssm_lookup_task_def_active(): + sb = _make_sandbox() + sb._ssm = MagicMock() + sb._ssm.get_parameter.return_value = {"Parameter": {"Value": "arn:cached-td"}} + sb._ecs = MagicMock() + sb._ecs.describe_task_definition.return_value = {"taskDefinition": {"status": "ACTIVE"}} + assert sb._ssm_lookup_task_def("hash1") == "arn:cached-td" + + +def test_ssm_lookup_task_def_not_found(): + sb = _make_sandbox() + sb._ssm = MagicMock() + sb._ssm.get_parameter.side_effect = _client_error("ParameterNotFound", "GetParameter") + assert sb._ssm_lookup_task_def("hash1") is None + + +def test_ssm_lookup_task_def_other_param_error_returns_none(): + sb = _make_sandbox() + sb._ssm = MagicMock() + sb._ssm.get_parameter.side_effect = _client_error("ThrottlingException", "GetParameter") + assert sb._ssm_lookup_task_def("hash1") is None + + +def test_ssm_lookup_task_def_describe_fails_returns_none(): + sb = _make_sandbox() + sb._ssm = MagicMock() + sb._ssm.get_parameter.return_value = {"Parameter": {"Value": "arn:gone"}} + sb._ecs = MagicMock() + sb._ecs.describe_task_definition.side_effect = _client_error("ClientException", "DescribeTaskDefinition") + assert sb._ssm_lookup_task_def("hash1") is None + + +def test_ssm_lookup_task_def_inactive_returns_none(): + sb = _make_sandbox() + sb._ssm = MagicMock() + sb._ssm.get_parameter.return_value = {"Parameter": {"Value": "arn:old"}} + sb._ecs = MagicMock() + sb._ecs.describe_task_definition.return_value = {"taskDefinition": {"status": "INACTIVE"}} + assert sb._ssm_lookup_task_def("hash1") is None + + +def test_ssm_write_task_def_ok_and_error_swallowed(): + sb = _make_sandbox() + sb._ssm = MagicMock() + sb._ssm_write_task_def("hash1", "arn:new") + sb._ssm.put_parameter.assert_called_once() + # A failing PutParameter must not raise (cache write is best-effort). + sb._ssm.put_parameter.side_effect = _client_error("AccessDenied", "PutParameter") + sb._ssm_write_task_def("hash2", "arn:new2") + + +def test_do_register_cache_hit(): + sb = _make_sandbox() + payload = {"family": "f", "containerDefinitions": [{"name": "main"}]} + h = engine._compute_task_def_hash(payload) + engine._task_def_cache[h] = "arn:cached" + assert sb._do_register(payload) == "arn:cached" + + +def test_do_register_fresh_path_writes_cache(): + sb = _make_sandbox() + payload = {"family": "f", "containerDefinitions": [{"name": "main"}]} + h = engine._compute_task_def_hash(payload) + with ( + patch.object(engine.EcsFargateSandbox, "_ssm_lookup_task_def", return_value=None), + patch.object(engine.EcsFargateSandbox, "_register_task_def_fresh", return_value="arn:fresh") as fresh, + ): + arn = sb._do_register(payload) + assert arn == "arn:fresh" + assert engine._task_def_cache[h] == "arn:fresh" + fresh.assert_called_once() + assert h not in engine._task_def_inflight # inflight released + + +def test_do_register_uses_ssm_cache_entry(): + sb = _make_sandbox() + payload = {"family": "f", "containerDefinitions": [{"name": "main"}]} + with ( + patch.object(engine.EcsFargateSandbox, "_ssm_lookup_task_def", return_value="arn:ssm"), + patch.object(engine.EcsFargateSandbox, "_register_task_def_fresh") as fresh, + ): + assert sb._do_register(payload) == "arn:ssm" + fresh.assert_not_called() + + +def test_do_register_waits_for_inflight_then_reads_cache(): + sb = _make_sandbox() + payload = {"family": "f", "containerDefinitions": [{"name": "main"}]} + h = engine._compute_task_def_hash(payload) + + class _WaitThenCache(threading.Event): + def wait(self, timeout=None): + # Simulate the in-flight builder finishing and publishing the arn. + engine._task_def_cache[h] = "arn:by-other" + return True + + engine._task_def_inflight[h] = _WaitThenCache() + with patch.object(engine.EcsFargateSandbox, "_register_task_def_fresh") as fresh: + assert sb._do_register(payload) == "arn:by-other" + fresh.assert_not_called() + + +def test_register_task_def_fresh_calls_ecs_and_writes_ssm(): + sb = _make_sandbox() + sb._ecs = MagicMock() + sb._ecs.register_task_definition.return_value = {"taskDefinition": {"taskDefinitionArn": "arn:reg"}} + with patch.object(engine.EcsFargateSandbox, "_ssm_write_task_def") as write: + arn = sb._register_task_def_fresh({"family": "f"}, "hash1") + assert arn == "arn:reg" + write.assert_called_once_with("hash1", "arn:reg") + + +# ── EcsFargateSandbox: _run_task ────────────────────────────────────── + + +def test_run_task_success_network_config(): + sb = _make_sandbox(assign_public_ip=True, container_name="main") + sb._task_arn = None + sb._ecs = MagicMock() + sb._ecs.run_task.return_value = {"tasks": [{"taskArn": "arn:task-1"}]} + arn = sb._run_task("arn:td") + assert arn == "arn:task-1" + kwargs = sb._ecs.run_task.call_args.kwargs + vpc = kwargs["networkConfiguration"]["awsvpcConfiguration"] + assert vpc["assignPublicIp"] == "ENABLED" + assert vpc["subnets"] == ["subnet-a"] and vpc["securityGroups"] == ["sg-a"] + assert "overrides" not in kwargs # no per-invocation env + + +def test_run_task_includes_runtime_overrides_and_disabled_public_ip(): + sb = _make_sandbox(assign_public_ip=False, container_name="main", platform_version="1.5.0") + sb._runtime_container_env = {"_NEL_EFS_SESSION": "sess-1"} + sb._ecs = MagicMock() + sb._ecs.run_task.return_value = {"tasks": [{"taskArn": "arn:task-2"}]} + sb._run_task("arn:td") + kwargs = sb._ecs.run_task.call_args.kwargs + assert kwargs["networkConfiguration"]["awsvpcConfiguration"]["assignPublicIp"] == "DISABLED" + assert kwargs["platformVersion"] == "1.5.0" + overrides = kwargs["overrides"]["containerOverrides"][0] + assert overrides["name"] == "main" + assert overrides["environment"] == [{"name": "_NEL_EFS_SESSION", "value": "sess-1"}] + + +def test_run_task_efs_forces_platform_version(): + spec = engine.SandboxSpec(image="img", volumes=[engine.VolumeMount(container_path="/m", efs_filesystem_id="fs-1")]) + sb = _make_sandbox(spec=spec) + sb._ecs = MagicMock() + sb._ecs.run_task.return_value = {"tasks": [{"taskArn": "arn:task-3"}]} + sb._run_task("arn:td") + assert sb._ecs.run_task.call_args.kwargs["platformVersion"] == "1.4.0" + + +def test_run_task_no_tasks_raises(): + sb = _make_sandbox() + sb._ecs = MagicMock() + sb._ecs.run_task.return_value = {"tasks": []} + with pytest.raises(RuntimeError, match="no tasks"): + sb._run_task("arn:td") + + +def test_run_task_non_retryable_exception_reraised(): + sb = _make_sandbox() + sb._ecs = MagicMock() + sb._ecs.run_task.side_effect = ValueError("permanent") + with patch(f"{_ENG}.time.sleep"): + with pytest.raises(ValueError, match="permanent"): + sb._run_task("arn:td") + + +def test_run_task_retryable_exception_then_success(): + sb = _make_sandbox(run_task_max_retries=5) + sb._ecs = MagicMock() + sb._ecs.run_task.side_effect = [ + _client_error("ThrottlingException", "RunTask"), + {"tasks": [{"taskArn": "arn:task-ok"}]}, + ] + # Isolate the outer retry loop from the inner per-call backoff helper. + with ( + patch(f"{_ENG}._retry_with_backoff", side_effect=lambda func, **_kw: func()), + patch(f"{_ENG}.time.sleep"), + ): + assert sb._run_task("arn:td") == "arn:task-ok" + + +def test_run_task_non_retryable_failures_raise(): + sb = _make_sandbox(run_task_max_retries=3) + sb._ecs = MagicMock() + sb._ecs.run_task.return_value = {"failures": [{"reason": "AccessDenied on subnet"}]} + with patch(f"{_ENG}.time.sleep"): + with pytest.raises(RuntimeError, match="run_task failures"): + sb._run_task("arn:td") + + +def test_run_task_retryable_failures_exhaust(): + sb = _make_sandbox(run_task_max_retries=2) + sb._ecs = MagicMock() + sb._ecs.run_task.return_value = {"failures": [{"reason": "Capacity is unavailable right now"}]} + with patch(f"{_ENG}.time.sleep"): + with pytest.raises(RuntimeError, match="run_task failures"): + sb._run_task("arn:td") + assert sb._ecs.run_task.call_count == 2 + + +def test_run_task_zero_retries_raises_final(): + sb = _make_sandbox(run_task_max_retries=0) + sb._ecs = MagicMock() + with pytest.raises(RuntimeError, match="failed after 0 retries"): + sb._run_task("arn:td") + sb._ecs.run_task.assert_not_called() + + +# ── EcsFargateSandbox: _wait_for_running ────────────────────────────── + + +def test_wait_for_running_returns_when_running(): + sb = _make_sandbox() + sb._task_arn = "arn:task" + sb._ecs = MagicMock() + sb._ecs.describe_tasks.return_value = {"tasks": [{"lastStatus": "RUNNING"}]} + with patch(f"{_ENG}.time.monotonic", return_value=0.0): + sb._wait_for_running() + + +def test_wait_for_running_transitions_then_running(): + sb = _make_sandbox() + sb._task_arn = "arn:task" + sb._ecs = MagicMock() + sb._ecs.describe_tasks.side_effect = [ + {"tasks": [{"lastStatus": "PROVISIONING"}]}, + {"tasks": [{"lastStatus": "RUNNING"}]}, + ] + with ( + patch(f"{_ENG}.time.monotonic", return_value=0.0), + patch(f"{_ENG}.time.sleep"), + patch(f"{_ENG}.random.random", return_value=0.0), + ): + sb._wait_for_running() + assert sb._ecs.describe_tasks.call_count == 2 + + +def test_wait_for_running_stopped_raises(): + sb = _make_sandbox() + sb._task_arn = "arn:task" + sb._ecs = MagicMock() + sb._ecs.describe_tasks.return_value = {"tasks": [{"lastStatus": "STOPPED", "stoppedReason": "OOM"}]} + with patch(f"{_ENG}.time.monotonic", return_value=0.0): + with pytest.raises(RuntimeError, match="ECS task stopped: OOM"): + sb._wait_for_running() + + +def test_wait_for_running_no_tasks_raises(): + sb = _make_sandbox() + sb._task_arn = "arn:task" + sb._ecs = MagicMock() + sb._ecs.describe_tasks.return_value = {"tasks": []} + with patch(f"{_ENG}.time.monotonic", return_value=0.0): + with pytest.raises(RuntimeError, match="disappeared"): + sb._wait_for_running() + + +def test_wait_for_running_timeout(): + sb = _make_sandbox(startup_timeout_sec=300.0) + sb._task_arn = "arn:task" + sb._ecs = MagicMock() + with patch(f"{_ENG}.time.monotonic", side_effect=[0.0, 999.0]): + with pytest.raises(TimeoutError, match="not RUNNING"): + sb._wait_for_running() + + +def test_wait_for_running_retryable_describe_then_running(): + sb = _make_sandbox() + sb._task_arn = "arn:task" + sb._ecs = MagicMock() + sb._ecs.describe_tasks.side_effect = [ + _client_error("ThrottlingException", "DescribeTasks"), + {"tasks": [{"lastStatus": "RUNNING"}]}, + ] + with ( + patch(f"{_ENG}.time.monotonic", return_value=0.0), + patch(f"{_ENG}.time.sleep"), + patch(f"{_ENG}.random.random", return_value=0.0), + ): + sb._wait_for_running() + + +def test_wait_for_running_non_retryable_describe_raises(): + sb = _make_sandbox() + sb._task_arn = "arn:task" + sb._ecs = MagicMock() + sb._ecs.describe_tasks.side_effect = ValueError("fatal") + with patch(f"{_ENG}.time.monotonic", return_value=0.0): + with pytest.raises(ValueError, match="fatal"): + sb._wait_for_running() + + +# ── EcsFargateSandbox: _get_task_public_ip ──────────────────────────── + + +def _eni_task(detail_name: str, value: str): + return { + "tasks": [ + {"attachments": [{"type": "ElasticNetworkInterface", "details": [{"name": detail_name, "value": value}]}]} + ] + } + + +def test_get_task_public_ip_returns_public(): + sb = _make_sandbox() + sb._task_arn = "arn:task" + sb._ecs = MagicMock() + sb._ecs.describe_tasks.return_value = _eni_task("networkInterfaceId", "eni-1") + sb._ec2 = MagicMock() + sb._ec2.describe_network_interfaces.return_value = { + "NetworkInterfaces": [{"Association": {"PublicIp": "1.2.3.4"}}] + } + assert sb._get_task_public_ip() == "1.2.3.4" + + +def test_get_task_public_ip_returns_private_when_no_public(): + sb = _make_sandbox() + sb._task_arn = "arn:task" + sb._ecs = MagicMock() + sb._ecs.describe_tasks.return_value = _eni_task("networkInterfaceId", "eni-1") + sb._ec2 = MagicMock() + sb._ec2.describe_network_interfaces.return_value = {"NetworkInterfaces": [{"PrivateIpAddress": "10.0.0.5"}]} + assert sb._get_task_public_ip() == "10.0.0.5" + + +def test_get_task_public_ip_uses_private_detail_without_eni(): + sb = _make_sandbox() + sb._task_arn = "arn:task" + sb._ecs = MagicMock() + sb._ecs.describe_tasks.return_value = _eni_task("privateIPv4Address", "10.9.9.9") + sb._ec2 = MagicMock() + assert sb._get_task_public_ip() == "10.9.9.9" + sb._ec2.describe_network_interfaces.assert_not_called() + + +def test_get_task_public_ip_no_ip_details_then_succeeds(): + sb = _make_sandbox() + sb._task_arn = "arn:task" + sb._ecs = MagicMock() + # First poll: an attachment with neither ENI id nor a usable private IP. + no_ip = {"tasks": [{"attachments": [{"type": "Other", "details": [{"name": "subnetId", "value": "sn"}]}]}]} + sb._ecs.describe_tasks.side_effect = [no_ip, _eni_task("privateIPv4Address", "10.2.2.2")] + sb._ec2 = MagicMock() + with patch(f"{_ENG}.time.sleep"), patch(f"{_ENG}.random.random", return_value=0.0): + assert sb._get_task_public_ip() == "10.2.2.2" + + +def test_get_task_public_ip_retries_on_missing_ip_then_succeeds(): + sb = _make_sandbox() + sb._task_arn = "arn:task" + sb._ecs = MagicMock() + sb._ecs.describe_tasks.return_value = _eni_task("networkInterfaceId", "eni-1") + sb._ec2 = MagicMock() + sb._ec2.describe_network_interfaces.side_effect = [ + {"NetworkInterfaces": [{}]}, # no IP yet + {"NetworkInterfaces": [{"Association": {"PublicIp": "5.6.7.8"}}]}, + ] + with patch(f"{_ENG}.time.sleep"), patch(f"{_ENG}.random.random", return_value=0.0): + assert sb._get_task_public_ip() == "5.6.7.8" + + +def test_get_task_public_ip_retryable_error_then_succeeds(): + sb = _make_sandbox() + sb._task_arn = "arn:task" + sb._ecs = MagicMock() + sb._ecs.describe_tasks.side_effect = [ + _client_error("ThrottlingException", "DescribeTasks"), + _eni_task("privateIPv4Address", "10.1.1.1"), + ] + sb._ec2 = MagicMock() + with patch(f"{_ENG}.time.sleep"), patch(f"{_ENG}.random.random", return_value=0.0): + assert sb._get_task_public_ip() == "10.1.1.1" + + +def test_get_task_public_ip_exhausts_retries(): + sb = _make_sandbox() + sb._task_arn = "arn:task" + sb._ecs = MagicMock() + sb._ecs.describe_tasks.return_value = {"tasks": []} + sb._ec2 = MagicMock() + with patch(f"{_ENG}.time.sleep"), patch(f"{_ENG}.random.random", return_value=0.0): + with pytest.raises(RuntimeError, match="Task not found"): + sb._get_task_public_ip() + + +# ── EcsFargateSandbox: _wait_for_ssh_ready ──────────────────────────── + + +def test_wait_for_ssh_ready_success(): + with ( + patch(f"{_ENG}.socket.socket", return_value=_connect_sock(recv=b"SSH-2.0-OpenSSH")), + patch(f"{_ENG}.time.monotonic", side_effect=[0.0, 0.1]), + patch(f"{_ENG}.time.sleep"), + ): + engine.EcsFargateSandbox._wait_for_ssh_ready("1.2.3.4", 2222, timeout=10.0) + + +def test_wait_for_ssh_ready_timeout(): + with ( + patch(f"{_ENG}.socket.socket", return_value=_connect_sock(raise_oserror=True)), + patch(f"{_ENG}.time.monotonic", side_effect=[0.0, 0.1, 100.0]), + patch(f"{_ENG}.time.sleep"), + ): + with pytest.raises(TimeoutError, match="SSH not ready"): + engine.EcsFargateSandbox._wait_for_ssh_ready("1.2.3.4", 2222, timeout=10.0) + + +# ── EcsFargateSandbox: _open_tunnel ─────────────────────────────────── + + +def test_open_tunnel_exec_server_mode(): + sb = _make_sandbox() + sb._task_ip = "1.2.3.4" + sb._ssh_key_file = "/tmp/key" + sb._outside_endpoint_routing = engine._OutsideEndpointRouting.for_exec_server( + [engine.OutsideEndpoint("http://10.0.0.1:4000/v1", "MODEL")], _exec_sidecar() + ) + tunnel = MagicMock() + with patch(f"{_ENG}.SshTunnel", return_value=tunnel) as ctor: + sb._open_tunnel(_exec_sidecar(exec_server_port=5000)) + kwargs = ctor.call_args.kwargs + assert kwargs["host"] == "1.2.3.4" and kwargs["port"] == 2222 + assert kwargs["forward_port"] == 5000 + assert any(s.endswith(":10.0.0.1:4000") for s in kwargs["reverses"]) + tunnel.open.assert_called_once() + + +def test_open_tunnel_agent_server_mode(): + sb = _make_sandbox(container_port=8080) + sb._task_ip = "1.2.3.4" + sb._ssh_key_file = "/tmp/key" + sb._ssh_tunnel_port = 7000 + sb._outside_endpoint_routing = engine._OutsideEndpointRouting.for_agent_server( + [engine.OutsideEndpoint("http://10.0.0.2:9000/v1", "MODEL")] + ) + tunnel = MagicMock() + with patch(f"{_ENG}.SshTunnel", return_value=tunnel) as ctor, patch(f"{_ENG}._free_port", return_value=15000): + sb._open_tunnel(_agent_sidecar()) + kwargs = ctor.call_args.kwargs + assert kwargs["forwards"] == ["15000:localhost:8080"] + assert kwargs["reverses"] == ["7000:10.0.0.2:9000"] + assert kwargs["local_port_override"] == 15000 + tunnel.open.assert_called_once() + + +def test_open_tunnel_agent_server_requires_container_port(): + sb = _make_sandbox(container_port=None) + sb._task_ip = "1.2.3.4" + sb._ssh_key_file = "/tmp/key" + sb._ssh_tunnel_port = 7000 + sb._outside_endpoint_routing = engine._OutsideEndpointRouting.for_agent_server( + [engine.OutsideEndpoint("http://10.0.0.2:9000/v1", "MODEL")] + ) + with patch(f"{_ENG}._free_port", return_value=15000): + with pytest.raises(ValueError, match="container_port is required"): + sb._open_tunnel(_agent_sidecar()) + + +# ── EcsFargateSandbox: cleanup ──────────────────────────────────────── + + +def test_cleanup_closes_tunnel_stops_task_removes_key(): + sb = _make_sandbox() + sb._ssh_tunnel = MagicMock() + sb._task_arn = "arn:task" + sb._ecs = MagicMock() + sb._ssh_key_file = "/tmp/key" + with patch(f"{_ENG}.os.remove") as rm: + sb._cleanup() + assert sb._ssh_tunnel is None + sb._ecs.stop_task.assert_called_once() + rm.assert_called_once_with("/tmp/key") + assert sb._ssh_key_file is None + + +def test_cleanup_swallows_all_errors(): + sb = _make_sandbox() + tunnel = MagicMock() + tunnel.close.side_effect = RuntimeError("close boom") + sb._ssh_tunnel = tunnel + sb._task_arn = "arn:task" + sb._ecs = MagicMock() + sb._ecs.stop_task.side_effect = RuntimeError("stop boom") + sb._ssh_key_file = "/tmp/key" + with patch(f"{_ENG}.os.remove", side_effect=OSError("rm boom")): + sb._cleanup() # must not raise + assert sb._ssh_tunnel is None + + +def test_sync_stop_idempotent(): + sb = _make_sandbox() + with patch.object(engine.EcsFargateSandbox, "_cleanup") as cleanup: + sb._sync_stop() + assert sb._stopped is True + sb._sync_stop() # second call short-circuits + cleanup.assert_called_once() + + +def test_require_exec_client_raises_in_agent_mode(): + sb = _make_sandbox() + with pytest.raises(RuntimeError, match="require exec-server mode"): + sb._require_exec_client() + _attach_exec_client(sb) + sb._require_exec_client() # no error once set + + +def test_register_and_unregister_for_cleanup(): + sb = _make_sandbox() + with patch(f"{_ENG}.atexit.register") as reg: + sb._register_for_cleanup() + assert id(sb) in engine._active_sandboxes + reg.assert_called_once() + sb._unregister_from_cleanup() + assert id(sb) not in engine._active_sandboxes + + +def test_emergency_cleanup_invokes_sync_stop(): + good = MagicMock() + bad = MagicMock() + bad._sync_stop.side_effect = RuntimeError("boom") + engine._active_sandboxes[1] = good + engine._active_sandboxes[2] = bad + engine._emergency_cleanup() # bad's error is swallowed + good._sync_stop.assert_called_once() + bad._sync_stop.assert_called_once() + + +async def test_upload_via_s3_packs_and_extracts(tmp_path): + f = tmp_path / "payload.txt" + f.write_text("data") + sb = _make_sandbox(s3_bucket="bkt", s3_prefix="pfx") + ec = _attach_exec_client(sb, engine.ExecResult("ok\n", "", 0)) + s3 = MagicMock() + s3.generate_presigned_url.return_value = "https://signed/url" + with _patch_aws(s3=s3): + await sb._upload_via_s3([f], "/dest") + s3.put_object.assert_called_once() + dl_cmd = ec.exec.await_args.args[0] + assert "tar xzf" in dl_cmd and "/dest" in dl_cmd + + +async def test_upload_via_s3_packs_a_directory(tmp_path): + d = tmp_path / "dir" + (d / "sub").mkdir(parents=True) + (d / "sub" / "f.txt").write_text("nested") + sb = _make_sandbox(s3_bucket="bkt") + _attach_exec_client(sb, engine.ExecResult("ok\n", "", 0)) + s3 = MagicMock() + s3.generate_presigned_url.return_value = "https://signed/url" + with _patch_aws(s3=s3): + await sb._upload_via_s3([d], "/dest") + s3.put_object.assert_called_once() + # The directory branch tars children with paths relative to the directory. + import tarfile + + body = s3.put_object.call_args.kwargs["Body"] + with tarfile.open(fileobj=io.BytesIO(body), mode="r:gz") as tar: + assert "sub/f.txt" in tar.getnames() + + +async def test_upload_via_s3_requires_bucket(): + sb = _make_sandbox(s3_bucket=None) + _attach_exec_client(sb) + with pytest.raises(ValueError, match="s3_bucket is required"): + await sb._upload_via_s3([Path("/x")], "/dest") + + +async def test_upload_via_s3_raises_when_extraction_fails(tmp_path): + f = tmp_path / "payload.txt" + f.write_text("data") + sb = _make_sandbox(s3_bucket="bkt") + _attach_exec_client(sb, engine.ExecResult("nope", "tar error", 2)) + s3 = MagicMock() + s3.generate_presigned_url.return_value = "https://signed/url" + with _patch_aws(s3=s3): + with pytest.raises(RuntimeError, match="S3 upload extraction failed"): + await sb._upload_via_s3([f], "/dest") + + +# ── EcsFargateSandbox: properties ───────────────────────────────────── + + +def test_properties_reflect_state(): + spec = engine.SandboxSpec(image="img:1") + sb = _make_sandbox(spec=spec) + assert sb.spec is spec + assert sb.is_running is False + sb._started = True + assert sb.is_running is True + sb._stopped = True + assert sb.is_running is False + + sb._task_arn = "arn:task" + assert sb.task_arn == "arn:task" + sb._task_ip = "1.2.3.4" + assert sb.container_ip == "1.2.3.4" + sb._ssh_tunnel_port = 7000 + assert sb.model_tunnel_port == 7000 + + ec = _attach_exec_client(sb) + assert sb.exec_client is ec + + +def test_local_port_property_paths(): + sb = _make_sandbox() + assert sb.local_port is None # no tunnel + tunnel = MagicMock() + tunnel.local_port = 18000 + sb._ssh_tunnel = tunnel + assert sb.ssh_tunnel is tunnel + assert sb.local_port == 18000 + # A not-yet-open tunnel raises RuntimeError, surfaced as None. + type(tunnel).local_port = property(lambda self: (_ for _ in ()).throw(RuntimeError("not open"))) + assert sb.local_port is None + + +def test_resolved_endpoint_url_and_resolve_outside_endpoint(): + sb = _make_sandbox() + sb._outside_endpoint_routing = engine._OutsideEndpointRouting.for_agent_server( + [engine.OutsideEndpoint("http://model:7000/v1", "MODEL")] + ) + assert sb.resolved_endpoint_url("MODEL") == "http://127.0.0.1:7000/v1" + assert sb.resolve_outside_endpoint("http://anything/v1").startswith("http://127.0.0.1:7000") + + +# ── EcsFargateSandbox: async start/stop ─────────────────────────────── + + +async def test_start_runs_do_start_and_marks_started(): + sb = _make_sandbox(ssh_sidecar=_exec_sidecar()) + with patch.object(engine.EcsFargateSandbox, "_do_start") as do_start: + await sb.start() + do_start.assert_called_once() + assert sb._started is True + + +async def test_start_is_idempotent(): + sb = _make_sandbox(ssh_sidecar=_exec_sidecar()) + sb._started = True + with patch.object(engine.EcsFargateSandbox, "_do_start") as do_start: + await sb.start() + do_start.assert_not_called() + + +async def test_start_agent_mode_validates_endpoints(): + sb = _make_sandbox(ssh_sidecar=_agent_sidecar()) + with patch.object(engine.EcsFargateSandbox, "_do_start") as do_start: + with pytest.raises(ValueError, match="requires OutsideEndpoint"): + await sb.start(outside_endpoints=[]) + do_start.assert_not_called() + + +async def test_start_failure_cleans_up_and_reraises(): + sb = _make_sandbox(ssh_sidecar=_exec_sidecar()) + ec = _attach_exec_client(sb) + with ( + patch.object(engine.EcsFargateSandbox, "_do_start", side_effect=RuntimeError("startup failed")), + patch.object(engine.EcsFargateSandbox, "_cleanup") as cleanup, + ): + with pytest.raises(RuntimeError, match="startup failed"): + await sb.start() + ec.close.assert_awaited_once() + cleanup.assert_called_once() + assert sb._started is False + + +async def test_stop_closes_client_and_cleans_up(): + sb = _make_sandbox(ssh_sidecar=_exec_sidecar()) + ec = _attach_exec_client(sb) + with ( + patch.object(engine.EcsFargateSandbox, "_cleanup") as cleanup, + patch.object(engine.EcsFargateSandbox, "_unregister_from_cleanup") as unreg, + ): + await sb.stop() + assert sb._stopped is True + ec.close.assert_awaited_once() + cleanup.assert_called_once() + unreg.assert_called_once() + + +async def test_stop_is_idempotent(): + sb = _make_sandbox() + sb._stopped = True + with patch.object(engine.EcsFargateSandbox, "_cleanup") as cleanup: + await sb.stop() + cleanup.assert_not_called() + + +async def test_aenter_aexit_delegate_to_start_stop(): + sb = _make_sandbox(ssh_sidecar=_exec_sidecar()) + with ( + patch.object(engine.EcsFargateSandbox, "start", AsyncMock()) as start, + patch.object(engine.EcsFargateSandbox, "stop", AsyncMock()) as stop, + ): + async with sb as ctx: + assert ctx is sb + start.assert_awaited_once() + stop.assert_awaited_once() + + +# ── EcsFargateSandbox: async exec ───────────────────────────────────── + + +async def test_exec_passes_timeout_as_int(): + sb = _make_sandbox(ssh_sidecar=_exec_sidecar()) + ec = _attach_exec_client(sb, engine.ExecResult("hi", "", 0)) + out = await sb.exec("echo hi", timeout_sec=42.7) + assert out.return_code == 0 + assert ec.exec.await_args.args[0] == "echo hi" + assert ec.exec.await_args.kwargs["timeout"] == 42 + + +async def test_exec_wraps_env_cwd_and_user_string(): + sb = _make_sandbox(ssh_sidecar=_exec_sidecar()) + ec = _attach_exec_client(sb) + await sb.exec("run", env={"A": "1"}, cwd="/work", user="appuser") + shell_cmd = ec.exec.await_args.args[0] + assert shell_cmd.startswith("su -s /bin/bash appuser -c ") + assert "cd /work &&" in shell_cmd + assert "export A=1 &&" in shell_cmd + + +async def test_exec_user_int_uses_getent(): + sb = _make_sandbox(ssh_sidecar=_exec_sidecar()) + ec = _attach_exec_client(sb) + await sb.exec("run", user=1000) + assert "getent passwd 1000" in ec.exec.await_args.args[0] + + +async def test_exec_connection_error_with_live_tunnel_reraises(): + sb = _make_sandbox(ssh_sidecar=_exec_sidecar()) + ec = _attach_exec_client(sb) + ec.exec = AsyncMock(side_effect=ConnectionError("dead")) + sb._ssh_tunnel = MagicMock(is_open=True) + with pytest.raises(ConnectionError): + await sb.exec("run") + + +async def test_exec_reconnects_tunnel_then_succeeds(): + sb = _make_sandbox(ssh_sidecar=_exec_sidecar()) + ec = _attach_exec_client(sb) + ec.exec = AsyncMock(side_effect=ConnectionError("dead")) + tunnel = MagicMock(is_open=False, local_port=18000) + sb._ssh_tunnel = tunnel + new_client = MagicMock() + new_client.exec = AsyncMock(return_value=engine.ExecResult("recovered", "", 0)) + with ( + patch.object(engine.EcsFargateSandbox, "reconnect_tunnel", AsyncMock()) as reconnect, + patch(f"{_ENG}.ExecClient", return_value=new_client), + ): + out = await sb.exec("run") + assert out.stdout == "recovered" + reconnect.assert_awaited_once() + ec.close.assert_awaited_once() # old client closed + tunnel.wait_ready.assert_called_once() + + +async def test_exec_reconnect_failure_reraises_original(): + sb = _make_sandbox(ssh_sidecar=_exec_sidecar()) + ec = _attach_exec_client(sb) + ec.exec = AsyncMock(side_effect=ConnectionError("dead")) + sb._ssh_tunnel = MagicMock(is_open=False) + with patch.object(engine.EcsFargateSandbox, "reconnect_tunnel", AsyncMock(side_effect=RuntimeError("no luck"))): + with pytest.raises(ConnectionError): + await sb.exec("run") + + +async def test_exec_requires_exec_client(): + sb = _make_sandbox() + with pytest.raises(RuntimeError, match="require exec-server mode"): + await sb.exec("run") + + +# ── EcsFargateSandbox: async upload/download ────────────────────────── + + +async def test_upload_small_file_uses_exec_client(tmp_path): + f = tmp_path / "small.txt" + f.write_text("tiny") + sb = _make_sandbox(ssh_sidecar=_exec_sidecar()) + ec = _attach_exec_client(sb) + await sb.upload(f, "/remote/small.txt") + ec.upload.assert_awaited_once() + assert ec.upload.await_args.args[0] == "/remote/small.txt" + + +async def test_upload_directory_recurses(tmp_path): + d = tmp_path / "dir" + (d / "sub").mkdir(parents=True) + (d / "a.txt").write_text("a") + (d / "sub" / "b.txt").write_text("b") + sb = _make_sandbox(ssh_sidecar=_exec_sidecar()) + ec = _attach_exec_client(sb) + await sb.upload(d, "/remote") + uploaded = {call.args[0] for call in ec.upload.await_args_list} + assert uploaded == {"/remote/a.txt", "/remote/sub/b.txt"} + + +async def test_upload_large_file_uses_s3(tmp_path): + f = tmp_path / "big.bin" + f.write_bytes(b"x" * (512 * 1024 + 1)) + sb = _make_sandbox(ssh_sidecar=_exec_sidecar(), s3_bucket="bkt") + _attach_exec_client(sb) + with patch.object(engine.EcsFargateSandbox, "_upload_via_s3", AsyncMock()) as via_s3: + await sb.upload(f, "/remote/big.bin") + via_s3.assert_awaited_once() + assert via_s3.await_args.args[1] == "/remote" + + +async def test_download_writes_bytes(tmp_path): + sb = _make_sandbox(ssh_sidecar=_exec_sidecar()) + _attach_exec_client(sb) + dest = tmp_path / "nested" / "out.bin" + await sb.download("/remote/out.bin", dest) + assert dest.read_bytes() == b"payload" + + +# ── EcsFargateSandbox: reconnect_tunnel ─────────────────────────────── + + +async def test_reconnect_tunnel_rejects_stopped_or_unstarted(): + sb = _make_sandbox(ssh_sidecar=_exec_sidecar()) + with pytest.raises(RuntimeError, match="stopped/unstarted"): + await sb.reconnect_tunnel() # not started + sb._started = True + sb._stopped = True + with pytest.raises(RuntimeError, match="stopped/unstarted"): + await sb.reconnect_tunnel() + + +async def test_reconnect_tunnel_noop_without_sidecar(): + sb = _make_sandbox(ssh_sidecar=None) + sb._started = True + with patch.object(engine.EcsFargateSandbox, "_open_tunnel") as open_tunnel: + await sb.reconnect_tunnel() + open_tunnel.assert_not_called() + + +async def test_reconnect_tunnel_reopens(): + sb = _make_sandbox(ssh_sidecar=_exec_sidecar()) + sb._started = True + old_tunnel = MagicMock() + sb._ssh_tunnel = old_tunnel + with patch.object(engine.EcsFargateSandbox, "_open_tunnel") as open_tunnel: + await sb.reconnect_tunnel() + old_tunnel.close.assert_called_once() + open_tunnel.assert_called_once() + + +# ── EcsFargateSandbox: _do_start orchestration ──────────────────────── + + +@contextlib.contextmanager +def _do_start_seams(tunnel=None, exec_client=None): + """Patch every AWS/SSH seam so the *real* ``_do_start`` runs offline.""" + tunnel = tunnel or MagicMock() + tunnel.local_port = 19000 + ec = exec_client or MagicMock() + with ( + patch.object(engine.EcsFargateSandbox, "_init_aws_clients"), + patch.object(engine.EcsFargateSandbox, "_register_task_definition", return_value="td-arn"), + patch.object(engine.EcsFargateSandbox, "_run_task", return_value="task-arn"), + patch.object(engine.EcsFargateSandbox, "_register_for_cleanup"), + patch.object(engine.EcsFargateSandbox, "_wait_for_running"), + patch.object(engine.EcsFargateSandbox, "_get_task_public_ip", return_value="1.2.3.4"), + patch.object(engine.EcsFargateSandbox, "_wait_for_ssh_ready"), + patch(f"{_ENG}.download_secret_to_file", return_value="/tmp/key"), + patch(f"{_ENG}.download_secret_to_string", return_value="ssh-rsa fake"), + patch(f"{_ENG}.build_ssh_sidecar_container", return_value={"name": "ssh-tunnel"}), + patch(f"{_ENG}._free_port", return_value=19001), + patch(f"{_ENG}.SshTunnel", return_value=tunnel), + patch(f"{_ENG}.ExecClient", return_value=ec), + ): + yield tunnel, ec + + +def test_do_start_requires_sidecar(): + sb = _make_sandbox(ssh_sidecar=None) + with pytest.raises(ValueError, match="ssh_sidecar must be configured"): + sb._do_start() + + +def test_do_start_requires_key_arns(): + sb = _make_sandbox(ssh_sidecar=_exec_sidecar(private_key_secret_arn="")) + with _do_start_seams(): + with pytest.raises(ValueError, match="private_key_secret_arn and public_key_secret_arn"): + sb._do_start() + + +def test_do_start_exec_server_full_path(): + sb = _make_sandbox(spec=engine.SandboxSpec(image="python:3.12"), ssh_sidecar=_exec_sidecar(exec_server_port=5000)) + with _do_start_seams() as (tunnel, ec): + sb._do_start() + assert sb._task_arn == "task-arn" + assert sb._task_def_arn == "td-arn" + assert sb._task_ip == "1.2.3.4" + assert sb._ssh_tunnel is tunnel + assert sb._exec_client is ec # exec-server mode creates an ExecClient + + +def test_do_start_agent_server_full_path(): + sb = _make_sandbox( + spec=engine.SandboxSpec(image="python:3.12"), + ssh_sidecar=_agent_sidecar(), + container_port=9000, + ) + sb._outside_endpoints = [engine.OutsideEndpoint("http://10.0.0.9:9000/v1", "MODEL")] + with _do_start_seams() as (tunnel, _ec): + sb._do_start() + assert sb._ssh_tunnel is tunnel + assert sb._exec_client is None # agent-server mode does not create an ExecClient + assert sb._ssh_tunnel_port == 9000 # agent tunnel target port + + +def test_do_start_builds_image_when_ecr_and_env_dir(): + sb = _make_sandbox( + spec=engine.SandboxSpec(image="python:3.12", environment_dir="/env"), + ssh_sidecar=_exec_sidecar(exec_server_port=5000), + ecr_repository=_ecr_repo(), + ) + with ( + _do_start_seams(), + patch.object(engine.ImageBuilder, "ensure_image_built", return_value=f"{_ecr_repo()}:built") as build, + ): + sb._do_start() + build.assert_called_once() + assert sb._task_def_arn == "td-arn" diff --git a/tests/unit_tests/test_ecs_fargate_provider.py b/tests/unit_tests/test_ecs_fargate_provider.py index fe73607d5e..bad5731b8e 100644 --- a/tests/unit_tests/test_ecs_fargate_provider.py +++ b/tests/unit_tests/test_ecs_fargate_provider.py @@ -47,8 +47,8 @@ def _provider_config(**overrides): task_role_arn="arn:aws:iam::1234:role/ecsTask", ssh_sidecar={ "sshd_port": 2222, - "public_key_secret_arn": "arn:aws:secretsmanager:us-east-1:1234:secret:pub", - "private_key_secret_arn": "arn:aws:secretsmanager:us-east-1:1234:secret:priv", + "public_key_secret_arn": "arn:aws:secretsmanager:us-east-1:1234:secret:pub", # pragma: allowlist secret + "private_key_secret_arn": "arn:aws:secretsmanager:us-east-1:1234:secret:priv", # pragma: allowlist secret "exec_server_port": 5000, }, ) @@ -118,8 +118,8 @@ def test_config_ssm_autodiscovery_merges_and_yaml_wins(): "security_groups": ["sg-ssm"], "execution_role_arn": "arn:ssm:exec", "ssh_sidecar": { - "public_key_secret_arn": "arn:ssm:pub", - "private_key_secret_arn": "arn:ssm:priv", + "public_key_secret_arn": "arn:ssm:pub", # pragma: allowlist secret + "private_key_secret_arn": "arn:ssm:priv", # pragma: allowlist secret "sshd_port": 52222, }, } @@ -131,7 +131,7 @@ def test_config_ssm_autodiscovery_merges_and_yaml_wins(): assert cfg.cluster == "ssm-cluster" # filled from SSM assert cfg.subnets == ["subnet-override"] # explicit YAML wins assert cfg.execution_role_arn == "arn:ssm:exec" - assert cfg.ssh_sidecar.public_key_secret_arn == "arn:ssm:pub" + assert cfg.ssh_sidecar.public_key_secret_arn == "arn:ssm:pub" # pragma: allowlist secret assert cfg.ssh_sidecar.sshd_port == 52222 @@ -270,8 +270,13 @@ def _resolve_image_for(image, **cfg_overrides): def test_resolve_image_routes_bare_name_to_ecr_mirror(): # Bare/public names are mirrored to the ECR tag, never pulled directly. ecr = "463701203462.dkr.ecr.us-east-1.amazonaws.com/harbor-us-east-1" - resolved = _resolve_image_for("docker.io/swebench/sweb.eval.x86_64.astropy_1776_astropy-12907:latest", ecr_repository=ecr) - assert resolved == f"{ecr}:{engine._sanitize_id('docker.io/swebench/sweb.eval.x86_64.astropy_1776_astropy-12907:latest')}" + resolved = _resolve_image_for( + "docker.io/swebench/sweb.eval.x86_64.astropy_1776_astropy-12907:latest", ecr_repository=ecr + ) + assert ( + resolved + == f"{ecr}:{engine._sanitize_id('docker.io/swebench/sweb.eval.x86_64.astropy_1776_astropy-12907:latest')}" + ) assert "docker.io" not in resolved.split(":", 1)[1] # origin registry not used for the pull @@ -285,9 +290,7 @@ def test_resolve_image_passes_through_existing_ecr_ref(): def test_resolve_image_template_takes_precedence(): ecr = "463701203462.dkr.ecr.us-east-1.amazonaws.com/harbor-us-east-1" - resolved = _resolve_image_for( - "anything", ecr_repository=ecr, image_template="{task_id}-built" - ) + resolved = _resolve_image_for("anything", ecr_repository=ecr, image_template="{task_id}-built") assert resolved == "anything-built" @@ -301,7 +304,7 @@ def test_generate_mirror_buildspec_pulls_tags_and_pushes(): cfg = engine.EcsFargateConfig( region="us-east-1", ecr_repository="123.dkr.ecr.us-east-1.amazonaws.com/mirror", - dockerhub_secret_arn="arn:aws:secretsmanager:us-east-1:123:secret:dh", + dockerhub_secret_arn="arn:aws:secretsmanager:us-east-1:123:secret:dh", # pragma: allowlist secret ) src = "docker.io/swebench/sweb.eval.x86_64.astropy_1776_astropy-12907:latest" ecr_url = f"{cfg.ecr_repository}:{engine._sanitize_id(src)}" From 35c204b9b72da586c4ea482b664f05846ce64464 Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Wed, 24 Jun 2026 15:36:24 +0200 Subject: [PATCH 04/11] fix(sandbox): address review feedback on #1645 Per @hemildesai's review: - mini_swe_agent_2: reword the conda-activation comment to be provider-agnostic (the sandbox exec shell is non-login across apptainer/docker/ECS); the command itself was already generic, so no behavior change. - pyproject: fold boto3 into the `sandbox` extra and drop the separate `sandbox-ecs` extra; update the mini_swe_agent_2 requirement, test comments, and uv.lock to match. Signed-off-by: Michal Bien --- pyproject.toml | 2 -- responses_api_agents/mini_swe_agent_2/requirements.txt | 2 +- .../mini_swe_agent_2/sandbox_environment.py | 9 +++++---- tests/unit_tests/test_ecs_fargate_engine.py | 6 +++--- 4 files changed, 9 insertions(+), 10 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 1cb3ad4f6d..e9e074dbca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -248,9 +248,7 @@ sandbox = [ # Updated: Sat May 16, 2026 with opensandbox>=0.1.9 # License: Apache 2.0 "opensandbox>=0.1.9", -] -sandbox-ecs = [ # boto3: AWS SDK used by the ECS Fargate sandbox provider for ECS/EC2/ECR/ # CodeBuild/S3/SSM/Secrets Manager calls. # Updated: Tue Jun 03, 2026 with boto3>=1.34 diff --git a/responses_api_agents/mini_swe_agent_2/requirements.txt b/responses_api_agents/mini_swe_agent_2/requirements.txt index 22fb4a3220..2f82dbd4c7 100644 --- a/responses_api_agents/mini_swe_agent_2/requirements.txt +++ b/responses_api_agents/mini_swe_agent_2/requirements.txt @@ -1,3 +1,3 @@ --e nemo-gym[dev,sandbox,sandbox-ecs] @ ../../ +-e nemo-gym[dev,sandbox] @ ../../ mini-swe-agent==2.1.0 swebench==4.1.0 diff --git a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py index e18501eb32..d9d2f7e009 100644 --- a/responses_api_agents/mini_swe_agent_2/sandbox_environment.py +++ b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py @@ -124,10 +124,11 @@ def _command(self, command: str) -> str: if not self.config.activate_conda or not self.config.conda_env: return command quoted_env = shlex.quote(self.config.conda_env) - # Resolve conda from common install roots before activating: the ECS exec shell - # is non-login, so `conda` isn't on PATH and `conda info --base` can't be relied - # on. The grouped loop (not an `&&` chain) keeps a missing root from aborting the - # command; cwd is handled by exec(cwd=...), so we don't `cd` here. + # Resolve conda from common install roots before activating. The sandbox exec + # shell is non-login (apptainer/docker/ECS alike), so `conda` may not be on PATH + # and `conda info --base` can't be relied on. The grouped loop (not an `&&` chain) + # keeps a missing root from aborting the command; cwd is handled by exec(cwd=...), + # so we don't `cd` here. return ( '{ for __base in /opt/miniconda3 /opt/conda "$HOME/miniconda3" ' '"$(command -v conda >/dev/null 2>&1 && conda info --base 2>/dev/null)"; do ' diff --git a/tests/unit_tests/test_ecs_fargate_engine.py b/tests/unit_tests/test_ecs_fargate_engine.py index 79a83d95e4..941e843e27 100644 --- a/tests/unit_tests/test_ecs_fargate_engine.py +++ b/tests/unit_tests/test_ecs_fargate_engine.py @@ -46,7 +46,7 @@ class ClientError(Exception): """Local stand-in for ``botocore.exceptions.ClientError``. - CI's base venv has no boto3/botocore (it's the ``sandbox-ecs`` extra), so we + CI's base venv has no boto3/botocore (it's the ``sandbox`` extra), so we don't import it. engine.py only references ``ClientError`` via ``_require_aws_sdks()`` — which ``_patch_aws`` patches to return this class — so engine's ``except ClientError`` catches these instances unchanged. @@ -65,7 +65,7 @@ def __init__(self, error_response, operation_name="Operation"): def _clear_engine_module_state(request): """Reset process-global engine caches and stub the AWS SDK accessor. - CI's base venv has no boto3/botocore (it's the sandbox-ecs extra), so patch + CI's base venv has no boto3/botocore (it's the sandbox extra), so patch ``engine._require_aws_sdks`` to hand back mocks + the local fake ``ClientError`` for every test — so engine's ``except ClientError`` catches the fakes the tests raise without real boto3. Tests needing specific clients override it via @@ -168,7 +168,7 @@ def _attach_exec_client(sb: engine.EcsFargateSandbox, result: engine.ExecResult def test_require_aws_sdks_returns_triple(): - pytest.importorskip("boto3") # real boto3/botocore only with the sandbox-ecs extra + pytest.importorskip("boto3") # real boto3/botocore only with the sandbox extra from botocore.exceptions import ClientError as RealClientError boto3, config_cls, client_error = engine._require_aws_sdks() From 69985585c8640a46e69a23da8e7049f59987f36c Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Fri, 26 Jun 2026 11:00:42 +0200 Subject: [PATCH 05/11] fix(sandbox): address ECS review (phase A: API drift, build dedup, AWS bugs) P0 fixes from @ananthsub's review: - Sandbox API drift: config sandbox_spec.timeout_s -> ttl_s; drop stale sandbox_environment_kwargs.delete; provider.close() drops the removed `delete` kwarg; provider maps spec.resources (cpu/memory_mib/disk_gib) + spec.ttl_s onto the ECS config (rejects GPU, which Fargate lacks). - Image-build dedup: wait on the in-flight event OUTSIDE cls._lock (one waiter no longer serializes all builds) and verify via ECR on wake, so a failed build no longer returns a URL for an image that was never pushed. - Fargate ephemeral storage: omit the field unless >=21 GiB is requested and validate 21-200 (an explicit 20 is rejected by RegisterTaskDefinition). - Shell quoting: shlex.quote env values + cwd in exec (injection / breakage). - {task_ip}: fail loudly instead of baking a literal placeholder into container env (the IP is only known after the task starts). Tests: rewrite the build/mirror dedup tests to the corrected contract + add failure-path, ephemeral-range, task_ip-guard, and resource/ttl-mapping tests. Signed-off-by: Michal Bien --- .../sandbox/providers/ecs_fargate/engine.py | 167 ++++++++++-------- .../sandbox/providers/ecs_fargate/provider.py | 32 +++- .../configs/mini_swe_agent_ecs_fargate.yaml | 3 +- tests/unit_tests/test_ecs_fargate_engine.py | 108 +++++------ tests/unit_tests/test_ecs_fargate_provider.py | 54 +++++- 5 files changed, 233 insertions(+), 131 deletions(-) diff --git a/nemo_gym/sandbox/providers/ecs_fargate/engine.py b/nemo_gym/sandbox/providers/ecs_fargate/engine.py index 49d19609c1..d164b58511 100644 --- a/nemo_gym/sandbox/providers/ecs_fargate/engine.py +++ b/nemo_gym/sandbox/providers/ecs_fargate/engine.py @@ -481,6 +481,22 @@ def _free_port() -> int: return s.getsockname()[1] +def _ephemeral_storage_block(requested: Any) -> dict[str, int] | None: + """Task-def ``ephemeralStorage`` block for an explicit request, or ``None`` to take Fargate's + implicit 20 GiB default. + + Fargate accepts only 21-200 GiB; an explicit ``20`` is rejected by RegisterTaskDefinition (20 is + valid only as the omitted default), so we omit the field unless a larger size is requested and + validate the range. + """ + if not requested: + return None + size = int(requested) + if not 21 <= size <= 200: + raise ValueError(f"ephemeral storage must be between 21 and 200 GiB (got {size})") + return {"sizeInGiB": size} + + def download_secret_to_file(secret_arn: str, region: str | None = None) -> str: """Fetch a Secrets Manager secret → temp file (mode 0600).""" key_material = download_secret_to_string(secret_arn, region=region) @@ -1084,48 +1100,69 @@ def docker_push_to_ecr(local_image: str, ecr_repository: str, tag: str) -> str: return ecr_url @classmethod - def ensure_image_built(cls, *, cfg: EcsFargateConfig, environment_name: str, force_build: bool = False) -> str: - ecr_repo = cfg.ecr_repository - env_dir = cfg.environment_dir - if not ecr_repo or not env_dir: - raise ValueError("ecr_repository and environment_dir are required for image building") - tag = cls.get_ecr_image_tag(env_dir, environment_name) - image_url = f"{ecr_repo}:{tag}" + def _run_deduped_build( + cls, *, cfg: EcsFargateConfig, tag: str, image_url: str, force: bool, build: Callable[[], None] + ) -> str: + """Run ``build`` for ``tag`` at most once across threads, returning the ECR ``image_url``. - # Dedup: check if another thread is already building this tag - with cls._lock: - if tag in cls._inflight_builds: - event = cls._inflight_builds[tag] - event.wait() - return image_url - if not force_build and cls.image_exists_in_ecr(ecr_repo, tag, cfg.region): + Concurrent callers for the same tag block on a shared in-flight event **outside** + ``cls._lock`` (so one waiter can't stall every other tag's build) and then confirm the + result via ECR rather than the event -- a failed build sets the event too, so a waiter must + not return a URL for an image that was never pushed. + """ + ecr_repo = cfg.ecr_repository + if not force and cls.image_exists_in_ecr(ecr_repo, tag, cfg.region): logger.info("ECR cache hit — skipping build: %s", image_url) return image_url - # Register as builder - event = threading.Event() - with cls._lock: - if tag in cls._inflight_builds: - cls._inflight_builds[tag].wait() - return image_url - cls._inflight_builds[tag] = event with cls._lock: + event = cls._inflight_builds.get(tag) + is_builder = event is None + if is_builder: + event = threading.Event() + cls._inflight_builds[tag] = event if cls._build_semaphore is None or cls._build_semaphore_size != cfg.build_parallelism: cls._build_semaphore = threading.Semaphore(cfg.build_parallelism) cls._build_semaphore_size = cfg.build_parallelism + semaphore = cls._build_semaphore + + if not is_builder: + event.wait() # wait outside the lock so one waiter can't block other tags' builds + if cls.image_exists_in_ecr(ecr_repo, tag, cfg.region): + return image_url + raise RuntimeError(f"Concurrent build for {image_url} failed; image is not in ECR") + try: - cls._build_semaphore.acquire() # type: ignore[union-attr] + semaphore.acquire() try: - if not force_build and cls.image_exists_in_ecr(ecr_repo, tag, cfg.region): + if not force and cls.image_exists_in_ecr(ecr_repo, tag, cfg.region): return image_url - cls._build_and_push(cfg=cfg, environment_name=environment_name, tag=tag, image_url=image_url) + build() + return image_url finally: - cls._build_semaphore.release() # type: ignore[union-attr] + semaphore.release() finally: - event.set() with cls._lock: cls._inflight_builds.pop(tag, None) - return image_url + event.set() + + @classmethod + def ensure_image_built(cls, *, cfg: EcsFargateConfig, environment_name: str, force_build: bool = False) -> str: + ecr_repo = cfg.ecr_repository + env_dir = cfg.environment_dir + if not ecr_repo or not env_dir: + raise ValueError("ecr_repository and environment_dir are required for image building") + tag = cls.get_ecr_image_tag(env_dir, environment_name) + image_url = f"{ecr_repo}:{tag}" + return cls._run_deduped_build( + cfg=cfg, + tag=tag, + image_url=image_url, + force=force_build, + build=lambda: cls._build_and_push( + cfg=cfg, environment_name=environment_name, tag=tag, image_url=image_url + ), + ) @classmethod def ensure_mirrored(cls, *, cfg: EcsFargateConfig, src_image: str, force: bool = False) -> str: @@ -1145,45 +1182,18 @@ def ensure_mirrored(cls, *, cfg: EcsFargateConfig, src_image: str, force: bool = tag = _sanitize_id(src_image) image_url = f"{ecr_repo}:{tag}" - with cls._lock: - if tag in cls._inflight_builds: - cls._inflight_builds[tag].wait() - return image_url - if not force and cls.image_exists_in_ecr(ecr_repo, tag, cfg.region): - logger.info("ECR mirror hit — skipping pull: %s", image_url) - return image_url + def _mirror() -> None: + logger.info("Mirroring %s -> %s via CodeBuild ...", src_image, image_url) + buildspec = cls._generate_mirror_buildspec(cfg, src_image, image_url) + cls.run_buildspec_via_codebuild( + cfg=cfg, + buildspec=buildspec, + job_label=f"mirror::{tag}", + timeout_minutes=cfg.codebuild_build_timeout, + ) + logger.info("Mirrored OK: %s -> %s", src_image, image_url) - event = threading.Event() - with cls._lock: - if tag in cls._inflight_builds: - cls._inflight_builds[tag].wait() - return image_url - cls._inflight_builds[tag] = event - with cls._lock: - if cls._build_semaphore is None or cls._build_semaphore_size != cfg.build_parallelism: - cls._build_semaphore = threading.Semaphore(cfg.build_parallelism) - cls._build_semaphore_size = cfg.build_parallelism - try: - cls._build_semaphore.acquire() # type: ignore[union-attr] - try: - if not force and cls.image_exists_in_ecr(ecr_repo, tag, cfg.region): - return image_url - logger.info("Mirroring %s -> %s via CodeBuild ...", src_image, image_url) - buildspec = cls._generate_mirror_buildspec(cfg, src_image, image_url) - cls.run_buildspec_via_codebuild( - cfg=cfg, - buildspec=buildspec, - job_label=f"mirror::{tag}", - timeout_minutes=cfg.codebuild_build_timeout, - ) - logger.info("Mirrored OK: %s -> %s", src_image, image_url) - finally: - cls._build_semaphore.release() # type: ignore[union-attr] - finally: - event.set() - with cls._lock: - cls._inflight_builds.pop(tag, None) - return image_url + return cls._run_deduped_build(cfg=cfg, tag=tag, image_url=image_url, force=force, build=_mirror) @staticmethod def _generate_mirror_buildspec(cfg: EcsFargateConfig, src_image: str, ecr_url: str) -> str: @@ -1550,10 +1560,10 @@ async def exec( self._require_exec_client() shell_cmd = command if env: - exports = " ".join(f"{k}={v}" for k, v in env.items()) + exports = " ".join(f"{k}={shlex.quote(v)}" for k, v in env.items()) shell_cmd = f"export {exports} && {shell_cmd}" if cwd: - shell_cmd = f"cd {cwd} && {shell_cmd}" + shell_cmd = f"cd {shlex.quote(cwd)} && {shell_cmd}" if user is not None: if isinstance(user, int): shell_cmd = f'su -s /bin/bash "$(getent passwd {user} | cut -d: -f1)" -c {shlex.quote(shell_cmd)}' @@ -1809,6 +1819,14 @@ def _render_env_value(self, value: str) -> str: value = value.replace("{ssh_tunnel_port}", str(self._ssh_tunnel_port)) if self._task_ip: value = value.replace("{task_ip}", self._task_ip) + elif "{task_ip}" in value: + # _task_ip is only assigned after the task is RUNNING (_get_task_public_ip), which is + # after container env is baked into the task-def / RunTask overrides. Fail loudly rather + # than baking a literal "{task_ip}" that later breaks model egress with a confusing error. + raise ValueError( + "{task_ip} is not available in container env: the task IP is only known after the " + "sandbox starts. Resolve the container's own address at runtime instead." + ) value = value.replace("{image}", self._spec.image or "") return value @@ -1900,12 +1918,12 @@ def _register_from_base( "cpu": str(max(int(base.get("cpu") or "256"), int(cfg.cpu))), "memory": str(max(int(base.get("memory") or "512"), int(cfg.memory))), "containerDefinitions": containers, - "ephemeralStorage": { - "sizeInGiB": max( - (base.get("ephemeralStorage") or {}).get("sizeInGiB", 20), cfg.ephemeral_storage_gib or 20 - ) - }, } + ephemeral = _ephemeral_storage_block( + (base.get("ephemeralStorage") or {}).get("sizeInGiB") or cfg.ephemeral_storage_gib + ) + if ephemeral is not None: + payload["ephemeralStorage"] = ephemeral for k in ("taskRoleArn", "executionRoleArn", "runtimePlatform", "volumes"): if base.get(k) is not None: payload[k] = base[k] @@ -1957,8 +1975,9 @@ def _register_from_scratch( payload["volumes"] = task_volumes if cfg.task_role_arn: payload["taskRoleArn"] = cfg.task_role_arn - if cfg.ephemeral_storage_gib: - payload["ephemeralStorage"] = {"sizeInGiB": cfg.ephemeral_storage_gib} + ephemeral = _ephemeral_storage_block(cfg.ephemeral_storage_gib) + if ephemeral is not None: + payload["ephemeralStorage"] = ephemeral return self._do_register(payload) def _build_efs_volumes(self) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: diff --git a/nemo_gym/sandbox/providers/ecs_fargate/provider.py b/nemo_gym/sandbox/providers/ecs_fargate/provider.py index 6c4013702f..ab364c11db 100644 --- a/nemo_gym/sandbox/providers/ecs_fargate/provider.py +++ b/nemo_gym/sandbox/providers/ecs_fargate/provider.py @@ -34,6 +34,7 @@ SandboxCreateError, SandboxExecResult, SandboxHandle, + SandboxResources, SandboxSpec, SandboxStatus, ) @@ -77,6 +78,31 @@ def _engine_spec(spec: SandboxSpec) -> engine.SandboxSpec: ) +def _apply_spec_overrides(cfg: engine.EcsFargateConfig, spec: SandboxSpec) -> engine.EcsFargateConfig: + """Apply per-sandbox ``SandboxSpec`` requests (readiness/TTL/resources) onto the provider config. + + ``SandboxResources.cpu`` is in vCPUs; Fargate task CPU is in 1024-unit increments. GPU is not + supported on Fargate. + """ + overrides: dict[str, Any] = {} + if spec.ready_timeout_s is not None: + overrides["startup_timeout_sec"] = float(spec.ready_timeout_s) + if spec.ttl_s is not None: + overrides["max_task_lifetime_sec"] = int(spec.ttl_s) + resources = spec.resources + if not isinstance(resources, SandboxResources): + resources = SandboxResources.from_mapping(resources) + if resources.gpu: + raise SandboxCreateError("ECS Fargate does not support GPU sandboxes (spec.resources.gpu)") + if resources.cpu is not None: + overrides["cpu"] = str(int(resources.cpu * 1024)) + if resources.memory_mib is not None: + overrides["memory"] = str(int(resources.memory_mib)) + if resources.disk_gib is not None: + overrides["ephemeral_storage_gib"] = int(resources.disk_gib) + return replace(cfg, **overrides) if overrides else cfg + + class EcsFargateProvider: """Run sandboxes as AWS ECS Fargate tasks behind an SSH sidecar.""" @@ -86,9 +112,7 @@ def __init__(self, **config: Any) -> None: self._cfg = engine_config_from_mapping(config) async def create(self, spec: SandboxSpec) -> SandboxHandle: - cfg = self._cfg - if spec.ready_timeout_s is not None: - cfg = replace(cfg, startup_timeout_sec=float(spec.ready_timeout_s)) + cfg = _apply_spec_overrides(self._cfg, spec) sandbox = engine.EcsFargateSandbox(_engine_spec(spec), ecs_config=cfg) try: await sandbox.start(outside_endpoints=_outside_endpoints(spec)) @@ -136,7 +160,7 @@ async def status(self, handle: SandboxHandle) -> SandboxStatus: sandbox: engine.EcsFargateSandbox = handle.raw return SandboxStatus.RUNNING if sandbox.is_running else SandboxStatus.STOPPED - async def close(self, handle: SandboxHandle, *, delete: bool = False) -> None: + async def close(self, handle: SandboxHandle) -> None: sandbox: engine.EcsFargateSandbox = handle.raw await sandbox.stop() diff --git a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml index 89f8e3876c..e7179350bf 100644 --- a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml +++ b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml @@ -20,7 +20,7 @@ mini_swe_agent_2: memory: "8192" ephemeral_storage_gib: 50 sandbox_spec: - timeout_s: 18000 + ttl_s: 18000 ready_timeout_s: 1200 metadata: benchmark: swebench-verified @@ -31,7 +31,6 @@ mini_swe_agent_2: conda_env: testbed activate_conda: true user: root - delete: true run_golden: false step_timeout: 600 eval_timeout: 1800 diff --git a/tests/unit_tests/test_ecs_fargate_engine.py b/tests/unit_tests/test_ecs_fargate_engine.py index 941e843e27..da3eccaab3 100644 --- a/tests/unit_tests/test_ecs_fargate_engine.py +++ b/tests/unit_tests/test_ecs_fargate_engine.py @@ -1183,42 +1183,46 @@ def test_ensure_image_built_force_skips_cache_check(): def test_ensure_image_built_dedupes_on_inflight(): + # A peer build for this tag is in flight; once it finishes the image is in ECR, so the waiter + # returns the URL without building. The waiter verifies via ECR (not the event) on wake. cfg = _build_cfg() tag = "env__abcd1234" done = threading.Event() done.set() engine.ImageBuilder._inflight_builds[tag] = done - with ( - patch.object(engine.ImageBuilder, "get_ecr_image_tag", return_value=tag), - patch.object(engine.ImageBuilder, "image_exists_in_ecr") as exists, - patch.object(engine.ImageBuilder, "_build_and_push") as build, - ): - url = engine.ImageBuilder.ensure_image_built(cfg=cfg, environment_name="env") - assert url == f"{cfg.ecr_repository}:{tag}" - exists.assert_not_called() - build.assert_not_called() + try: + with ( + patch.object(engine.ImageBuilder, "get_ecr_image_tag", return_value=tag), + patch.object(engine.ImageBuilder, "image_exists_in_ecr", return_value=True), + patch.object(engine.ImageBuilder, "_build_and_push") as build, + ): + url = engine.ImageBuilder.ensure_image_built(cfg=cfg, environment_name="env") + assert url == f"{cfg.ecr_repository}:{tag}" + build.assert_not_called() + finally: + engine.ImageBuilder._inflight_builds.pop(tag, None) -def test_ensure_image_built_dedupes_on_second_check(): - # A concurrent builder registers the tag between the first cache check and - # the second in-flight check; we must wait on it rather than rebuild. +def test_ensure_image_built_waiter_raises_when_peer_build_failed(): + # A peer build for this tag is in flight but fails (image never lands in ECR). The waiter must + # raise rather than return a URL for an image that was never pushed (it would fail later at the + # ECS image pull). It verifies via ECR on wake instead of trusting the in-flight event. cfg = _build_cfg() tag = "env__deadbeef" - - def _register_inflight(*_a, **_k): - ev = threading.Event() - ev.set() - engine.ImageBuilder._inflight_builds[tag] = ev - return False - - with ( - patch.object(engine.ImageBuilder, "get_ecr_image_tag", return_value=tag), - patch.object(engine.ImageBuilder, "image_exists_in_ecr", side_effect=_register_inflight), - patch.object(engine.ImageBuilder, "_build_and_push") as build, - ): - url = engine.ImageBuilder.ensure_image_built(cfg=cfg, environment_name="env") - assert url == f"{cfg.ecr_repository}:{tag}" - build.assert_not_called() + done = threading.Event() + done.set() + engine.ImageBuilder._inflight_builds[tag] = done + try: + with ( + patch.object(engine.ImageBuilder, "get_ecr_image_tag", return_value=tag), + patch.object(engine.ImageBuilder, "image_exists_in_ecr", return_value=False), + patch.object(engine.ImageBuilder, "_build_and_push") as build, + pytest.raises(RuntimeError, match="not in ECR"), + ): + engine.ImageBuilder.ensure_image_built(cfg=cfg, environment_name="env") + build.assert_not_called() + finally: + engine.ImageBuilder._inflight_builds.pop(tag, None) def test_ensure_image_built_cache_filled_after_semaphore(): @@ -1241,38 +1245,42 @@ def test_ensure_mirrored_requires_repo(): def test_ensure_mirrored_dedupes_on_inflight(): + # Peer mirror in flight + image now in ECR -> waiter returns the URL without re-mirroring. cfg = _build_cfg() tag = engine._sanitize_id("ubuntu:24.04") done = threading.Event() done.set() engine.ImageBuilder._inflight_builds[tag] = done - with ( - patch.object(engine.ImageBuilder, "image_exists_in_ecr") as exists, - patch.object(engine.ImageBuilder, "run_buildspec_via_codebuild") as cb, - ): - url = engine.ImageBuilder.ensure_mirrored(cfg=cfg, src_image="ubuntu:24.04") - assert url.endswith(tag) - exists.assert_not_called() - cb.assert_not_called() + try: + with ( + patch.object(engine.ImageBuilder, "image_exists_in_ecr", return_value=True), + patch.object(engine.ImageBuilder, "run_buildspec_via_codebuild") as cb, + ): + url = engine.ImageBuilder.ensure_mirrored(cfg=cfg, src_image="ubuntu:24.04") + assert url.endswith(tag) + cb.assert_not_called() + finally: + engine.ImageBuilder._inflight_builds.pop(tag, None) -def test_ensure_mirrored_dedupes_on_second_check(): +def test_ensure_mirrored_waiter_raises_when_peer_build_failed(): + # Peer mirror in flight but failed (image absent from ECR) -> waiter raises instead of returning + # a URL for an image that was never pushed. cfg = _build_cfg() tag = engine._sanitize_id("ubuntu:24.04") - - def _register_inflight(*_a, **_k): - ev = threading.Event() - ev.set() - engine.ImageBuilder._inflight_builds[tag] = ev - return False - - with ( - patch.object(engine.ImageBuilder, "image_exists_in_ecr", side_effect=_register_inflight), - patch.object(engine.ImageBuilder, "run_buildspec_via_codebuild") as cb, - ): - url = engine.ImageBuilder.ensure_mirrored(cfg=cfg, src_image="ubuntu:24.04") - assert url.endswith(tag) - cb.assert_not_called() + done = threading.Event() + done.set() + engine.ImageBuilder._inflight_builds[tag] = done + try: + with ( + patch.object(engine.ImageBuilder, "image_exists_in_ecr", return_value=False), + patch.object(engine.ImageBuilder, "run_buildspec_via_codebuild") as cb, + pytest.raises(RuntimeError, match="not in ECR"), + ): + engine.ImageBuilder.ensure_mirrored(cfg=cfg, src_image="ubuntu:24.04") + cb.assert_not_called() + finally: + engine.ImageBuilder._inflight_builds.pop(tag, None) def test_ensure_mirrored_cache_filled_after_semaphore(): diff --git a/tests/unit_tests/test_ecs_fargate_provider.py b/tests/unit_tests/test_ecs_fargate_provider.py index bad5731b8e..9251df8ad7 100644 --- a/tests/unit_tests/test_ecs_fargate_provider.py +++ b/tests/unit_tests/test_ecs_fargate_provider.py @@ -30,7 +30,7 @@ list_providers, ) from nemo_gym.sandbox.providers.ecs_fargate import EcsFargateProvider, engine -from nemo_gym.sandbox.providers.ecs_fargate.provider import engine_config_from_mapping +from nemo_gym.sandbox.providers.ecs_fargate.provider import _apply_spec_overrides, engine_config_from_mapping _ENG = "nemo_gym.sandbox.providers.ecs_fargate.engine" @@ -373,3 +373,55 @@ def test_get_ecr_image_tag_is_content_addressed(tmp_path): assert tag1 == tag2 and tag1.startswith("env__") (tmp_path / "Dockerfile").write_text("FROM alpine") assert engine.ImageBuilder.get_ecr_image_tag(tmp_path, "env") != tag1 + + +# ── Phase A review fixes ────────────────────────────────────────────── + + +def test_ephemeral_storage_block_validates_fargate_range(): + # Omit the field (take Fargate's implicit 20 GiB) unless a larger size is requested; an explicit + # 20 -- or anything outside 21-200 -- is rejected by RegisterTaskDefinition. + assert engine._ephemeral_storage_block(None) is None + assert engine._ephemeral_storage_block(0) is None + assert engine._ephemeral_storage_block(50) == {"sizeInGiB": 50} + for bad in (20, 201): + with pytest.raises(ValueError, match="21 and 200"): + engine._ephemeral_storage_block(bad) + + +def test_render_env_value_rejects_unresolved_task_ip(): + cfg = engine.EcsFargateConfig(region="us-east-1") + sandbox = engine.EcsFargateSandbox(engine.SandboxSpec(image="img"), ecs_config=cfg) + # _task_ip is unknown until the task is running, so it cannot be baked into container env. + with pytest.raises(ValueError, match="task_ip"): + sandbox._render_env_value("BASE=http://{task_ip}:8000") + # once known it resolves; unrelated values pass through untouched. + sandbox._task_ip = "10.0.0.5" + assert sandbox._render_env_value("BASE=http://{task_ip}:8000") == "BASE=http://10.0.0.5:8000" + assert sandbox._render_env_value("PLAIN=1") == "PLAIN=1" + + +def test_apply_spec_overrides_maps_resources_and_ttl(): + cfg = engine.EcsFargateConfig(region="us-east-1", cpu="256", memory="512") + spec = SandboxSpec( + image="img", + ttl_s=3600, + ready_timeout_s=120, + resources={"cpu": 2, "memory_mib": 4096, "disk_gib": 50}, + ) + out = _apply_spec_overrides(cfg, spec) + assert out.max_task_lifetime_sec == 3600 + assert out.startup_timeout_sec == 120.0 + assert out.cpu == "2048" # vCPUs -> Fargate CPU units + assert out.memory == "4096" + assert out.ephemeral_storage_gib == 50 + # no per-sandbox requests -> config returned unchanged (same object) + assert _apply_spec_overrides(cfg, SandboxSpec(image="img")) is cfg + + +def test_apply_spec_overrides_rejects_gpu(): + from nemo_gym.sandbox.providers import SandboxCreateError + + cfg = engine.EcsFargateConfig(region="us-east-1") + with pytest.raises(SandboxCreateError, match="GPU"): + _apply_spec_overrides(cfg, SandboxSpec(image="img", resources={"gpu": 1})) From bb843f821ac1413a2f420c9ccd8a3ca2010cf32a Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Fri, 26 Jun 2026 11:28:17 +0200 Subject: [PATCH 06/11] fix(sandbox): ECS review phase B (concurrency, Fargate validation, leaks) P1 fixes from @ananthsub's review: - Build semaphore: initialize once; never replace a live semaphore when a later caller passes a different build_parallelism (a permit held against the old object would drift the cap). - Fargate cpu/memory: validate the pair (not independent maxes) before RegisterTaskDefinition -- cpu/memory derived as separate max-of-base-and-config could form an unsupported combination. - S3 artifacts: delete the CodeBuild context zip after the build (try/finally) and the per-rollout upload tarballs on cleanup. The process-cached exec-server script is left alone -- it is shared across sandboxes, so per-instance deletion would be unsafe. - _get_task_public_ip: when assign_public_ip is enabled, keep retrying for the PublicIp to propagate before falling back to the private IP on the last attempt (a private IP is unreachable from the orchestrator and otherwise yields a misleading SSH timeout later). Tests: cpu/memory validation, public-IP retry, and best-effort S3 delete. Signed-off-by: Michal Bien --- .../sandbox/providers/ecs_fargate/engine.py | 113 ++++++++++++++---- tests/unit_tests/test_ecs_fargate_provider.py | 46 +++++++ 2 files changed, 135 insertions(+), 24 deletions(-) diff --git a/nemo_gym/sandbox/providers/ecs_fargate/engine.py b/nemo_gym/sandbox/providers/ecs_fargate/engine.py index d164b58511..786a402675 100644 --- a/nemo_gym/sandbox/providers/ecs_fargate/engine.py +++ b/nemo_gym/sandbox/providers/ecs_fargate/engine.py @@ -497,6 +497,44 @@ def _ephemeral_storage_block(requested: Any) -> dict[str, int] | None: return {"sizeInGiB": size} +# Valid Fargate task cpu (units) -> inclusive memory range (MiB). cpu/memory are not independent: +# https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task-cpu-memory-error.html +_FARGATE_MEMORY_RANGE_MIB: dict[int, tuple[int, int]] = { + 256: (512, 2048), + 512: (1024, 4096), + 1024: (2048, 8192), + 2048: (4096, 16384), + 4096: (8192, 30720), + 8192: (16384, 61440), + 16384: (32768, 122880), +} + + +def _validate_fargate_cpu_memory(cpu: int, memory: int) -> None: + """Validate that a Fargate task cpu (units) + memory (MiB) form a supported combination. + + Fargate rejects independent cpu/memory values (e.g. cpu=4096 with memory=512), so callers that + derive the two separately (max-of-base-and-config) must check the pair before RegisterTaskDefinition. + """ + memory_range = _FARGATE_MEMORY_RANGE_MIB.get(cpu) + if memory_range is None: + raise ValueError(f"Fargate task cpu must be one of {sorted(_FARGATE_MEMORY_RANGE_MIB)} units (got {cpu})") + low, high = memory_range + if not low <= memory <= high: + raise ValueError(f"Fargate memory for cpu={cpu} must be {low}-{high} MiB (got {memory})") + + +def _delete_s3_object(cfg: EcsFargateConfig, key: str) -> None: + """Best-effort delete of a transient S3 artifact; never raises.""" + if not cfg.s3_bucket or not key: + return + try: + boto3, *_ = _require_aws_sdks() + boto3.client("s3", region_name=cfg.region).delete_object(Bucket=cfg.s3_bucket, Key=key) + except Exception: + logger.debug("Failed to delete S3 object s3://%s/%s", cfg.s3_bucket, key, exc_info=True) + + def download_secret_to_file(secret_arn: str, region: str | None = None) -> str: """Fetch a Secrets Manager secret → temp file (mode 0600).""" key_material = download_secret_to_string(secret_arn, region=region) @@ -1121,9 +1159,17 @@ def _run_deduped_build( if is_builder: event = threading.Event() cls._inflight_builds[tag] = event - if cls._build_semaphore is None or cls._build_semaphore_size != cfg.build_parallelism: + if cls._build_semaphore is None: cls._build_semaphore = threading.Semaphore(cfg.build_parallelism) cls._build_semaphore_size = cfg.build_parallelism + elif cls._build_semaphore_size != cfg.build_parallelism: + # Init once: never replace a live semaphore (a permit held against the old object + # would be released into the new one, drifting the cap). Honor the first caller's size. + logger.warning( + "build_parallelism=%d ignored; build semaphore already initialized at %d", + cfg.build_parallelism, + cls._build_semaphore_size, + ) semaphore = cls._build_semaphore if not is_builder: @@ -1342,27 +1388,31 @@ def _build_and_push(cls, *, cfg: EcsFargateConfig, environment_name: str, tag: s nonce = uuid.uuid4().hex[:8] logger.info("Building image via CodeBuild: %s", image_url) s3_key = cls._upload_build_context(cfg, environment_name, nonce) - cb = boto3.client("codebuild", region_name=cfg.region) - project_name = cls._resolve_codebuild_project(cfg, cb, nonce) - buildspec = cls._generate_buildspec(cfg, repo_name, tag, image_url) - resp = _retry_with_backoff( - lambda: cb.start_build( - projectName=project_name, - sourceTypeOverride="S3", - sourceLocationOverride=f"{cfg.s3_bucket}/{s3_key}", - buildspecOverride=buildspec, - timeoutInMinutesOverride=cfg.codebuild_build_timeout, - privilegedModeOverride=True, - environmentTypeOverride="LINUX_CONTAINER", - imageOverride="aws/codebuild/amazonlinux-x86_64-standard:5.0", - computeTypeOverride=cfg.codebuild_compute_type, - ), - operation_name="codebuild.start_build", - max_retries=5, - ) - build_id = resp["build"]["id"] - logger.info("CodeBuild started: %s", build_id) - cls._poll_codebuild(cb, build_id, image_url) + try: + cb = boto3.client("codebuild", region_name=cfg.region) + project_name = cls._resolve_codebuild_project(cfg, cb, nonce) + buildspec = cls._generate_buildspec(cfg, repo_name, tag, image_url) + resp = _retry_with_backoff( + lambda: cb.start_build( + projectName=project_name, + sourceTypeOverride="S3", + sourceLocationOverride=f"{cfg.s3_bucket}/{s3_key}", + buildspecOverride=buildspec, + timeoutInMinutesOverride=cfg.codebuild_build_timeout, + privilegedModeOverride=True, + environmentTypeOverride="LINUX_CONTAINER", + imageOverride="aws/codebuild/amazonlinux-x86_64-standard:5.0", + computeTypeOverride=cfg.codebuild_compute_type, + ), + operation_name="codebuild.start_build", + max_retries=5, + ) + build_id = resp["build"]["id"] + logger.info("CodeBuild started: %s", build_id) + cls._poll_codebuild(cb, build_id, image_url) + finally: + # The build-context zip is only needed during the CodeBuild job. + _delete_s3_object(cfg, s3_key) @classmethod def run_buildspec_via_codebuild( @@ -1470,6 +1520,7 @@ def __init__(self, spec: SandboxSpec, *, ecs_config: EcsFargateConfig) -> None: self._ec2: Any = None self._ssm: Any = None self._runtime_container_env: dict[str, str] = {} + self._s3_artifacts: list[str] = [] # transient S3 keys to delete on cleanup self._ssh_tunnel_port: int | None = None self._agent_forward_port: int | None = None self._outside_endpoints: list[OutsideEndpoint] = [] @@ -1911,12 +1962,15 @@ def _register_from_base( target["mountPoints"] = existing_mounts + mount_points family = self._make_family_name() + cpu = max(int(base.get("cpu") or "256"), int(cfg.cpu)) + memory = max(int(base.get("memory") or "512"), int(cfg.memory)) + _validate_fargate_cpu_memory(cpu, memory) payload: dict[str, Any] = { "family": family, "networkMode": base.get("networkMode", "awsvpc"), "requiresCompatibilities": base.get("requiresCompatibilities", ["FARGATE"]), - "cpu": str(max(int(base.get("cpu") or "256"), int(cfg.cpu))), - "memory": str(max(int(base.get("memory") or "512"), int(cfg.memory))), + "cpu": str(cpu), + "memory": str(memory), "containerDefinitions": containers, } ephemeral = _ephemeral_storage_block( @@ -1962,6 +2016,7 @@ def _register_from_scratch( if mount_points: container_def["mountPoints"] = mount_points + _validate_fargate_cpu_memory(int(cfg.cpu), int(cfg.memory)) payload: dict[str, Any] = { "family": self._make_family_name(), "networkMode": "awsvpc", @@ -2227,6 +2282,12 @@ def _get_task_public_ip(self) -> str: logger.info("Container public IP: %s", pub) return pub priv = iface.get("PrivateIpAddress") + # When a public IP is expected, keep retrying for it to propagate before falling back + # to the private IP on the last attempt — the orchestrator is outside the VPC and + # cannot reach the private address, so an early private fallback only yields a + # misleading SSH-timeout later. + if self._cfg.assign_public_ip and attempt < max_retries: + raise RuntimeError(f"ENI {eni_id} has no public IP yet") if priv: logger.info("Container private IP: %s", priv) return priv @@ -2316,6 +2377,9 @@ def _cleanup(self) -> None: except Exception: logger.debug("Failed to remove SSH key file %s", self._ssh_key_file, exc_info=True) self._ssh_key_file = None + for key in self._s3_artifacts: + _delete_s3_object(self._cfg, key) + self._s3_artifacts = [] def _sync_stop(self) -> None: """Synchronous stop for emergency cleanup (atexit handler).""" @@ -2361,6 +2425,7 @@ def _pack() -> bytes: operation_name="s3.put_object(upload)", max_retries=5, ) + self._s3_artifacts.append(key) # remove on cleanup (the container downloads it during the run) url = await asyncio.to_thread( s3.generate_presigned_url, "get_object", diff --git a/tests/unit_tests/test_ecs_fargate_provider.py b/tests/unit_tests/test_ecs_fargate_provider.py index 9251df8ad7..653c71399e 100644 --- a/tests/unit_tests/test_ecs_fargate_provider.py +++ b/tests/unit_tests/test_ecs_fargate_provider.py @@ -425,3 +425,49 @@ def test_apply_spec_overrides_rejects_gpu(): cfg = engine.EcsFargateConfig(region="us-east-1") with pytest.raises(SandboxCreateError, match="GPU"): _apply_spec_overrides(cfg, SandboxSpec(image="img", resources={"gpu": 1})) + + +# ── Phase B review fixes ────────────────────────────────────────────── + + +def test_validate_fargate_cpu_memory(): + engine._validate_fargate_cpu_memory(2048, 8192) # valid pair + engine._validate_fargate_cpu_memory(256, 512) # valid low edge + with pytest.raises(ValueError, match="cpu must be"): + engine._validate_fargate_cpu_memory(777, 2048) # unsupported cpu value + with pytest.raises(ValueError, match="memory for cpu"): + # the independent-max bug: cpu=256 cannot pair with 8192 MiB + engine._validate_fargate_cpu_memory(256, 8192) + + +def test_get_task_public_ip_waits_for_public_ip_when_enabled(): + # assign_public_ip=ENABLED: must keep retrying for the PublicIp to propagate instead of + # returning the (VPC-internal, unreachable) private IP on the first poll. + cfg = engine.EcsFargateConfig(region="us-east-1", cluster="c", assign_public_ip=True) + sandbox = engine.EcsFargateSandbox(engine.SandboxSpec(image="img"), ecs_config=cfg) + sandbox._task_arn = "task-arn" + sandbox._ecs = MagicMock() + sandbox._ecs.describe_tasks.return_value = { + "tasks": [ + {"attachments": [{"type": "ElasticNetworkInterface", "details": [{"name": "networkInterfaceId", "value": "eni-1"}]}]} + ] + } + sandbox._ec2 = MagicMock() + sandbox._ec2.describe_network_interfaces.side_effect = [ + {"NetworkInterfaces": [{"PrivateIpAddress": "10.0.0.1"}]}, # no PublicIp yet + {"NetworkInterfaces": [{"Association": {"PublicIp": "1.2.3.4"}, "PrivateIpAddress": "10.0.0.1"}]}, + ] + with patch(f"{_ENG}.time.sleep"): + assert sandbox._get_task_public_ip() == "1.2.3.4" + + +def test_delete_s3_object_best_effort(): + cfg = engine.EcsFargateConfig(region="us-east-1", s3_bucket="bucket") + fake_s3 = MagicMock() + fake_boto3 = MagicMock() + fake_boto3.client.return_value = fake_s3 + with patch(f"{_ENG}._require_aws_sdks", return_value=(fake_boto3, MagicMock(), MagicMock())): + engine._delete_s3_object(cfg, "some/key.zip") + fake_s3.delete_object.assert_called_once_with(Bucket="bucket", Key="some/key.zip") + # no bucket configured -> no-op (no client built) + engine._delete_s3_object(engine.EcsFargateConfig(region="us-east-1"), "k") From 5c87e103673ad49f9eb0e1c6e775244498371290 Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Fri, 26 Jun 2026 12:04:05 +0200 Subject: [PATCH 07/11] fix(sandbox): ECS review phase C (injection, async, EFS, RepoNotFound, CodeBuild) Final batch from @ananthsub's review: - Mirror injection: validate src_image (dataset-controlled) against a strict registry/repo/tag/ digest charset before it is interpolated into the privileged CodeBuild docker pull/tag commands. - Async hygiene: offload the blocking SSH wait_ready in the exec() reconnect path via asyncio.to_thread so it can't stall the event loop. - EFS: a volume may inherit the provider-level efs_filesystem_id / efs_access_point_id defaults (opt in via `efs: true`); a volume with no resolvable filesystem id fails fast. - RepositoryNotFound: image_exists_in_ecr fails fast with a clear error for a missing ECR repo instead of reporting "image absent" and building/pushing into a non-existent repo. - CodeBuild: reuse a single shared project (created idempotently) instead of one per build, so per-build projects don't accumulate. - _free_port: documented the TOCTOU window (mitigated by SshTunnel.open() retries). Tests: image-ref validation, EFS default inheritance + missing-fs error, RepositoryNotFound fail-fast, shared CodeBuild project; updated dedup/codebuild tests to the new contracts. Signed-off-by: Michal Bien --- .../sandbox/providers/ecs_fargate/engine.py | 69 +++++++++++++++---- tests/unit_tests/test_ecs_fargate_engine.py | 29 +++++--- tests/unit_tests/test_ecs_fargate_provider.py | 35 ++++++++++ 3 files changed, 111 insertions(+), 22 deletions(-) diff --git a/nemo_gym/sandbox/providers/ecs_fargate/engine.py b/nemo_gym/sandbox/providers/ecs_fargate/engine.py index 786a402675..410fb3bbf5 100644 --- a/nemo_gym/sandbox/providers/ecs_fargate/engine.py +++ b/nemo_gym/sandbox/providers/ecs_fargate/engine.py @@ -88,13 +88,15 @@ class VolumeMount: host_path: str = "" container_path: str = "" readonly: bool = False + efs: bool = False efs_filesystem_id: str | None = None efs_root_directory: str | None = None efs_access_point_id: str | None = None @property def is_efs(self) -> bool: - return self.efs_filesystem_id is not None + # Either names its own filesystem, or opts into the provider-level EFS default via `efs: true`. + return self.efs or self.efs_filesystem_id is not None @dataclass @@ -476,6 +478,9 @@ def _retry_with_backoff( def _free_port() -> int: + # TOCTOU: the port is free when probed but the socket is closed before ssh binds it, so a racing + # process can claim it under high concurrency. SshTunnel.open() retries on port-not-open, which + # mitigates this in practice; holding the socket and handing the fd to ssh would remove it. with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind(("127.0.0.1", 0)) return s.getsockname()[1] @@ -535,6 +540,21 @@ def _delete_s3_object(cfg: EcsFargateConfig, key: str) -> None: logger.debug("Failed to delete S3 object s3://%s/%s", cfg.s3_bucket, key, exc_info=True) +# Registry/repo/tag/digest characters only — used to gate dataset-controlled image refs before they +# are interpolated into shell `docker pull/tag` commands + the buildspec of a privileged CodeBuild job. +_IMAGE_REF_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:/@-]*$") + + +def _validate_image_ref(ref: str) -> None: + """Reject an image reference containing characters outside the safe registry/repo/tag/digest set. + + ``src_image`` comes from the dataset/benchmark and is interpolated into shell commands inside a + privileged build, so anything outside the safe set is a code-injection risk and is refused. + """ + if not ref or not _IMAGE_REF_RE.match(ref): + raise ValueError(f"Unsafe image reference {ref!r}: only [A-Za-z0-9._:/@-] characters are allowed") + + def download_secret_to_file(secret_arn: str, region: str | None = None) -> str: """Fetch a Secrets Manager secret → temp file (mode 0600).""" key_material = download_secret_to_string(secret_arn, region=region) @@ -1070,8 +1090,13 @@ def image_exists_in_ecr(ecr_repository: str, tag: str, region: str | None = None return True except ClientError as exc: code = exc.response.get("Error", {}).get("Code", "") - if code in ("ImageNotFoundException", "RepositoryNotFoundException"): + if code == "ImageNotFoundException": return False + if code == "RepositoryNotFoundException": + raise RuntimeError( + f"ECR repository {ecr_repository!r} does not exist; create it (e.g. via Terraform) " + f"before building or mirroring images" + ) from exc raise @staticmethod @@ -1225,6 +1250,7 @@ def ensure_mirrored(cls, *, cfg: EcsFargateConfig, src_image: str, force: bool = ecr_repo = cfg.ecr_repository if not ecr_repo: raise ValueError("ecr_repository is required to mirror images") + _validate_image_ref(src_image) tag = _sanitize_id(src_image) image_url = f"{ecr_repo}:{tag}" @@ -1296,13 +1322,16 @@ def _upload_build_context(cfg: EcsFargateConfig, environment_name: str, nonce: s return s3_key @staticmethod - def _resolve_codebuild_project(cfg: EcsFargateConfig, cb: Any, nonce: str) -> str: + def _resolve_codebuild_project(cfg: EcsFargateConfig, cb: Any) -> str: _, _, ClientError = _require_aws_sdks() if cfg.codebuild_project: return cfg.codebuild_project if not cfg.codebuild_service_role: raise RuntimeError("codebuild_project or codebuild_service_role is required") - project_name = f"ecs-sandbox-build-{nonce}" + # Reuse a single shared project (created once, idempotently) instead of one per build, so + # per-build CodeBuild projects don't accumulate. Source + buildspec are overridden per + # start_build, so one generic NO_SOURCE project serves every build. + project_name = "ecs-sandbox-build" try: _retry_with_backoff( lambda: cb.create_project( @@ -1390,7 +1419,7 @@ def _build_and_push(cls, *, cfg: EcsFargateConfig, environment_name: str, tag: s s3_key = cls._upload_build_context(cfg, environment_name, nonce) try: cb = boto3.client("codebuild", region_name=cfg.region) - project_name = cls._resolve_codebuild_project(cfg, cb, nonce) + project_name = cls._resolve_codebuild_project(cfg, cb) buildspec = cls._generate_buildspec(cfg, repo_name, tag, image_url) resp = _retry_with_backoff( lambda: cb.start_build( @@ -1431,9 +1460,8 @@ def run_buildspec_via_codebuild( internally). """ boto3, *_ = _require_aws_sdks() - nonce = uuid.uuid4().hex[:8] cb = boto3.client("codebuild", region_name=cfg.region) - project_name = cls._resolve_codebuild_project(cfg, cb, nonce) + project_name = cls._resolve_codebuild_project(cfg, cb) timeout = timeout_minutes or cfg.codebuild_build_timeout logger.info("Starting CodeBuild harness build: %s (timeout=%dm)", job_label, timeout) @@ -1630,7 +1658,12 @@ async def exec( sidecar = self._cfg.ssh_sidecar if sidecar and sidecar.exec_server_port is not None: health_url = f"http://127.0.0.1:{self._ssh_tunnel.local_port}/health" # type: ignore[union-attr] - self._ssh_tunnel.wait_ready(health_url=health_url, timeout=60.0) # type: ignore[union-attr] + # wait_ready blocks (sleep + urlopen); offload so it can't stall the event loop. + await asyncio.to_thread( + self._ssh_tunnel.wait_ready, # type: ignore[union-attr] + health_url=health_url, + timeout=60.0, + ) old_client = self._exec_client self._exec_client = ExecClient(port=self._ssh_tunnel.local_port) # type: ignore[union-attr] if old_client is not None: @@ -2036,20 +2069,32 @@ def _register_from_scratch( return self._do_register(payload) def _build_efs_volumes(self) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: - """Build EFS volume definitions and mount points from spec.volumes.""" + """Build EFS volume definitions and mount points from spec.volumes. + + A volume may name its own filesystem / access point, or inherit the provider-level + ``efs_filesystem_id`` / ``efs_access_point_id`` defaults (populated from YAML/SSM). + """ + cfg = self._cfg task_volumes: list[dict[str, Any]] = [] mount_points: list[dict[str, Any]] = [] for i, vol in enumerate(self._spec.volumes): if not vol.is_efs: continue + filesystem_id = vol.efs_filesystem_id or cfg.efs_filesystem_id + if not filesystem_id: + raise ValueError( + f"EFS volume at {vol.container_path!r} has no filesystem id; set it on the volume " + f"or configure the provider's efs_filesystem_id" + ) + access_point_id = vol.efs_access_point_id or cfg.efs_access_point_id vol_name = f"efs-{i}" efs_cfg: dict[str, Any] = { - "fileSystemId": vol.efs_filesystem_id, + "fileSystemId": filesystem_id, "transitEncryption": "ENABLED", } - if vol.efs_access_point_id: + if access_point_id: efs_cfg["authorizationConfig"] = { - "accessPointId": vol.efs_access_point_id, + "accessPointId": access_point_id, "iam": "ENABLED", } elif vol.efs_root_directory: diff --git a/tests/unit_tests/test_ecs_fargate_engine.py b/tests/unit_tests/test_ecs_fargate_engine.py index da3eccaab3..e1fdcde463 100644 --- a/tests/unit_tests/test_ecs_fargate_engine.py +++ b/tests/unit_tests/test_ecs_fargate_engine.py @@ -1034,14 +1034,23 @@ def test_image_exists_in_ecr_true(): ecr.describe_images.assert_called_once_with(repositoryName="sandbox", imageIds=[{"imageTag": "tag1"}]) -@pytest.mark.parametrize("code", ["ImageNotFoundException", "RepositoryNotFoundException"]) -def test_image_exists_in_ecr_false_on_missing(code): +def test_image_exists_in_ecr_false_on_image_not_found(): ecr = MagicMock() - ecr.describe_images.side_effect = _client_error(code, "DescribeImages") + ecr.describe_images.side_effect = _client_error("ImageNotFoundException", "DescribeImages") with _patch_aws(ecr=ecr): assert engine.ImageBuilder.image_exists_in_ecr(_ecr_repo(), "tag1") is False +def test_image_exists_in_ecr_fails_fast_on_missing_repo(): + # A missing repository fails fast with a clear error rather than reporting "image absent" and + # then building/pushing into a non-existent repo (which would only fail later at the push). + ecr = MagicMock() + ecr.describe_images.side_effect = _client_error("RepositoryNotFoundException", "DescribeImages") + with _patch_aws(ecr=ecr): + with pytest.raises(RuntimeError, match="does not exist"): + engine.ImageBuilder.image_exists_in_ecr(_ecr_repo(), "tag1") + + def test_image_exists_in_ecr_reraises_other(): ecr = MagicMock() ecr.describe_images.side_effect = _client_error("AccessDeniedException", "DescribeImages") @@ -1325,22 +1334,22 @@ def test_resolve_codebuild_project_uses_explicit(): cfg = _build_cfg(codebuild_project="my-project") with _patch_aws(): # No client calls expected when project is explicit. - assert engine.ImageBuilder._resolve_codebuild_project(cfg, MagicMock(), "nonce") == "my-project" + assert engine.ImageBuilder._resolve_codebuild_project(cfg, MagicMock()) == "my-project" def test_resolve_codebuild_project_requires_role(): cfg = _build_cfg(codebuild_service_role=None, codebuild_project=None) with _patch_aws(): with pytest.raises(RuntimeError, match="codebuild_project or codebuild_service_role"): - engine.ImageBuilder._resolve_codebuild_project(cfg, MagicMock(), "nonce") + engine.ImageBuilder._resolve_codebuild_project(cfg, MagicMock()) def test_resolve_codebuild_project_creates_project(): cfg = _build_cfg() cb = MagicMock() with _patch_aws(): - name = engine.ImageBuilder._resolve_codebuild_project(cfg, cb, "abc123") - assert name == "ecs-sandbox-build-abc123" + name = engine.ImageBuilder._resolve_codebuild_project(cfg, cb) + assert name == "ecs-sandbox-build" cb.create_project.assert_called_once() assert cb.create_project.call_args.kwargs["serviceRole"] == "arn:aws:iam::123:role/cb" @@ -1352,8 +1361,8 @@ def test_resolve_codebuild_project_tolerates_already_exists(): "ResourceAlreadyExistsException", "CreateProject", msg="already exists" ) with _patch_aws(): - name = engine.ImageBuilder._resolve_codebuild_project(cfg, cb, "abc123") - assert name == "ecs-sandbox-build-abc123" + name = engine.ImageBuilder._resolve_codebuild_project(cfg, cb) + assert name == "ecs-sandbox-build" def test_resolve_codebuild_project_reraises_other_error(): @@ -1362,7 +1371,7 @@ def test_resolve_codebuild_project_reraises_other_error(): cb.create_project.side_effect = _client_error("AccessDenied", "CreateProject", msg="forbidden") with _patch_aws(): with pytest.raises(ClientError): - engine.ImageBuilder._resolve_codebuild_project(cfg, cb, "abc123") + engine.ImageBuilder._resolve_codebuild_project(cfg, cb) def test_poll_codebuild_success(): diff --git a/tests/unit_tests/test_ecs_fargate_provider.py b/tests/unit_tests/test_ecs_fargate_provider.py index 653c71399e..7382a38889 100644 --- a/tests/unit_tests/test_ecs_fargate_provider.py +++ b/tests/unit_tests/test_ecs_fargate_provider.py @@ -471,3 +471,38 @@ def test_delete_s3_object_best_effort(): fake_s3.delete_object.assert_called_once_with(Bucket="bucket", Key="some/key.zip") # no bucket configured -> no-op (no client built) engine._delete_s3_object(engine.EcsFargateConfig(region="us-east-1"), "k") + + +# ── Phase C review fixes ────────────────────────────────────────────── + + +def test_validate_image_ref_blocks_injection(): + # Legit references pass (registry/repo/tag and digests). + engine._validate_image_ref("docker.io/swebench/sweb.eval.x86_64.astropy_1776-12907:latest") + engine._validate_image_ref("123.dkr.ecr.us-east-1.amazonaws.com/mirror:tag@sha256:abc123") + engine._validate_image_ref("ubuntu:24.04") + # Anything with shell metacharacters (the privileged-build injection surface) is refused. + for bad in ["ubuntu:24.04; rm -rf /", "img$(whoami)", "img`id`", "a && curl evil", "x\nFROM y", ""]: + with pytest.raises(ValueError, match="Unsafe image reference"): + engine._validate_image_ref(bad) + + +def test_build_efs_volumes_inherits_provider_defaults(): + # A volume opts into EFS (efs=True) without naming a filesystem; it inherits the provider-level + # efs_filesystem_id / efs_access_point_id (from YAML/SSM). + cfg = engine.EcsFargateConfig(region="us-east-1", efs_filesystem_id="fs-123", efs_access_point_id="fsap-9") + spec = engine.SandboxSpec(image="img", volumes=[engine.VolumeMount(container_path="/data", efs=True)]) + sandbox = engine.EcsFargateSandbox(spec, ecs_config=cfg) + task_volumes, mount_points = sandbox._build_efs_volumes() + efs = task_volumes[0]["efsVolumeConfiguration"] + assert efs["fileSystemId"] == "fs-123" + assert efs["authorizationConfig"]["accessPointId"] == "fsap-9" + assert mount_points[0]["containerPath"] == "/data" + + +def test_build_efs_volumes_requires_a_filesystem_id(): + cfg = engine.EcsFargateConfig(region="us-east-1") # no provider-level EFS default + spec = engine.SandboxSpec(image="img", volumes=[engine.VolumeMount(container_path="/data", efs=True)]) + sandbox = engine.EcsFargateSandbox(spec, ecs_config=cfg) + with pytest.raises(ValueError, match="no filesystem id"): + sandbox._build_efs_volumes() From ea0a92a7b7ed5d0752769121676d3b760f57d03b Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Fri, 26 Jun 2026 12:24:55 +0200 Subject: [PATCH 08/11] style(sandbox): ruff-format ECS provider test (line wrap) Signed-off-by: Michal Bien --- tests/unit_tests/test_ecs_fargate_provider.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/test_ecs_fargate_provider.py b/tests/unit_tests/test_ecs_fargate_provider.py index 7382a38889..5832c1b62f 100644 --- a/tests/unit_tests/test_ecs_fargate_provider.py +++ b/tests/unit_tests/test_ecs_fargate_provider.py @@ -449,7 +449,11 @@ def test_get_task_public_ip_waits_for_public_ip_when_enabled(): sandbox._ecs = MagicMock() sandbox._ecs.describe_tasks.return_value = { "tasks": [ - {"attachments": [{"type": "ElasticNetworkInterface", "details": [{"name": "networkInterfaceId", "value": "eni-1"}]}]} + { + "attachments": [ + {"type": "ElasticNetworkInterface", "details": [{"name": "networkInterfaceId", "value": "eni-1"}]} + ] + } ] } sandbox._ec2 = MagicMock() From 8a5d8c8cb881067ddca344942a82b283ec5cb0fc Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Fri, 26 Jun 2026 13:53:34 +0200 Subject: [PATCH 09/11] fix(sandbox): correct ECS large-upload path; de-narrate engine/provider - Large S3 uploads (>512 KiB) now land at the requested remote path; the tar was keyed by the local temp basename and extracted into the dest dir, so the file previously landed under the wrong name. - Surface S3 staging cleanup gaps once (missing s3:DeleteObject) instead of silently swallowing them; cleanup stays best-effort and never raises. - Strip nemo-evaluator-next/FEP/NEL provenance and condense verbose docstrings and comments to terse one-liners (no logic changes). Signed-off-by: Michal Bien --- .../sandbox/providers/ecs_fargate/engine.py | 176 ++++++++---------- .../sandbox/providers/ecs_fargate/provider.py | 17 +- .../configs/mini_swe_agent_ecs_fargate.yaml | 5 +- tests/unit_tests/test_ecs_fargate_engine.py | 42 ++++- 4 files changed, 122 insertions(+), 118 deletions(-) diff --git a/nemo_gym/sandbox/providers/ecs_fargate/engine.py b/nemo_gym/sandbox/providers/ecs_fargate/engine.py index 410fb3bbf5..863b3e18b2 100644 --- a/nemo_gym/sandbox/providers/ecs_fargate/engine.py +++ b/nemo_gym/sandbox/providers/ecs_fargate/engine.py @@ -14,11 +14,9 @@ """ECS Fargate sandbox engine. -Lifted from nemo-evaluator-next (``nemo_evaluator/sandbox/ecs_fargate.py``). -The orchestration is unchanged; only the three host-protocol types it relied on -(``ExecResult``, ``OutsideEndpoint``, ``SandboxSpec``) and the SSM/port -constants are vendored here so the engine stands alone inside the provider. -The Gym-facing adapter lives in ``provider.py``. +Runs each sandbox as an ECS Fargate task behind an SSH sidecar. Defines the host-protocol +types (``ExecResult``, ``OutsideEndpoint``, ``SandboxSpec``) and the engine; the Gym-facing +provider adapter lives in ``provider.py``. """ from __future__ import annotations @@ -264,10 +262,7 @@ class EcsFargateConfig: codebuild_build_timeout: int = 60 dockerhub_secret_arn: str | None = None build_parallelism: int = 50 - # When a public/bare image is routed to the ECR mirror and is not yet - # present, pull it into ECR on demand (via CodeBuild) during create. The - # mirror normally runs ahead of time, but this keeps create self-healing so - # callers never have to pre-stage images manually. Set False to require a + # Mirror a missing public/bare image into ECR on demand during create. Set False to require a # pre-populated mirror and fail fast on a miss. auto_mirror: bool = True efs_filesystem_id: str | None = None @@ -478,21 +473,17 @@ def _retry_with_backoff( def _free_port() -> int: - # TOCTOU: the port is free when probed but the socket is closed before ssh binds it, so a racing - # process can claim it under high concurrency. SshTunnel.open() retries on port-not-open, which - # mitigates this in practice; holding the socket and handing the fd to ssh would remove it. + # TOCTOU: the probed port can be claimed before ssh binds it; SshTunnel.open() retries on + # port-not-open to mitigate this under concurrency. with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind(("127.0.0.1", 0)) return s.getsockname()[1] def _ephemeral_storage_block(requested: Any) -> dict[str, int] | None: - """Task-def ``ephemeralStorage`` block for an explicit request, or ``None`` to take Fargate's - implicit 20 GiB default. + """Task-def ``ephemeralStorage`` block, or ``None`` for Fargate's implicit 20 GiB default. - Fargate accepts only 21-200 GiB; an explicit ``20`` is rejected by RegisterTaskDefinition (20 is - valid only as the omitted default), so we omit the field unless a larger size is requested and - validate the range. + Fargate accepts only 21-200 GiB explicitly (20 is valid only as the omitted default). """ if not requested: return None @@ -516,11 +507,7 @@ def _ephemeral_storage_block(requested: Any) -> dict[str, int] | None: def _validate_fargate_cpu_memory(cpu: int, memory: int) -> None: - """Validate that a Fargate task cpu (units) + memory (MiB) form a supported combination. - - Fargate rejects independent cpu/memory values (e.g. cpu=4096 with memory=512), so callers that - derive the two separately (max-of-base-and-config) must check the pair before RegisterTaskDefinition. - """ + """Validate a Fargate cpu (units) + memory (MiB) pair; Fargate rejects unsupported combinations.""" memory_range = _FARGATE_MEMORY_RANGE_MIB.get(cpu) if memory_range is None: raise ValueError(f"Fargate task cpu must be one of {sorted(_FARGATE_MEMORY_RANGE_MIB)} units (got {cpu})") @@ -529,28 +516,41 @@ def _validate_fargate_cpu_memory(cpu: int, memory: int) -> None: raise ValueError(f"Fargate memory for cpu={cpu} must be {low}-{high} MiB (got {memory})") +_s3_delete_warned = False + + def _delete_s3_object(cfg: EcsFargateConfig, key: str) -> None: - """Best-effort delete of a transient S3 artifact; never raises.""" + """Best-effort delete of a transient S3 staging artifact; never raises. + + Warns once if the role lacks ``s3:DeleteObject`` (objects then rely on a bucket lifecycle policy). + """ + global _s3_delete_warned if not cfg.s3_bucket or not key: return try: boto3, *_ = _require_aws_sdks() boto3.client("s3", region_name=cfg.region).delete_object(Bucket=cfg.s3_bucket, Key=key) except Exception: - logger.debug("Failed to delete S3 object s3://%s/%s", cfg.s3_bucket, key, exc_info=True) + if not _s3_delete_warned: + _s3_delete_warned = True + logger.warning( + "Could not delete S3 staging object s3://%s/%s (further failures silenced); grant " + "s3:DeleteObject or set a bucket lifecycle policy to avoid leaking staging artifacts", + cfg.s3_bucket, + key, + exc_info=True, + ) + else: + logger.debug("Failed to delete S3 object s3://%s/%s", cfg.s3_bucket, key, exc_info=True) -# Registry/repo/tag/digest characters only — used to gate dataset-controlled image refs before they -# are interpolated into shell `docker pull/tag` commands + the buildspec of a privileged CodeBuild job. +# Safe registry/repo/tag/digest characters. Gates dataset-controlled image refs before they are +# interpolated into shell commands inside a privileged CodeBuild job (injection guard). _IMAGE_REF_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:/@-]*$") def _validate_image_ref(ref: str) -> None: - """Reject an image reference containing characters outside the safe registry/repo/tag/digest set. - - ``src_image`` comes from the dataset/benchmark and is interpolated into shell commands inside a - privileged build, so anything outside the safe set is a code-injection risk and is refused. - """ + """Reject an image reference with characters outside the safe set (shell-injection guard).""" if not ref or not _IMAGE_REF_RE.match(ref): raise ValueError(f"Unsafe image reference {ref!r}: only [A-Za-z0-9._:/@-] characters are allowed") @@ -927,17 +927,15 @@ def _err(self, code, msg): class ExecClient: """Async HTTP client for the exec server (through the SSH tunnel). - Methods are coroutines so each in-flight request occupies an event-loop - slot, not an executor thread — required to scale past asyncio's default - thread-pool cap when many long agent commands run concurrently (FEP-886). + Methods are coroutines so concurrent requests occupy event-loop slots rather than executor + threads, scaling past asyncio's default thread-pool cap. """ def __init__(self, *, port: int, connect_timeout: float = 30.0) -> None: self._base = f"http://127.0.0.1:{port}" self._timeout = connect_timeout - # Lazy: aiohttp.ClientSession must be created inside a running event - # loop, but ExecClient is constructed from `_do_start` (which runs in - # asyncio.to_thread). Defer creation to the first request. + # Lazy: aiohttp.ClientSession needs a running loop, but ExecClient is built inside a + # to_thread; create it on the first request. self._session: aiohttp.ClientSession | None = None async def _ensure_session(self) -> aiohttp.ClientSession: @@ -1109,12 +1107,7 @@ def _ecr_region(ecr_repository: str, fallback: str | None = None) -> str | None: @staticmethod def list_ecr_tags(ecr_repository: str, region: str | None = None) -> set[str]: - """Return all image tags present in an ECR repository. - - Uses paginated ``list_images`` to fetch every tagged image in - a handful of API calls, rather than one ``describe_images`` - call per tag. - """ + """Return all image tags present in an ECR repository (paginated ``list_images``).""" boto3, _, ClientError = _require_aws_sdks() ecr_region = ImageBuilder._ecr_region(ecr_repository, fallback=region) ecr = boto3.client("ecr", region_name=ecr_region) @@ -1168,10 +1161,8 @@ def _run_deduped_build( ) -> str: """Run ``build`` for ``tag`` at most once across threads, returning the ECR ``image_url``. - Concurrent callers for the same tag block on a shared in-flight event **outside** - ``cls._lock`` (so one waiter can't stall every other tag's build) and then confirm the - result via ECR rather than the event -- a failed build sets the event too, so a waiter must - not return a URL for an image that was never pushed. + Concurrent callers wait on a shared in-flight event outside ``cls._lock`` (so one waiter + can't stall other tags' builds), then confirm via ECR -- a failed build also sets the event. """ ecr_repo = cfg.ecr_repository if not force and cls.image_exists_in_ecr(ecr_repo, tag, cfg.region): @@ -1188,8 +1179,7 @@ def _run_deduped_build( cls._build_semaphore = threading.Semaphore(cfg.build_parallelism) cls._build_semaphore_size = cfg.build_parallelism elif cls._build_semaphore_size != cfg.build_parallelism: - # Init once: never replace a live semaphore (a permit held against the old object - # would be released into the new one, drifting the cap). Honor the first caller's size. + # Init once: never replace a live semaphore (would drift the permit cap). logger.warning( "build_parallelism=%d ignored; build semaphore already initialized at %d", cfg.build_parallelism, @@ -1237,15 +1227,10 @@ def ensure_image_built(cls, *, cfg: EcsFargateConfig, environment_name: str, for @classmethod def ensure_mirrored(cls, *, cfg: EcsFargateConfig, src_image: str, force: bool = False) -> str: - """Ensure ``src_image`` is present in the ECR mirror, pulling it if not. - - Public/bare image references are served from the ECR mirror tag - (``{ecr_repository}:{sanitize(src_image)}``) rather than pulled from - their origin registry at task-launch time. This mirrors the image into - ECR on demand via a self-contained CodeBuild job (privileged DinD that - logs into Docker Hub + ECR, pulls, retags, and pushes), so callers never - have to pre-stage images. Concurrent callers for the same tag dedup on a - shared in-flight event, exactly like :meth:`ensure_image_built`. + """Ensure ``src_image`` is present in the ECR mirror tag, mirroring it on demand if not. + + Mirrors via a self-contained CodeBuild job (privileged DinD: login, pull, retag, push). + Concurrent callers dedup on a shared in-flight event, like :meth:`ensure_image_built`. """ ecr_repo = cfg.ecr_repository if not ecr_repo: @@ -1328,9 +1313,8 @@ def _resolve_codebuild_project(cfg: EcsFargateConfig, cb: Any) -> str: return cfg.codebuild_project if not cfg.codebuild_service_role: raise RuntimeError("codebuild_project or codebuild_service_role is required") - # Reuse a single shared project (created once, idempotently) instead of one per build, so - # per-build CodeBuild projects don't accumulate. Source + buildspec are overridden per - # start_build, so one generic NO_SOURCE project serves every build. + # One shared project (created idempotently); source + buildspec are overridden per build, + # so a generic NO_SOURCE project serves every build instead of accumulating per-build ones. project_name = "ecs-sandbox-build" try: _retry_with_backoff( @@ -1452,13 +1436,7 @@ def run_buildspec_via_codebuild( job_label: str = "harness-build", timeout_minutes: int | None = None, ) -> None: - """Run an arbitrary buildspec via CodeBuild (privileged mode for DinD). - - Unlike :meth:`_build_and_push` which uploads a Dockerfile context to S3, - this method uses ``NO_SOURCE`` — the buildspec is fully self-contained - (e.g. it installs packages and runs a harness that builds Docker images - internally). - """ + """Run a self-contained buildspec via CodeBuild (privileged DinD, ``NO_SOURCE``).""" boto3, *_ = _require_aws_sdks() cb = boto3.client("codebuild", region_name=cfg.region) project_name = cls._resolve_codebuild_project(cfg, cb) @@ -1496,21 +1474,15 @@ def run_buildspec_via_codebuild( _task_def_cache_lock = threading.Lock() _task_def_inflight: dict[str, threading.Event] = {} -# Env vars whose values vary per sandbox invocation (workspace session id, etc). -# Routed via RunTask containerOverrides so the underlying task definition stays -# content-stable and the FEP-866 hash cache hits across invocations of the same -# task. Per-invocation env keys discovered dynamically from OutsideEndpoint -# routing (e.g. MODEL_BASE_URL with session-scoped URLs) are merged in at call -# time; this set holds the keys that are NOT visible to OutsideEndpoint routing. +# Env keys whose values vary per invocation. Routed via RunTask containerOverrides (not baked into +# the task def) so the task-def hash cache stays stable across invocations. OutsideEndpoint routing +# keys are merged in dynamically at call time; this set holds the rest. _PER_INVOCATION_ENV_KEYS: frozenset[str] = frozenset({"_NEL_EFS_SESSION"}) def _compute_task_def_hash(payload: dict[str, Any]) -> str: - # Strip logConfiguration from every container definition before hashing. - # Log config (group, stream-prefix, region) is a visibility annotation — it has - # no effect on what the sandbox does. Two task defs that differ only in - # log_stream_prefix or log_group are functionally identical and should share - # a cache entry so cross-run SSM cache hits work across differently-named runs. + # Strip logConfiguration before hashing: it's a visibility annotation with no behavioral + # effect, so task defs differing only in log group/prefix should share a cache entry. def _strip_log_cfg(containers: list) -> list: return [{k: v for k, v in c.items() if k != "logConfiguration"} for c in containers] @@ -1682,7 +1654,13 @@ async def upload(self, local_path: Path, remote_path: str) -> None: await self.upload(child, f"{remote_path}/{child.relative_to(local)}") return if local.stat().st_size > 512 * 1024 and self._cfg.s3_bucket: - await self._upload_via_s3([local], os.path.dirname(remote_path) or "/tmp") + # Map the local temp name to the requested remote filename: the S3 tar is keyed by + # basename and extracted into dest_dir, so otherwise it lands under the local basename. + await self._upload_via_s3( + [local], + os.path.dirname(remote_path) or "/tmp", + arcnames={local: os.path.basename(remote_path)}, + ) else: await self._exec_client.upload(remote_path, local) # type: ignore[union-attr] @@ -1739,9 +1717,7 @@ def _do_start(self) -> None: and not cfg.image_template and not _is_ecr_image_ref(self._spec.image) ): - # The bare/public image is served from the ECR mirror tag; pull it - # into ECR on demand so a missing mirror entry self-heals instead of - # failing the task's image pull. + # Pull the bare/public image into the ECR mirror on demand so a missing entry self-heals. ImageBuilder.ensure_mirrored(cfg=cfg, src_image=self._spec.image) image = self._resolve_image(built_image) @@ -1815,15 +1791,10 @@ def _resolve_image(self, built_image: str | None = None) -> str: f"placeholders like {{task_id}}." ) from exc if self._spec.image: - # A reference that already points at an ECR registry (e.g. the - # configured mirror) is used verbatim — re-prefixing it under - # ecr_repository and re-sanitizing would corrupt the tag and make - # an existing image unresolvable. + # An existing ECR reference is used verbatim; re-prefixing it would corrupt the tag. if _is_ecr_image_ref(self._spec.image): return self._spec.image - # Bare / public names are routed to the ECR mirror tag rather than - # pulled directly from their origin registry (avoids Docker Hub - # rate limits when many tasks start concurrently). + # Bare/public names route to the ECR mirror tag (avoids origin-registry rate limits). if cfg.ecr_repository: return f"{cfg.ecr_repository}:{_sanitize_id(self._spec.image)}" return self._spec.image @@ -1904,9 +1875,8 @@ def _render_env_value(self, value: str) -> str: if self._task_ip: value = value.replace("{task_ip}", self._task_ip) elif "{task_ip}" in value: - # _task_ip is only assigned after the task is RUNNING (_get_task_public_ip), which is - # after container env is baked into the task-def / RunTask overrides. Fail loudly rather - # than baking a literal "{task_ip}" that later breaks model egress with a confusing error. + # _task_ip is only known after the task is RUNNING, i.e. after env is baked into the + # task-def; fail loudly rather than baking a literal "{task_ip}". raise ValueError( "{task_ip} is not available in container env: the task IP is only known after the " "sandbox starts. Resolve the container's own address at runtime instead." @@ -2069,10 +2039,10 @@ def _register_from_scratch( return self._do_register(payload) def _build_efs_volumes(self) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: - """Build EFS volume definitions and mount points from spec.volumes. + """Build EFS volume defs + mount points from spec.volumes. - A volume may name its own filesystem / access point, or inherit the provider-level - ``efs_filesystem_id`` / ``efs_access_point_id`` defaults (populated from YAML/SSM). + A volume may name its own filesystem/access point or inherit the provider-level + ``efs_filesystem_id`` / ``efs_access_point_id`` defaults. """ cfg = self._cfg task_volumes: list[dict[str, Any]] = [] @@ -2327,10 +2297,8 @@ def _get_task_public_ip(self) -> str: logger.info("Container public IP: %s", pub) return pub priv = iface.get("PrivateIpAddress") - # When a public IP is expected, keep retrying for it to propagate before falling back - # to the private IP on the last attempt — the orchestrator is outside the VPC and - # cannot reach the private address, so an early private fallback only yields a - # misleading SSH-timeout later. + # Keep retrying for the public IP to propagate before falling back to the private IP + # on the last attempt (the orchestrator is outside the VPC and can't reach private). if self._cfg.assign_public_ip and attempt < max_retries: raise RuntimeError(f"ENI {eni_id} has no public IP yet") if priv: @@ -2440,17 +2408,21 @@ def _require_exec_client(self) -> None: "(ssh_sidecar.exec_server_port). In agent-server mode use sandbox.ssh_tunnel." ) - async def _upload_via_s3(self, paths: list[Path], dest_dir: str) -> None: + async def _upload_via_s3( + self, paths: list[Path], dest_dir: str, *, arcnames: dict[Path, str] | None = None + ) -> None: cfg = self._cfg if not cfg.s3_bucket: raise ValueError("s3_bucket is required for S3 staging") + names = arcnames or {} + def _pack() -> bytes: buf = io.BytesIO() with tarfile.open(fileobj=buf, mode="w:gz") as tar: for p in paths: if p.is_file(): - tar.add(str(p), arcname=p.name) + tar.add(str(p), arcname=names.get(p, p.name)) elif p.is_dir(): for child in p.rglob("*"): if child.is_file(): diff --git a/nemo_gym/sandbox/providers/ecs_fargate/provider.py b/nemo_gym/sandbox/providers/ecs_fargate/provider.py index ab364c11db..b586363940 100644 --- a/nemo_gym/sandbox/providers/ecs_fargate/provider.py +++ b/nemo_gym/sandbox/providers/ecs_fargate/provider.py @@ -14,14 +14,10 @@ """ECS Fargate sandbox provider. -Adapts the lifted :mod:`engine` (one stateful ``EcsFargateSandbox`` per -sandbox) to Gym's stateless ``SandboxProvider`` contract: per-sandbox engine -state lives in ``SandboxHandle.raw``; the provider methods delegate to it. - -The model endpoint is assumed reachable from inside the sandbox directly, or -routed via the SSH reverse tunnel when ``outside_endpoints`` are supplied -through ``spec.provider_options``. See the package README for the planned -Teleport follow-up. +Adapts :mod:`engine` (one stateful ``EcsFargateSandbox`` per sandbox) to Gym's ``SandboxProvider`` +contract: per-sandbox engine state lives in ``SandboxHandle.raw`` and the provider methods delegate +to it. The model endpoint is reached directly from the sandbox, or via the SSH reverse tunnel when +``outside_endpoints`` are supplied through ``spec.provider_options``. """ from __future__ import annotations @@ -171,9 +167,8 @@ async def aclose(self) -> None: def engine_config_from_mapping(config: dict[str, Any]) -> engine.EcsFargateConfig: """Build an ``EcsFargateConfig`` from provider kwargs. - When ``region`` is given but ``cluster`` is omitted, infrastructure is - auto-discovered from SSM (written by the reference Terraform); explicit - kwargs always win over SSM. Mirrors NEL's orchestrator resolution. + When ``region`` is given but ``cluster`` is omitted, infrastructure is auto-discovered from SSM; + explicit kwargs always win over SSM. """ config = dict(config) region = config.get("region") diff --git a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml index e7179350bf..f8ba83b15c 100644 --- a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml +++ b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml @@ -10,9 +10,8 @@ mini_swe_agent_2: name: policy_model concurrency: 8 env: sandbox - # Region-only config: the provider auto-discovers cluster/subnets/SGs/roles - # and the SSH-sidecar key ARNs from SSM (//ecs-sandbox/config, - # ssm_project defaults to "harbor"), exactly like NEL. + # Region-only config: the provider auto-discovers cluster/subnets/SGs/roles and the + # SSH-sidecar key ARNs from SSM (//ecs-sandbox/config, ssm_project="harbor"). sandbox_provider: ecs_fargate: region: ${oc.env:AWS_REGION} diff --git a/tests/unit_tests/test_ecs_fargate_engine.py b/tests/unit_tests/test_ecs_fargate_engine.py index e1fdcde463..3bf416ce80 100644 --- a/tests/unit_tests/test_ecs_fargate_engine.py +++ b/tests/unit_tests/test_ecs_fargate_engine.py @@ -28,6 +28,7 @@ import contextlib import io import json +import logging import threading import types from pathlib import Path @@ -2321,6 +2322,23 @@ async def test_upload_via_s3_packs_a_directory(tmp_path): assert "sub/f.txt" in tar.getnames() +async def test_upload_via_s3_honors_arcname_for_single_file(tmp_path): + # A large single file must land under the requested remote name, not the local temp basename. + f = tmp_path / "local-temp-xyz.bin" + f.write_text("data") + sb = _make_sandbox(s3_bucket="bkt") + _attach_exec_client(sb, engine.ExecResult("ok\n", "", 0)) + s3 = MagicMock() + s3.generate_presigned_url.return_value = "https://signed/url" + import tarfile + + with _patch_aws(s3=s3): + await sb._upload_via_s3([f], "/dest", arcnames={f: "final-name.bin"}) + body = s3.put_object.call_args.kwargs["Body"] + with tarfile.open(fileobj=io.BytesIO(body), mode="r:gz") as tar: + assert tar.getnames() == ["final-name.bin"] + + async def test_upload_via_s3_requires_bucket(): sb = _make_sandbox(s3_bucket=None) _attach_exec_client(sb) @@ -2328,6 +2346,23 @@ async def test_upload_via_s3_requires_bucket(): await sb._upload_via_s3([Path("/x")], "/dest") +def test_delete_s3_object_warns_once_then_silent(caplog): + cfg = _make_sandbox(s3_bucket="bkt")._cfg + s3 = MagicMock() + s3.delete_object.side_effect = Exception("AccessDenied") + engine._s3_delete_warned = False + try: + with _patch_aws(s3=s3): + with caplog.at_level(logging.WARNING, logger="nemo_gym.sandbox.providers.ecs_fargate.engine"): + engine._delete_s3_object(cfg, "k1") # first failure -> one WARNING + engine._delete_s3_object(cfg, "k2") # subsequent failures silenced at WARNING + finally: + engine._s3_delete_warned = False + warnings = [r for r in caplog.records if r.levelno == logging.WARNING and "s3:DeleteObject" in r.message] + assert len(warnings) == 1 + assert s3.delete_object.call_count == 2 # still attempted both times, never raised + + async def test_upload_via_s3_raises_when_extraction_fails(tmp_path): f = tmp_path / "payload.txt" f.write_text("data") @@ -2560,14 +2595,17 @@ async def test_upload_directory_recurses(tmp_path): async def test_upload_large_file_uses_s3(tmp_path): - f = tmp_path / "big.bin" + # Local basename intentionally differs from the requested remote name to catch the case where + # the S3 path drops the target filename. + f = tmp_path / "local-temp-name.bin" f.write_bytes(b"x" * (512 * 1024 + 1)) sb = _make_sandbox(ssh_sidecar=_exec_sidecar(), s3_bucket="bkt") _attach_exec_client(sb) with patch.object(engine.EcsFargateSandbox, "_upload_via_s3", AsyncMock()) as via_s3: - await sb.upload(f, "/remote/big.bin") + await sb.upload(f, "/remote/renamed.bin") via_s3.assert_awaited_once() assert via_s3.await_args.args[1] == "/remote" + assert via_s3.await_args.kwargs["arcnames"] == {f: "renamed.bin"} async def test_download_writes_bytes(tmp_path): From 5259409292a2fb47e0a7f697a8a346ec1cd11b70 Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Fri, 26 Jun 2026 13:53:44 +0200 Subject: [PATCH 10/11] docs(sandbox): add ECS Fargate guide under infrastructure/sandbox - Replace the provider README with a Fern page following the sandbox provider-page shape (setup / config / provider_options / resource mapping + isolation / first-run example / troubleshooting). - Add a minimal sandbox/index.mdx so the section renders standalone; complements the sandbox API docs structure without depending on that PR. Signed-off-by: Michal Bien --- .../sandbox/adding-a-provider.mdx | 2 +- .../infrastructure/sandbox/ecs-fargate.mdx | 125 ++++++++++++++++++ .../pages/infrastructure/sandbox/index.mdx | 6 + .../sandbox/providers/ecs_fargate/README.md | 79 ----------- 4 files changed, 132 insertions(+), 80 deletions(-) create mode 100644 fern/versions/latest/pages/infrastructure/sandbox/ecs-fargate.mdx delete mode 100644 nemo_gym/sandbox/providers/ecs_fargate/README.md diff --git a/fern/versions/latest/pages/infrastructure/sandbox/adding-a-provider.mdx b/fern/versions/latest/pages/infrastructure/sandbox/adding-a-provider.mdx index 5ae9f2b546..3f358b936c 100644 --- a/fern/versions/latest/pages/infrastructure/sandbox/adding-a-provider.mdx +++ b/fern/versions/latest/pages/infrastructure/sandbox/adding-a-provider.mdx @@ -1,7 +1,7 @@ --- title: "Adding a Sandbox Provider" description: "Implement and register a sandbox runtime backend for NeMo Gym." -position: 3 +position: 4 --- Add a provider when NeMo Gym needs to create sandboxes through a new runtime backend, such as a container service, HPC isolation layer, or in-house execution platform. The public `AsyncSandbox` and `Sandbox` facades stay the same; the provider owns runtime-specific create, command, file transfer, status, and cleanup behavior. diff --git a/fern/versions/latest/pages/infrastructure/sandbox/ecs-fargate.mdx b/fern/versions/latest/pages/infrastructure/sandbox/ecs-fargate.mdx new file mode 100644 index 0000000000..8f46c465ce --- /dev/null +++ b/fern/versions/latest/pages/infrastructure/sandbox/ecs-fargate.mdx @@ -0,0 +1,125 @@ +--- +title: "ECS Fargate" +description: "Run sandboxes as AWS ECS Fargate tasks reached over an SSH sidecar." +position: 3 +--- + +The `ecs_fargate` provider runs each sandbox as an AWS ECS Fargate task reached over an SSH sidecar. It implements the provider-neutral `SandboxProvider` contract, so any sandbox-backed agent or resources server can use it by selecting `ecs_fargate` in its sandbox config. + +## Setup + +ECS Fargate needs the AWS SDK and a provisioned account/region: + +```bash +uv pip install boto3 +``` + +- **Infrastructure.** The reference Terraform stack publishes its outputs to SSM at `//ecs-sandbox/config` (`ssm_project` defaults to `harbor`): cluster, subnets, security groups, task/execution roles, ECR mirror, EFS, and the SSH-sidecar key ARNs. A missing parameter raises an actionable error. +- **Credentials.** `AWS_ACCESS_KEY_ID` / `AWS_SECRET_ACCESS_KEY` (or an instance role) plus `AWS_REGION`. +- **Network.** The host must reach each task's SSH sidecar port (`52222`). Run inside the sandbox VPC/peered network, or allow the host IP on the sidecar security group. Exec, file transfer, and model reverse-tunnels all ride this SSH connection. + +## Provider Configuration + +Region-only is enough; everything else auto-discovers from SSM (explicit YAML always wins): + +```yaml +sandbox_provider: + ecs_fargate: + region: ${oc.env:AWS_REGION} + cpu: "2048" + memory: "8192" + ephemeral_storage_gib: 50 +``` + +| Field | Default | Purpose | +| --- | --- | --- | +| `region` | — | AWS region; enables SSM auto-discovery when `cluster` is omitted | +| `cpu` / `memory` | `"4096"` / `"8192"` | Fargate task size (CPU units / MiB) when set directly | +| `ephemeral_storage_gib` | 20 (implicit) | Task scratch disk; explicit values must be 21–200 | +| `auto_mirror` | `true` | Mirror a missing public image into the ECR mirror on demand | +| `ssm_project` | `harbor` | SSM namespace for auto-discovery | +| `environment_dir` | — | Build the task image from a Dockerfile directory via CodeBuild instead of using a prebuilt image | + +Cluster, subnets, security groups, roles, ECR mirror, EFS defaults, and the SSH-sidecar key ARNs are filled in from SSM when omitted. + +## Provider Options + +Per-sandbox options go in `SandboxSpec.provider_options`: + +| Key | Purpose | +| --- | --- | +| `volumes` | EFS mounts: a list of `{"container_path": "/mnt/efs", "efs": true}` entries. Each inherits the provider's `efs_filesystem_id` / `efs_access_point_id` unless it names its own. | +| `outside_endpoints` | Host URLs exposed inside the sandbox via an SSH reverse tunnel, as `{"url": ..., "env_var": ...}` (e.g. a model server). | +| `environment_dir` | Dockerfile directory to build the task image from via CodeBuild. | + +Common `SandboxSpec` fields map onto the task as follows: `ttl_s` → sidecar watchdog that stops the task, `ready_timeout_s` → task-startup timeout, `env` / `files` / `workdir` → container environment, seed files, and working directory. + +## Resource Mapping and Isolation + +`SandboxResources` maps onto the Fargate task definition: + +| `SandboxResources` | Fargate | +| --- | --- | +| `cpu` (vCPU) | task CPU units (`cpu * 1024`), validated against Fargate's CPU/memory pairs | +| `memory_mib` | task memory (MiB) | +| `disk_gib` | task ephemeral storage (21–200 GiB) | +| `gpu` | unsupported — raises `SandboxCreateError` | + +Each sandbox is a dedicated Fargate task with its own kernel, network namespace, and microVM boundary — there is no host sharing between sandboxes. The orchestrator holds a per-task SSH connection for exec and file transfer, and TTL is enforced by a sidecar watchdog that stops the task. + +## Images and On-Demand Mirroring + +ECS pulls task images from the account ECR mirror rather than their origin registry. A bare/public image (e.g. `docker.io/swebench/sweb.eval.x86_64.:latest`) resolves to the mirror tag `:`. Resolution order: + +1. `environment_dir` set → build the image via CodeBuild and use it. +2. Image is already an ECR reference → use verbatim (never re-mirrored). +3. Bare/public name + `auto_mirror=true` → mirror into ECR on demand (CodeBuild pull → retag → push) during `create`, then launch. + +The first task for a new image waits on a one-time build (~1–3 min for typical SWE-bench images); later tasks hit the ECR cache, and concurrent tasks for the same image de-duplicate onto a single build. Set `auto_mirror: false` to require a pre-populated mirror and fail fast on a miss. + +## First-Run Example + +Drive a sandbox directly through the provider-neutral API: + +```python +import asyncio + +from nemo_gym.sandbox import AsyncSandbox, SandboxResources, SandboxSpec + + +provider_config = {"ecs_fargate": {"region": "us-east-1"}} +spec = SandboxSpec( + image="python:3.12-slim", + ttl_s=1800, + ready_timeout_s=300, + resources=SandboxResources(cpu=2, memory_mib=8192), +) + + +async def main() -> None: + async with AsyncSandbox(provider_config, spec) as sandbox: + await sandbox.start() + result = await sandbox.exec("python3 --version", timeout_s=60) + print(result.stdout or result.stderr) + + +asyncio.run(main()) +``` + +For an end-to-end agent run, `mini_swe_agent_2` ships a ready config at `responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml` that selects this provider with region-only auto-discovery. + +## Troubleshooting + +| Symptom | Cause / fix | +| --- | --- | +| `SSM parameter '/…/ecs-sandbox/config' not found` | Infrastructure is not provisioned in that region. Run the Terraform stack, or set the ECS fields explicitly in YAML. | +| Exec/SSH times out after the task reaches RUNNING | The host cannot reach the sidecar port `52222`. Run inside the sandbox VPC or allow the host IP on the sidecar security group. | +| `Fargate memory for cpu=… must be …` | The CPU/memory pair is not a supported Fargate combination. Pick a valid pair. | +| `ephemeral storage must be between 21 and 200 GiB` | `disk_gib` is out of range. Omit it for the implicit 20 GiB default. | +| First task is slow or the image pull fails | The first task mirrors the image via CodeBuild (~1–3 min); later tasks hit the ECR cache. Set `auto_mirror: false` to require a pre-staged mirror. | +| S3 staging objects accumulate | The orchestrator role lacks `s3:DeleteObject`. Grant it, or set a bucket lifecycle policy to reap `*/ecs-sandbox/*` staging artifacts. | +| `ECS Fargate does not support GPU sandboxes` | Fargate has no GPU; use a GPU-capable provider instead. | + + +The reference sidecar security group allows `0.0.0.0/0` on `52222` for convenience. Restrict it to the orchestrator's egress IP (or move to a private/peered path) before non-smoke use. + diff --git a/fern/versions/latest/pages/infrastructure/sandbox/index.mdx b/fern/versions/latest/pages/infrastructure/sandbox/index.mdx index 8aaefd8740..4ba8f0a4e5 100644 --- a/fern/versions/latest/pages/infrastructure/sandbox/index.mdx +++ b/fern/versions/latest/pages/infrastructure/sandbox/index.mdx @@ -30,6 +30,12 @@ Run sandboxes as local Apptainer instances on a host or HPC node. provider apptainer + +Run sandboxes as AWS ECS Fargate tasks reached over an SSH sidecar. + +provider ecs-fargate + + Implement the `SandboxProvider` protocol and register a new runtime backend. diff --git a/nemo_gym/sandbox/providers/ecs_fargate/README.md b/nemo_gym/sandbox/providers/ecs_fargate/README.md deleted file mode 100644 index 5010060028..0000000000 --- a/nemo_gym/sandbox/providers/ecs_fargate/README.md +++ /dev/null @@ -1,79 +0,0 @@ -# ECS Fargate sandbox provider - -Runs each `nemo_gym.sandbox` sandbox as an AWS ECS Fargate task behind an SSH -sidecar. It implements the provider-neutral `SandboxProvider` contract, so any -sandbox-backed agent (not just mini-swe-agent) can use it by setting -`sandbox_provider.ecs_fargate` in its config. - -## Prerequisites - -- **Infrastructure** provisioned in the target account/region. The reference - Terraform stack publishes its outputs to SSM at - `//ecs-sandbox/config` (`ssm_project` defaults to `harbor`): - cluster, subnets, security groups, task/execution roles, ECR mirror, EFS, and - the SSH-sidecar key ARNs. A missing parameter raises an actionable error. -- **Credentials**: `AWS_ACCESS_KEY_ID`/`AWS_SECRET_ACCESS_KEY` (or an instance - role) plus `AWS_REGION`. -- **Network**: the host must reach each task's SSH sidecar port (`52222`) — run - inside the sandbox VPC/peered network, or allow the host IP on the sidecar - security group. Exec, file transfer, and model reverse-tunnels ride this SSH - connection. - -## Configuration - -Region-only is enough; everything else auto-discovers from SSM (explicit YAML -always wins): - -```yaml -sandbox_provider: - ecs_fargate: - region: ${oc.env:AWS_REGION} - cpu: "2048" - memory: "8192" - ephemeral_storage_gib: 50 -``` - -Common fields: - -| Field | Default | Purpose | -| --- | --- | --- | -| `region` | — | AWS region; enables SSM auto-discovery when `cluster` is omitted | -| `cpu` / `memory` | `"4096"` / `"8192"` | Fargate task size (CPU units / MiB) | -| `ephemeral_storage_gib` | task default | Task ephemeral disk | -| `auto_mirror` | `true` | Pull a missing public image into the ECR mirror on demand (see below) | -| `ssm_project` | `harbor` | SSM namespace for auto-discovery | -| `environment_dir` | — | Build a task image from a Dockerfile dir via CodeBuild instead of using a prebuilt image | - -Per-sandbox `ready_timeout_s`, `env`, `files`, `metadata`, and -`provider_options` (e.g. `platform`, `outside_endpoints`) come from the -`SandboxSpec`. - -## Images and on-demand mirroring - -ECS pulls task images from the account ECR mirror, not their origin registry. A -bare/public image (e.g. `docker.io/swebench/sweb.eval.x86_64.:latest`) -resolves to the mirror tag `:`. Resolution order: - -1. `environment_dir` set → build the image via CodeBuild and use it. -2. Image is already an ECR reference → use verbatim (never re-mirrored). -3. Bare/public name + `auto_mirror=true` → mirror into ECR on demand (CodeBuild - pull → retag → push) during `create`, then launch. Concurrent tasks for the - same image de-duplicate onto one build. - -Set `auto_mirror: false` to require a pre-populated mirror and fail fast on a -miss. The first task for a new image waits on a one-time build (~1–3 min for -typical SWE-bench images); later tasks hit the ECR cache. - -## Lifecycle - -`create` launches the task + SSH sidecar and returns once the exec server is -healthy. `exec`, `upload_file`, `download_file`, `status`, and `close` delegate -to the per-sandbox engine over the SSH tunnel. `outside_endpoints` (via -`spec.provider_options`) open reverse tunnels so an in-sandbox process can reach -a host-side endpoint (e.g. a model server). - -## Security - -The sidecar security group in the reference stack allows `0.0.0.0/0` on `52222` -for convenience. Restrict it to the orchestrator's egress IP (or move to a -private/Teleport path) before non-smoke use. From dbeb031fa3836e66e7140355d505b1403f34ee51 Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Mon, 6 Jul 2026 11:10:28 +0200 Subject: [PATCH 11/11] fix(sandbox): migrate ECS config to the #1708 decoupled provider pattern Rebased onto main, which merged #1708 (decouple sandbox provider config from the agent config). Align the ECS Fargate provider with that pattern: - Add a standalone provider config nemo_gym/sandbox/providers/ecs_fargate/configs/ecs_fargate.yaml -- a top-level `sandbox` block selected via `sandbox_provider: sandbox`, matching opensandbox/apptainer. Region-only; task cpu/memory/disk come from sandbox_spec.resources. - Drop the agent-coupled configs/mini_swe_agent_ecs_fargate.yaml; the generic mini_swe_agent_2.yaml now runs on ECS by adding the provider config path. - Bump the generic config's sandbox_spec.resources.disk_gib 20 -> 30 so it is valid on every provider (Fargate rejects an explicit 20; 21-200 only). - Update the ECS Fargate docs (Fern page + agent README) to the config-paths recipe. - Reconcile uv.lock (boto3 / botocore, sandbox extra) after the rebase. Signed-off-by: Michal Bien --- .../infrastructure/sandbox/ecs-fargate.mdx | 15 +++--- .../ecs_fargate/configs/ecs_fargate.yaml | 26 ++++++++++ .../mini_swe_agent_2/README.md | 12 ++--- .../configs/mini_swe_agent_2.yaml | 4 +- .../configs/mini_swe_agent_ecs_fargate.yaml | 37 ------------- uv.lock | 52 +++++++++++++++++++ 6 files changed, 96 insertions(+), 50 deletions(-) create mode 100644 nemo_gym/sandbox/providers/ecs_fargate/configs/ecs_fargate.yaml delete mode 100644 responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml diff --git a/fern/versions/latest/pages/infrastructure/sandbox/ecs-fargate.mdx b/fern/versions/latest/pages/infrastructure/sandbox/ecs-fargate.mdx index 8f46c465ce..838d85692d 100644 --- a/fern/versions/latest/pages/infrastructure/sandbox/ecs-fargate.mdx +++ b/fern/versions/latest/pages/infrastructure/sandbox/ecs-fargate.mdx @@ -20,15 +20,14 @@ uv pip install boto3 ## Provider Configuration -Region-only is enough; everything else auto-discovers from SSM (explicit YAML always wins): +NeMo Gym ships this provider config at `nemo_gym/sandbox/providers/ecs_fargate/configs/ecs_fargate.yaml`. It defines a top-level `sandbox` block that an agent selects with `sandbox_provider: sandbox`. Region-only is enough; everything else auto-discovers from SSM (explicit YAML always wins), and task cpu/memory/disk come from the agent's `sandbox_spec.resources`: ```yaml -sandbox_provider: +sandbox: + default_metadata: + sandbox-api: ecs-fargate ecs_fargate: region: ${oc.env:AWS_REGION} - cpu: "2048" - memory: "8192" - ephemeral_storage_gib: 50 ``` | Field | Default | Purpose | @@ -106,7 +105,11 @@ async def main() -> None: asyncio.run(main()) ``` -For an end-to-end agent run, `mini_swe_agent_2` ships a ready config at `responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml` that selects this provider with region-only auto-discovery. +For an end-to-end agent run, add this provider config to `mini_swe_agent_2`'s config paths — the agent config is unchanged: + +```bash +ng_run "+config_paths=[responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_2.yaml, nemo_gym/sandbox/providers/ecs_fargate/configs/ecs_fargate.yaml, ]" +``` ## Troubleshooting diff --git a/nemo_gym/sandbox/providers/ecs_fargate/configs/ecs_fargate.yaml b/nemo_gym/sandbox/providers/ecs_fargate/configs/ecs_fargate.yaml new file mode 100644 index 0000000000..84e9046961 --- /dev/null +++ b/nemo_gym/sandbox/providers/ecs_fargate/configs/ecs_fargate.yaml @@ -0,0 +1,26 @@ +# ECS Fargate sandbox provider config. +# +# `sandbox` is the instance name an agent references via `sandbox_provider: sandbox`; +# the child key `ecs_fargate` selects the provider and its value is passed to the +# provider constructor. Every shipped provider config binds the same name `sandbox`, +# so swapping providers is swapping this config path in `+config_paths` (no agent edit): +# +# AGENT=responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_2.yaml +# MODEL=responses_api_models/vllm_model/configs/vllm_model.yaml +# ng_run "+config_paths=[$AGENT, nemo_gym/sandbox/providers/ecs_fargate/configs/ecs_fargate.yaml, $MODEL]" +# +# Requires AWS credentials (AWS_REGION + access keys or an instance role) and the `sandbox` +# extra (boto3). See the ECS Fargate provider page under infrastructure/sandbox for setup +# (SSM infrastructure, the :52222 network requirement, and on-demand image mirroring). +# +# `default_metadata` (optional) is merged into every sandbox's spec metadata +# (SandboxSpec.metadata); an agent's own sandbox_spec.metadata overrides it. +sandbox: + default_metadata: + sandbox-api: ecs-fargate + ecs_fargate: + # Region-only: cluster / subnets / security-groups / roles and the SSH-sidecar key ARNs + # auto-discover from SSM (//ecs-sandbox/config; ssm_project defaults to harbor). + # Task cpu / memory / disk come from the agent's sandbox_spec.resources (mapped onto the + # Fargate task); set cpu / memory / ephemeral_storage_gib here only to change the defaults. + region: ${oc.env:AWS_REGION} diff --git a/responses_api_agents/mini_swe_agent_2/README.md b/responses_api_agents/mini_swe_agent_2/README.md index ac64da39c5..20109c255b 100644 --- a/responses_api_agents/mini_swe_agent_2/README.md +++ b/responses_api_agents/mini_swe_agent_2/README.md @@ -487,11 +487,11 @@ Kubernetes wrapper, add any outer per-sample guard there. ## Running SWE-bench on ECS Fargate -Swap OpenSandbox for ECS Fargate by using -`configs/mini_swe_agent_ecs_fargate.yaml` — the agent loop and SWE-bench -verifier are unchanged. For provider setup (AWS infra/SSM, credentials, the -`:52222` network requirement, and automatic image mirroring via `auto_mirror`) -see the [ECS Fargate provider README](../../nemo_gym/sandbox/providers/ecs_fargate/README.md). +Swap OpenSandbox for ECS Fargate by adding the ECS provider config to your +`+config_paths` — the agent config (`configs/mini_swe_agent_2.yaml`), agent loop, +and SWE-bench verifier are unchanged. For provider setup (AWS infra/SSM, +credentials, the `:52222` network requirement, and automatic image mirroring via +`auto_mirror`) see the ECS Fargate provider page under `infrastructure/sandbox`. SWE-bench task images are pulled into the ECR mirror on first use, so no manual image staging is needed. @@ -503,7 +503,7 @@ Each input row needs the SWE-bench instance fields shown under the verifier in-sandbox, so the model is never called: ```bash -CONFIG_PATHS="responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml,responses_api_models/vllm_model/configs/vllm_model.yaml" +CONFIG_PATHS="responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_2.yaml,nemo_gym/sandbox/providers/ecs_fargate/configs/ecs_fargate.yaml,responses_api_models/vllm_model/configs/vllm_model.yaml" AWS_REGION=us-east-1 ng_run "+config_paths=[$CONFIG_PATHS]" \ ++mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.run_golden=true diff --git a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_2.yaml b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_2.yaml index 9f225012aa..4c3b6ce505 100644 --- a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_2.yaml +++ b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_2.yaml @@ -27,7 +27,9 @@ mini_swe_agent_2: resources: cpu: 2 memory_mib: 8192 - disk_gib: 20 + # 30 GiB works on every provider: within Fargate's 21-200 GiB ephemeral range + # (an explicit 20 is rejected there) and fine as an ephemeral request elsewhere. + disk_gib: 30 provider_options: platform: os: linux diff --git a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml b/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml deleted file mode 100644 index f8ba83b15c..0000000000 --- a/responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_ecs_fargate.yaml +++ /dev/null @@ -1,37 +0,0 @@ -mini_swe_agent_2: - responses_api_agents: - mini_swe_agent_2: - entrypoint: app.py - domain: coding - description: Software engineering tasks driven by mini-swe-agent harness on ECS Fargate. - value: Improve agentic software engineering capabilities. - model_server: - type: responses_api_models - name: policy_model - concurrency: 8 - env: sandbox - # Region-only config: the provider auto-discovers cluster/subnets/SGs/roles and the - # SSH-sidecar key ARNs from SSM (//ecs-sandbox/config, ssm_project="harbor"). - sandbox_provider: - ecs_fargate: - region: ${oc.env:AWS_REGION} - cpu: "2048" - memory: "8192" - ephemeral_storage_gib: 50 - sandbox_spec: - ttl_s: 18000 - ready_timeout_s: 1200 - metadata: - benchmark: swebench-verified - harness: mini-swe-agent - sandbox-api: ecs-fargate - sandbox_environment_kwargs: - cwd: /testbed - conda_env: testbed - activate_conda: true - user: root - run_golden: false - step_timeout: 600 - eval_timeout: 1800 - skip_if_exists: false - step_limit: 250 diff --git a/uv.lock b/uv.lock index 0536f51bfb..778dae4a1e 100644 --- a/uv.lock +++ b/uv.lock @@ -297,6 +297,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/04/eb/f4151e0c7377a6e08a38108609ba5cede57986802757848688aeedd1b9e8/beautifulsoup4-4.13.5-py3-none-any.whl", hash = "sha256:642085eaa22233aceadff9c69651bc51e8bf3f874fb6d7104ece2beb24b47c4a", size = 105113, upload-time = "2025-08-24T14:06:14.884Z" }, ] +[[package]] +name = "boto3" +version = "1.43.40" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "botocore" }, + { name = "jmespath" }, + { name = "s3transfer" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3b/43/38cdaf9c7bb50f0922ef87e6b6696c6aa589d978e4311a8287537e84ecf7/boto3-1.43.40.tar.gz", hash = "sha256:a7108b9ce25b8f92d1ce96b9e35090794d5c58b219ff2f6a7a65c8986a4ed6f4", size = 112655, upload-time = "2026-07-03T00:28:26.776Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d7/3b/2abcc7f47e955a8cb45a6378ec69cd8ed2dcb3d14ceb67717dea020a794a/boto3-1.43.40-py3-none-any.whl", hash = "sha256:5fa80de4b4b7bab383dd8c563e235d51fcc7df4502662af56f0a2472c3c15651", size = 140029, upload-time = "2026-07-03T00:28:25.028Z" }, +] + +[[package]] +name = "botocore" +version = "1.43.40" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jmespath" }, + { name = "python-dateutil" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d0/50/269986277f852cc83029bccbdcdc0b343a685cfd570599e58029792808d8/botocore-1.43.40.tar.gz", hash = "sha256:2085a4314cfd2c8bc1d08ab8039f76c92e99278db0d2a0e2437010526d5d5d70", size = 15639899, upload-time = "2026-07-03T00:28:16.125Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/8c/5f7e73fd66b28f0705bc55d7060d41ef72328b656b86ca53e75765b3ba2c/botocore-1.43.40-py3-none-any.whl", hash = "sha256:0bc9d352267c9e48415c5d7bb61ff05c3f193eac2fc7e69cfd229a05fbab67d6", size = 15323870, upload-time = "2026-07-03T00:28:12.56Z" }, +] + [[package]] name = "cachetools" version = "5.5.2" @@ -1143,6 +1171,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b3/4a/4175a563579e884192ba6e81725fc0448b042024419be8d83aa8a80a3f44/jiter-0.10.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3aa96f2abba33dc77f79b4cf791840230375f9534e5fac927ccceb58c5e604a5", size = 354213, upload-time = "2025-05-18T19:04:41.894Z" }, ] +[[package]] +name = "jmespath" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d3/59/322338183ecda247fb5d1763a6cbe46eff7222eaeebafd9fa65d4bf5cb11/jmespath-1.1.0.tar.gz", hash = "sha256:472c87d80f36026ae83c6ddd0f1d05d4e510134ed462851fd5f754c8c3cbb88d", size = 27377, upload-time = "2026-01-22T16:35:26.279Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl", hash = "sha256:a5663118de4908c91729bea0acadca56526eb2698e83de10cd116ae0f4e97c64", size = 20419, upload-time = "2026-01-22T16:35:24.919Z" }, +] + [[package]] name = "joblib" version = "1.5.3" @@ -1714,6 +1751,7 @@ dependencies = [ [package.optional-dependencies] all = [ + { name = "boto3" }, { name = "coverage" }, { name = "mypy" }, { name = "opensandbox" }, @@ -1738,6 +1776,7 @@ dev = [ { name = "ruff" }, ] sandbox = [ + { name = "boto3" }, { name = "opensandbox" }, { name = "tenacity" }, ] @@ -1760,6 +1799,7 @@ docs = [ requires-dist = [ { name = "aiohttp", specifier = ">=3.14.1" }, { name = "anthropic", specifier = "<=0.109.2" }, + { name = "boto3", marker = "extra == 'sandbox'", specifier = ">=1.34" }, { name = "coverage", extras = ["toml"], marker = "extra == 'dev'" }, { name = "datasets" }, { name = "devtools" }, @@ -2950,6 +2990,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/84/a8/001d4a7c2b37623a3fd7463208267fb906df40ff31db496157549cfd6e72/ruff-0.12.11-py3-none-win_arm64.whl", hash = "sha256:bae4d6e6a2676f8fb0f98b74594a048bae1b944aab17e9f5d504062303c6dbea", size = 12135290, upload-time = "2025-08-28T13:59:06.933Z" }, ] +[[package]] +name = "s3transfer" +version = "0.19.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "botocore" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f6/94/dcdaeb1713cab9c84def276cfac7388b17c7d9855bbcfe88d77e4dbafd44/s3transfer-0.19.0.tar.gz", hash = "sha256:ce436931687addc4c1712d52d40b32f53e88315723f107ffa20ba82b05a0f685", size = 165171, upload-time = "2026-06-16T19:44:51.599Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/46/5f/4c174edad94f82de888ac00a5ddd8d07b35609b6c94f0bdf4d74af57703e/s3transfer-0.19.0-py3-none-any.whl", hash = "sha256:777cc2415536f1debadb5c2ef7779275d0fc0fe0e042411cdd6caebeb2685262", size = 90101, upload-time = "2026-06-16T19:44:50.439Z" }, +] + [[package]] name = "scikit-learn" version = "1.9.0"