diff --git a/pyproject.toml b/pyproject.toml index 940668da..9146812a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,7 @@ classifiers = [ dependencies = [ "anthropic>=0.97.0", "anyio>=4.13.0", + "boto3>=1.40.11", "botocore>=1.43.11", "claude-agent-sdk>=0.1.17", "daytona>=0.121.0", diff --git a/scripts/e2e_sqs_kubernetes.py b/scripts/e2e_sqs_kubernetes.py new file mode 100644 index 00000000..a25440f4 --- /dev/null +++ b/scripts/e2e_sqs_kubernetes.py @@ -0,0 +1,936 @@ +#!/usr/bin/env python3 +"""End-to-end test for pier's SQSKubernetesEnvironment against a real +sandbox-eks cluster. + +Drives ``pier.environments.sqs_kubernetes.SQSKubernetesEnvironment`` against a +live ``docker_k8s_consumer`` running on an EKS cluster (see the agent-dist +repo for the consumer-side chart). All infrastructure is shared/multi-tenant, +so the script: + +- caps concurrency (``--concurrency``, default 20) +- caps total lifecycle count (``--iterations`` * concurrency, hard-capped at 100 + by default via ``--max-total`` so the tests can't run away) +- labels every sandbox with ``slurm_user`` (default the pier PR marker below) +- runs a final orphan sweep by listing pods matching the label and issuing + ``DELETE`` for any survivors + +Bypasses ``BaseEnvironment.__init__`` the same way ``tests/environments/ +test_sqs_kubernetes.py::_bare_env`` does so the script does not need a real +pier TrialPaths / EnvironmentConfig context — it only needs the SQS-Kubernetes +wire protocol. + +Usage (from the pier repo root, inside its uv venv):: + + python scripts/e2e_sqs_kubernetes.py \\ + --queue-url https://sqs.eu-west-1.amazonaws.com/475108760152/docker-requests \\ + --s3-bucket sqs-message-queue-large-objects-475108760152-eu-west-1-an \\ + --registry-url 475108760152.dkr.ecr.eu-west-1.amazonaws.com/sandbox + +Requires: +- valid AWS credentials in the boto3 credential chain +- kubectl configured for the prod EKS context (only used for the orphan sweep + fallback; the tests themselves do not touch kubectl) +""" + +from __future__ import annotations + +import argparse +import asyncio +import hashlib +import json +import logging +import os +import random +import secrets +import shutil +import statistics +import subprocess +import sys +import tempfile +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock + +import boto3 + +# Make ``src/`` importable without ``pip install -e .`` when the script is +# executed directly from the repo root. +_REPO_ROOT = Path(__file__).resolve().parent.parent +if (_REPO_ROOT / "src").is_dir(): + sys.path.insert(0, str(_REPO_ROOT / "src")) + +from pier.environments.sqs_kubernetes import ( # noqa: E402 + K8sInfraError, + S3PayloadError, + SQSKubernetesEnvironment, + SqsConsumerError, +) +from pier.models.task.config import TaskOS # noqa: E402 + +# Marker label stamped on every test sandbox for orphan cleanup + accounting. +DEFAULT_SLURM_USER = "pier-e2e-test-2026-08-21" + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(name)s: %(message)s", + datefmt="%H:%M:%S", +) +log = logging.getLogger("pier-e2e") + + +def _make_env( + *, + sqs_url: str, + s3_bucket: str, + registry_url: str, + sqs_region: str, + s3_region: str, + slurm_user: str, + docker_image: str, + session_id: str, +) -> SQSKubernetesEnvironment: + """Instantiate SQSKubernetesEnvironment WITHOUT running Base.__init__. + + Same pattern as tests/environments/test_sqs_kubernetes.py::_bare_env — + the SQS-Kubernetes wire protocol is a self-contained subset of pier's + environment API and doesn't need a real TrialPaths / EnvironmentConfig + to exercise start/exec/stop/upload/download. + """ + env = SQSKubernetesEnvironment.__new__(SQSKubernetesEnvironment) + + # SQS-K8s specific config + env._sqs_queue_url = sqs_url + env._s3_bucket = s3_bucket + env._registry_url = registry_url + env._sqs_region = sqs_region + env._s3_region = s3_region + env._s3_access_key_id = None + env._s3_secret_access_key = None + env._first_chunk_timeout = 120 + env._s3_threshold = 200 * 1024 + env._create_max_attempts = 5 + env._pre_install_commands = [] + env._ephemeral_storage_limit = "" + env._slurm_user = slurm_user + env._slurm_job_id = "pier-e2e" + env._repo_url = "" + env._repo_commit = "" + env._repo_dest = "repo" + env._inject_prebuilt_env_files = False # skip Dockerfile parsing + env._explicit_volume_mounts = None + env._volume_mounts = [] # no volume mounts for the smoke test + env._egress_proxy_url = "http://egress-proxy.sandbox-proxy.svc.cluster.local:3128" + env._egress_token = None + env._metrics_bridge_queue = "" + env._metrics_bridge_url = None + env._metrics_push_task = None + env._proxy_id = secrets.token_hex(8) + env._pre_generated_cid = None + env._all_pre_generated_cids = [] + env._sandbox_container_id = None + env._timing = {} + env._exec_count = 0 + env._exec_total_secs = 0.0 + env._start_time = None + + # boto3 clients — sender is process-wide (mirrors real __init__ path). + if SQSKubernetesEnvironment._shared_sender_sqs_client is None: + SQSKubernetesEnvironment._shared_sender_sqs_client = boto3.client( + "sqs", region_name=sqs_region + ) + env._sqs = SQSKubernetesEnvironment._shared_sender_sqs_client + env._s3 = boto3.client("s3", region_name=s3_region) + + # BaseEnvironment fields + env.environment_dir = Path(tempfile.gettempdir()) # unused when docker_image set + env.environment_name = f"pier-e2e-{session_id[:8]}" + env.session_id = session_id + env.trial_paths = MagicMock() + env.default_user = None + env._persistent_env = {} + env.logger = logging.getLogger(f"pier-e2e.env.{session_id[:8]}") + + # Minimum viable task_env_config + task_env_config = MagicMock() + task_env_config.docker_image = docker_image + task_env_config.workdir = None + task_env_config.allow_internet = True + task_env_config.os = TaskOS.LINUX + task_env_config.env = {} + task_env_config.cpus = None + task_env_config.memory_mb = None + task_env_config.storage_mb = None + task_env_config.gpus = None + env.task_env_config = task_env_config + + # Default network_allowlist (empty) — test 4 overrides this before start(). + from pier.models.agent.network import NetworkAllowlist + + env.network_allowlist = NetworkAllowlist() + + # Resource-mode plumbing used by start() + from pier.models.trial.config import ResourceMode + + env._cpu_resource_mode = ResourceMode.AUTO + env._memory_resource_mode = ResourceMode.AUTO + env._override_cpus = None + env._override_memory_mb = None + env._override_storage_mb = None + env._override_gpus = None + + return env + + +# ── result types ───────────────────────────────────────────────────────────── + + +@dataclass +class TestResult: + name: str + passed: bool + duration_s: float + evidence: dict[str, Any] = field(default_factory=dict) + error: str | None = None + + +@dataclass +class LatencySample: + op: str + duration_ms: float + sandbox_id: str + + +# ── Test 1: E2E happy path ─────────────────────────────────────────────────── + + +async def test_1_happy_path(cfg: argparse.Namespace) -> TestResult: + t0 = time.time() + ev: dict[str, Any] = {} + env = _make_env( + sqs_url=cfg.queue_url, + s3_bucket=cfg.s3_bucket, + registry_url=cfg.registry_url, + sqs_region=cfg.region, + s3_region=cfg.region, + slurm_user=cfg.slurm_user, + docker_image=cfg.image, + session_id=f"t1-{secrets.token_hex(4)}", + ) + try: + t_start = time.time() + await env.start() + ev["create_latency_s"] = round(time.time() - t_start, 3) + ev["sandbox_id"] = (env._sandbox_container_id or "")[:12] + + t_exec = time.time() + r1 = await env.exec("echo hello world") + ev["exec1_latency_s"] = round(time.time() - t_exec, 3) + ev["exec1_stdout"] = r1.stdout.strip()[:60] + ev["exec1_rc"] = r1.return_code + + t_exec2 = time.time() + r2 = await env.exec("cat /etc/os-release 2>/dev/null || uname -a") + ev["exec2_latency_s"] = round(time.time() - t_exec2, 3) + ev["exec2_stdout_head"] = r2.stdout.strip().split("\n")[0][:80] + + # Upload a small file via archive PUT + with tempfile.NamedTemporaryFile( + "wb", delete=False, suffix=".bin" + ) as f: + payload_up = secrets.token_bytes(1024) + f.write(payload_up) + up_path = Path(f.name) + try: + t_up = time.time() + await env.upload_file(up_path, "/tmp/pier_e2e_smoke.bin") + ev["upload_latency_s"] = round(time.time() - t_up, 3) + + # sha256 verify via download_file + with tempfile.TemporaryDirectory() as td: + dst = Path(td) / "downloaded.bin" + t_dl = time.time() + await env.download_file("/tmp/pier_e2e_smoke.bin", dst) + ev["download_latency_s"] = round(time.time() - t_dl, 3) + got = dst.read_bytes() + match = hashlib.sha256(got).digest() == hashlib.sha256(payload_up).digest() + ev["upload_download_roundtrip_ok"] = match + if not match: + raise AssertionError( + f"roundtrip mismatch: sent {len(payload_up)}B, got {len(got)}B" + ) + finally: + up_path.unlink(missing_ok=True) + + r3 = await env.exec("printenv OPENAI_BASE_URL || true") + ev["openai_base_url"] = r3.stdout.strip() or "" + r4 = await env.exec("printenv HTTP_PROXY || true") + ev["http_proxy"] = r4.stdout.strip() or "" + + # Strict pass rule: exec must return rc=0 and stdout must contain the + # echoed marker. Prevents images without `bash` (rc=127) from being + # counted as a pass. + strict_ok = ( + r1.return_code == 0 + and "hello world" in r1.stdout + and r2.return_code == 0 + and ev["upload_download_roundtrip_ok"] + ) + ev["strict_pass"] = strict_ok + return TestResult("1_happy_path", strict_ok, time.time() - t0, ev) + except Exception as exc: # noqa: BLE001 - surface every failure with type + return TestResult( + "1_happy_path", + False, + time.time() - t0, + ev, + error=f"{type(exc).__name__}: {exc}", + ) + finally: + try: + await env.stop(delete=True) + except Exception as exc: # noqa: BLE001 - best-effort cleanup + log.warning("test_1 stop() failed: %s", exc) + + +# ── Test 2: Large-payload roundtrip ────────────────────────────────────────── + + +async def test_2_large_payload(cfg: argparse.Namespace) -> TestResult: + t0 = time.time() + ev: dict[str, Any] = {} + env = _make_env( + sqs_url=cfg.queue_url, + s3_bucket=cfg.s3_bucket, + registry_url=cfg.registry_url, + sqs_region=cfg.region, + s3_region=cfg.region, + slurm_user=cfg.slurm_user, + docker_image=cfg.image, + session_id=f"t2-{secrets.token_hex(4)}", + ) + s3 = boto3.client("s3", region_name=cfg.region) + + try: + await env.start() + ev["sandbox_id"] = (env._sandbox_container_id or "")[:12] + + # 2 MiB upload — must go through S3 fallback (threshold is 200 KiB). + # Instrument _upload_to_s3 with a call-counter so we can prove S3 was + # actually used regardless of how many other payloads land in the + # bucket during the test. + payload = secrets.token_bytes(2 * 1024 * 1024) + s3_uploads: list[str] = [] + real_upload = env._upload_to_s3 + + async def _spy(body: bytes, prefix: str = "sqs-k8s-payloads") -> str: + key = await real_upload(body, prefix=prefix) + s3_uploads.append(key) + return key + + env._upload_to_s3 = _spy # type: ignore[method-assign] + + with tempfile.NamedTemporaryFile("wb", delete=False, suffix=".bin") as f: + f.write(payload) + up_path = Path(f.name) + try: + t_up = time.time() + await env.upload_file(up_path, "/tmp/pier_e2e_big.bin") + ev["upload_2mib_s"] = round(time.time() - t_up, 3) + finally: + up_path.unlink(missing_ok=True) + ev["s3_uploads_count"] = len(s3_uploads) + ev["s3_uploads_sample"] = s3_uploads[:2] + ev["s3_fallback_observed"] = len(s3_uploads) >= 1 + + # Download and byte-compare. + with tempfile.TemporaryDirectory() as td: + dst = Path(td) / "big.bin" + t_dl = time.time() + await env.download_file("/tmp/pier_e2e_big.bin", dst) + ev["download_2mib_s"] = round(time.time() - t_dl, 3) + got = dst.read_bytes() + ev["download_size_match"] = len(got) == len(payload) + ev["download_sha256_match"] = ( + hashlib.sha256(got).hexdigest() == hashlib.sha256(payload).hexdigest() + ) + if not ev["download_sha256_match"]: + raise AssertionError("2MiB roundtrip sha mismatch") + + # 3 MiB dd — pod-generated, verify size + envelope roundtrip. + r = await env.exec( + "dd if=/dev/urandom of=/tmp/pier_dd_big bs=1M count=3 2>/dev/null " + "&& stat -c%s /tmp/pier_dd_big 2>/dev/null || wc -c int: + """Cheap count-only pass — Bounded by pagination first-page + short window.""" + try: + resp = s3.list_objects_v2(Bucket=bucket, Prefix=prefix, MaxKeys=1000) + except Exception: # noqa: BLE001 - listing is opportunistic + return -1 + contents = resp.get("Contents", []) + cutoff = time.time() - minutes * 60 + return sum( + 1 for o in contents if o.get("LastModified").timestamp() >= cutoff + ) + + +# ── Test 3: Concurrency stress ─────────────────────────────────────────────── + + +async def _one_lifecycle(cfg: argparse.Namespace, idx: int) -> dict[str, Any]: + env = _make_env( + sqs_url=cfg.queue_url, + s3_bucket=cfg.s3_bucket, + registry_url=cfg.registry_url, + sqs_region=cfg.region, + s3_region=cfg.region, + slurm_user=cfg.slurm_user, + docker_image=cfg.image, + session_id=f"stress-{idx}-{secrets.token_hex(3)}", + ) + latencies: dict[str, float] = {} + error: str | None = None + error_code: str | None = None + try: + t0 = time.time() + await env.start() + latencies["create"] = time.time() - t0 + + for k in range(3): + t1 = time.time() + r = await env.exec(f"echo stress-{idx}-{k} && sleep 0.1 && echo done") + latencies[f"exec{k}"] = time.time() - t1 + if r.return_code != 0: + error = f"exec{k} rc={r.return_code} stderr={r.stderr[:100]}" + break + return { + "idx": idx, + "sandbox_id": (env._sandbox_container_id or "")[:12], + "latencies": latencies, + "error": error, + "error_code": error_code, + } + except SqsConsumerError as exc: + return { + "idx": idx, + "sandbox_id": (env._sandbox_container_id or "")[:12] if env._sandbox_container_id else "", + "latencies": latencies, + "error": f"{type(exc).__name__}: {exc.message}", + "error_code": exc.error_code, + } + except K8sInfraError as exc: + return { + "idx": idx, + "sandbox_id": (env._sandbox_container_id or "")[:12] if env._sandbox_container_id else "", + "latencies": latencies, + "error": f"K8sInfraError: {exc.message}", + "error_code": exc.error_code, + } + except Exception as exc: # noqa: BLE001 + return { + "idx": idx, + "sandbox_id": (env._sandbox_container_id or "")[:12] if env._sandbox_container_id else "", + "latencies": latencies, + "error": f"{type(exc).__name__}: {exc}", + "error_code": None, + } + finally: + try: + await env.stop(delete=True) + except Exception as exc: # noqa: BLE001 + log.warning("stress[%d] stop failed: %s", idx, exc) + + +async def test_3_stress(cfg: argparse.Namespace) -> TestResult: + t0 = time.time() + concurrency = cfg.concurrency + total = min(cfg.concurrency * cfg.iterations, cfg.max_total) + if total < 1: + return TestResult( + "3_stress", False, 0.0, error="max_total < 1, nothing to do" + ) + rounds = (total + concurrency - 1) // concurrency + + log.info( + "Test 3: %d rounds x %d concurrent = %d total lifecycles", + rounds, concurrency, total, + ) + + all_results: list[dict[str, Any]] = [] + idx = 0 + for r in range(rounds): + n = min(concurrency, total - idx) + if n <= 0: + break + log.info("stress round %d/%d: %d concurrent", r + 1, rounds, n) + batch = await asyncio.gather( + *(_one_lifecycle(cfg, idx + i) for i in range(n)), + return_exceptions=False, + ) + all_results.extend(batch) + idx += n + + create_lats = [r["latencies"]["create"] for r in all_results if "create" in r["latencies"]] + exec_lats = [ + v + for r in all_results + for k, v in r["latencies"].items() + if k.startswith("exec") + ] + errors = [r for r in all_results if r["error"]] + error_codes: dict[str, int] = {} + for e in errors: + c = e.get("error_code") or "no_error_code" + error_codes[c] = error_codes.get(c, 0) + 1 + + def _pct(xs: list[float], p: float) -> float: + if not xs: + return 0.0 + xs2 = sorted(xs) + k = max(0, min(len(xs2) - 1, int(round((p / 100.0) * (len(xs2) - 1))))) + return round(xs2[k], 3) + + ev = { + "total_lifecycles": len(all_results), + "concurrency": concurrency, + "rounds": rounds, + "wall_clock_s": round(time.time() - t0, 2), + "successes": len(all_results) - len(errors), + "failures": len(errors), + "success_rate": ( + round((len(all_results) - len(errors)) / max(1, len(all_results)) * 100, 1) + ), + "error_code_breakdown": error_codes, + "sample_errors": [e["error"] for e in errors[:5]], + "create_p50_s": _pct(create_lats, 50), + "create_p95_s": _pct(create_lats, 95), + "create_p99_s": _pct(create_lats, 99), + "exec_p50_s": _pct(exec_lats, 50), + "exec_p95_s": _pct(exec_lats, 95), + "exec_p99_s": _pct(exec_lats, 99), + } + # Pass rule: ≥90% success rate (prod is noisy — some SQS/pod-scheduling + # transients are expected). Absolute failures cap at 10% before we mark it + # a regression. + passed = ev["success_rate"] >= 90.0 + return TestResult("3_stress", passed, time.time() - t0, ev) + + +# ── Test 4: network_allowlist plumbing check ───────────────────────────────── + + +async def test_4_network_allowlist(cfg: argparse.Namespace) -> TestResult: + """Verify per-agent egress allowlist enforcement end-to-end. + + Requires the target cluster to be running the server-side changes from + agent-dist PR #216 (docker_k8s_consumer mints an ``egress_token`` and + publishes to the ``egress-allowlist`` ConfigMap; egress-proxy checks it + via the ``Proxy-Authorization`` header). On a cluster still running the + pre-PR-216 consumer, the create response omits ``egress_token`` — this + test detects that (``env._egress_token is None``), emits a WARN, and + falls back to the pre-PR-216 "documented as no-op" behaviour so the + test still passes as long as the pod's egress is at least reachable. + """ + from pier.models.agent.network import NetworkAllowlist + + t0 = time.time() + ev: dict[str, Any] = {} + env = _make_env( + sqs_url=cfg.queue_url, + s3_bucket=cfg.s3_bucket, + registry_url=cfg.registry_url, + sqs_region=cfg.region, + s3_region=cfg.region, + slurm_user=cfg.slurm_user, + docker_image=cfg.image, + session_id=f"t4-{secrets.token_hex(4)}", + ) + # Deny general internet, allow only httpbin.org. If the server-side + # allowlist is wired up, https://example.com must 403. + env.network_allowlist = NetworkAllowlist(domains=["httpbin.org"]) + env.task_env_config.allow_internet = False + try: + await env.start() + ev["sandbox_id"] = (env._sandbox_container_id or "")[:12] + + ev["egress_token_minted"] = bool(env._egress_token) + if env._egress_token: + # Mask by design — never log the whole secret to summary JSON. + ev["egress_token_prefix"] = env._egress_token[:8] + "..." + + # Compute the env dict the agent would actually see. + agent_env = env.agent_process_env({}) or {} + ev["agent_http_proxy"] = agent_env.get("HTTP_PROXY", "") + ev["agent_no_proxy_head"] = (agent_env.get("NO_PROXY") or "")[:80] + + if not env._egress_token: + ev["finding"] = ( + "WARN: consumer response did not include egress_token. Target " + "cluster is running pre-PR-216 consumer code — per-agent " + "allowlist enforcement is a no-op here. Falling back to " + "documented-as-no-op behaviour for the runtime assertion " + "portion of the test." + ) + log.warning(ev["finding"]) + # Documented-as-no-op fallback: just verify the sandbox is up. + r = await env.exec("echo hello && exit 0") + ev["fallback_exec_rc"] = r.return_code + passed = r.return_code == 0 + return TestResult("4_network_allowlist", passed, time.time() - t0, ev) + + # Server-side allowlist enforcement path. + # Build shell environment string so curl sees the proxy vars. + env_prefix = " ".join( + f"{k}={v}" for k, v in agent_env.items() if k in ( + "HTTP_PROXY", "HTTPS_PROXY", "NO_PROXY", + "http_proxy", "https_proxy", "no_proxy", + ) + ) + + # httpbin.org is on the allowlist → expect 200. + r_allowed = await env.exec( + f"{env_prefix} curl -sS -o /dev/null -w '%{{http_code}}' " + "--max-time 20 https://httpbin.org/get" + ) + ev["httpbin_code"] = r_allowed.stdout.strip() + + # example.com is NOT on the allowlist → expect 403 (egress-proxy denies). + r_denied = await env.exec( + f"{env_prefix} curl -sS -o /dev/null -w '%{{http_code}}' " + "--max-time 20 https://example.com" + ) + ev["example_com_code"] = r_denied.stdout.strip() + + # Pass criteria: httpbin OK (2xx), example.com denied (403). + # We accept any 2xx for the allowed host to tolerate httpbin's + # occasional 301/redirect noise. + allowed_ok = ev["httpbin_code"].startswith("2") + denied_ok = ev["example_com_code"] == "403" + passed = allowed_ok and denied_ok + ev["allowlist_enforced"] = passed + ev["finding"] = ( + "PASS: allowlist enforcement observed end-to-end" + if passed + else ( + f"FAIL: httpbin={ev['httpbin_code']} (want 2xx), " + f"example.com={ev['example_com_code']} (want 403)" + ) + ) + return TestResult("4_network_allowlist", passed, time.time() - t0, ev) + except Exception as exc: # noqa: BLE001 + return TestResult( + "4_network_allowlist", + False, + time.time() - t0, + ev, + error=f"{type(exc).__name__}: {exc}", + ) + finally: + try: + await env.stop(delete=True) + except Exception as exc: # noqa: BLE001 + log.warning("test_4 stop() failed: %s", exc) + + +# ── orphan cleanup ─────────────────────────────────────────────────────────── + + +def cleanup_orphans(cfg: argparse.Namespace) -> dict[str, Any]: + """Post-run: list any surviving pods with slurm_user=, delete them. + + Uses kubectl (rather than the K8s Python client) to avoid pulling in an + extra dependency. Assumes the caller already has a working kubectl context + for the target cluster. + """ + out: dict[str, Any] = {"orphans_found": 0, "orphans_deleted": 0, "sample": []} + if not shutil.which("kubectl"): + out["error"] = "kubectl not on PATH — skipping orphan sweep" + return out + try: + j = subprocess.run( + [ + "kubectl", "--context", cfg.kubectl_context, "-n", cfg.k8s_namespace, + "get", "pods", "-l", f"slurm_user={cfg.slurm_user}", + "-o", "json", + ], + capture_output=True, text=True, check=True, timeout=30, + ) + except subprocess.CalledProcessError as exc: + out["error"] = f"kubectl list failed: {exc.stderr.strip()[:200]}" + return out + except subprocess.TimeoutExpired: + out["error"] = "kubectl list timed out" + return out + try: + data = json.loads(j.stdout) + except json.JSONDecodeError as exc: + out["error"] = f"kubectl json parse failed: {exc}" + return out + pods = data.get("items", []) + out["orphans_found"] = len(pods) + out["sample"] = [p["metadata"]["name"] for p in pods[:10]] + if not pods: + return out + log.warning("orphan sweep: %d pods still labelled %s — deleting", + len(pods), cfg.slurm_user) + try: + subprocess.run( + [ + "kubectl", "--context", cfg.kubectl_context, "-n", cfg.k8s_namespace, + "delete", "pods", "-l", f"slurm_user={cfg.slurm_user}", + "--wait=false", "--grace-period=5", + ], + capture_output=True, text=True, check=True, timeout=60, + ) + out["orphans_deleted"] = len(pods) + except subprocess.CalledProcessError as exc: + out["delete_error"] = exc.stderr.strip()[:200] + return out + + +def cleanup_response_queues(cfg: argparse.Namespace) -> dict[str, Any]: + """Delete any SQS response queues we left behind (belt-and-braces).""" + sqs = boto3.client("sqs", region_name=cfg.region) + out: dict[str, Any] = {"queues_deleted": 0, "sample": []} + try: + resp = sqs.list_queues(QueueNamePrefix="pier-sqs-k8s-resp-", MaxResults=1000) + except Exception as exc: # noqa: BLE001 + out["error"] = f"list_queues failed: {exc}" + return out + urls = resp.get("QueueUrls", []) + for url in urls: + try: + sqs.delete_queue(QueueUrl=url) + out["queues_deleted"] += 1 + if len(out["sample"]) < 5: + out["sample"].append(url.split("/")[-1]) + except Exception as exc: # noqa: BLE001 + out.setdefault("delete_errors", []).append(f"{url}: {exc}") + return out + + +# ── main ───────────────────────────────────────────────────────────────────── + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument( + "--queue-url", + default=os.environ.get( + "PIER_E2E_QUEUE_URL", + "https://sqs.eu-west-1.amazonaws.com/475108760152/docker-requests", + ), + ) + p.add_argument( + "--s3-bucket", + default=os.environ.get( + "PIER_E2E_S3_BUCKET", + "sqs-message-queue-large-objects-475108760152-eu-west-1-an", + ), + ) + p.add_argument( + "--registry-url", + default=os.environ.get( + "PIER_E2E_REGISTRY_URL", + "475108760152.dkr.ecr.eu-west-1.amazonaws.com/sandbox", + ), + ) + p.add_argument("--region", default=os.environ.get("PIER_E2E_REGION", "eu-west-1")) + p.add_argument( + "--image", + default=os.environ.get("PIER_E2E_IMAGE", "bash:latest"), + help="Default is bash:latest — pier hardcodes `bash -c` in _sqs_exec_run " + "so busybox/alpine (no bash) fails with rc=127.", + ) + p.add_argument("--slurm-user", default=DEFAULT_SLURM_USER, + help="Label applied to every test sandbox for cleanup + accounting") + p.add_argument("--concurrency", type=int, default=20, + help="Max concurrent sandboxes in test 3 (HARD CAP: 20)") + p.add_argument("--iterations", type=int, default=5, + help="Rounds of --concurrency in test 3") + p.add_argument("--max-total", type=int, default=100, + help="Absolute cap on lifecycles across test 3 (HARD CAP: 100)") + p.add_argument("--kubectl-context", + default="arn:aws:eks:eu-west-1:475108760152:cluster/sandbox-eks") + p.add_argument("--k8s-namespace", default="sandbox-proxy") + p.add_argument("--skip", nargs="+", default=[], + choices=["1", "2", "3", "4"], + help="Skip individual tests by number") + p.add_argument("--json-out", type=Path, default=None, + help="Write a machine-readable summary JSON here") + args = p.parse_args() + if args.concurrency > 20: + p.error("--concurrency must be ≤ 20 (shared prod cluster)") + if args.max_total > 100: + p.error("--max-total must be ≤ 100 (shared prod cluster)") + return args + + +async def _amain(cfg: argparse.Namespace) -> int: + all_tests = [ + ("1", test_1_happy_path), + ("2", test_2_large_payload), + ("3", test_3_stress), + ("4", test_4_network_allowlist), + ] + results: list[TestResult] = [] + try: + for tag, fn in all_tests: + if tag in cfg.skip: + log.info("skipping test %s (--skip)", tag) + continue + log.info("=== running test %s: %s ===", tag, fn.__name__) + try: + r = await fn(cfg) + except Exception as exc: # noqa: BLE001 + r = TestResult(fn.__name__, False, 0.0, error=f"UNCAUGHT: {exc}") + results.append(r) + status = "PASS" if r.passed else "FAIL" + log.info( + "=== %s: %s (%.1fs) ===", + r.name, status, r.duration_s, + ) + if r.error: + log.warning(" error: %s", r.error) + if r.evidence: + for k, v in r.evidence.items(): + log.info(" %s: %s", k, v) + finally: + log.info("=== cleanup: orphan pod sweep ===") + orphans = cleanup_orphans(cfg) + for k, v in orphans.items(): + log.info(" %s: %s", k, v) + log.info("=== cleanup: SQS response-queue sweep ===") + queues = cleanup_response_queues(cfg) + for k, v in queues.items(): + log.info(" %s: %s", k, v) + + summary = { + "timestamp": time.time(), + "config": { + "queue_url": cfg.queue_url, + "s3_bucket": cfg.s3_bucket, + "registry_url": cfg.registry_url, + "region": cfg.region, + "image": cfg.image, + "slurm_user": cfg.slurm_user, + "concurrency": cfg.concurrency, + "iterations": cfg.iterations, + "max_total": cfg.max_total, + }, + "results": [ + { + "name": r.name, + "passed": r.passed, + "duration_s": round(r.duration_s, 2), + "evidence": r.evidence, + "error": r.error, + } + for r in results + ], + "orphans": orphans, + "queue_cleanup": queues, + } + if cfg.json_out: + cfg.json_out.write_text(json.dumps(summary, indent=2, default=str)) + log.info("wrote %s", cfg.json_out) + + print("\n" + "=" * 78) + print(f"{'test':<28} {'result':<8} {'duration':>10}") + print("-" * 78) + for r in results: + status = "PASS" if r.passed else "FAIL" + print(f"{r.name:<28} {status:<8} {r.duration_s:>8.1f}s") + print("-" * 78) + print(f"orphans found: {orphans.get('orphans_found', '?')} " + f"deleted: {orphans.get('orphans_deleted', '?')}") + print(f"response queues deleted: {queues.get('queues_deleted', '?')}") + print("=" * 78) + + return 0 if all(r.passed for r in results) else 1 + + +def main() -> None: + cfg = parse_args() + try: + rc = asyncio.run(_amain(cfg)) + except KeyboardInterrupt: + log.warning("interrupted — running orphan sweep before exit") + cleanup_orphans(cfg) + cleanup_response_queues(cfg) + rc = 130 + sys.exit(rc) + + +if __name__ == "__main__": + main() diff --git a/src/pier/environments/factory.py b/src/pier/environments/factory.py index 7f80ba7b..0c8d78b2 100644 --- a/src/pier/environments/factory.py +++ b/src/pier/environments/factory.py @@ -45,6 +45,11 @@ class _EnvEntry(NamedTuple): "DaytonaEnvironment", "daytona", ), + EnvironmentType.SQS_KUBERNETES: _EnvEntry( + "pier.environments.sqs_kubernetes", + "SQSKubernetesEnvironment", + "sqs-kubernetes", + ), } diff --git a/src/pier/environments/sqs_kubernetes.py b/src/pier/environments/sqs_kubernetes.py new file mode 100644 index 00000000..613e435c --- /dev/null +++ b/src/pier/environments/sqs_kubernetes.py @@ -0,0 +1,1792 @@ +"""SQSKubernetesEnvironment — pier BaseEnvironment for remote EKS via SQS. + +Routes exec / file / lifecycle operations to a ``docker_k8s_consumer`` +Deployment running on EKS through an AWS SQS request queue plus a +process-shared response queue. Large payloads (post-compression) fall back +to S3. + +This class is a pier-native port of Harbor's +``harbor.environments.sqs_kubernetes.SQSKubernetesEnvironment`` (see the +LLM360/agent-dist repo for the consumer-side infrastructure). It does not +depend on the ``harbor`` runtime — the on-the-wire protocol is duplicated +here so pier can consume the same EKS infrastructure without pulling in +Harbor's environment abstractions. + +Configuration kwargs (passed through ``TrialEnvironmentConfig.kwargs`` or +directly to the constructor): + + sqs_queue_url: Full SQS URL for the docker-requests queue (required) + sqs_region: AWS region for SQS (default: eu-west-1) + s3_bucket: S3 bucket for large payloads (required) + s3_region: AWS region for S3 (default: same as sqs_region) + s3_access_key_id / s3_secret_access_key: static S3 creds (optional; + falls back to the boto3 credential chain when unset) + registry_url: ECR registry prefix, e.g. + ``123456.dkr.ecr.eu-west-1.amazonaws.com/sandbox`` (required) + first_chunk_timeout: seconds to wait for the first SQS response chunk + (default: 60) + create_max_attempts: max ``containers/create`` attempts on retryable + consumer errors (default: 8; set to 1 to fail-fast) + s3_threshold: bytes above which bodies are uploaded to S3 (default: + 204800). Applied to the post-compression body. + volume_mounts: list of ``(container_path, local_path)`` tuples for + bind-mount emulation. Defaults to + ``(env_paths.agent_dir, trial.agent_dir)`` etc. + metrics_bridge_queue: SQS queue name for pushing trial-level metrics + (default: "" — disabled; falls back to + ``METRICS_BRIDGE_QUEUE_NAME`` env var). + pre_install_commands: list of shell strings run once inside the pod + at create-time (idempotent — SQS is at-least-once). + ephemeral_storage_limit: per-task override for the sandbox pod's + ``ephemeral-storage`` limit (K8s quantity, e.g. ``20Gi``). + egress_proxy_url: base URL of the cluster's egress-proxy Service — used + for the ``HTTP_PROXY`` / ``HTTPS_PROXY`` env injection + performed by :meth:`agent_process_env` when the + consumer mints an ``egress_token`` for a per-agent + ``network_allowlist``. Default: + ``http://egress-proxy.sandbox-proxy.svc.cluster.local:3128`` + (matches the ``sandbox-proxy`` helm chart). + slurm_user / slurm_job_id: stamped as pod labels for provenance. + repo_url / repo_commit / repo_dest: optional server-side git clone + for BuildKit (see agent-dist BuildKit init container). + inject_prebuilt_env_files: when ``True`` (default) and a pre-built + ``docker_image`` is used, reproduce the Dockerfile's + ``COPY`` semantics by uploading matching files from + ``environment_dir`` after create. + +AWS credentials come from the boto3 credential chain unless explicitly +passed. +""" + +from __future__ import annotations + +import asyncio +import base64 +import binascii +import io +import itertools +import json +import logging +import os +import random +import re +import struct +import tarfile +import time +import zlib +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from uuid import uuid4 + +try: + import boto3 + from botocore.config import Config as BotoConfig + from botocore.exceptions import BotoCoreError, ClientError + + _HAS_BOTO3 = True +except ImportError: # pragma: no cover - handled at __init__ time + _HAS_BOTO3 = False + boto3 = None # type: ignore[assignment] + BotoConfig = None # type: ignore[assignment] + BotoCoreError = Exception # type: ignore[assignment] + ClientError = Exception # type: ignore[assignment] + +from pier.environments.base import BaseEnvironment, ExecResult +from pier.environments.capabilities import ( + EnvironmentCapabilities, + EnvironmentResourceCapabilities, +) +from pier.models.environment_type import EnvironmentType +from pier.models.task.config import TaskOS +from pier.models.trial.config import ResourceMode +from pier.models.trial.paths import EnvironmentPaths + +# ── image-name sanitiser (mirrors DockerEnvironment) ───────────────────────── + +_TAG_SANITISE_RE = re.compile(r"[^a-z0-9._-]") + + +def _sanitize_image_name(name: str) -> str: + name = name.lower() + if not re.match(r"^[a-z0-9]", name): + name = "0" + name + return _TAG_SANITISE_RE.sub("-", name) + + +def _split_family(task_slug: str) -> tuple[str | None, str]: + """Split ``task_slug`` into ``(family, task)`` on the first ``__``. + + Mirrors Harbor's family split so ECR repos stay under the 100k cap. + ``openswe__akuli__mantaray-165`` → ``("openswe", "akuli__mantaray-165")``. + """ + parts = task_slug.split("__", 1) + if len(parts) == 2: + return parts[0], parts[1] + return None, task_slug + + +# K8s resource-quantity validator for `ephemeral_storage_limit`. +_EPHEMERAL_STORAGE_RE = re.compile(r"^\d+(\.\d+)?(Ki|Mi|Gi|Ti|Pi|Ei|K|M|G|T|P|E)?$") + +# Threshold above which _encode_body switches to zlib+b64. Kept below the +# 256 KiB SQS hard limit so the compressed body still fits inline when it +# happens to compress well; the caller (_build_sqs_message) then decides +# whether the *encoded* body still exceeds `s3_threshold` and needs S3. +_COMPRESSION_THRESHOLD = 250 * 1024 + +# ── helpers ─────────────────────────────────────────────────────────────────── + + +def _b64compress(data: bytes) -> str: + return base64.b64encode(zlib.compress(data)).decode() + + +def _b64decompress(s: str) -> bytes: + return zlib.decompress(base64.b64decode(s)) + + +def _decode_stream_chunk(raw, ctype: str, req_id: str = "") -> bytes: + """Decode a STREAM/FULL chunk based on its ``content_type``. + + Text-shaped bodies (``text/plain``, ``application/json``, ``text/*``, + ``*+json``) arrive as UTF-8 strings; binary bodies + (``application/octet-stream``, ``application/x-tar``, + ``application/vnd.docker.raw-stream``) are base64-encoded by the + consumer. Unknown types default to the base64 path with a defensive + fallback so a single bad chunk doesn't crash the trial. + """ + if isinstance(raw, bytes): + return raw + if ( + ctype == "text/plain" + or ctype == "application/json" + or ctype.startswith("text/") + or ctype.endswith("+json") + ): + return raw.encode("utf-8") if isinstance(raw, str) else raw + if isinstance(raw, str): + try: + return base64.b64decode(raw) + except (binascii.Error, ValueError) as exc: + logging.getLogger(__name__).warning( + "[sqs-k8s] b64 STREAM chunk decode failed " + "(ctype=%r, %d chars, req_id=%s): %s; returning empty", + ctype, + len(raw), + (req_id or "")[:8], + exc, + ) + return b"" + return raw + + +def _decode_mux_frames(data: bytes) -> str: + """Strip Docker multiplex headers (8-byte header per frame).""" + out: list[bytes] = [] + i = 0 + while i + 8 <= len(data): + size = struct.unpack_from(">I", data, i + 4)[0] + i += 8 + if i + size <= len(data): + out.append(data[i : i + size]) + i += size + else: + break + return b"".join(out).decode("utf-8", errors="replace") + + +def _parse_dockerfile_copies(dockerfile_text: str) -> dict[str, str]: + """Extract ``COPY src dest`` pairs from a Dockerfile for env-file injection. + + Handles ``--chown`` / ``--chmod`` flags, line continuations, and + multi-source COPYs. Skips ``COPY --from=stage`` (source is another + image layer, not the host context). + """ + joined = re.sub(r"\\\s*\n", " ", dockerfile_text) + copy_map: dict[str, str] = {} + for line in joined.splitlines(): + if not re.match(r"\s*COPY\s", line, re.IGNORECASE): + continue + tokens = line.split()[1:] + from_stage = False + while tokens and tokens[0].startswith("--"): + from_stage = from_stage or tokens[0].lower().startswith("--from") + tokens.pop(0) + if from_stage or len(tokens) < 2 or tokens[0].startswith("["): + continue + dest, srcs = tokens[-1], tokens[:-1] + if len(srcs) > 1 and not dest.endswith("/"): + dest += "/" + for src in srcs: + copy_map[src.rstrip("/")] = dest + return copy_map + + +# ── error classes ───────────────────────────────────────────────────────────── + + +class SqsConsumerError(RuntimeError): + """Structured error returned by the EKS consumer. + + Carries the parsed body so callers (and retry loops) can route on + ``error_code`` / ``retryable`` instead of the status alone. + """ + + def __init__( + self, + status: int, + error_code: str, + retryable: bool, + message: str, + details: dict, + raw_body: bytes, + op_label: str = "", + ) -> None: + self.status = status + self.error_code = error_code + self.retryable = retryable + self.message = message + self.details = details + self.raw_body = raw_body + self.op_label = op_label + super().__init__( + f"[sqs-k8s] {op_label or 'request'} failed: " + f"status={status} error_code={error_code} " + f"retryable={retryable} message={message!r}" + ) + + +class K8sInfraError(RuntimeError): + """Infra-side exec failure that retry cannot fix. + + Raised by ``exec()`` when the consumer marks an error non-retryable + OR the retry budget is exhausted. Workload errors (OOM, disk-full) + take a different path: they are synthesised into an ``ExecResult`` and + returned normally, so the agent loop sees a "user command failed". + """ + + def __init__(self, error_code: str, message: str, details: dict) -> None: + self.error_code = error_code + self.message = message + self.details = details + super().__init__(f"K8sInfraError[{error_code}]: {message}") + + +class S3PayloadError(RuntimeError): + """Payload could not be uploaded to / downloaded from the S3 fallback. + + Surface as a distinct exception so callers don't confuse an S3 failure + with an SQS or consumer failure. Never silently fall back to inline — + that would risk ``MessageTooLong`` from SQS when the body is above the + 256 KiB hard limit. + """ + + def __init__(self, direction: str, key: str, cause: BaseException) -> None: + self.direction = direction + self.key = key + self.cause = cause + super().__init__( + f"[sqs-k8s] S3 {direction} failed for key={key!r}: " + f"{type(cause).__name__}: {cause}" + ) + + +# ── exec()-level retry knobs ───────────────────────────────────────────────── + +_EXEC_MAX_ATTEMPTS = 4 +_EXEC_BASE_DELAY = 2.0 +_EXEC_MAX_DELAY = 30.0 +_KEEPALIVE_ROLLING_TIMEOUT_SEC = 60 + +# Codes the consumer flags as workload failures. Kept small — we trust +# `details.is_workload` as the single source of truth; this map only +# supplies the synthesised exit code. +_WORKLOAD_EXIT_CODES: dict[str, int] = { + "OOM_KILLED": 137, + "DISK_FULL": 1, +} + + +def _is_workload_error(err: SqsConsumerError) -> bool: + return err.details.get("is_workload") is True + + +def _synthesize_workload_exec_result(err: SqsConsumerError) -> ExecResult: + parts: list[str] = [f"(killed: {err.error_code})"] + if err.message: + parts.append(err.message) + if err.details: + detail_bits = [ + f"{k}={v}" + for k, v in err.details.items() + if k not in ("is_workload", "retryable", "operation") + ] + if detail_bits: + parts.append("| " + " ".join(detail_bits)) + return ExecResult( + stdout="", + stderr=" ".join(parts), + return_code=_WORKLOAD_EXIT_CODES.get(err.error_code, 1), + ) + + +def _parse_error_body( + resp_bytes: bytes, status: int, op_label: str = "" +) -> SqsConsumerError: + """Parse a non-2xx response body into a ``SqsConsumerError``. + + Trusts the consumer's ``retryable`` field. Missing / malformed bodies + default to ``retryable=False`` (fail-closed: better to fail loudly on a + truly-unknown response than to retry-forever). + """ + body: dict = {} + if resp_bytes: + try: + parsed = json.loads(resp_bytes.decode("utf-8", errors="replace")) + if isinstance(parsed, dict): + body = parsed + except (json.JSONDecodeError, UnicodeDecodeError): + pass + return SqsConsumerError( + status=status, + error_code=body.get("error_code", "UNKNOWN"), + retryable=bool(body.get("retryable", False)), + message=body.get("message", ""), + details=body.get("details", {}) + if isinstance(body.get("details"), dict) + else {}, + raw_body=resp_bytes[:500] if resp_bytes else b"", + op_label=op_label, + ) + + +def _maybe_raise_terminal_stream_end(chunk: dict, error_code: str | None) -> None: + """Raise on STREAM_END chunks carrying a terminal infra classification.""" + if not error_code: + return + details = chunk.get("details") or {} + if bool(details.get("is_workload")): + return + raise SqsConsumerError( + status=500, + error_code=error_code, + retryable=bool(chunk.get("retryable", False)), + message=chunk.get("message") or f"exec stream ended with {error_code}", + details=details, + raw_body=b"", + op_label="exec_run", + ) + + +# ── main class ──────────────────────────────────────────────────────────────── + + +class SQSKubernetesEnvironment(BaseEnvironment): + """Pier environment that runs sandboxes on EKS via SQS. + + Routes every ``exec`` / ``upload_*`` / ``download_*`` call over an SQS + request queue to a ``docker_k8s_consumer`` Deployment running in the + ``sandbox-proxy`` namespace on EKS. See the LLM360/agent-dist repo for + the consumer-side chart and its wire protocol. + """ + + # Process-wide S3 build-context upload cache. + _build_context_uploaded: set[str] = set() + _build_context_locks: dict[str, asyncio.Lock] = {} + + # Process-wide shared response queue. Every instance in the same + # process shares one SQS response queue + poller pool; messages are + # routed by ``req_id`` via ``_shared_pending``. + _shared_response_queue_url: str | None = None + _shared_poller_tasks: list[asyncio.Task] = [] + _shared_pending: dict[str, asyncio.PriorityQueue] = {} + _shared_refcount: int = 0 + _shared_lock: asyncio.Lock | None = None + _shared_sqs_client = None # type: ignore[assignment] + _shared_sender_sqs_client = None # type: ignore[assignment] + _delete_queue: "asyncio.Queue | None" = None + _poller_executor: "ThreadPoolExecutor | None" = None + _sender_executor: "ThreadPoolExecutor | None" = None + _heap_tiebreaker: itertools.count = itertools.count() + + _NUM_POLLERS = 50 + + # ── constructor / plain properties ──────────────────────────────────────── + + def __init__( + self, + *args, + sqs_queue_url: str, + s3_bucket: str, + registry_url: str, + sqs_region: str = "eu-west-1", + s3_region: str | None = None, + s3_access_key_id: str | None = None, + s3_secret_access_key: str | None = None, + first_chunk_timeout: int = 60, + s3_threshold: int = 200 * 1024, + volume_mounts: list[tuple[str, str]] | None = None, + metrics_bridge_queue: str = "", + pre_install_commands: list[str] | None = None, + ephemeral_storage_limit: str = "", + slurm_user: str = "", + slurm_job_id: str = "", + repo_url: str = "", + repo_commit: str = "", + repo_dest: str = "repo", + inject_prebuilt_env_files: bool = True, + create_max_attempts: int = 8, + egress_proxy_url: str = ( + "http://egress-proxy.sandbox-proxy.svc.cluster.local:3128" + ), + **kwargs, + ) -> None: + if not _HAS_BOTO3: + raise ImportError( + "SQSKubernetesEnvironment requires boto3. Install it with " + "`pip install boto3` (already declared as a base dependency " + "of pier)." + ) + if create_max_attempts < 1: + raise ValueError("create_max_attempts must be at least 1") + if ephemeral_storage_limit and not _EPHEMERAL_STORAGE_RE.match( + ephemeral_storage_limit + ): + raise ValueError( + f"ephemeral_storage_limit={ephemeral_storage_limit!r} is not a " + f"valid K8s resource quantity (e.g. '20Gi')." + ) + + # Store config on self BEFORE calling super().__init__ so that + # _validate_definition (invoked from Base.__init__) can read it. + self._sqs_queue_url = sqs_queue_url + self._s3_bucket = s3_bucket + self._registry_url = registry_url.rstrip("/") if registry_url else "" + self._sqs_region = sqs_region + self._s3_region = s3_region or sqs_region + self._s3_access_key_id = s3_access_key_id or os.environ.get( + "S3_ACCESS_KEY_ID" + ) + self._s3_secret_access_key = s3_secret_access_key or os.environ.get( + "S3_SECRET_ACCESS_KEY" + ) + self._first_chunk_timeout = first_chunk_timeout + self._s3_threshold = s3_threshold + self._create_max_attempts = create_max_attempts + self._pre_install_commands = pre_install_commands or [] + self._ephemeral_storage_limit = ephemeral_storage_limit + self._slurm_user = ( + slurm_user + or os.environ.get("SLURM_JOB_USER") + or os.environ.get("USER") + or "UNKNOWN_USER" + ) + self._slurm_job_id = ( + slurm_job_id or os.environ.get("SLURM_JOB_ID") or "UNKNOWN_JOB_ID" + ) + self._repo_url = repo_url + self._repo_commit = repo_commit + self._repo_dest = repo_dest + self._inject_prebuilt_env_files = inject_prebuilt_env_files + self._explicit_volume_mounts = volume_mounts + self._metrics_bridge_queue = metrics_bridge_queue or os.environ.get( + "METRICS_BRIDGE_QUEUE_NAME", "" + ) + self._egress_proxy_url = egress_proxy_url.rstrip("/") if egress_proxy_url else "" + # Token minted by the consumer on ``containers/create`` when a + # per-agent ``network_allowlist`` is passed. Absent (``None``) means + # either the caller did not request allowlist enforcement OR the + # target cluster is running a pre-PR-216 consumer that doesn't mint + # the token yet. Either way :meth:`agent_process_env` is a no-op + # unless this attribute is set to a non-empty string. + self._egress_token: str | None = None + + super().__init__(*args, **kwargs) + + # Share one SQS sender client process-wide to avoid a TLS + # thundering-herd at startup (see incident notes in agent-dist). + if SQSKubernetesEnvironment._shared_sender_sqs_client is None: + SQSKubernetesEnvironment._shared_sender_sqs_client = boto3.client( + "sqs", + region_name=self._sqs_region, + config=BotoConfig( + max_pool_connections=1000, + retries={"max_attempts": 5, "mode": "standard"}, + ), + ) + self._sqs = SQSKubernetesEnvironment._shared_sender_sqs_client + self._s3 = boto3.client( + "s3", + region_name=self._s3_region, + aws_access_key_id=self._s3_access_key_id, + aws_secret_access_key=self._s3_secret_access_key, + ) + + self._proxy_id = uuid4().hex + + # Container-id lifecycle mirrors Harbor's semantics: pre-generate + # a uuid per attempt so a 409 Conflict can't happen on retry, but + # keep every id the consumer *may* have acted on so ``stop()`` can + # issue delayed DELETEs. + self._pre_generated_cid: str | None = None + self._all_pre_generated_cids: list[str] = [] + self._sandbox_container_id: str | None = None + + # Volume mounts default to the pier env_paths ↔ trial_paths pairing. + if self._explicit_volume_mounts is not None: + self._volume_mounts = self._explicit_volume_mounts + else: + self._volume_mounts = [ + (str(EnvironmentPaths.agent_dir), str(self.trial_paths.agent_dir)), + (str(EnvironmentPaths.verifier_dir), str(self.trial_paths.verifier_dir)), + ( + str(EnvironmentPaths.artifacts_dir), + str(self.trial_paths.artifacts_dir), + ), + ] + + # Metrics bridge (optional; Harbor pushes trial-level timing). + self._metrics_bridge_url: str | None = None + self._metrics_push_task: asyncio.Task | None = None + self._timing: dict[str, float] = {} + self._exec_count = 0 + self._exec_total_secs = 0.0 + self._start_time: float | None = None + + @staticmethod + def type() -> EnvironmentType: + return EnvironmentType.SQS_KUBERNETES + + @property + def capabilities(self) -> EnvironmentCapabilities: + # ``mounted=True`` because ``stop()`` pulls each volume mount from + # the pod via ``archive GET`` before deletion — the caller sees + # host-side files as if they were bind-mounted. + # ``filtered_egress=True`` because :meth:`start` forwards the + # per-agent ``network_allowlist`` to the consumer, which mints a + # token published to the shared ``egress-allowlist`` ConfigMap and + # :meth:`agent_process_env` injects ``HTTP(S)_PROXY`` pointing at + # the cluster's egress-proxy. Declared statically — clusters + # running a pre-PR-216 consumer simply won't return a token and + # ``agent_process_env`` degrades to a no-op (the create body still + # carries the field; the old consumer ignores unknown keys). + return EnvironmentCapabilities( + gpus=False, + disable_internet=False, + filtered_egress=True, + preinstall_agents=False, + windows=False, + mounted=True, + docker_compose=False, + ) + + def agent_process_env( + self, env: dict[str, str] | None + ) -> dict[str, str] | None: + """Inject egress-proxy env vars for installed-agent commands. + + No-op unless :meth:`start` captured an ``egress_token`` from the + consumer response (per agent-dist PR #216) — without the token the + egress-proxy has no allowlist entry keyed to this sandbox, so + setting proxy vars would just route traffic through the shared + rate-limit path with no per-agent scoping. Caller-supplied env + wins on key collision, mirroring the docker / modal backends. + """ + if not self._egress_token: + return env + # Accept an ``egress_proxy_url`` either with an explicit scheme + # (``http://host:port``) or bare (``host:port``); either way emit + # ``http://agent:@host:port`` so HTTP libraries send + # ``Proxy-Authorization: Basic b64("agent:")``. The server's + # ``parseProxyAuthToken`` (agent-dist egress-proxy/tokenstore.go) + # splits on the first ``:`` and takes the password half as the + # token; the ``agent`` username is convention only. + host_port = self._egress_proxy_url + for scheme in ("http://", "https://"): + if host_port.startswith(scheme): + host_port = host_port[len(scheme) :] + break + proxy_url = f"http://agent:{self._egress_token}@{host_port}" + no_proxy = ".svc,.cluster.local,10.0.0.0/8,127.0.0.1,localhost" + merged: dict[str, str] = { + "HTTP_PROXY": proxy_url, + "http_proxy": proxy_url, + "HTTPS_PROXY": proxy_url, + "https_proxy": proxy_url, + "NO_PROXY": no_proxy, + "no_proxy": no_proxy, + } + if env: + merged.update(env) # caller-supplied env wins on collision + return merged + + @classmethod + def resource_capabilities(cls) -> EnvironmentResourceCapabilities: + # EKS supports both request and limit for CPU / memory. + return EnvironmentResourceCapabilities( + cpu_limit=True, + cpu_request=True, + memory_limit=True, + memory_request=True, + ) + + def _validate_definition(self) -> None: + if not self._sqs_queue_url: + raise ValueError("SQSKubernetesEnvironment: sqs_queue_url is required") + if not self._sqs_queue_url.startswith(("https://sqs.", "http://")): + raise ValueError( + f"SQSKubernetesEnvironment: sqs_queue_url does not look like " + f"an SQS URL: {self._sqs_queue_url!r}" + ) + if not self._s3_bucket: + raise ValueError("SQSKubernetesEnvironment: s3_bucket is required") + if not self._registry_url: + raise ValueError("SQSKubernetesEnvironment: registry_url is required") + if not self.environment_dir.is_dir(): + raise FileNotFoundError( + f"SQSKubernetesEnvironment: environment_dir does not exist: " + f"{self.environment_dir}" + ) + # SQS-Kubernetes runs Linux containers on EKS Linux nodegroups. + # Windows tasks are rejected at Base validation time via + # capabilities.windows=False; enforce here as well for a clearer + # message when someone constructs the env directly. + if self.task_env_config.os == TaskOS.WINDOWS: + raise RuntimeError( + "SQSKubernetesEnvironment does not support Windows containers" + ) + + # ── shared response-queue infrastructure ───────────────────────────────── + + @classmethod + async def _ensure_shared_queue(cls, sqs_client, logger: logging.Logger) -> str: + if cls._shared_lock is None: + cls._shared_lock = asyncio.Lock() + async with cls._shared_lock: + cls._shared_refcount += 1 + if cls._shared_response_queue_url is not None: + return cls._shared_response_queue_url + + cls._shared_sqs_client = sqs_client + proxy_id = uuid4().hex + resp_name = f"pier-sqs-k8s-resp-{proxy_id}" + dlq_name = f"pier-sqs-k8s-resp-{proxy_id}-dlq" + + await asyncio.to_thread(sqs_client.create_queue, QueueName=dlq_name) + dlq_url = await asyncio.to_thread( + lambda: sqs_client.get_queue_url(QueueName=dlq_name)["QueueUrl"] + ) + dlq_arn = await asyncio.to_thread( + lambda: sqs_client.get_queue_attributes( + QueueUrl=dlq_url, AttributeNames=["QueueArn"] + )["Attributes"]["QueueArn"] + ) + resp = await asyncio.to_thread( + lambda: sqs_client.create_queue( + QueueName=resp_name, + Attributes={ + "RedrivePolicy": json.dumps( + {"deadLetterTargetArn": dlq_arn, "maxReceiveCount": "3"} + ), + }, + ) + ) + cls._shared_response_queue_url = resp["QueueUrl"] + + cls._poller_executor = ThreadPoolExecutor( + max_workers=cls._NUM_POLLERS + 10, + thread_name_prefix="pier-sqs-poller", + ) + cls._sender_executor = ThreadPoolExecutor( + max_workers=1000, + thread_name_prefix="pier-sqs-sender", + ) + + cls._delete_queue = asyncio.Queue() + delete_executor = ThreadPoolExecutor( + max_workers=50, thread_name_prefix="pier-sqs-deleter" + ) + + async def _delete_worker() -> None: + loop = asyncio.get_event_loop() + delete_queue = cls._delete_queue + assert delete_queue is not None + while True: + try: + client, entries = await delete_queue.get() + await loop.run_in_executor( + delete_executor, + lambda c=client, e=entries: c.delete_message_batch( + QueueUrl=cls._shared_response_queue_url, + Entries=e, + ), + ) + except asyncio.CancelledError: + return + except Exception: # noqa: BLE001 - best-effort deleter + pass + + for _ in range(20): + cls._shared_poller_tasks.append( + asyncio.create_task(_delete_worker()) + ) + + region = sqs_client.meta.region_name + for _ in range(cls._NUM_POLLERS): + poller_client = boto3.client("sqs", region_name=region) + cls._shared_poller_tasks.append( + asyncio.create_task(cls._shared_poll_loop(poller_client, logger)) + ) + logger.info( + "[sqs-k8s] created shared response queue with %d pollers: %s", + cls._NUM_POLLERS, + cls._shared_response_queue_url, + ) + return cls._shared_response_queue_url + + @classmethod + async def _shared_poll_loop(cls, sqs_client, logger: logging.Logger) -> None: + loop = asyncio.get_event_loop() + while True: + if not cls._shared_response_queue_url: + return + try: + response = await loop.run_in_executor( + cls._poller_executor, + lambda: sqs_client.receive_message( + QueueUrl=cls._shared_response_queue_url, + MaxNumberOfMessages=10, + WaitTimeSeconds=2, + ), + ) + messages = response.get("Messages", []) + if not messages: + continue + to_delete = [] + t_poller_recv = time.time() + for msg in messages: + body = json.loads(msg["Body"]) + body["_poller_recv_ts"] = t_poller_recv + req_id = body.get("req_id") + seq_num = body.get("seq_num", 0) + if req_id in cls._shared_pending: + await cls._shared_pending[req_id].put( + (seq_num, next(cls._heap_tiebreaker), body) + ) + to_delete.append( + {"Id": msg["MessageId"], "ReceiptHandle": msg["ReceiptHandle"]} + ) + if to_delete and cls._delete_queue: + await cls._delete_queue.put((sqs_client, to_delete)) + except asyncio.CancelledError: + return + except Exception as e: # noqa: BLE001 - poller must not die + if "NonExistentQueue" in str(e): + return + logger.error("[sqs-k8s] poller error: %s", e, exc_info=True) + await asyncio.sleep(5) + + @classmethod + async def _release_shared_queue(cls, logger: logging.Logger) -> None: + if cls._shared_lock is None: + return + async with cls._shared_lock: + cls._shared_refcount -= 1 + if cls._shared_refcount > 0: + return + if cls._shared_poller_tasks: + for t in cls._shared_poller_tasks: + t.cancel() + await asyncio.gather(*cls._shared_poller_tasks, return_exceptions=True) + cls._shared_poller_tasks = [] + if cls._shared_response_queue_url and cls._shared_sqs_client: + try: + await asyncio.to_thread( + cls._shared_sqs_client.delete_queue, + QueueUrl=cls._shared_response_queue_url, + ) + except Exception: # noqa: BLE001 - best-effort teardown + pass + cls._shared_response_queue_url = None + if cls._poller_executor: + cls._poller_executor.shutdown(wait=False) + cls._poller_executor = None + if cls._sender_executor: + cls._sender_executor.shutdown(wait=False) + cls._sender_executor = None + + async def _startup(self) -> None: + await self._ensure_shared_queue(self._sqs, self.logger) + + async def _teardown(self) -> None: + await self._release_shared_queue(self.logger) + + # ── body encoding / S3 fallback ────────────────────────────────────────── + + async def _upload_to_s3(self, body: bytes, prefix: str = "sqs-k8s-payloads") -> str: + key = f"{prefix}/{uuid4().hex}" + try: + await asyncio.to_thread( + self._s3.put_object, Bucket=self._s3_bucket, Key=key, Body=body + ) + except (BotoCoreError, ClientError) as exc: + raise S3PayloadError("upload", key, exc) from exc + return key + + async def _download_from_s3(self, key: str) -> bytes: + try: + resp = await asyncio.to_thread( + self._s3.get_object, Bucket=self._s3_bucket, Key=key + ) + return resp["Body"].read() + except (BotoCoreError, ClientError) as exc: + raise S3PayloadError("download", key, exc) from exc + + def _encode_body(self, body: bytes) -> tuple[str, bool]: + """Encode a body for the SQS ``content`` field. + + Returns ``(content_str, compress_flag)``. ``s3_key`` fallback is + the caller's concern (see ``_build_sqs_message``). + """ + if len(body) > _COMPRESSION_THRESHOLD: + return _b64compress(body), True + if b"\x00" in body: + return base64.b64encode(body).decode(), False + try: + return body.decode("utf-8"), False + except UnicodeDecodeError: + return base64.b64encode(body).decode(), False + + async def _send_sqs( + self, + queue_url: str, + message_body: str, + _max_retries: int = 5, + ) -> tuple[float, float]: + t0 = time.time() + + def _timed_send() -> float: + t_api = time.time() + self._sqs.send_message(QueueUrl=queue_url, MessageBody=message_body) + return time.time() - t_api + + loop = asyncio.get_event_loop() + executor = SQSKubernetesEnvironment._sender_executor + for attempt in range(_max_retries): + try: + if executor: + api_time = await loop.run_in_executor(executor, _timed_send) + else: + api_time = await asyncio.to_thread(_timed_send) + thread_wait = (time.time() - t0) - api_time + return thread_wait, api_time + except Exception as send_err: # noqa: BLE001 - retry all failures + if attempt < _max_retries - 1: + delay = min(2 * (2**attempt), 15) * (0.5 + random.random()) + self.logger.warning( + "[sqs-send-retry] %s (attempt %d/%d, retry in %.1fs): %s", + type(send_err).__name__, + attempt + 1, + _max_retries, + delay, + send_err, + ) + await asyncio.sleep(delay) + else: + raise + raise RuntimeError("_send_sqs exhausted retries without returning") + + async def _send_sqs_message( + self, + method: str, + path: str, + query: str = "", + body: bytes = b"", + delay_seconds: int = 0, + ) -> None: + """Fire-and-forget SQS message (used for orphan-cleanup deletes).""" + msg = await self._build_sqs_message(method, path, query, body=body) + message_body = json.dumps(msg) + + def _send() -> None: + self._sqs.send_message( + QueueUrl=self._sqs_queue_url, + MessageBody=message_body, + DelaySeconds=delay_seconds, + ) + + loop = asyncio.get_event_loop() + executor = SQSKubernetesEnvironment._sender_executor + if executor: + await loop.run_in_executor(executor, _send) + else: + await asyncio.to_thread(_send) + + async def _build_sqs_message( + self, + method: str, + path: str, + query: str = "", + headers: dict | None = None, + body: bytes = b"", + extra_fields: dict | None = None, + ) -> dict: + """Assemble an SQS request; upload to S3 when the encoded body is large. + + Three-tier body handling: + + 1. Empty body → ``content=""``, ``s3_key`` absent. + 2. Small body → inline ``content`` (optionally zlib+base64 encoded + when ``len(body) > _COMPRESSION_THRESHOLD``). + 3. Large body (encoded size > ``s3_threshold``) → uploaded to S3; + the message carries ``s3_key`` instead of ``content``. + + Never silently falls back to inline when S3 upload fails — a + raised ``S3PayloadError`` is preferable to a downstream + ``MessageTooLong`` from SQS (256 KiB hard limit). + """ + req_id = uuid4().hex + msg: dict = { + "req_id": req_id, + "channel": SQSKubernetesEnvironment._shared_response_queue_url, + "method": method, + "path": path, + "query": query, + "headers": headers or {}, + "content": "", + "compress": False, + "sent_at": time.time(), + "task_id": self.session_id, + } + if extra_fields: + msg.update(extra_fields) + + if not body: + return msg + + # Encode first; the encoded size is what matters for SQS. + encoded, compress = self._encode_body(body) + if len(encoded.encode("utf-8")) > self._s3_threshold: + msg["s3_key"] = await self._upload_to_s3(body) + else: + msg["content"] = encoded + msg["compress"] = compress + return msg + + def _decode_response_body(self, resp: dict) -> bytes: + """Decode the body from a FULL response message (no S3 in this path).""" + if resp.get("s3_key"): + raise RuntimeError( + "Caller must handle s3_key async before calling _decode_response_body" + ) + if resp.get("compress"): + return _b64decompress(resp.get("content", "")) + content = resp.get("content", "") + if resp.get("content_type", "").startswith( + ("application/x-tar", "application/octet-stream") + ): + try: + return base64.b64decode(content) + except Exception: # noqa: BLE001 - fall through to str path + pass + if isinstance(content, str): + return content.encode("utf-8") + return content + + # ── SQS round-trip ──────────────────────────────────────────────────────── + + async def _sqs_round_trip( + self, + method: str, + path: str, + query: str = "", + headers: dict | None = None, + body: bytes = b"", + extra_fields: dict | None = None, + on_send: Callable[[], None] | None = None, + wants_keepalive: bool = False, + ) -> tuple[int, bytes]: + """Send one request over SQS and wait for a FULL response. + + ``wants_keepalive=True`` tolerates zero-or-more empty STREAM + "still working" beacons before the terminal FULL response — used + by ``containers/create`` to cover slow image pulls / inline builds. + Non-empty STREAM frames on this path are a protocol violation. + """ + msg = await self._build_sqs_message( + method, path, query, headers, body, extra_fields + ) + req_id = msg["req_id"] + t_send = time.time() + q: asyncio.PriorityQueue = asyncio.PriorityQueue() + SQSKubernetesEnvironment._shared_pending[req_id] = q + try: + await self._send_sqs(self._sqs_queue_url, json.dumps(msg)) + if on_send is not None: + on_send() + + _, _, resp = await asyncio.wait_for( + q.get(), timeout=self._first_chunk_timeout + ) + while wants_keepalive and resp.get("type") == "STREAM": + if resp.get("content", "") != "": + raise RuntimeError( + f"[sqs-k8s] _sqs_round_trip({path}): unexpected " + f"non-empty STREAM frame on FULL-response path " + f"req_id={req_id[:8]}" + ) + _, _, resp = await asyncio.wait_for( + q.get(), timeout=_KEEPALIVE_ROLLING_TIMEOUT_SEC + ) + + if resp.get("s3_key"): + resp_bytes = await self._download_from_s3(resp["s3_key"]) + if resp.get("s3_encoding") == "base64": + resp_bytes = base64.b64decode(resp_bytes) + else: + resp_bytes = self._decode_response_body(resp) + + self.logger.debug( + "[sqs-k8s] round_trip %s %s elapsed_s=%.3f req_id=%s", + method, + path, + time.time() - t_send, + req_id[:8], + ) + return int(resp.get("status_code", 200)), resp_bytes + finally: + SQSKubernetesEnvironment._shared_pending.pop(req_id, None) + + async def _round_trip_with_retry( + self, + method: str, + path: str, + *, + query: str = "", + headers: dict | None = None, + body: bytes = b"", + extra_fields: dict | None = None, + op_label: str = "", + max_attempts: int = 4, + base_delay: float = 2.0, + max_delay: float = 30.0, + treat_404_as_success: bool = False, + ) -> tuple[int, bytes]: + last_err: SqsConsumerError | None = None + for attempt in range(1, max_attempts + 1): + status, resp_bytes = await self._sqs_round_trip( + method, + path, + query=query, + headers=headers, + body=body, + extra_fields=extra_fields, + ) + if status in (200, 201, 204): + return status, resp_bytes + if treat_404_as_success and status == 404: + return status, resp_bytes + err = _parse_error_body(resp_bytes, status, op_label) + last_err = err + if not err.retryable or attempt == max_attempts: + raise err + delay = min(base_delay * (2 ** (attempt - 1)), max_delay) + delay *= 0.5 + random.random() + self.logger.warning( + "[sqs-k8s] %s retryable %s status=%d attempt=%d/%d sleeping %.1fs", + op_label, + err.error_code, + status, + attempt, + max_attempts, + delay, + ) + await asyncio.sleep(delay) + assert last_err is not None + raise last_err + + # ── exec-run streaming ──────────────────────────────────────────────────── + + async def _sqs_exec_run( + self, + container_id: str, + cmd: list[str], + user: str = "", + env: dict[str, str] | None = None, + timeout_sec: int | None = None, + ) -> ExecResult: + exec_body: dict = { + "Cmd": cmd, + "AttachStdout": True, + "AttachStderr": True, + "Detach": False, + "User": user, + "client": "sqs_env", + "exec_id": uuid4().hex, + "keepalive": True, + } + if env: + exec_body["Env"] = [f"{k}={v}" for k, v in env.items()] + if timeout_sec and timeout_sec > 0: + exec_body["timeout_sec"] = timeout_sec + + msg = await self._build_sqs_message( + method="POST", + path=f"v1.43/exec_run/{container_id}", + body=json.dumps(exec_body).encode(), + ) + req_id = msg["req_id"] + t_send = time.time() + q: asyncio.PriorityQueue = asyncio.PriorityQueue() + SQSKubernetesEnvironment._shared_pending[req_id] = q + + effective_timeout = timeout_sec or self._first_chunk_timeout + stdout_chunks: list[bytes] = [] + exit_code = -1 + is_mux = False + + try: + await self._send_sqs(self._sqs_queue_url, json.dumps(msg)) + _, _, first = await asyncio.wait_for( + q.get(), timeout=self._first_chunk_timeout + ) + + if first.get("type") == "STREAM_END": + exit_code = first.get("exit_code", -1) + if exit_code is None: + exit_code = -1 + exec_error_code = first.get("error_code") or None + _maybe_raise_terminal_stream_end(first, exec_error_code) + stdout = first.get("content", "") or "" + stderr = first.get("stderr", "") or "" + return ExecResult( + stdout=stdout, stderr=stderr, return_code=exit_code + ) + + current_seq = 0 + exec_error_code: str | None = None + buffer: dict[int, tuple[str, str]] = {} + consumed_seqs: set[int] = set() + seq = first.get("seq_num", 0) + first_content = first.get("content", "") + keepalive_active = first_content == "" + rolling_timeout = ( + _KEEPALIVE_ROLLING_TIMEOUT_SEC + if keepalive_active + else effective_timeout + ) + if not keepalive_active: + buffer[seq] = (first_content, first.get("content_type", "")) + else: + consumed_seqs.add(seq) + current_seq = seq + 1 + + stderr = "" + while True: + while current_seq in buffer: + raw_chunk, ctype = buffer.pop(current_seq) + chunk_bytes = _decode_stream_chunk( + raw_chunk, ctype, req_id=req_id + ) + if ( + not is_mux + and stdout_chunks == [] + and len(chunk_bytes) >= 8 + and chunk_bytes[0] in (1, 2) + ): + is_mux = True + stdout_chunks.append(chunk_bytes) + consumed_seqs.add(current_seq) + current_seq += 1 + + try: + _, _, chunk = await asyncio.wait_for( + q.get(), timeout=rolling_timeout + ) + except asyncio.TimeoutError: + self.logger.warning( + "[sqs-k8s] exec_run stream timeout after %ds req_id=%s", + rolling_timeout, + req_id[:8], + ) + break + + if chunk.get("type") == "STREAM_END": + exit_code = chunk.get("exit_code", -1) + if exit_code is None: + exit_code = -1 + exec_error_code = chunk.get("error_code") or None + _maybe_raise_terminal_stream_end(chunk, exec_error_code) + stderr = chunk.get("stderr", "") or "" + end_content = chunk.get("content", "") + if end_content and end_content != "[EXIT]" and not stdout_chunks: + stdout_chunks.append(end_content.encode("utf-8")) + break + + chunk_seq = chunk.get("seq_num", current_seq) + chunk_content = chunk.get("content", "") + if keepalive_active and chunk_content == "": + consumed_seqs.add(chunk_seq) + if chunk_seq == current_seq: + current_seq += 1 + continue + if chunk_seq in consumed_seqs: + continue + buffer[chunk_seq] = (chunk_content, chunk.get("content_type", "")) + + for k in sorted(buffer.keys()): + if k in consumed_seqs: + continue + raw_chunk, ctype = buffer[k] + stdout_chunks.append(_decode_stream_chunk(raw_chunk, ctype, req_id)) + + raw_output = b"".join(stdout_chunks) + stdout = ( + _decode_mux_frames(raw_output) + if is_mux + else raw_output.decode("utf-8", errors="replace") + ) + + elapsed = time.time() - t_send + self._exec_count += 1 + self._exec_total_secs += elapsed + return ExecResult(stdout=stdout, stderr=stderr, return_code=exit_code) + finally: + SQSKubernetesEnvironment._shared_pending.pop(req_id, None) + + # ── tar helpers ─────────────────────────────────────────────────────────── + + def _make_tar(self, source: Path, arcname: str | None = None) -> bytes: + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w") as tf: + tf.add(str(source), arcname=arcname or source.name) + return buf.getvalue() + + def _extract_tar(self, tar_bytes: bytes, dest: Path) -> None: + dest.mkdir(parents=True, exist_ok=True) + with tarfile.open(fileobj=io.BytesIO(tar_bytes), mode="r:*") as tf: + members = tf.getmembers() + top = members[0].name.split("/")[0] if members else None + for member in members: + if member.name == top: + continue + if top and member.name.startswith(top + "/"): + member.name = member.name[len(top) + 1 :] + if not member.name: + continue + tf.extract(member, str(dest)) + + async def _pull_volume_mounts(self) -> None: + if self._sandbox_container_id is None: + return + for container_path, host_path in self._volume_mounts: + max_retries = 3 + for attempt in range(max_retries): + try: + status, tar_bytes = await self._sqs_round_trip( + "GET", + f"v1.43/containers/{self._sandbox_container_id}/archive", + query=f"path={container_path}", + ) + if status != 200 or not tar_bytes: + break + self._extract_tar(tar_bytes, Path(host_path)) + break + except tarfile.ReadError: + if attempt < max_retries - 1: + await asyncio.sleep(1) + else: + self.logger.error( + "[sqs-k8s] volume pull failed %s (all attempts)", + container_path, + ) + except Exception as e: # noqa: BLE001 - best-effort pull + self.logger.error( + "[sqs-k8s] volume pull error %s: %s", + container_path, + e, + ) + break + + # ── orphan-cid cleanup ──────────────────────────────────────────────────── + + async def _cleanup_orphan_cids(self, cids: list[str]) -> None: + for cid in cids: + try: + await self._send_sqs_message( + "DELETE", + f"v1.43/containers/{cid}", + query="force=true", + delay_seconds=180, + ) + except Exception as e: # noqa: BLE001 - best-effort orphan cleanup + self.logger.warning( + "[sqs-k8s] orphan cleanup failed for cid=%s: %s", + cid[:12], + e, + ) + + # ── BaseEnvironment interface ───────────────────────────────────────────── + + async def start(self, force_build: bool = False) -> None: + await self._startup() + self._start_time = time.time() + + # Resolve metrics-bridge queue URL (optional; missing queue is not fatal). + if self._metrics_bridge_queue and not self._metrics_bridge_url: + try: + resp = await asyncio.to_thread( + lambda: self._sqs.get_queue_url( + QueueName=self._metrics_bridge_queue + ) + ) + self._metrics_bridge_url = resp["QueueUrl"] + except Exception as e: # noqa: BLE001 - optional + self.logger.debug( + "[sqs-k8s] metrics bridge queue not found: %s", e + ) + + prebuilt_image = self.task_env_config.docker_image or "" + tag = _sanitize_image_name(self.environment_name) + build_tag = prebuilt_image + ctx_key = "" + + if not prebuilt_image or force_build: + prebuilt_image = "" + family, task = _split_family(tag) + build_tag = f"{family}:{task}-latest" if family else f"{tag}:latest" + + import hashlib + + ctx_hash = hashlib.sha256(build_tag.encode()).hexdigest()[:16] + ctx_key = f"build-contexts/{build_tag}/{ctx_hash}.tar.gz" + + if ctx_key not in SQSKubernetesEnvironment._build_context_uploaded: + lock = SQSKubernetesEnvironment._build_context_locks.setdefault( + ctx_key, asyncio.Lock() + ) + async with lock: + if ctx_key not in SQSKubernetesEnvironment._build_context_uploaded: + env_dir = self.environment_dir + + def _make_ctx_tar() -> bytes: + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tf: + for p in sorted(env_dir.rglob("*")): + arcname = str(p.relative_to(env_dir)) + tf.add(str(p), arcname=arcname, recursive=False) + return buf.getvalue() + + build_context = await asyncio.to_thread(_make_ctx_tar) + try: + await asyncio.to_thread( + self._s3.put_object, + Bucket=self._s3_bucket, + Key=ctx_key, + Body=build_context, + ) + except (BotoCoreError, ClientError) as exc: + raise S3PayloadError("upload", ctx_key, exc) from exc + SQSKubernetesEnvironment._build_context_uploaded.add(ctx_key) + + # Resource limits / requests plumbed into the create body so the + # consumer's build_pod_spec can honour them. The consumer today + # only wires ephemeral_storage_limit; cpu / memory here are for + # forward-compatibility with the consumer extension noted in the + # docs. Empty / None values are omitted so we don't clobber + # cluster defaults. + create_body: dict = { + "Image": build_tag, + "Cmd": ["/bin/sh", "-c", "sleep infinity"], + "AttachStdout": False, + "AttachStderr": False, + "HostConfig": { + "Binds": [ + f"{host_path}:{container_path}" + for container_path, host_path in self._volume_mounts + ] + }, + "keepalive": True, + } + if self._pre_install_commands: + create_body["PreInstallCommands"] = self._pre_install_commands + cpu_limit = self._resource_limit_value("cpu", auto_mode=ResourceMode.LIMIT) + cpu_request = self._resource_request_value( + "cpu", auto_mode=ResourceMode.LIMIT + ) + mem_limit = self._resource_limit_value("memory", auto_mode=ResourceMode.LIMIT) + mem_request = self._resource_request_value( + "memory", auto_mode=ResourceMode.LIMIT + ) + if cpu_limit is not None: + create_body["Cpu"] = cpu_limit + if cpu_request is not None: + create_body["CpuRequest"] = cpu_request + if mem_limit is not None: + create_body["Memory"] = mem_limit + if mem_request is not None: + create_body["MemoryRequest"] = mem_request + + extra: dict = {} + if ctx_key: + extra["build_tag"] = build_tag + extra["build_context_s3_key"] = ctx_key + extra["force_build"] = force_build + if self._repo_url and self._repo_commit: + extra["repo_clone"] = { + "url": self._repo_url, + "commit": self._repo_commit, + "dest": self._repo_dest, + } + if self._ephemeral_storage_limit: + extra["ephemeral_storage_limit"] = self._ephemeral_storage_limit + extra["slurm_user"] = self._slurm_user + extra["slurm_job_id"] = self._slurm_job_id + # Per-agent egress allowlist enforcement — see agent-dist PR #216 + # (server-side implementation) and #215 (design doc). Only sent when + # the task explicitly denies internet AND supplies a non-empty + # allowlist; both preconditions guarantee zero behaviour change on + # tasks that don't opt in and on clusters running an older + # consumer that doesn't understand the field. + allowlist = getattr(self, "network_allowlist", None) + allowlist_domains = list(allowlist.domains) if allowlist else [] + if not self.task_env_config.allow_internet and allowlist_domains: + extra["network_allowlist"] = {"domains": allowlist_domains} + + create_attempts = 0 + t_start = time.time() + while True: + self._pre_generated_cid = uuid4().hex + create_body["ContainerId"] = self._pre_generated_cid + cid_for_attempt = self._pre_generated_cid + try: + status, resp_bytes = await self._sqs_round_trip( + "POST", + "v1.43/containers/create", + body=json.dumps(create_body).encode(), + extra_fields=extra, + on_send=lambda: self._all_pre_generated_cids.append( + cid_for_attempt + ), + wants_keepalive=True, + ) + except S3PayloadError: + raise + except Exception as conn_err: # noqa: BLE001 - connection retry + create_attempts += 1 + if create_attempts >= 5 or time.time() - t_start >= 580: + raise + delay = min(5 * (2 ** min(create_attempts - 1, 3)), 30) * ( + 0.5 + random.random() + ) + self.logger.warning( + "[sqs-k8s] create connection error (attempt %d): %s. " + "Retrying in %.1fs...", + create_attempts, + conn_err, + delay, + ) + await asyncio.sleep(delay) + continue + if status in (200, 201): + break + err = _parse_error_body(resp_bytes, status, op_label="containers/create") + create_attempts += 1 + if not err.retryable or create_attempts >= self._create_max_attempts: + raise err + base_delay = min(10 * (2 ** min(create_attempts - 1, 4)), 120) + delay = base_delay * (0.5 + random.random()) + self.logger.warning( + "[sqs-k8s] containers/create %s status=%d attempt=%d/%d " + "sleeping %.1fs", + err.error_code, + status, + create_attempts, + self._create_max_attempts, + delay, + ) + await asyncio.sleep(delay) + + create_response = json.loads(resp_bytes) + cid = create_response.get("Id", "") + if not cid: + raise RuntimeError("[sqs-k8s] containers/create: no Id in response") + self._sandbox_container_id = cid + # egress_token is minted by the consumer iff network_allowlist was + # in the request AND the consumer is running PR-216 code. Absent + # (empty / missing) → agent_process_env stays a no-op, matching + # the behaviour on older clusters. + egress_token = create_response.get("egress_token") or None + self._egress_token = egress_token if isinstance(egress_token, str) else None + self.logger.info( + "[sqs-k8s] sandbox created: %s%s", + cid[:12], + " (egress-token minted)" if self._egress_token else "", + ) + + # Reproduce the Dockerfile's `COPY` semantics for pre-built images so + # task-specific environment/ files land where the Dockerfile would + # have put them. + if prebuilt_image and self._inject_prebuilt_env_files: + dockerfile = self.environment_dir / "Dockerfile" + copy_map = ( + _parse_dockerfile_copies(dockerfile.read_text()) + if dockerfile.exists() + else {} + ) + for item in sorted(self.environment_dir.iterdir()): + name = item.name.lower() + if item.name == "Dockerfile" or name.startswith("solution"): + continue + dest = copy_map.get(item.name) + if dest is None: + continue + if item.is_dir(): + dest_dir = dest.rstrip("/") + for child in sorted(item.iterdir()): + target = f"{dest_dir}/{child.name}" + if child.is_dir(): + await self.upload_dir(child, target) + else: + await self.upload_file(child, target) + elif dest.endswith("/"): + await self.upload_file(item, f"{dest.rstrip('/')}/{item.name}") + else: + await self.upload_file(item, dest) + + if self._metrics_bridge_url: + self._metrics_push_task = asyncio.create_task(self._periodic_metrics_push()) + + def _build_metrics_payload(self, final: bool = False) -> dict: + total = time.time() - self._start_time if self._start_time else 0 + return { + "task_id": self.session_id, + "total_duration_seconds": round(total, 3), + "final": final, + "timing_phases": { + k: {"duration_seconds": round(v, 3)} for k, v in self._timing.items() + }, + "sqs_timing": { + "exec_count": self._exec_count, + "exec_total_seconds": round(self._exec_total_secs, 3), + "exec_avg_seconds": ( + round(self._exec_total_secs / self._exec_count, 3) + if self._exec_count + else 0 + ), + }, + } + + async def _push_metrics(self, final: bool = False) -> None: + if not self._metrics_bridge_url: + return + payload = self._build_metrics_payload(final=final) + try: + await self._send_sqs(self._metrics_bridge_url, json.dumps(payload)) + except Exception as e: # noqa: BLE001 - optional + self.logger.warning("[sqs-k8s] metrics bridge push failed: %s", e) + + async def _periodic_metrics_push(self, interval: float = 60.0) -> None: + try: + while True: + await asyncio.sleep(interval) + await self._push_metrics(final=False) + except asyncio.CancelledError: + return + + async def stop(self, delete: bool = True) -> None: + orphan_cids = [ + cid + for cid in self._all_pre_generated_cids + if cid != self._sandbox_container_id + ] + if self._sandbox_container_id is None: + if delete and orphan_cids: + await self._cleanup_orphan_cids(orphan_cids) + self._pre_generated_cid = None + self._all_pre_generated_cids = [] + self._egress_token = None + await self._teardown() + return + if delete and orphan_cids: + await self._cleanup_orphan_cids(orphan_cids) + self._all_pre_generated_cids = [] + if self._metrics_push_task: + self._metrics_push_task.cancel() + await asyncio.gather(self._metrics_push_task, return_exceptions=True) + self._metrics_push_task = None + t_stop = time.time() + try: + await self._pull_volume_mounts() + except Exception as e: # noqa: BLE001 - best-effort pull on stop + self.logger.warning("[sqs-k8s] volume pull error during stop: %s", e) + self._timing["volume_pull"] = time.time() - t_stop + + if delete: + try: + t_delete = time.time() + await self._sqs_round_trip( + "DELETE", + f"v1.43/containers/{self._sandbox_container_id}", + query="force=true", + ) + self._timing["delete"] = time.time() - t_delete + except Exception as e: # noqa: BLE001 - best-effort delete + self.logger.warning("[sqs-k8s] containers/delete error: %s", e) + self._sandbox_container_id = None + + # Drop the egress token so any post-stop call to agent_process_env + # returns env unchanged — the consumer removes the corresponding + # ConfigMap entry on containers/delete, so the token wouldn't work + # anyway. Reset happens unconditionally (even when delete=False) + # because the semantic invariant is "no live sandbox = no proxy". + self._egress_token = None + + await self._push_metrics(final=True) + await self._teardown() + + async def exec( + self, + command: str, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_sec: int | None = None, + user: str | int | None = None, + ) -> ExecResult: + """Run a shell command inside the sandbox pod over SQS. + + Workload failures (OOM, disk-full) are converted into a normal + ``ExecResult`` (non-zero exit) so the agent loop can proceed. + Retryable infra failures are retried with backoff; non-retryable + or budget-exhausted infra failures raise ``K8sInfraError``. + """ + if self._sandbox_container_id is None: + raise K8sInfraError( + error_code="CONTAINER_NOT_FOUND", + message="sandbox container id is None at exec start", + details={"operation": "exec"}, + ) + + effective_user = self._resolve_user(user) + merged_env = self._merge_env(env) + effective_cwd = cwd or self.task_env_config.workdir + if effective_cwd: + command = f"cd {effective_cwd} && {command}" + + last_err: SqsConsumerError | None = None + for attempt in range(1, _EXEC_MAX_ATTEMPTS + 1): + try: + return await self._sqs_exec_run( + container_id=self._sandbox_container_id, + cmd=["bash", "-c", command], + user=str(effective_user) if effective_user is not None else "", + env=merged_env, + timeout_sec=timeout_sec, + ) + except SqsConsumerError as err: + last_err = err + if _is_workload_error(err): + return _synthesize_workload_exec_result(err) + if not err.retryable or attempt == _EXEC_MAX_ATTEMPTS: + raise K8sInfraError( + error_code=err.error_code, + message=err.message, + details=err.details, + ) from err + delay = min(_EXEC_BASE_DELAY * (2 ** (attempt - 1)), _EXEC_MAX_DELAY) + delay *= 0.5 + random.random() + self.logger.warning( + "[sqs-k8s] exec retryable %s attempt=%d/%d sleeping %.1fs", + err.error_code, + attempt, + _EXEC_MAX_ATTEMPTS, + delay, + ) + await asyncio.sleep(delay) + + assert last_err is not None + raise K8sInfraError( + error_code=last_err.error_code, + message=last_err.message, + details=last_err.details, + ) + + async def upload_file(self, source_path: Path | str, target_path: str) -> None: + if self._sandbox_container_id is None: + return + source = Path(source_path) + tar_bytes = self._make_tar(source, arcname=os.path.basename(target_path)) + target_dir = os.path.dirname(target_path).rstrip("/") or "/" + await self._round_trip_with_retry( + "PUT", + f"v1.43/containers/{self._sandbox_container_id}/archive", + query=f"path={target_dir}", + headers={"Content-Type": "application/x-tar"}, + body=tar_bytes, + op_label="upload_file", + ) + + async def upload_dir(self, source_dir: Path | str, target_dir: str) -> None: + if self._sandbox_container_id is None: + return + source = Path(source_dir) + tar_bytes = self._make_tar(source, arcname=source.name) + parent = os.path.dirname(target_dir.rstrip("/")) or "/" + await self._round_trip_with_retry( + "PUT", + f"v1.43/containers/{self._sandbox_container_id}/archive", + query=f"path={parent}", + headers={"Content-Type": "application/x-tar"}, + body=tar_bytes, + op_label="upload_dir", + ) + + async def download_file(self, source_path: str, target_path: Path | str) -> None: + if self._sandbox_container_id is None: + return + status, tar_bytes = await self._round_trip_with_retry( + "GET", + f"v1.43/containers/{self._sandbox_container_id}/archive", + query=f"path={source_path}", + op_label="download_file", + ) + if not tar_bytes: + raise RuntimeError( + f"[sqs-k8s] download_file {source_path}: status={status} empty body" + ) + dest = Path(target_path) + dest.parent.mkdir(parents=True, exist_ok=True) + with tarfile.open(fileobj=io.BytesIO(tar_bytes), mode="r:*") as tf: + members = [m for m in tf.getmembers() if not m.isdir()] + if members: + f = tf.extractfile(members[0]) + if f: + dest.write_bytes(f.read()) + + async def download_dir(self, source_dir: str, target_dir: Path | str) -> None: + if self._sandbox_container_id is None: + return + try: + status, tar_bytes = await self._round_trip_with_retry( + "GET", + f"v1.43/containers/{self._sandbox_container_id}/archive", + query=f"path={source_dir}", + op_label="download_dir", + ) + except SqsConsumerError as e: + self.logger.warning( + "[sqs-k8s] download_dir %s: %s status=%d — skipping", + source_dir, + e.error_code, + e.status, + ) + return + if not tar_bytes: + self.logger.warning( + "[sqs-k8s] download_dir %s: status=%d empty body — skipping", + source_dir, + status, + ) + return + self._extract_tar(tar_bytes, Path(target_dir)) diff --git a/src/pier/models/environment_type.py b/src/pier/models/environment_type.py index cc378e44..07afb94d 100644 --- a/src/pier/models/environment_type.py +++ b/src/pier/models/environment_type.py @@ -5,3 +5,4 @@ class EnvironmentType(str, Enum): DOCKER = "docker" MODAL = "modal" DAYTONA = "daytona" + SQS_KUBERNETES = "sqs_kubernetes" diff --git a/tests/environments/__init__.py b/tests/environments/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/environments/test_sqs_kubernetes.py b/tests/environments/test_sqs_kubernetes.py new file mode 100644 index 00000000..7ac16381 --- /dev/null +++ b/tests/environments/test_sqs_kubernetes.py @@ -0,0 +1,856 @@ +"""Unit tests for the pier-native SQSKubernetesEnvironment port. + +Focus areas: +* ``_validate_definition`` — required fields, URL shape, Windows rejection. +* Body encoding + S3 fallback — the load-bearing large-payload path. +* Response s3_key handling on the receive side. +* Consumer-error parsing (retryable vs non-retryable, workload synthesis). +* ``capabilities`` shape (mounted=True, gpus=False, filtered_egress=False). +* Signature parity between ``exec`` and pier's ``BaseEnvironment.exec``. +* Factory registration. + +No integration tests here — they need a real EKS + SQS + S3 setup. The PR +body documents the manual smoke-test procedure. +""" + +from __future__ import annotations + +import asyncio +import base64 +import inspect +import io +import json +import zlib +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +from pier.environments.base import BaseEnvironment, ExecResult +from pier.environments.capabilities import EnvironmentCapabilities +from pier.environments.factory import _ENVIRONMENT_REGISTRY +from pier.environments.sqs_kubernetes import ( + K8sInfraError, + S3PayloadError, + SQSKubernetesEnvironment, + SqsConsumerError, + _b64compress, + _b64decompress, + _decode_stream_chunk, + _is_workload_error, + _parse_error_body, + _parse_dockerfile_copies, + _sanitize_image_name, + _split_family, + _synthesize_workload_exec_result, +) +from pier.models.environment_type import EnvironmentType + + +# ── helpers to build a stub environment without running __init__ ───────────── + + +def _bare_env( + *, + sqs_url: str = "https://sqs.eu-west-1.amazonaws.com/123456/docker-requests", + s3_bucket: str = "sqs-bucket", + registry_url: str = "123.dkr.ecr.eu-west-1.amazonaws.com/sandbox", + s3_threshold: int = 200 * 1024, + environment_dir: Path | None = None, + task_workdir: str | None = None, + task_os_windows: bool = False, + sqs_client: MagicMock | None = None, + s3_client: MagicMock | None = None, +) -> SQSKubernetesEnvironment: + """Build an SQSKubernetesEnvironment without invoking Base.__init__. + + Mirrors the pattern used by ``tests/test_daytona_network_limits.py``: + ``__new__`` + explicit attribute setup lets us exercise internal + helpers without needing a real TrialPaths, boto3 client, event loop, + or shared response queue. + """ + env = SQSKubernetesEnvironment.__new__(SQSKubernetesEnvironment) + env._sqs_queue_url = sqs_url + env._s3_bucket = s3_bucket + env._registry_url = registry_url + env._sqs_region = "eu-west-1" + env._s3_region = "eu-west-1" + env._s3_access_key_id = None + env._s3_secret_access_key = None + env._first_chunk_timeout = 60 + env._s3_threshold = s3_threshold + env._create_max_attempts = 8 + env._pre_install_commands = [] + env._ephemeral_storage_limit = "" + env._slurm_user = "test-user" + env._slurm_job_id = "test-job" + env._repo_url = "" + env._repo_commit = "" + env._repo_dest = "repo" + env._inject_prebuilt_env_files = True + env._explicit_volume_mounts = None + env._metrics_bridge_queue = "" + env._metrics_bridge_url = None + env._metrics_push_task = None + env._proxy_id = "test-proxy-id" + env._pre_generated_cid = None + env._all_pre_generated_cids = [] + env._sandbox_container_id = None + env._volume_mounts = [] + env._timing = {} + env._exec_count = 0 + env._exec_total_secs = 0.0 + env._start_time = None + env._sqs = sqs_client or MagicMock() + env._s3 = s3_client or MagicMock() + env._egress_proxy_url = "http://egress-proxy.sandbox-proxy.svc.cluster.local:3128" + env._egress_token = None + + env.environment_dir = environment_dir or Path(".") + env.environment_name = "test-env" + env.session_id = "test-session" + env.trial_paths = MagicMock() + env.default_user = None + env._persistent_env = {} + env.logger = MagicMock() + + from pier.models.agent.network import NetworkAllowlist + + env.network_allowlist = NetworkAllowlist() + + task_env_config = MagicMock() + task_env_config.docker_image = None + task_env_config.workdir = task_workdir + task_env_config.allow_internet = True + from pier.models.task.config import TaskOS + + task_env_config.os = TaskOS.WINDOWS if task_os_windows else TaskOS.LINUX + env.task_env_config = task_env_config + return env + + +# ── factory + type registration ────────────────────────────────────────────── + + +def test_environment_type_enum_has_sqs_kubernetes(): + assert EnvironmentType.SQS_KUBERNETES.value == "sqs_kubernetes" + + +def test_factory_registry_contains_sqs_kubernetes(): + assert EnvironmentType.SQS_KUBERNETES in _ENVIRONMENT_REGISTRY + entry = _ENVIRONMENT_REGISTRY[EnvironmentType.SQS_KUBERNETES] + assert entry.module == "pier.environments.sqs_kubernetes" + assert entry.class_name == "SQSKubernetesEnvironment" + + +def test_type_returns_enum_member(): + assert SQSKubernetesEnvironment.type() == EnvironmentType.SQS_KUBERNETES + + +# ── capabilities + resource shapes ─────────────────────────────────────────── + + +def test_capabilities_shape_matches_infra_reality(): + env = _bare_env() + caps = env.capabilities + assert isinstance(caps, EnvironmentCapabilities) + # Mounted so BaseEnvironment.stop() semantics match: volume-mount + # emulation via archive-pull happens during stop(). + assert caps.mounted is True + # EKS Linux nodegroups; no GPU nodegroup wired up here. + assert caps.gpus is False + # Per-agent egress allowlist enforcement wired through + # ``containers/create`` (agent-dist PR #216) — declared statically so + # older clusters still accept the create body (unknown field ignored) + # and simply return no ``egress_token``. + assert caps.filtered_egress is True + # No image-time agent preinstall; sandbox pods pull cached images. + assert caps.preinstall_agents is False + # EKS Linux nodegroups only. + assert caps.windows is False + + +def test_resource_capabilities_flag_both_request_and_limit(): + caps = SQSKubernetesEnvironment.resource_capabilities() + assert caps.cpu_request is True + assert caps.cpu_limit is True + assert caps.memory_request is True + assert caps.memory_limit is True + + +# ── _validate_definition ───────────────────────────────────────────────────── + + +def test_validate_definition_missing_sqs_url_rejected(tmp_path): + env = _bare_env(sqs_url="", environment_dir=tmp_path) + with pytest.raises(ValueError, match="sqs_queue_url is required"): + env._validate_definition() + + +def test_validate_definition_bad_sqs_url_rejected(tmp_path): + env = _bare_env(sqs_url="not-a-url", environment_dir=tmp_path) + with pytest.raises(ValueError, match="does not look like"): + env._validate_definition() + + +def test_validate_definition_missing_s3_bucket_rejected(tmp_path): + env = _bare_env(s3_bucket="", environment_dir=tmp_path) + with pytest.raises(ValueError, match="s3_bucket is required"): + env._validate_definition() + + +def test_validate_definition_missing_registry_rejected(tmp_path): + env = _bare_env(registry_url="", environment_dir=tmp_path) + with pytest.raises(ValueError, match="registry_url is required"): + env._validate_definition() + + +def test_validate_definition_missing_env_dir_rejected(tmp_path): + env = _bare_env(environment_dir=tmp_path / "does-not-exist") + with pytest.raises(FileNotFoundError, match="environment_dir does not exist"): + env._validate_definition() + + +def test_validate_definition_rejects_windows_task(tmp_path): + env = _bare_env(environment_dir=tmp_path, task_os_windows=True) + with pytest.raises(RuntimeError, match="Windows"): + env._validate_definition() + + +def test_validate_definition_happy_path(tmp_path): + env = _bare_env(environment_dir=tmp_path) + env._validate_definition() # does not raise + + +# ── body encoding ──────────────────────────────────────────────────────────── + + +def test_encode_body_small_utf8_is_inline_uncompressed(): + env = _bare_env() + content, compress = env._encode_body(b'{"hello":"world"}') + assert content == '{"hello":"world"}' + assert compress is False + + +def test_encode_body_binary_falls_back_to_base64_uncompressed(): + env = _bare_env() + binary = bytes(range(256)) * 10 # 2560 bytes with NULs + content, compress = env._encode_body(binary) + assert compress is False + assert base64.b64decode(content) == binary + + +def test_encode_body_above_compression_threshold_zlib_b64(): + env = _bare_env() + body = b"A" * (300 * 1024) # 300 KiB — highly compressible + content, compress = env._encode_body(body) + assert compress is True + assert _b64decompress(content) == body + + +def test_b64compress_roundtrip(): + body = b"hello world" * 500 + assert _b64decompress(_b64compress(body)) == body + + +# ── _build_sqs_message inline vs S3 ────────────────────────────────────────── + + +def test_build_sqs_message_no_body_is_inline_empty(): + env = _bare_env() + msg = asyncio.run( + env._build_sqs_message("GET", "v1.43/containers/abc", query="path=/x") + ) + assert msg["method"] == "GET" + assert msg["path"] == "v1.43/containers/abc" + assert msg["query"] == "path=/x" + assert msg["content"] == "" + assert msg["compress"] is False + assert "s3_key" not in msg + + +def test_build_sqs_message_below_threshold_inline(): + # 150 KiB non-compressible body → post-encode still under 200 KiB → inline. + import os as _os + + body = _os.urandom(150 * 1024) + env = _bare_env(s3_threshold=200 * 1024) + msg = asyncio.run(env._build_sqs_message("POST", "v1.43/x", body=body)) + assert "s3_key" not in msg + assert msg["content"] != "" + env._s3.put_object.assert_not_called() + + +def test_build_sqs_message_above_threshold_uploads_to_s3(): + # 250 KiB uncompressible body → encoded body > 200 KiB → S3. + import os as _os + + body = _os.urandom(250 * 1024) + env = _bare_env(s3_threshold=200 * 1024) + msg = asyncio.run(env._build_sqs_message("POST", "v1.43/x", body=body)) + assert msg["content"] == "" + assert msg["s3_key"].startswith("sqs-k8s-payloads/") + env._s3.put_object.assert_called_once() + call_kwargs = env._s3.put_object.call_args.kwargs + assert call_kwargs["Bucket"] == env._s3_bucket + assert call_kwargs["Key"] == msg["s3_key"] + assert call_kwargs["Body"] == body + + +def test_build_sqs_message_two_mib_body_uses_s3(): + import os as _os + + body = _os.urandom(2 * 1024 * 1024) + env = _bare_env(s3_threshold=200 * 1024) + msg = asyncio.run(env._build_sqs_message("PUT", "v1.43/x/archive", body=body)) + assert "s3_key" in msg + assert msg["content"] == "" + env._s3.put_object.assert_called_once() + + +def test_build_sqs_message_extra_fields_merged(): + env = _bare_env() + msg = asyncio.run( + env._build_sqs_message( + "POST", + "v1.43/containers/create", + body=b"{}", + extra_fields={"slurm_user": "u", "slurm_job_id": "1"}, + ) + ) + assert msg["slurm_user"] == "u" + assert msg["slurm_job_id"] == "1" + + +# ── S3 failure modes ───────────────────────────────────────────────────────── + + +def test_s3_upload_failure_surfaces_S3PayloadError(): + from botocore.exceptions import BotoCoreError + + env = _bare_env(s3_threshold=100) + env._s3.put_object.side_effect = BotoCoreError() + with pytest.raises(S3PayloadError) as ei: + asyncio.run(env._build_sqs_message("PUT", "v1.43/x", body=b"A" * 500)) + assert ei.value.direction == "upload" + + +def test_s3_download_failure_surfaces_S3PayloadError(): + from botocore.exceptions import BotoCoreError + + env = _bare_env() + env._s3.get_object.side_effect = BotoCoreError() + with pytest.raises(S3PayloadError) as ei: + asyncio.run(env._download_from_s3("some/key")) + assert ei.value.direction == "download" + assert ei.value.key == "some/key" + + +# ── response-side S3 download ──────────────────────────────────────────────── + + +def test_response_decode_s3_key_raises_if_not_downloaded_first(): + env = _bare_env() + with pytest.raises(RuntimeError, match="Caller must handle s3_key"): + env._decode_response_body({"s3_key": "abc"}) + + +def test_response_decode_inline_utf8(): + env = _bare_env() + out = env._decode_response_body( + {"content": "hello world", "compress": False, "content_type": "text/plain"} + ) + assert out == b"hello world" + + +def test_response_decode_inline_compressed(): + env = _bare_env() + payload = b"payload " * 100 + encoded = base64.b64encode(zlib.compress(payload)).decode() + out = env._decode_response_body({"content": encoded, "compress": True}) + assert out == payload + + +def test_response_decode_binary_content_type_base64(): + env = _bare_env() + payload = b"\x00\x01\x02\x03" + encoded = base64.b64encode(payload).decode() + out = env._decode_response_body( + { + "content": encoded, + "compress": False, + "content_type": "application/x-tar", + } + ) + assert out == payload + + +# ── roundtrip: request via S3, response via S3 ─────────────────────────────── + + +def test_full_roundtrip_large_request_and_large_response(monkeypatch): + """2 MiB request → S3; consumer response with s3_key → downloaded from S3.""" + import os as _os + + request_body = _os.urandom(2 * 1024 * 1024) + response_body = _os.urandom(3 * 1024 * 1024) + + env = _bare_env(s3_threshold=200 * 1024) + + # Simulate S3: put_object records, get_object returns the response payload. + s3_state: dict[str, bytes] = {} + + def _put(Bucket, Key, Body): + assert Bucket == env._s3_bucket + s3_state[Key] = Body + + def _get(Bucket, Key): + return {"Body": io.BytesIO(s3_state[Key])} + + env._s3.put_object.side_effect = _put + env._s3.get_object.side_effect = _get + + # Simulate the consumer's response by patching _shared_pending after send. + async def _run(): + # Manually build request message → verifies request-side S3 upload. + msg = await env._build_sqs_message( + "POST", "v1.43/containers/create", body=request_body + ) + assert msg["s3_key"] in s3_state + assert s3_state[msg["s3_key"]] == request_body + + # Consumer would put the response bytes into S3 under `resp_key` and + # send an SQS response referring to it. Simulate the download side. + resp_key = "sandbox-response/xyz" + s3_state[resp_key] = response_body + downloaded = await env._download_from_s3(resp_key) + assert downloaded == response_body + + asyncio.run(_run()) + + +# ── consumer error classification ──────────────────────────────────────────── + + +def test_parse_error_body_trusts_retryable_field(): + body = json.dumps( + { + "message": "cluster is full", + "error_code": "CLUSTER_FULL", + "retryable": True, + "details": {"reason": "quota_saturated"}, + } + ).encode() + err = _parse_error_body(body, 503, op_label="containers/create") + assert err.error_code == "CLUSTER_FULL" + assert err.retryable is True + assert err.details == {"reason": "quota_saturated"} + assert err.status == 503 + + +def test_parse_error_body_missing_body_fail_closed(): + err = _parse_error_body(b"", 500, op_label="x") + assert err.retryable is False + assert err.error_code == "UNKNOWN" + + +def test_parse_error_body_bad_json_fail_closed(): + err = _parse_error_body(b"not-json", 500, op_label="x") + assert err.retryable is False + assert err.error_code == "UNKNOWN" + + +def test_parse_error_body_non_retryable_build_failure(): + body = json.dumps( + { + "message": "Dockerfile broken", + "error_code": "BUILD_FAILED", + "retryable": False, + } + ).encode() + err = _parse_error_body(body, 400, op_label="containers/create") + assert err.retryable is False + + +# ── workload synthesis ────────────────────────────────────────────────────── + + +def test_workload_error_synthesised_as_ExecResult_137_for_oom(): + err = SqsConsumerError( + status=500, + error_code="OOM_KILLED", + retryable=False, + message="memory cgroup OOM", + details={"is_workload": True, "memory_used": "8Gi"}, + raw_body=b"", + ) + assert _is_workload_error(err) is True + result = _synthesize_workload_exec_result(err) + assert result.return_code == 137 + assert "OOM_KILLED" in result.stderr + assert "memory_used=8Gi" in result.stderr + assert result.stdout == "" + + +def test_workload_error_synthesised_as_ExecResult_1_for_unknown(): + err = SqsConsumerError( + status=500, + error_code="DISK_FULL", + retryable=False, + message="", + details={"is_workload": True}, + raw_body=b"", + ) + result = _synthesize_workload_exec_result(err) + assert result.return_code == 1 # DISK_FULL maps to 1 + + +def test_non_workload_error_not_flagged(): + err = SqsConsumerError( + status=500, + error_code="INTERNAL_ERROR", + retryable=True, + message="oops", + details={}, + raw_body=b"", + ) + assert _is_workload_error(err) is False + + +# ── decode helpers ─────────────────────────────────────────────────────────── + + +def test_decode_stream_chunk_text_json_no_b64(): + assert _decode_stream_chunk('{"a":1}', "application/json") == b'{"a":1}' + + +def test_decode_stream_chunk_text_plain_no_b64(): + assert _decode_stream_chunk("hello", "text/plain") == b"hello" + + +def test_decode_stream_chunk_binary_base64(): + payload = b"\xff\xfe\xfd" + encoded = base64.b64encode(payload).decode() + assert _decode_stream_chunk(encoded, "application/octet-stream") == payload + + +def test_decode_stream_chunk_bad_b64_returns_empty(): + # Bad base64 should degrade to empty, not raise. + assert _decode_stream_chunk("!!not-b64!!", "application/octet-stream") == b"" + + +# ── image-name / family sanitiser ──────────────────────────────────────────── + + +def test_sanitize_image_name_lowercases_and_prefixes_digit(): + assert _sanitize_image_name("MyImage.Bad!") == "myimage.bad-" + assert _sanitize_image_name(".hidden") == "0.hidden" + + +def test_split_family(): + assert _split_family("openswe__akuli__mantaray-165") == ( + "openswe", + "akuli__mantaray-165", + ) + assert _split_family("scale-swe__foo-bar") == ("scale-swe", "foo-bar") + assert _split_family("no-double-underscore") == (None, "no-double-underscore") + + +# ── Dockerfile COPY parsing ────────────────────────────────────────────────── + + +def test_parse_dockerfile_copies_handles_flags_and_multi_source(): + text = """ + FROM python:3.11 + COPY --chown=root:root src/ /app/src/ + COPY a b c /dest + COPY --from=builder /out/bin /usr/local/bin/prog + COPY solution.sh /solution/solve.sh + """ + m = _parse_dockerfile_copies(text) + assert m["src"] == "/app/src/" + assert m["a"] == "/dest/" + assert m["b"] == "/dest/" + assert m["c"] == "/dest/" + # --from copies are skipped (source is another image layer). + assert "/out/bin" not in m + assert m["solution.sh"] == "/solution/solve.sh" + + +# ── signature parity with pier's BaseEnvironment.exec ──────────────────────── + + +def test_exec_signature_matches_base_environment(): + base_sig = inspect.signature(BaseEnvironment.exec) + sub_sig = inspect.signature(SQSKubernetesEnvironment.exec) + assert list(base_sig.parameters.keys()) == list(sub_sig.parameters.keys()) + for name in ("command", "cwd", "env", "timeout_sec", "user"): + assert name in sub_sig.parameters, f"missing parameter: {name}" + + +# ── K8sInfraError propagation ──────────────────────────────────────────────── + + +def test_exec_with_no_container_raises_K8sInfraError(): + env = _bare_env() + env._sandbox_container_id = None + with pytest.raises(K8sInfraError) as ei: + asyncio.run(env.exec("echo hi")) + assert ei.value.error_code == "CONTAINER_NOT_FOUND" + + +# ── volume mount default sourcing ──────────────────────────────────────────── + + +def test_default_volume_mounts_from_env_paths(tmp_path): + # Bypass __init__ but exercise the default-mounts branch by hand. + env = _bare_env() + trial_paths = MagicMock() + trial_paths.agent_dir = tmp_path / "agent" + trial_paths.verifier_dir = tmp_path / "verifier" + trial_paths.artifacts_dir = tmp_path / "artifacts" + env.trial_paths = trial_paths + + # Simulate the constructor's default-mount branch. + env._explicit_volume_mounts = None + from pier.models.trial.paths import EnvironmentPaths + + default_mounts = [ + (str(EnvironmentPaths.agent_dir), str(env.trial_paths.agent_dir)), + (str(EnvironmentPaths.verifier_dir), str(env.trial_paths.verifier_dir)), + (str(EnvironmentPaths.artifacts_dir), str(env.trial_paths.artifacts_dir)), + ] + assert default_mounts[0][0] == "/logs/agent" + assert default_mounts[1][0] == "/logs/verifier" + assert default_mounts[2][0] == "/logs/artifacts" + + +# ── ExecResult shape sanity ───────────────────────────────────────────────── + + +def test_execresult_used_by_synth_matches_pier_shape(): + # Pier's ExecResult has stdout/stderr/return_code only — no error_code. + err = SqsConsumerError( + status=500, + error_code="OOM_KILLED", + retryable=False, + message="oom", + details={"is_workload": True}, + raw_body=b"", + ) + result = _synthesize_workload_exec_result(err) + assert isinstance(result, ExecResult) + # Must not carry a rogue `error_code` field (pier's ExecResult lacks it). + assert not hasattr(result, "error_code") or result.__class__.model_fields.get( + "error_code" + ) is None + + +# ── agent_process_env: egress-proxy injection ─────────────────────────────── + + +def test_agent_process_env_returns_env_unchanged_without_token(): + """No token minted (default) → env passthrough (identity for None).""" + env = _bare_env() + env._egress_token = None + assert env.agent_process_env(None) is None + caller = {"FOO": "bar"} + assert env.agent_process_env(caller) is caller + + +def test_agent_process_env_injects_proxy_when_token_set(): + """Token → HTTP(S)_PROXY / NO_PROXY populated in ``http://agent:@host:port`` form. + + The ``agent`` username is convention (matches agent-dist egress-proxy); + the server takes the password half via ``parseProxyAuthToken``. + """ + env = _bare_env() + env._egress_token = "deadbeef" * 8 + out = env.agent_process_env(None) + assert out is not None + for key in ("HTTP_PROXY", "http_proxy", "HTTPS_PROXY", "https_proxy"): + assert ( + out[key] + == "http://agent:deadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef" + "@egress-proxy.sandbox-proxy.svc.cluster.local:3128" + ), key + for key in ("NO_PROXY", "no_proxy"): + assert ".svc" in out[key] + assert ".cluster.local" in out[key] + assert "10.0.0.0/8" in out[key] + assert "127.0.0.1" in out[key] + assert "localhost" in out[key] + + +def test_agent_process_env_caller_env_wins_on_key_collision(): + """Caller-supplied HTTP_PROXY (or any key) survives the merge.""" + env = _bare_env() + env._egress_token = "tok-xyz" + caller = { + "HTTP_PROXY": "http://caller-set:1/", + "NO_PROXY": "caller-no-proxy", + "CUSTOM": "keep-me", + } + out = env.agent_process_env(caller) + assert out["HTTP_PROXY"] == "http://caller-set:1/" + assert out["NO_PROXY"] == "caller-no-proxy" + assert out["CUSTOM"] == "keep-me" + # Untouched keys still get our default (proves we're actually merging). + assert out["HTTPS_PROXY"].startswith("http://agent:tok-xyz@") + assert out["https_proxy"].startswith("http://agent:tok-xyz@") + + +def test_agent_process_env_respects_custom_egress_proxy_url(): + """Override kwarg replaces both scheme+host+port in the emitted URL.""" + env = _bare_env() + env._egress_proxy_url = "http://other-proxy.example.svc:9000" + env._egress_token = "tok" + out = env.agent_process_env(None) + assert out is not None + assert out["HTTP_PROXY"] == "http://agent:tok@other-proxy.example.svc:9000" + + +def test_agent_process_env_accepts_bare_host_port(): + """``egress_proxy_url`` without a scheme still emits a valid http:// URL.""" + env = _bare_env() + env._egress_proxy_url = "internal-proxy.svc:3128" + env._egress_token = "tok" + out = env.agent_process_env(None) + assert out is not None + assert out["HTTP_PROXY"] == "http://agent:tok@internal-proxy.svc:3128" + + +# ── containers/create body wiring for network_allowlist ───────────────────── + + +def _fake_create_roundtrip(env, captured: dict, egress_token: str | None = None): + """Patch _sqs_round_trip to record the create call and return 201 + Id.""" + from pier.models.agent.network import NetworkAllowlist # noqa: F401 + + async def _rt(method, path, *, query="", headers=None, body=b"", + extra_fields=None, on_send=None, wants_keepalive=False): + captured["method"] = method + captured["path"] = path + captured["body"] = json.loads(body.decode()) if body else {} + captured["extra_fields"] = dict(extra_fields or {}) + if on_send is not None: + on_send() + resp: dict = {"Id": "container-abc123"} + if egress_token is not None: + resp["egress_token"] = egress_token + return 201, json.dumps(resp).encode() + + env._sqs_round_trip = _rt # type: ignore[method-assign] + + +def _prep_env_for_start(env, *, environment_dir: Path): + """Stub the parts of start() that require real infra.""" + from pier.models.trial.config import ResourceMode + + env.environment_dir = environment_dir + env._cpu_resource_mode = ResourceMode.AUTO + env._memory_resource_mode = ResourceMode.AUTO + env._override_cpus = None + env._override_memory_mb = None + env._override_storage_mb = None + env._override_gpus = None + env.task_env_config.cpus = None + env.task_env_config.memory_mb = None + env.task_env_config.storage_mb = None + env.task_env_config.gpus = None + # Pre-built image path → skips build-context upload branch entirely. + env.task_env_config.docker_image = "prebuilt:tag" + + async def _noop_startup(): + return + + env._startup = _noop_startup # type: ignore[method-assign] + # Skip file-injection branch (no Dockerfile in tmp). + env._inject_prebuilt_env_files = False + env._metrics_bridge_queue = "" + env._metrics_bridge_url = None + + +def test_start_includes_network_allowlist_when_domains_set_and_internet_denied( + tmp_path, +): + from pier.models.agent.network import NetworkAllowlist + + env = _bare_env(environment_dir=tmp_path) + env.network_allowlist = NetworkAllowlist(domains=["httpbin.org", ".anthropic.com"]) + env.task_env_config.allow_internet = False + _prep_env_for_start(env, environment_dir=tmp_path) + + captured: dict = {} + _fake_create_roundtrip(env, captured, egress_token="minted-tok-123") + + asyncio.run(env.start()) + + extra = captured["extra_fields"] + assert "network_allowlist" in extra + assert extra["network_allowlist"] == { + "domains": [".anthropic.com", "httpbin.org"] # NetworkAllowlist sorts + } + # Consumer response's egress_token is captured on self. + assert env._egress_token == "minted-tok-123" + + +def test_start_skips_network_allowlist_when_allow_internet_true(tmp_path): + from pier.models.agent.network import NetworkAllowlist + + env = _bare_env(environment_dir=tmp_path) + env.network_allowlist = NetworkAllowlist(domains=["httpbin.org"]) + env.task_env_config.allow_internet = True # internet allowed → no allowlist + _prep_env_for_start(env, environment_dir=tmp_path) + + captured: dict = {} + _fake_create_roundtrip(env, captured, egress_token=None) + + asyncio.run(env.start()) + + assert "network_allowlist" not in captured["extra_fields"] + assert env._egress_token is None + + +def test_start_skips_network_allowlist_when_domains_empty(tmp_path): + from pier.models.agent.network import NetworkAllowlist + + env = _bare_env(environment_dir=tmp_path) + env.network_allowlist = NetworkAllowlist(domains=[]) + env.task_env_config.allow_internet = False + _prep_env_for_start(env, environment_dir=tmp_path) + + captured: dict = {} + _fake_create_roundtrip(env, captured, egress_token=None) + + asyncio.run(env.start()) + + assert "network_allowlist" not in captured["extra_fields"] + assert env._egress_token is None + + +def test_start_missing_egress_token_leaves_agent_process_env_a_noop(tmp_path): + """Older cluster: allowlist sent but consumer doesn't understand → no token → no injection.""" + from pier.models.agent.network import NetworkAllowlist + + env = _bare_env(environment_dir=tmp_path) + env.network_allowlist = NetworkAllowlist(domains=["httpbin.org"]) + env.task_env_config.allow_internet = False + _prep_env_for_start(env, environment_dir=tmp_path) + + captured: dict = {} + _fake_create_roundtrip(env, captured, egress_token=None) # server omits token + + asyncio.run(env.start()) + + assert env._egress_token is None + assert env.agent_process_env({"X": "y"}) == {"X": "y"} + + +# ── capabilities flag ─────────────────────────────────────────────────────── + + +def test_capabilities_filtered_egress_true(): + """Filtered-egress is declared statically; see comment on the property.""" + env = _bare_env() + assert env.capabilities.filtered_egress is True diff --git a/uv.lock b/uv.lock index c479fa26..67d308e5 100644 --- a/uv.lock +++ b/uv.lock @@ -197,18 +197,32 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309", size = 67548, upload-time = "2026-03-19T14:22:23.645Z" }, ] +[[package]] +name = "boto3" +version = "1.43.77" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "botocore" }, + { name = "jmespath" }, + { name = "s3transfer" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3c/6b/45cbf87e918e8a970b83feca75148787e084e60cad1de791cd2da089c926/boto3-1.43.77.tar.gz", hash = "sha256:790082cca710ce8b49b69960731df49fa7cef709823c162e4697412f35d6c813", size = 112639, upload-time = "2026-08-21T03:46:43.044Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9f/e1/d0ddad94091eae29b14d974d02fd6f30e7dfca71e36bc0484d64f921dc0f/boto3-1.43.77-py3-none-any.whl", hash = "sha256:3be161c9edf7645ba09651a442899b5fda9a7e5529f94eb4d9e2a276f1fc1208", size = 140028, upload-time = "2026-08-21T03:46:40.987Z" }, +] + [[package]] name = "botocore" -version = "1.43.11" +version = "1.43.77" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jmespath" }, { name = "python-dateutil" }, { name = "urllib3" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e9/fa/4bec16fa5a4cde7b593e549238bfeb8ed1bdba9d427888a18c460a1f2352/botocore-1.43.11.tar.gz", hash = "sha256:d7d479cc2809ec2728f2898521003adfb79bfe6a4615c59dfd222ec52b0cee6b", size = 15364020, upload-time = "2026-05-19T19:39:58.317Z" } +sdist = { url = "https://files.pythonhosted.org/packages/56/e4/21c88d018b11fd9999d492c45ed010aca4015f4d3bdec0fcda9d069b0770/botocore-1.43.77.tar.gz", hash = "sha256:b33300395e2fb994ec8baafff76a600184e426ec7ea462cc59a0e8d1469d557e", size = 15980152, upload-time = "2026-08-21T03:46:37.742Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e1/9a/9f1d955c2eebefb6bd20de740ae7a05e7b015c63f0f01dba338dcf29cc68/botocore-1.43.11-py3-none-any.whl", hash = "sha256:0108b5604df5a26918936c845e1e761866ee9ea8d1c1f9358ed3c69afdc37436", size = 15043467, upload-time = "2026-05-19T19:39:53.176Z" }, + { url = "https://files.pythonhosted.org/packages/4b/aa/978683172e66f87b0217b47a743a1435f69e2da05856e9796d8a09428401/botocore-1.43.77-py3-none-any.whl", hash = "sha256:e88d86643daa1d85e1c4e9644187aca60325955c596ea2eea646730857998631", size = 15674302, upload-time = "2026-08-21T03:46:33.636Z" }, ] [[package]] @@ -488,6 +502,7 @@ source = { editable = "." } dependencies = [ { name = "anthropic" }, { name = "anyio" }, + { name = "boto3" }, { name = "botocore" }, { name = "claude-agent-sdk" }, { name = "daytona" }, @@ -520,6 +535,7 @@ dev = [ requires-dist = [ { name = "anthropic", specifier = ">=0.97.0" }, { name = "anyio", specifier = ">=4.13.0" }, + { name = "boto3", specifier = ">=1.40.11" }, { name = "botocore", specifier = ">=1.43.11" }, { name = "claude-agent-sdk", specifier = ">=0.1.17" }, { name = "daytona", specifier = ">=0.121.0" }, @@ -2716,6 +2732,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c0/98/6beb4b351e472e5f4c4613f7c35a5290b8be2497e183825310c4c3a3984b/ruff-0.15.12-py3-none-win_arm64.whl", hash = "sha256:a538f7a82d061cee7be55542aca1d86d1393d55d81d4fcc314370f4340930d4f", size = 11120821, upload-time = "2026-04-24T18:16:57.979Z" }, ] +[[package]] +name = "s3transfer" +version = "0.19.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "botocore" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/76/43/35e4d8aa320bffe8287fe8f65f578fa2d2db0a64212f0e710dce58267854/s3transfer-0.19.2.tar.gz", hash = "sha256:ba0309fd86be3c27dbf78cdd813c13c5e1df16e5874b99d2535ebbdfb9892993", size = 165592, upload-time = "2026-07-22T19:30:44.432Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bc/e7/5c595c75e9f41a44f30e526eda465ea0b4eec93470e074e4a111b253f13a/s3transfer-0.19.2-py3-none-any.whl", hash = "sha256:d8168eccca828cbb2cd573675333f3bddd254313a9c42494b84c76b539e8ba25", size = 90216, upload-time = "2026-07-22T19:30:43.251Z" }, +] + [[package]] name = "scantree" version = "0.0.4"