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..838d85692d --- /dev/null +++ b/fern/versions/latest/pages/infrastructure/sandbox/ecs-fargate.mdx @@ -0,0 +1,128 @@ +--- +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 + +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: + default_metadata: + sandbox-api: ecs-fargate + ecs_fargate: + region: ${oc.env:AWS_REGION} +``` + +| 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, 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 + +| 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/__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/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/nemo_gym/sandbox/providers/ecs_fargate/engine.py b/nemo_gym/sandbox/providers/ecs_fargate/engine.py new file mode 100644 index 0000000000..863b3e18b2 --- /dev/null +++ b/nemo_gym/sandbox/providers/ecs_fargate/engine.py @@ -0,0 +1,2477 @@ +# 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. + +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 + +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: 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: + # 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 +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" + + +_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. + + 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 + # 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 + 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: + # 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, or ``None`` for Fargate's implicit 20 GiB default. + + Fargate accepts only 21-200 GiB explicitly (20 is valid only as the omitted default). + """ + 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} + + +# 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 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})") + low, high = memory_range + if not low <= memory <= high: + 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 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: + 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) + + +# 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 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") + + +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 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 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: + 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 == "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 + 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 (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) + 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 _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``. + + 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): + logger.info("ECR cache hit — skipping build: %s", image_url) + return image_url + + 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: + 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 (would drift the permit cap). + 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: + 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: + semaphore.acquire() + try: + if not force and cls.image_exists_in_ecr(ecr_repo, tag, cfg.region): + return image_url + build() + return image_url + finally: + semaphore.release() + finally: + with cls._lock: + cls._inflight_builds.pop(tag, None) + 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: + """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: + 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}" + + 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) + + 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: + 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() + 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) -> 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") + # 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( + 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) + try: + cb = boto3.client("codebuild", region_name=cfg.region) + 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( + 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( + cls, + *, + cfg: EcsFargateConfig, + buildspec: str, + job_label: str = "harness-build", + timeout_minutes: int | None = None, + ) -> None: + """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) + 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 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 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] + + 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._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] = [] + 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}={shlex.quote(v)}" for k, v in env.items()) + shell_cmd = f"export {exports} && {shell_cmd}" + if cwd: + 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)}' + 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] + # 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: + 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: + # 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] + + 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") + ) + 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) + ): + # 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) + + 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 self._spec.image: + # 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 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 + 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) + elif "{task_ip}" in value: + # _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." + ) + 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() + 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(cpu), + "memory": str(memory), + "containerDefinitions": containers, + } + 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] + 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 + + _validate_fargate_cpu_memory(int(cfg.cpu), int(cfg.memory)) + 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 + 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]]]: + """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. + """ + 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": filesystem_id, + "transitEncryption": "ENABLED", + } + if access_point_id: + efs_cfg["authorizationConfig"] = { + "accessPointId": 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") + # 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: + 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 + 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).""" + 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, *, 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=names.get(p, 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, + ) + 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", + 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..b586363940 --- /dev/null +++ b/nemo_gym/sandbox/providers/ecs_fargate/provider.py @@ -0,0 +1,256 @@ +# 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 :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 + +from dataclasses import replace +from pathlib import Path +from typing import Any + +from nemo_gym.sandbox.providers.base import ( + SandboxCreateError, + SandboxExecResult, + SandboxHandle, + SandboxResources, + 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"), + ) + + +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.""" + + name = "ecs_fargate" + + def __init__(self, **config: Any) -> None: + self._cfg = engine_config_from_mapping(config) + + async def create(self, spec: SandboxSpec) -> SandboxHandle: + 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)) + 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) -> 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; + explicit kwargs always win over SSM. + """ + 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, + 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"), + 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..e9e074dbca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -248,6 +248,12 @@ sandbox = [ # Updated: Sat May 16, 2026 with opensandbox>=0.1.9 # License: Apache 2.0 "opensandbox>=0.1.9", + + # 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). diff --git a/responses_api_agents/mini_swe_agent_2/README.md b/responses_api_agents/mini_swe_agent_2/README.md index ac285143ac..20109c255b 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 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. + +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_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 + +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/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/sandbox_environment.py b/responses_api_agents/mini_swe_agent_2/sandbox_environment.py index 704122b9ce..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,7 +124,18 @@ 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 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 ' + '[ -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_engine.py b/tests/unit_tests/test_ecs_fargate_engine.py new file mode 100644 index 0000000000..3bf416ce80 --- /dev/null +++ b/tests/unit_tests/test_ecs_fargate_engine.py @@ -0,0 +1,2728 @@ +# 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 logging +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`` 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 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 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"}]) + + +def test_image_exists_in_ecr_false_on_image_not_found(): + ecr = MagicMock() + 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") + 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(): + # 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 + 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_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" + 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(): + # 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(): + # 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 + 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_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") + 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(): + 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()) == "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()) + + +def test_resolve_codebuild_project_creates_project(): + cfg = _build_cfg() + cb = MagicMock() + with _patch_aws(): + 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" + + +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) + assert name == "ecs-sandbox-build" + + +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) + + +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_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) + with pytest.raises(ValueError, match="s3_bucket is required"): + 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") + 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): + # 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/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): + 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 new file mode 100644 index 0000000000..5832c1b62f --- /dev/null +++ b/tests/unit_tests/test_ecs_fargate_provider.py @@ -0,0 +1,512 @@ +# 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 _apply_spec_overrides, 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", # pragma: allowlist secret + "private_key_secret_arn": "arn:aws:secretsmanager:us-east-1:1234:secret:priv", # pragma: allowlist secret + "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", # pragma: allowlist secret + "private_key_secret_arn": "arn:ssm:priv", # pragma: allowlist secret + "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" # pragma: allowlist secret + 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 _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", # 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)}" + 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") + 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 + + +# ── 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})) + + +# ── 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") + + +# ── 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() 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"