From 5569fd0580262bc0f2c547d9ee56e9bf6684612b Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Tue, 10 Mar 2026 13:34:09 -0700 Subject: [PATCH 01/31] moved perf/enroot fixes from internal Gym to hypotest --- src/hypotest/env/interpreter_env.py | 974 ++++++++++++++++++++++- src/hypotest/env/kernel_server.py | 112 ++- src/hypotest/env/tools/filesystem.py | 18 +- src/hypotest/env/utils/__init__.py | 2 +- src/hypotest/env/utils/notebook_utils.py | 2 +- 5 files changed, 1033 insertions(+), 75 deletions(-) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index deb936d..a2865b8 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -4,6 +4,9 @@ code in Jupyter kernels. It focuses on direct code execution via run_cell(). """ +import sys +import ray +import signal import argparse import asyncio import json @@ -11,11 +14,15 @@ import os import shutil import socket +import statistics +import threading import time +import random from pathlib import Path from tempfile import mkdtemp from typing import TYPE_CHECKING, Any, cast from uuid import UUID +import uuid import aiodocker import httpx @@ -36,6 +43,9 @@ from nbformat import NotebookNode from pydantic import BaseModel, Field, JsonValue +import subprocess +from textwrap import dedent + from . import config as cfg from .config import ExecutionConfig from .interpreter import ExecutionResult, Interpreter @@ -49,21 +59,159 @@ # Port management for Docker containers _USED_PORTS: set[int] = set() +used_ports_lock = asyncio.Lock() +# container launch semaphore to limit concurrency +CONTAINER_LAUNCH_SEM = asyncio.Semaphore(128) -def get_free_port() -> int: +async def get_free_port() -> int: """Get a free port for the kernel server container.""" while True: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind(("", 0)) port = s.getsockname()[1] - if port not in _USED_PORTS: - _USED_PORTS.add(port) - return port + async with used_ports_lock: + if port not in _USED_PORTS: + _USED_PORTS.add(port) + return port logger = logging.getLogger(__name__) +# Container lifecycle log level: root logger defaults to WARNING and we +# cannot reconfigure it, so use WARNING for all container diagnostics to +# ensure they are visible in production logs. +_CONTAINER_LOG_LEVEL = logging.WARNING + + +def _kill_process_group(proc: subprocess.Popen, label: str = "enroot", sigterm_timeout: float = 15) -> None: + """Safely terminate a process group, escalating from SIGTERM to SIGKILL. + + Handles all edge cases: already-dead process, missing process group, etc. + """ + if proc.poll() is not None: + logger.log( + _CONTAINER_LOG_LEVEL, + "[%s] Process pid=%d already exited with returncode=%d", + label, proc.pid, proc.returncode, + ) + return + + pgid = None + try: + pgid = os.getpgid(proc.pid) + except ProcessLookupError: + logger.log(_CONTAINER_LOG_LEVEL, "[%s] Process pid=%d vanished before we could get pgid", label, proc.pid) + return + + # SIGTERM the whole group + try: + logger.log(_CONTAINER_LOG_LEVEL, "[%s] Sending SIGTERM to pgid=%d (pid=%d)", label, pgid, proc.pid) + os.killpg(pgid, signal.SIGTERM) + except ProcessLookupError: + logger.log(_CONTAINER_LOG_LEVEL, "[%s] Process group pgid=%d already gone after SIGTERM", label, pgid) + return + + try: + # proc.wait(timeout=sigterm_timeout) + proc.communicate(timeout=sigterm_timeout) + logger.log( + _CONTAINER_LOG_LEVEL, + "[%s] Process pid=%d exited after SIGTERM with returncode=%d", + label, proc.pid, proc.returncode, + ) + return + except subprocess.TimeoutExpired: + logger.warning( + "[%s] Process pid=%d did not exit within %.1fs of SIGTERM, sending SIGKILL", + label, proc.pid, sigterm_timeout, + ) + + # SIGKILL the whole group + try: + os.killpg(pgid, signal.SIGKILL) + except ProcessLookupError: + logger.log(_CONTAINER_LOG_LEVEL, "[%s] Process group pgid=%d already gone before SIGKILL", label, pgid) + return + + try: + # proc.wait(timeout=5) + proc.communicate(timeout=5) + logger.log( + _CONTAINER_LOG_LEVEL, + "[%s] Process pid=%d exited after SIGKILL with returncode=%d", + label, proc.pid, proc.returncode, + ) + except subprocess.TimeoutExpired: + logger.error("[%s] Process pid=%d still alive after SIGKILL — possible zombie", label, proc.pid) + +PERF_METRICS_ENABLED = os.getenv("NEMO_GYM_PERF_METRICS", "0") == "1" +PERF_DETAIL_SAMPLE_RATE = float(os.getenv("NEMO_GYM_PERF_DETAIL_SAMPLE_RATE", "0.2")) +PERF_SLOW_STEP_MS = float(os.getenv("NEMO_GYM_PERF_SLOW_STEP_MS", "5000")) +PERF_AGGREGATE_INTERVAL_S = float(os.getenv("NEMO_GYM_PERF_AGGREGATE_INTERVAL_S", "30")) +PERF_LOG_LEVEL = getattr(logging, os.getenv("NEMO_GYM_PERF_LOG_LEVEL", "WARNING").upper(), logging.WARNING) + +_PERF_METRIC_SAMPLES: dict[str, list[float]] = {} +_PERF_LAST_FLUSH_TS = time.perf_counter() +_PERF_LOCK = threading.Lock() + + +def _log_perf(payload: dict[str, Any]) -> None: + logger.log(PERF_LOG_LEVEL, json.dumps(payload)) + + +def _percentile(sorted_values: list[float], pct: float) -> float: + if not sorted_values: + return 0.0 + if len(sorted_values) == 1: + return sorted_values[0] + idx = (len(sorted_values) - 1) * pct + lo = int(idx) + hi = min(lo + 1, len(sorted_values) - 1) + frac = idx - lo + return sorted_values[lo] * (1.0 - frac) + sorted_values[hi] * frac + + +def _maybe_flush_perf_aggregates() -> None: + global _PERF_LAST_FLUSH_TS, _PERF_METRIC_SAMPLES + now = time.perf_counter() + if now - _PERF_LAST_FLUSH_TS < PERF_AGGREGATE_INTERVAL_S: + return + + with _PERF_LOCK: + now = time.perf_counter() + if now - _PERF_LAST_FLUSH_TS < PERF_AGGREGATE_INTERVAL_S: + return + snapshots = _PERF_METRIC_SAMPLES + _PERF_METRIC_SAMPLES = {} + _PERF_LAST_FLUSH_TS = now + + for metric_name, values in snapshots.items(): + if not values: + continue + ordered = sorted(values) + payload = { + "component": "aviary_hypotest.interpreter_env", + "event": "perf_aggregate", + "metric": metric_name, + "count": len(values), + "mean_ms": round(statistics.fmean(values), 3), + "p50_ms": round(_percentile(ordered, 0.50), 3), + "p90_ms": round(_percentile(ordered, 0.90), 3), + "p95_ms": round(_percentile(ordered, 0.95), 3), + "p99_ms": round(_percentile(ordered, 0.99), 3), + "max_ms": round(ordered[-1], 3), + } + _log_perf(payload) + + +def _record_perf_metric(metric_name: str, value_ms: float) -> None: + if not PERF_METRICS_ENABLED: + return + with _PERF_LOCK: + _PERF_METRIC_SAMPLES.setdefault(metric_name, []).append(value_ms) + _maybe_flush_perf_aggregates() + class ProblemInstance(BaseModel): uuid: UUID @@ -73,6 +221,376 @@ class ProblemInstance(BaseModel): rubric: str max_score: int = Field(alias="max_points") metadata: dict[str, JsonValue] = Field(default_factory=dict) + nb_primary_language: str + +def _prep_workspace_dir(work_dir: str, workspace_path: str = "/data_workspace") -> None: + wd = Path(work_dir) + (wd / "pydeps").mkdir(parents=True, exist_ok=True) + (wd / "pip-cache").mkdir(parents=True, exist_ok=True) + + (wd / "pip.conf").write_text( + "[global]\n" + "disable-pip-version-check = true\n" + "no-input = true\n" + f"cache-dir = {workspace_path}/pip-cache\n" + f"target = {workspace_path}/pydeps\n" + ) + +@ray.remote( + scheduling_strategy="SPREAD", + runtime_env={ + "py_executable": sys.executable, + }, +) +class EnrootKernelServer: + def __init__(self, container_sqsh_path: Path, execution_timeout: float): + self.container_sqsh_path = container_sqsh_path + self.execution_timeout = execution_timeout + self._enroot_proc: subprocess.Popen[str] | None = None + self._http_client: httpx.AsyncClient | None = None + self._container_port: int | None = None + self._container_log_path: Path | None = None + self._container_log_file: Any | None = None + + def _proc_label(self) -> str: + """Short label for log messages identifying this container.""" + port = self._container_port or "?" + pid = self._enroot_proc.pid if self._enroot_proc else "?" + return f"enroot(port={port}, pid={pid})" + + async def initialize(self, work_dir: Path, language: NBLanguage) -> None: + _prep_workspace_dir(work_dir) + + online = False + attempt = 0 + while not online: + attempt += 1 + self._container_port = await get_free_port() + startup_token = str(uuid.uuid4()) + + bash = dedent(f"""\ + set -euo pipefail + + WORKDIR="$(mktemp -d /dev/shm/data_workspace.XXXXXX)" + trap 'rm -rf "$WORKDIR"' EXIT + + mkdir -p /dev/shm/$WORKDIR + cp -a /data_workspace/. $WORKDIR/ + + cd $WORKDIR + + mkdir -p "$WORKDIR/r_libs" + cat >"$WORKDIR/Rprofile" <<'EOF' + .local_lib <- Sys.getenv("R_LIBS_USER") + if (nzchar(.local_lib)) .libPaths(unique(c(.local_lib, .libPaths()))) + EOF + + mkdir -p $WORKDIR/r_libs + export R_LIBS_USER="$WORKDIR/r_libs" + export R_PROFILE_USER="$WORKDIR/Rprofile" + + export PYTHONPATH="$WORKDIR/pydeps:${{PYTHONPATH}}" + export PIP_CONFIG_FILE=$WORKDIR/pip.conf + export target_platform=${{target_platform:-linux-64}} + + source activate /app/kernel_env + exec /app/kernel_env/bin/python /envs/kernel_server.py \\ + --work_dir $WORKDIR \\ + --language {language.value} \\ + --port {self._container_port} \\ + --startup-token {startup_token} + """).strip() + + env_dir = Path(__file__).parent + kernel_server_path = env_dir / "kernel_server.py" + assert kernel_server_path.is_file(), f"kernel server must be a valid path, found {kernel_server_path}" + + cmd = [ + "env", "-i", "PATH=/usr/sbin:/usr/bin:/sbin:/bin", 'HOME="$HOME"', 'USER="$USER"', + "enroot", "start", + "--mount", f"{work_dir}:/data_workspace", + "--mount", f"{kernel_server_path.resolve()}:/envs/kernel_server.py", + str(self.container_sqsh_path.resolve()), + "/bin/bash", "-lc", bash, + ] + + async with CONTAINER_LAUNCH_SEM: + launch_t0 = time.perf_counter() + # Redirect container output to a log file instead of + # subprocess.PIPE to avoid pipe-buffer deadlock (the 64KB + # OS pipe buffer fills up when the kernel server produces + # verbose DEBUG / uvicorn access logs, blocking the + # container process on write() and freezing the kernel). + log_dir = work_dir / ".container_logs" + log_dir.mkdir(exist_ok=True) + self._container_log_path = log_dir / "container.log" + self._container_log_file = open(self._container_log_path, "w") # noqa: SIM115 + self._enroot_proc = subprocess.Popen( + cmd, text=True, start_new_session=True, + stdout=self._container_log_file, stderr=subprocess.STDOUT, + ) + logger.log( + _CONTAINER_LOG_LEVEL, + "[%s] Container launch attempt #%d started (work_dir=%s, token=%s)", + self._proc_label(), attempt, work_dir, startup_token[:8], + ) + + # Create HTTP client + self._http_client = httpx.AsyncClient( + base_url=f"http://localhost:{self._container_port}", + timeout=httpx.Timeout(self.execution_timeout + 10, connect=30.0), + ) + + # Wait for health check + try: + await self._wait_for_health(expected_startup_token=startup_token) + launch_ms = (time.perf_counter() - launch_t0) * 1000.0 + logger.log( + _CONTAINER_LOG_LEVEL, + "[%s] Container online after %.1fms (attempt #%d)", + self._proc_label(), launch_ms, attempt, + ) + online = True + except Exception as e: + launch_ms = (time.perf_counter() - launch_t0) * 1000.0 + self._log_container_failure(attempt, launch_ms, e) + await self._cleanup_failed_startup() + await asyncio.sleep(10) + + def _log_container_failure(self, attempt: int, launch_ms: float, error: Exception) -> None: + """Log detailed diagnostics when a container fails to start.""" + label = self._proc_label() + proc = self._enroot_proc + + diag_parts = [ + f"attempt=#{attempt}", + f"elapsed={launch_ms:.0f}ms", + f"error={error!r}", + ] + + if proc is not None: + rc = proc.poll() + diag_parts.append(f"process_alive={rc is None}") + if rc is not None: + diag_parts.append(f"returncode={rc}") + + # Read tail of the container log file for diagnostics + log_tail = self._read_container_log_tail() + if log_tail: + diag_parts.append(f"log_tail={log_tail!r}") + + logger.warning("[%s] Container startup FAILED: %s", label, ", ".join(diag_parts)) + + def _read_container_log_tail(self, max_chars: int = 2000) -> str: + """Read the tail of the container log file for diagnostics.""" + if self._container_log_path is None or not self._container_log_path.exists(): + return "" + try: + # Flush so any buffered output is written to disk + if self._container_log_file and not self._container_log_file.closed: + self._container_log_file.flush() + text = self._container_log_path.read_text() + return text[-max_chars:] if len(text) > max_chars else text + except Exception: + return "" + + def _close_container_log(self) -> None: + """Close the container log file handle.""" + if self._container_log_file is not None: + try: + self._container_log_file.close() + except Exception: + pass + self._container_log_file = None + + async def _cleanup_failed_startup(self) -> None: + """Best-effort cleanup for failed startup attempts before retrying.""" + label = self._proc_label() + + if self._container_port is not None: + async with used_ports_lock: + _USED_PORTS.discard(self._container_port) + self._container_port = None + + if self._http_client is not None: + try: + await self._http_client.aclose() + except Exception: + pass + self._http_client = None + + if self._enroot_proc is not None: + _kill_process_group(self._enroot_proc, label=label) + self._enroot_proc = None + + self._close_container_log() + + async def _wait_for_health(self, expected_startup_token: str | None = None) -> None: + """Wait for the kernel server to become healthy.""" + start_time = time.perf_counter() + poll_count = 0 + last_status = "no_attempt" + # Use a short per-request timeout for health checks so that a single + # poll can never block longer than a few seconds. Without this, the + # client's default timeout (execution_timeout + 10 ≈ 190s) means one + # health poll can exceed the entire KERNEL_SERVER_STARTUP_TIMEOUT, + # e.g. when a port collision causes us to connect to the wrong server. + health_timeout = httpx.Timeout(5.0, connect=3.0) + + while time.perf_counter() - start_time < cfg.KERNEL_SERVER_STARTUP_TIMEOUT: + poll_count += 1 + elapsed = time.perf_counter() - start_time + + # Check if the enroot process has died before we even get an HTTP response + if self._enroot_proc is not None and self._enroot_proc.poll() is not None: + rc = self._enroot_proc.returncode + log_tail = self._read_container_log_tail(1000) + raise RuntimeError( + f"Enroot process exited prematurely with returncode={rc} " + f"after {elapsed:.1f}s. log: {log_tail!r}" + ) + + try: + assert self._http_client is not None + response = await self._http_client.get("/health", timeout=health_timeout) + if response.status_code == 200: + if expected_startup_token is None: + logger.log( + _CONTAINER_LOG_LEVEL, + "[%s] Kernel server healthy after %.1fs (%d polls)", + self._proc_label(), elapsed, poll_count, + ) + return + payload = response.json() + if payload.get("startup_token") == expected_startup_token: + logger.log( + _CONTAINER_LOG_LEVEL, + "[%s] Kernel server healthy (token matched) after %.1fs (%d polls)", + self._proc_label(), elapsed, poll_count, + ) + return + else: + last_status = f"token_mismatch(got={payload.get('startup_token', '?')[:8]})" + else: + last_status = f"http_{response.status_code}" + except httpx.ConnectError: + last_status = "connect_error" + except httpx.ReadError: + last_status = "read_error" + except httpx.TimeoutException: + last_status = "timeout" + except httpx.RemoteProtocolError: + last_status = "protocol_error" + + # Log progress every 5s + if poll_count % 10 == 0: + proc_alive = self._enroot_proc.poll() is None if self._enroot_proc else False + logger.log( + _CONTAINER_LOG_LEVEL, + "[%s] Health poll #%d at %.1fs: last_status=%s, process_alive=%s", + self._proc_label(), poll_count, elapsed, last_status, proc_alive, + ) + + await asyncio.sleep(0.5) + + total_elapsed = time.perf_counter() - start_time + raise TimeoutError( + f"Kernel server did not become healthy within {cfg.KERNEL_SERVER_STARTUP_TIMEOUT}s " + f"({poll_count} polls, last_status={last_status}, elapsed={total_elapsed:.1f}s)" + ) + + # @tenacity.retry( + # retry=tenacity.retry_if_exception_type((httpx.ConnectError, httpx.ReadError)), + # stop=tenacity.stop_after_attempt(3), + # wait=tenacity.wait_exponential(multiplier=1, min=1, max=10), + # reraise=True, + # ) + async def _execute_via_http(self, code: str, timeout: float | None = None) -> ExecutionResult: # noqa: ASYNC109 + """Execute code via HTTP to the containerized kernel server. + + Handles httpx.TimeoutException (including ReadTimeout) by converting to + an error ExecutionResult. This happens when the kernel server's internal + asyncio.timeout doesn't cancel the ZMQ recv promptly, causing the HTTP + read to time out before the kernel server responds. + """ + assert self._http_client is not None + + try: + response = await self._http_client.post( + "/execute", + json={"code": code, "timeout": timeout}, + ) + response.raise_for_status() + except httpx.TimeoutException as e: + effective_timeout = timeout if timeout is not None else self.execution_timeout + logger.warning( + "[%s] HTTP %s during /execute (requested kernel timeout=%.1fs): %s", + self._proc_label(), type(e).__name__, effective_timeout, e, + ) + timeout_output = nbformat.v4.new_output( + output_type="error", + ename="TimeoutError", + evalue=f"Code execution timed out after {effective_timeout}s (HTTP layer)", + traceback=[f"TimeoutError: Code execution timed out after {effective_timeout}s (HTTP layer)"], + ) + return ExecutionResult( + notebook_outputs=[timeout_output], + error_occurred=True, + execution_time=effective_timeout, + ) + + data = response.json() + + # Convert serialized outputs back to NotebookNode + notebook_outputs = [nbformat.from_dict(o) for o in data["notebook_outputs"]] + + return ExecutionResult( + notebook_outputs=notebook_outputs, + error_occurred=data["error_occurred"], + execution_time=data.get("execution_time"), + ) + + async def _reset_via_http(self) -> None: + """Reset the kernel via HTTP.""" + assert self._http_client is not None + try: + response = await self._http_client.post("/reset") + response.raise_for_status() + except httpx.TimeoutException as e: + logger.warning( + "[%s] HTTP %s during /reset: %s", + self._proc_label(), type(e).__name__, e, + ) + raise RuntimeError(f"Kernel reset timed out: {e}") from e + + async def close(self): + label = self._proc_label() + logger.log(_CONTAINER_LOG_LEVEL, "[%s] Closing EnrootKernelServer", label) + + if self._container_port is not None: + async with used_ports_lock: + _USED_PORTS.discard(self._container_port) + + if self._http_client is not None: + try: + response = await self._http_client.post("/close") + response.raise_for_status() + except (httpx.ConnectError, httpx.ReadError, httpx.RemoteProtocolError, httpx.TimeoutException): + logger.log(_CONTAINER_LOG_LEVEL, "[%s] Graceful /close request failed (container may already be down)", label) + except Exception: + logger.warning("[%s] Unexpected error on /close request", label, exc_info=True) + try: + await self._http_client.aclose() + except Exception: + pass + self._http_client = None + + if self._enroot_proc is not None: + _kill_process_group(self._enroot_proc, label=label) + self._enroot_proc = None + + self._close_container_log() + logger.log(_CONTAINER_LOG_LEVEL, "[%s] EnrootKernelServer closed", label) class InterpreterEnvState: @@ -90,6 +608,9 @@ def __init__( use_host_env_vars: bool = False, extra_envs: dict[str, str] | None = None, use_docker: bool = cfg.USE_DOCKER, + use_enroot: bool = False, + use_ray: bool = True, + container_sqsh_path: Path | None = None, save_dir: Path | None = None, ): self.work_dir = work_dir @@ -102,11 +623,14 @@ def __init__( self.actions: list[str] = [] self.done = False self.use_docker = use_docker + self.use_enroot = use_enroot + self.container_sqsh_path = container_sqsh_path self.save_dir = save_dir + self.use_ray = use_ray # Local interpreter (only used when use_docker=False) self.interpreter: Interpreter | None = None - if not use_docker: + if not use_docker and not use_enroot: self.interpreter = Interpreter( work_dir=work_dir, language=language, @@ -115,11 +639,14 @@ def __init__( extra_envs=extra_envs, ) - # Docker container state (only used when use_docker=True) + # Docker/Enroot container state (only used when use_docker=True or use_enroot=True) self._docker_client: aiodocker.Docker | None = None self._container: DockerContainer | None = None + self._enroot_proc = None self._container_port: int | None = None self._http_client: httpx.AsyncClient | None = None + self._container_log_path: Path | None = None + self._container_log_file: Any | None = None # Initialize notebook structure for state tracking self.nb: NotebookNode = nbformat.v4.new_notebook() @@ -133,13 +660,157 @@ def __init__( async def start(self): """Start the interpreter (local or Docker-based).""" - if self.use_docker: - await self._start_container() + if self.use_enroot: + assert self.container_sqsh_path is not None, "container_sqsh_path must be set in config to use enroot" + if self.use_ray: + await self._start_ray_enroot_container() + else: + await self._start_enroot_container() + elif self.use_docker: + await self._start_docker_container() else: assert self.interpreter is not None await self.interpreter.start() - async def _start_container(self) -> None: + async def _start_ray_enroot_container(self) -> None: + self.kernel_container = EnrootKernelServer.remote(self.container_sqsh_path, self.execution_timeout) + init_ref = self.kernel_container.initialize.remote(self.work_dir, self.language) + await init_ref + + def _enroot_label(self) -> str: + port = self._container_port or "?" + pid = self._enroot_proc.pid if self._enroot_proc else "?" + return f"enroot-state(port={port}, pid={pid})" + + async def _start_enroot_container(self) -> None: + _prep_workspace_dir(self.work_dir) + + online = False + attempt = 0 + while not online: + attempt += 1 + self._container_port = await get_free_port() + startup_token = str(uuid.uuid4()) + + bash = dedent(f"""\ + set -euo pipefail + cd /data_workspace + + export PYTHONPATH="/data_workspace/pydeps:${{PYTHONPATH}}" + export PIP_CONFIG_FILE=/data_workspace/pip.conf + exec /app/kernel_env/bin/python /envs/kernel_server.py \\ + --work_dir /data_workspace \\ + --language {self.language.value} \\ + --port {self._container_port} \\ + --startup-token {startup_token} + """).strip() + + env_dir = Path(__file__).parent + kernel_server_path = env_dir / "kernel_server.py" + assert kernel_server_path.is_file(), f"kernel server must be a valid path, found {kernel_server_path}" + + cmd = [ + "env", "-i", "PATH=/usr/sbin:/usr/bin:/sbin:/bin", 'HOME="$HOME"', 'USER="$USER"', + "enroot", "start", + "--mount", f"{self.work_dir}:/data_workspace", + "--mount", f"{kernel_server_path.resolve()}:/envs/kernel_server.py", + str(self.container_sqsh_path.resolve()), + "/bin/bash", "-lc", bash, + ] + + async with CONTAINER_LAUNCH_SEM: + launch_t0 = time.perf_counter() + # Redirect container output to a log file (see EnrootKernelServer + # for detailed rationale on why we avoid subprocess.PIPE here). + log_dir = self.work_dir / ".container_logs" + log_dir.mkdir(exist_ok=True) + self._container_log_path = log_dir / "container.log" + self._container_log_file = open(self._container_log_path, "w") # noqa: SIM115 + self._enroot_proc = subprocess.Popen( + cmd, text=True, start_new_session=True, + stdout=self._container_log_file, stderr=subprocess.STDOUT, + ) + logger.log( + _CONTAINER_LOG_LEVEL, + "[%s] Container launch attempt #%d (work_dir=%s, token=%s)", + self._enroot_label(), attempt, self.work_dir, startup_token[:8], + ) + + # Create HTTP client + self._http_client = httpx.AsyncClient( + base_url=f"http://localhost:{self._container_port}", + timeout=httpx.Timeout(self.execution_timeout + 10, connect=30.0), + ) + + # Wait for health check + try: + await self._wait_for_health(expected_startup_token=startup_token) + launch_ms = (time.perf_counter() - launch_t0) * 1000.0 + logger.log( + _CONTAINER_LOG_LEVEL, + "[%s] Container online after %.1fms (attempt #%d)", + self._enroot_label(), launch_ms, attempt, + ) + online = True + except Exception as e: + launch_ms = (time.perf_counter() - launch_t0) * 1000.0 + diag = [f"attempt=#{attempt}", f"elapsed={launch_ms:.0f}ms", f"error={e!r}"] + if self._enroot_proc is not None: + rc = self._enroot_proc.poll() + diag.append(f"process_alive={rc is None}") + if rc is not None: + diag.append(f"returncode={rc}") + log_tail = self._read_container_log_tail() + if log_tail: + diag.append(f"log_tail={log_tail!r}") + logger.warning("[%s] Container startup FAILED: %s", self._enroot_label(), ", ".join(diag)) + await self._cleanup_failed_startup() + await asyncio.sleep(10) + + def _read_container_log_tail(self, max_chars: int = 2000) -> str: + """Read the tail of the container log file for diagnostics.""" + if self._container_log_path is None or not self._container_log_path.exists(): + return "" + try: + if self._container_log_file and not self._container_log_file.closed: + self._container_log_file.flush() + text = self._container_log_path.read_text() + return text[-max_chars:] if len(text) > max_chars else text + except Exception: + return "" + + def _close_container_log(self) -> None: + """Close the container log file handle.""" + if self._container_log_file is not None: + try: + self._container_log_file.close() + except Exception: + pass + self._container_log_file = None + + async def _cleanup_failed_startup(self) -> None: + """Best-effort cleanup for failed startup attempts before retrying.""" + label = self._enroot_label() + + if self._container_port is not None: + async with used_ports_lock: + _USED_PORTS.discard(self._container_port) + self._container_port = None + + if self._http_client is not None: + try: + await self._http_client.aclose() + except Exception: + pass + self._http_client = None + + if self._enroot_proc is not None: + _kill_process_group(self._enroot_proc, label=label) + self._enroot_proc = None + + self._close_container_log() + + async def _start_docker_container(self) -> None: """Start a Docker container with the kernel server.""" self._docker_client = aiodocker.Docker() self._container_port = get_free_port() @@ -150,21 +821,21 @@ async def _start_container(self) -> None: "/app/kernel_env/bin/python", "/envs/kernel_server.py", "--work_dir", - "/workspace", + "/data_workspace", "--language", self.language.value, ], "HostConfig": { - "Binds": [f"{self.work_dir}:/workspace"], + "Binds": [f"{self.work_dir}:/data_workspace"], "PortBindings": {f"{cfg.KERNEL_SERVER_PORT}/tcp": [{"HostPort": str(self._container_port)}]}, }, - "WorkingDir": "/workspace", + "WorkingDir": "/data_workspace", "Tty": True, "ExposedPorts": {f"{cfg.KERNEL_SERVER_PORT}/tcp": {}}, } self._container = await self._docker_client.containers.run(config=cast(dict[str, Any], docker_config)) - logger.info(f"Started container on port {self._container_port}") + logger.log(_CONTAINER_LOG_LEVEL, "Started docker container on port %s", self._container_port) # Create HTTP client self._http_client = httpx.AsyncClient( @@ -175,20 +846,76 @@ async def _start_container(self) -> None: # Wait for health check await self._wait_for_health() - async def _wait_for_health(self) -> None: + async def _wait_for_health(self, expected_startup_token: str | None = None) -> None: """Wait for the kernel server to become healthy.""" start_time = time.perf_counter() + poll_count = 0 + last_status = "no_attempt" + # Short per-request timeout (see EnrootKernelServer._wait_for_health + # for rationale). + health_timeout = httpx.Timeout(5.0, connect=3.0) + while time.perf_counter() - start_time < cfg.KERNEL_SERVER_STARTUP_TIMEOUT: + poll_count += 1 + elapsed = time.perf_counter() - start_time + + # Check if the enroot process has died + if self._enroot_proc is not None and self._enroot_proc.poll() is not None: + rc = self._enroot_proc.returncode + log_tail = self._read_container_log_tail(1000) + raise RuntimeError( + f"Enroot process exited prematurely with returncode={rc} " + f"after {elapsed:.1f}s. log: {log_tail!r}" + ) + try: assert self._http_client is not None - response = await self._http_client.get("/health") + response = await self._http_client.get("/health", timeout=health_timeout) if response.status_code == 200: - logger.info("Kernel server is healthy") - return - except (httpx.ConnectError, httpx.ReadError, httpx.RemoteProtocolError): - pass + if expected_startup_token is None: + logger.log( + _CONTAINER_LOG_LEVEL, + "[%s] Kernel server healthy after %.1fs (%d polls)", + self._enroot_label(), elapsed, poll_count, + ) + return + payload = response.json() + if payload.get("startup_token") == expected_startup_token: + logger.log( + _CONTAINER_LOG_LEVEL, + "[%s] Kernel server healthy (token matched) after %.1fs (%d polls)", + self._enroot_label(), elapsed, poll_count, + ) + return + else: + last_status = f"token_mismatch(got={payload.get('startup_token', '?')[:8]})" + else: + last_status = f"http_{response.status_code}" + except httpx.ConnectError: + last_status = "connect_error" + except httpx.ReadError: + last_status = "read_error" + except httpx.TimeoutException: + last_status = "timeout" + except httpx.RemoteProtocolError: + last_status = "protocol_error" + + # Log progress every 5s + if poll_count % 10 == 0: + proc_alive = self._enroot_proc.poll() is None if self._enroot_proc else False + logger.log( + _CONTAINER_LOG_LEVEL, + "[%s] Health poll #%d at %.1fs: last_status=%s, process_alive=%s", + self._enroot_label(), poll_count, elapsed, last_status, proc_alive, + ) + await asyncio.sleep(0.5) - raise TimeoutError(f"Kernel server did not become healthy within {cfg.KERNEL_SERVER_STARTUP_TIMEOUT}s") + + total_elapsed = time.perf_counter() - start_time + raise TimeoutError( + f"Kernel server did not become healthy within {cfg.KERNEL_SERVER_STARTUP_TIMEOUT}s " + f"({poll_count} polls, last_status={last_status}, elapsed={total_elapsed:.1f}s)" + ) @tenacity.retry( retry=tenacity.retry_if_exception_type((httpx.ConnectError, httpx.ReadError)), @@ -197,14 +924,38 @@ async def _wait_for_health(self) -> None: reraise=True, ) async def _execute_via_http(self, code: str, timeout: float | None = None) -> ExecutionResult: # noqa: ASYNC109 - """Execute code via HTTP to the containerized kernel server.""" + """Execute code via HTTP to the containerized kernel server. + + Handles httpx.TimeoutException (including ReadTimeout) by converting to + an error ExecutionResult, since the kernel server's asyncio.timeout may + not cancel the ZMQ recv promptly. + """ assert self._http_client is not None - response = await self._http_client.post( - "/execute", - json={"code": code, "timeout": timeout}, - ) - response.raise_for_status() + try: + response = await self._http_client.post( + "/execute", + json={"code": code, "timeout": timeout}, + ) + response.raise_for_status() + except httpx.TimeoutException as e: + effective_timeout = timeout if timeout is not None else self.execution_timeout + logger.warning( + "[%s] HTTP %s during /execute (requested kernel timeout=%.1fs): %s", + self._enroot_label(), type(e).__name__, effective_timeout, e, + ) + timeout_output = nbformat.v4.new_output( + output_type="error", + ename="TimeoutError", + evalue=f"Code execution timed out after {effective_timeout}s (HTTP layer)", + traceback=[f"TimeoutError: Code execution timed out after {effective_timeout}s (HTTP layer)"], + ) + return ExecutionResult( + notebook_outputs=[timeout_output], + error_occurred=True, + execution_time=effective_timeout, + ) + data = response.json() # Convert serialized outputs back to NotebookNode @@ -219,8 +970,14 @@ async def _execute_via_http(self, code: str, timeout: float | None = None) -> Ex async def _reset_via_http(self) -> None: """Reset the kernel via HTTP.""" assert self._http_client is not None - response = await self._http_client.post("/reset") - response.raise_for_status() + try: + response = await self._http_client.post("/reset") + response.raise_for_status() + except httpx.TimeoutException as e: + logger.warning( + "HTTP %s during /reset: %s", type(e).__name__, e, + ) + raise RuntimeError(f"Kernel reset timed out: {e}") from e async def close(self): """Save the notebook and close the interpreter or container.""" @@ -231,12 +988,22 @@ async def close(self): self.save_dir.parent.mkdir(parents=True, exist_ok=True) shutil.move(self.work_dir, self.save_dir) - if self.use_docker: + if self.use_ray and self.use_enroot: + if self.kernel_container is not None: + close_ref = self.kernel_container.close.remote() + await close_ref + self.kernel_container = None + + elif self.use_docker or self.use_enroot: if self._container_port is not None: - _USED_PORTS.discard(self._container_port) + async with used_ports_lock: + _USED_PORTS.discard(self._container_port) if self._http_client is not None: - await self._http_client.aclose() + try: + await self._http_client.aclose() + except Exception: + pass self._http_client = None if self._container is not None: @@ -250,6 +1017,13 @@ async def close(self): if self._docker_client is not None: await self._docker_client.close() self._docker_client = None + + if self._enroot_proc is not None: + _kill_process_group(self._enroot_proc, label=self._enroot_label()) + self._enroot_proc = None + + self._close_container_log() + elif self.interpreter is not None: await self.interpreter.close() @@ -342,11 +1116,30 @@ async def execute_and_add_cell( Returns: Tuple of (ExecutionResult, actual_cell_index) """ - if self.use_docker: + execute_start = time.perf_counter() + if self.use_ray and self.use_enroot: + rpc_wait_start = time.perf_counter() + result_ref = self.kernel_container._execute_via_http.remote(code, timeout) + result = await result_ref + _record_perf_metric( + "execute_path_ms.ray_enroot_rpc_wait", + (time.perf_counter() - rpc_wait_start) * 1000.0, + ) + elif self.use_docker or self.use_enroot: + http_exec_start = time.perf_counter() result = await self._execute_via_http(code, timeout) + _record_perf_metric( + "execute_path_ms.http_kernel_execute", + (time.perf_counter() - http_exec_start) * 1000.0, + ) else: assert self.interpreter is not None + local_exec_start = time.perf_counter() result = await self.interpreter.execute_code(code, timeout) + _record_perf_metric( + "execute_path_ms.local_interpreter_execute", + (time.perf_counter() - local_exec_start) * 1000.0, + ) if cell_idx is None or cell_idx >= len(self.nb.cells): actual_idx = self._add_cell(code, result) @@ -354,6 +1147,7 @@ async def execute_and_add_cell( self._update_cell(cell_idx, code, result) actual_idx = cell_idx + _record_perf_metric("execute_and_add_cell_total_ms", (time.perf_counter() - execute_start) * 1000.0) return result, actual_idx @@ -365,6 +1159,8 @@ class InterpreterEnvConfig(BaseModel): execution_config: ExecutionConfig = Field(default_factory=ExecutionConfig) max_steps: int = cfg.AGENT_MAX_STEPS use_docker: bool = cfg.USE_DOCKER + use_enroot: bool = False + container_sqsh_path: Path | None = None normalize_reward: bool = True @@ -403,6 +1199,8 @@ def __init__( self.execution_timeout = self.execution_config.cell_execution_timeout self.max_steps = self.config.max_steps + logger.error(f"EXECUTION TIMEOUT SET TO {self.execution_timeout} SECONDS") + self.input_data = input_data self.output_data: list[dict[str, str | int]] = [] self.logger = logger @@ -412,6 +1210,19 @@ def __init__( self.state: InterpreterEnvState # prompting_config is set during reset() after language resolution self.prompting_config: PromptingConfig + self._perf_sample_step = False + + def _emit_perf_event(self, event: str, **fields: Any) -> None: + if not PERF_METRICS_ENABLED: + return + payload: dict[str, Any] = { + "component": "aviary_hypotest.interpreter_env", + "event": event, + "env_id": getattr(self, "_nemo_env_id", None), + "step_idx": self.step_count, + } + payload.update(fields) + _log_perf(payload) @property def language(self) -> NBLanguage: @@ -455,6 +1266,8 @@ async def reset(self) -> tuple[Messages, list[Tool]]: | self.extra_envs, save_dir=self.save_dir, use_docker=self.config.use_docker, + use_enroot=self.config.use_enroot, + container_sqsh_path=self.config.container_sqsh_path ) await self.state.start() @@ -470,7 +1283,7 @@ async def reset(self) -> tuple[Messages, list[Tool]]: Tool.from_function(self.run_cell), Tool.from_function(self.reset_kernel), Tool.from_function(self.submit_answer), - Tool.from_function(self._filesystem_tool.list_dir), + Tool.from_function(self.list_dir), ] messages.append( @@ -487,17 +1300,22 @@ async def reset(self) -> tuple[Messages, list[Tool]]: messages.append(self.get_env_state_msg()) # Always show initial directory listing (with truncation protection) - messages.append(Message(content=self._filesystem_tool.list_dir())) + messages.append(Message(content=self.list_dir())) return messages, self.tools async def step(self, action: ToolRequestMessage) -> tuple[Messages, float, bool, bool]: """Execute a step in the environment.""" self.step_count += 1 + self._perf_sample_step = PERF_METRICS_ENABLED and random.random() < PERF_DETAIL_SAMPLE_RATE + step_start = time.perf_counter() + exec_tool_calls_start = time.perf_counter() obs = cast( Messages, await self.exec_tool_calls(action, concurrency=False, handle_tool_exc=True), ) + exec_tool_calls_ms = (time.perf_counter() - exec_tool_calls_start) * 1000.0 + _record_perf_metric("exec_tool_calls_ms", exec_tool_calls_ms) obs = [*obs] if self.include_env_state_msg: @@ -512,6 +1330,20 @@ async def step(self, action: ToolRequestMessage) -> tuple[Messages, float, bool, self.state.actions.append(str(action)) reward = self.state.score if self.state.done else 0.0 + step_total_ms = (time.perf_counter() - step_start) * 1000.0 + _record_perf_metric("env_step_ms", step_total_ms) + should_log_detail = self._perf_sample_step or step_total_ms >= PERF_SLOW_STEP_MS + if should_log_detail: + obs_payload_chars = sum(len(str(getattr(m, "content", ""))) for m in obs) + self._emit_perf_event( + "env_step_end", + duration_ms=round(step_total_ms, 3), + exec_tool_calls_ms=round(exec_tool_calls_ms, 3), + action_tool_calls_count=len(action.tool_calls), + obs_messages_count=len(obs), + obs_payload_chars=obs_payload_chars, + slow=step_total_ms >= PERF_SLOW_STEP_MS, + ) return obs, reward, self.state.done, False # ========== Tools ========== @@ -526,15 +1358,9 @@ async def run_cell( This method allows running code in a new cell (append) or re-running an existing cell with updated code. - Usage Examples: - run_cell("print('Hello, world!')") # Run code in new cell - run_cell("print('Hello, world!')", idx=0) # Run code in existing cell at index 0 - Error Recovery: When a cell fails with an error, you MUST fix it by calling run_cell - with the corrected code and the SAME idx as the failed cell: - - run_cell("corrected_code", idx=3) # Fix error in Cell #3 + with the corrected code and the SAME idx as the failed cell. The cell number is shown in the output prefix (e.g., "[Cell #3]"). Do NOT create a new cell to fix an error - always edit the failed cell. @@ -549,6 +1375,7 @@ async def run_cell( Message with multimodal content if images present, otherwise string. The response includes the cell number (e.g., "[Cell #0] output..."). """ + start = time.perf_counter() remaining_seconds = self.get_remaining_time() if remaining_seconds <= self.execution_config.force_submit_threshold: @@ -575,9 +1402,12 @@ async def run_cell( cell_idx = None # Execute code and update notebook atomically + exec_start = time.perf_counter() result, actual_cell_idx = await self.state.execute_and_add_cell( code, cell_idx=cell_idx, timeout=effective_timeout ) + execute_and_add_cell_ms = (time.perf_counter() - exec_start) * 1000.0 + _record_perf_metric("execute_and_add_cell_ms", execute_and_add_cell_ms) # Build response with cell number cell_info = f"[Cell #{actual_cell_idx}] " @@ -592,14 +1422,27 @@ async def run_cell( images=cast(list[np.ndarray | str], image_urls), ) - return cell_info + result.get_truncated_text() + result = cell_info + result.get_truncated_text() + duration_ms = (time.perf_counter() - start) * 1000.0 + _record_perf_metric("tool_call_ms.run_cell", duration_ms) + if self._perf_sample_step: + self._emit_perf_event( + "tool_call_end", + tool_name="run_cell", + duration_ms=round(duration_ms, 3), + ) + return result async def reset_kernel(self) -> str: """Reset the kernel to a clean state. This clears all variables and execution state. """ - if self.state.use_docker: + start = time.perf_counter() + if self.state.use_ray: + reset_ref = self.state.kernel_container._reset_via_http.remote() + await reset_ref + elif self.state.use_docker or self.state.use_enroot: await self.state._reset_via_http() else: assert self.state.interpreter is not None @@ -611,6 +1454,14 @@ async def reset_kernel(self) -> str: self.state.notebook_runtime_errors = [] self.state._execution_count = 0 + duration_ms = (time.perf_counter() - start) * 1000.0 + _record_perf_metric("tool_call_ms.reset_kernel", duration_ms) + if self._perf_sample_step: + self._emit_perf_event( + "tool_call_end", + tool_name="reset_kernel", + duration_ms=round(duration_ms, 3), + ) return "Kernel reset successfully." @property @@ -674,7 +1525,10 @@ async def submit_answer(self, answer: str) -> str: Args: answer: Your final response to the research question """ + start = time.perf_counter() if self.state.done: + duration_ms = (time.perf_counter() - start) * 1000.0 + _record_perf_metric("tool_call_ms.submit_answer", duration_ms) return "Episode already finished." self.state.answer = answer @@ -682,11 +1536,45 @@ async def submit_answer(self, answer: str) -> str: if self.rubric_model is None: self.logger.warning("No rubric_model configured, skipping scoring") + duration_ms = (time.perf_counter() - start) * 1000.0 + _record_perf_metric("tool_call_ms.submit_answer", duration_ms) return answer correct = await self._score_solution(answer) + duration_ms = (time.perf_counter() - start) * 1000.0 + _record_perf_metric("tool_call_ms.submit_answer", duration_ms) + if self._perf_sample_step: + self._emit_perf_event( + "tool_call_end", + tool_name="submit_answer", + duration_ms=round(duration_ms, 3), + ) return CORRECT_MSG if correct else INCORRECT_MSG + def list_dir(self) -> str: + """List contents of a directory with truncation protection. + + Recursively lists files in a directory, with built-in protection against + overwhelming the context with too many files. Use this tool instead of + writing code to list directories to avoid context bloat. + + Args: + directory: Directory path to list (default: current working directory) + max_files: Maximum number of files to display (default: 20) + show_hidden: Whether to show hidden files starting with '.' (default: False) + """ + start = time.perf_counter() + result = self._filesystem_tool.list_dir() + duration_ms = (time.perf_counter() - start) * 1000.0 + _record_perf_metric("tool_call_ms.list_dir", duration_ms) + if self._perf_sample_step: + self._emit_perf_event( + "tool_call_end", + tool_name="list_dir", + duration_ms=round(duration_ms, 3), + ) + return result + # ========== Time Management ========== def get_remaining_time(self) -> int: diff --git a/src/hypotest/env/kernel_server.py b/src/hypotest/env/kernel_server.py index c2c1f2b..b62f4e3 100644 --- a/src/hypotest/env/kernel_server.py +++ b/src/hypotest/env/kernel_server.py @@ -9,12 +9,17 @@ from __future__ import annotations +import uuid import argparse import asyncio +import json import logging +import os +import random import time from enum import StrEnum, auto from pathlib import Path +from queue import Empty from typing import Any, assert_never import nbformat @@ -25,7 +30,19 @@ from nbformat import NotebookNode from pydantic import BaseModel + +class _DeadlineExceeded(Exception): + """Raised when a cooperative deadline check expires.""" + logger = logging.getLogger(__name__) +PERF_METRICS_ENABLED = os.getenv("NEMO_GYM_PERF_METRICS", "0") == "1" +PERF_DETAIL_SAMPLE_RATE = float(os.getenv("NEMO_GYM_PERF_DETAIL_SAMPLE_RATE", "0.2")) +PERF_SLOW_KERNEL_EXEC_MS = float(os.getenv("NEMO_GYM_PERF_SLOW_KERNEL_EXEC_MS", "5000")) +PERF_LOG_LEVEL = getattr(logging, os.getenv("NEMO_GYM_PERF_LOG_LEVEL", "WARNING").upper(), logging.WARNING) + + +def _log_perf(payload: dict[str, Any]) -> None: + logger.log(PERF_LOG_LEVEL, json.dumps(payload)) # ============================================================================= @@ -152,6 +169,14 @@ class ResetResponse(BaseModel): success: bool +class HealthResponse(BaseModel): + """Response model for /health endpoint.""" + + status: str + startup_token: str + kernel_ready: bool + + class KernelServer: """Manages a persistent Jupyter kernel and exposes it via HTTP.""" @@ -160,10 +185,12 @@ def __init__( work_dir: Path, language: NBLanguage, default_timeout: float = 600, + startup_token: str = "", ): self.work_dir = work_dir self.language = language self.default_timeout = default_timeout + self.startup_token = startup_token self._kernel_manager: AsyncKernelManager | None = None self._client: AsyncKernelClient | None = None @@ -175,7 +202,13 @@ async def start(self) -> None: return kernel_name = self.language.make_kernelspec()["name"] - self._kernel_manager = AsyncKernelManager(kernel_name=kernel_name) + kernel_runtime_dir = self.work_dir / ".jupyter_runtime" + kernel_runtime_dir.mkdir(exist_ok=True) + + conn_uuid = uuid.uuid4() + kernel_connect_file = (kernel_runtime_dir / f"conn_{conn_uuid}.json").resolve() + + self._kernel_manager = AsyncKernelManager(kernel_name=kernel_name, transport='ipc', connection_file=str(kernel_connect_file)) await self._kernel_manager.start_kernel(cwd=str(self.work_dir)) self._client = self._kernel_manager.client() @@ -195,11 +228,15 @@ async def execute(self, code: str, timeout: float | None = None) -> ExecuteRespo effective_timeout = timeout if timeout is not None else self.default_timeout start_time = time.perf_counter() + should_sample = PERF_METRICS_ENABLED and random.random() < PERF_DETAIL_SAMPLE_RATE + timed_out = False try: - async with asyncio.timeout(effective_timeout): - result = await self._execute_code(code) - except TimeoutError: + result = await self._execute_code( + code, deadline=start_time + effective_timeout, log_detail=should_sample, + ) + except _DeadlineExceeded: + timed_out = True timeout_output = MessageType.ERROR.to_notebook_output({ "ename": "TimeoutError", "evalue": f"Code execution timed out after {effective_timeout} seconds", @@ -222,26 +259,60 @@ async def execute(self, code: str, timeout: float | None = None) -> ExecuteRespo execution_time=time.perf_counter() - start_time, ) + total_ms = (time.perf_counter() - start_time) * 1000.0 + if PERF_METRICS_ENABLED and (should_sample or total_ms >= PERF_SLOW_KERNEL_EXEC_MS): + _log_perf( + { + "component": "aviary_hypotest.kernel_server", + "event": "kernel_execute_end", + "duration_ms": round(total_ms, 3), + "timed_out": timed_out, + "error_occurred": result.error_occurred, + "notebook_outputs_count": len(result.notebook_outputs), + "slow": total_ms >= PERF_SLOW_KERNEL_EXEC_MS, + } + ) + return result - async def _execute_code(self, code: str) -> ExecuteResponse: - """Internal method to execute code and collect outputs.""" + async def _execute_code(self, code: str, deadline: float, log_detail: bool = False) -> ExecuteResponse: + """Internal method to execute code and collect outputs. + + Uses cooperative deadline checking instead of asyncio.timeout, because + ZMQ socket operations may not respond to asyncio cancellation promptly. + Each get_iopub_msg call uses a short poll timeout so we can check the + deadline between messages. + """ if not self._client: raise RuntimeError("Kernel client not initialized") + # How long each ZMQ poll waits before we re-check the deadline. + # Shorter = more responsive timeout, slightly more overhead. + _POLL_INTERVAL_S = 2.0 + start_time = time.perf_counter() msg_id = self._client.execute(code) - logger.debug(f"Executing code with message ID: {msg_id}") notebook_outputs: list[dict[str, Any]] = [] error_occurred = False + iopub_msgs_count = 0 while True: - msg = await self._client.get_iopub_msg() - logger.debug(f"Received message type: {msg['msg_type']}") + # Check deadline before each poll + if time.perf_counter() >= deadline: + raise _DeadlineExceeded + + # Use a bounded poll so we never block longer than _POLL_INTERVAL_S. + # get_iopub_msg(timeout=T) raises queue.Empty if no message arrives + # within T seconds. + try: + msg = await self._client.get_iopub_msg(timeout=_POLL_INTERVAL_S) + except Empty: + continue if msg["parent_header"].get("msg_id") != msg_id: continue + iopub_msgs_count += 1 msg_type = MessageType.from_string(msg["msg_type"]) if msg_type is None: @@ -259,6 +330,17 @@ async def _execute_code(self, code: str) -> ExecuteResponse: error_occurred = True execution_time = time.perf_counter() - start_time + if PERF_METRICS_ENABLED and log_detail: + _log_perf( + { + "component": "aviary_hypotest.kernel_server", + "event": "kernel_iopub_summary", + "duration_ms": round(execution_time * 1000.0, 3), + "kernel_iopub_msgs_count": iopub_msgs_count, + "notebook_outputs_count": len(notebook_outputs), + "error_occurred": error_occurred, + } + ) return ExecuteResponse( notebook_outputs=notebook_outputs, @@ -287,6 +369,7 @@ async def close(self) -> None: self._client = None await self._kernel_manager.shutdown_kernel(now=True) + await self._kernel_manager.cleanup_resources(restart=False) self._is_ready = False logger.info("Kernel shutdown complete") @@ -304,8 +387,8 @@ async def reset() -> ResetResponse: return await server.reset() @app.get("/health") - async def health() -> str: - return "OK" + async def health() -> HealthResponse: + return HealthResponse(status="OK", startup_token=server.startup_token, kernel_ready=server._is_ready) @app.post("/close") async def close() -> dict[str, bool]: @@ -315,9 +398,9 @@ async def close() -> dict[str, bool]: return app -async def run_server(work_dir: Path, language: NBLanguage, port: int = 8000) -> None: +async def run_server(work_dir: Path, language: NBLanguage, port: int = 8000, startup_token: str = "") -> None: """Start the kernel server.""" - server = KernelServer(work_dir, language) + server = KernelServer(work_dir, language, startup_token=startup_token) await server.start() app = create_app(server) @@ -337,7 +420,8 @@ async def run_server(work_dir: Path, language: NBLanguage, port: int = 8000) -> parser.add_argument("--work_dir", type=Path, default=Path("/workspace")) parser.add_argument("--language", type=str, default="python", choices=["python", "r"]) parser.add_argument("--port", type=int, default=8000) + parser.add_argument("--startup-token", type=str, default="") args = parser.parse_args() language = NBLanguage.PYTHON if args.language == "python" else NBLanguage.R - asyncio.run(run_server(args.work_dir, language, args.port)) + asyncio.run(run_server(args.work_dir, language, args.port, args.startup_token)) diff --git a/src/hypotest/env/tools/filesystem.py b/src/hypotest/env/tools/filesystem.py index 8e57805..642177d 100644 --- a/src/hypotest/env/tools/filesystem.py +++ b/src/hypotest/env/tools/filesystem.py @@ -23,8 +23,8 @@ import numpy as np from aviary.core import Message, Tool -from hypotest.env.utils import img_utils -from hypotest.env.utils.notebook_utils import view_notebook +from resources_servers.aviary_hypotest.env.utils import img_utils +from resources_servers.aviary_hypotest.env.utils.notebook_utils import view_notebook # ============================================================================= # CONSTANTS @@ -343,25 +343,11 @@ def list_dir_tool( overwhelming the context with too many files. Use this tool instead of writing code to list directories to avoid context bloat. - Usage Examples: - list_dir() # List working directory - list_dir("data/") # List specific folder - list_dir(max_files=50) # Show more files - list_dir(show_hidden=True) # Include hidden files - Args: directory: Directory path to list (default: current working directory) max_files: Maximum number of files to display (default: 20) show_hidden: Whether to show hidden files starting with '.' (default: False) """ - # LLM can sometimes generate string instead of int; fallback to default when typecast fails. - try: - max_files = int(max_files) - except ValueError: - max_files = 20 - - show_hidden = bool(show_hidden) - try: dir_path = Path(directory).resolve() if not dir_path.exists() or not dir_path.is_dir(): diff --git a/src/hypotest/env/utils/__init__.py b/src/hypotest/env/utils/__init__.py index d1f3c1a..b7451c1 100644 --- a/src/hypotest/env/utils/__init__.py +++ b/src/hypotest/env/utils/__init__.py @@ -19,7 +19,7 @@ "view_notebook", ] -from hypotest.env.kernel_server import NBLanguage +from resources_servers.aviary_hypotest.env.kernel_server import NBLanguage from .core import extract_code_from_markdown, extract_xml_content from .img_utils import ( diff --git a/src/hypotest/env/utils/notebook_utils.py b/src/hypotest/env/utils/notebook_utils.py index e727aab..d2ec722 100644 --- a/src/hypotest/env/utils/notebook_utils.py +++ b/src/hypotest/env/utils/notebook_utils.py @@ -8,7 +8,7 @@ import nbformat from aiodocker.containers import DockerContainer -from hypotest.env import config as cfg +from resources_servers.aviary_hypotest.env import config as cfg from .img_utils import encode_image_to_base64 From 3ab76ffe0f8ff85cfe4dd04ea15c7893a2b57b18 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Tue, 10 Mar 2026 13:39:35 -0700 Subject: [PATCH 02/31] adjusted dataset config to support new enroot fields --- src/hypotest/dataset_server.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/src/hypotest/dataset_server.py b/src/hypotest/dataset_server.py index b6e323a..8401b32 100644 --- a/src/hypotest/dataset_server.py +++ b/src/hypotest/dataset_server.py @@ -19,7 +19,7 @@ from hypotest.env.kernel_server import NBLanguage -class DatasetConfig(BaseModel): +class HypotestDatasetConfig(BaseModel): model_config = ConfigDict(extra="forbid") problem_jsonl: FilePath @@ -30,11 +30,15 @@ class DatasetConfig(BaseModel): ) work_dir: Path | None = None - use_docker: bool = True + use_docker: bool = False + use_enroot: bool = True + container_sqsh_path: str | None = None force_python: bool = True normalize_reward: bool = True save_dir: Path | None = None + execution_config: dict[str, Any] + @model_validator(mode="after") def make_dirs(self) -> Self: for d in (self.work_dir, self.save_dir): @@ -43,8 +47,8 @@ def make_dirs(self) -> Self: return self -class Dataset(TaskDataset[InterpreterEnv]): - def __init__(self, config: DatasetConfig): +class HypotestDataset(TaskDataset[InterpreterEnv]): + def __init__(self, config: HypotestDatasetConfig): self.config = config self.problems = [ @@ -71,7 +75,7 @@ def get_new_env_by_idx(self, idx: int) -> InterpreterEnv: language = ( NBLanguage.PYTHON if self.config.force_python - else NBLanguage.from_string(cast(str, problem.metadata["nb_primary_language"]).upper()) + else NBLanguage.from_string(cast(str, problem.nb_primary_language).upper()) ) return InterpreterEnv( From facd076845651476d4b8a06f6bcce7838727a348 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Tue, 10 Mar 2026 13:41:35 -0700 Subject: [PATCH 03/31] added use ray field to config --- src/hypotest/dataset_server.py | 1 + src/hypotest/env/interpreter_env.py | 1 + 2 files changed, 2 insertions(+) diff --git a/src/hypotest/dataset_server.py b/src/hypotest/dataset_server.py index 8401b32..8ac9585 100644 --- a/src/hypotest/dataset_server.py +++ b/src/hypotest/dataset_server.py @@ -30,6 +30,7 @@ class HypotestDatasetConfig(BaseModel): ) work_dir: Path | None = None + use_ray: bool = True use_docker: bool = False use_enroot: bool = True container_sqsh_path: str | None = None diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index a2865b8..c347e59 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -1265,6 +1265,7 @@ async def reset(self) -> tuple[Messages, list[Tool]]: } | self.extra_envs, save_dir=self.save_dir, + use_ray=self.config.use_ray, use_docker=self.config.use_docker, use_enroot=self.config.use_enroot, container_sqsh_path=self.config.container_sqsh_path From 688173421c274858de7ae0c219fa3ec264606547 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Wed, 11 Mar 2026 09:54:52 -0700 Subject: [PATCH 04/31] updated interpreter env config to support use_ray --- src/hypotest/env/interpreter_env.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index c347e59..d3a90e8 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -1160,6 +1160,7 @@ class InterpreterEnvConfig(BaseModel): max_steps: int = cfg.AGENT_MAX_STEPS use_docker: bool = cfg.USE_DOCKER use_enroot: bool = False + use_ray: bool = False container_sqsh_path: Path | None = None normalize_reward: bool = True From a9e1f83c70fe5e38ae95f987382d253a7acb0cd8 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Thu, 12 Mar 2026 20:44:50 +0000 Subject: [PATCH 05/31] added code safety, perf fixes on enroot container launch to prevent hangs --- src/hypotest/env/config.py | 2 +- src/hypotest/env/interpreter_env.py | 237 +++++++++++++++++++--------- src/hypotest/env/kernel_server.py | 55 ++++++- 3 files changed, 220 insertions(+), 74 deletions(-) diff --git a/src/hypotest/env/config.py b/src/hypotest/env/config.py index e43cc8e..5d661e8 100644 --- a/src/hypotest/env/config.py +++ b/src/hypotest/env/config.py @@ -15,7 +15,7 @@ # Kernel server settings (for Docker-based execution) KERNEL_SERVER_PORT = 8000 -KERNEL_SERVER_STARTUP_TIMEOUT = 30.0 # seconds to wait for health check +KERNEL_SERVER_STARTUP_TIMEOUT = float(os.getenv("KERNEL_SERVER_STARTUP_TIMEOUT", "120.0")) # seconds to wait for health check MAX_FILES_TO_UPLOAD = int(os.getenv("MAX_FILES_TO_UPLOAD", "100")) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index d3a90e8..d822298 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -63,6 +63,9 @@ # container launch semaphore to limit concurrency CONTAINER_LAUNCH_SEM = asyncio.Semaphore(128) +MAX_CONTAINER_LAUNCH_RETRIES = int(os.getenv("MAX_CONTAINER_LAUNCH_RETRIES", "5")) +_RETRY_BASE_SLEEP = 1.0 +_RETRY_MAX_SLEEP = 16.0 async def get_free_port() -> int: """Get a free port for the kernel server container.""" @@ -84,6 +87,10 @@ async def get_free_port() -> int: _CONTAINER_LOG_LEVEL = logging.WARNING +class _PortCollisionError(Exception): + """Port already in use by another server — retry with a new port.""" + + def _kill_process_group(proc: subprocess.Popen, label: str = "enroot", sigterm_timeout: float = 15) -> None: """Safely terminate a process group, escalating from SIGTERM to SIGKILL. @@ -212,7 +219,6 @@ def _record_perf_metric(metric_name: str, value_ms: float) -> None: _PERF_METRIC_SAMPLES.setdefault(metric_name, []).append(value_ms) _maybe_flush_perf_aggregates() - class ProblemInstance(BaseModel): uuid: UUID hypothesis: str @@ -221,7 +227,7 @@ class ProblemInstance(BaseModel): rubric: str max_score: int = Field(alias="max_points") metadata: dict[str, JsonValue] = Field(default_factory=dict) - nb_primary_language: str + nb_primary_language: str = Field(default=str(NBLanguage.PYTHON)) def _prep_workspace_dir(work_dir: str, workspace_path: str = "/data_workspace") -> None: wd = Path(work_dir) @@ -259,22 +265,32 @@ def _proc_label(self) -> str: return f"enroot(port={port}, pid={pid})" async def initialize(self, work_dir: Path, language: NBLanguage) -> None: - _prep_workspace_dir(work_dir) + startup_token = str(uuid.uuid4()) + node_workdir = f"/tmp/data_workspace.{startup_token.split('-')[0]}" + + _prep_workspace_dir(work_dir, workspace_path=node_workdir) online = False attempt = 0 + last_err: Exception | None = None while not online: attempt += 1 + if attempt > MAX_CONTAINER_LAUNCH_RETRIES: + log_tail = self._read_container_log_tail(500) + raise RuntimeError( + f"Container failed to start after {MAX_CONTAINER_LAUNCH_RETRIES} attempts " + f"(last_error={last_err!r})" + f"{f' log_tail={log_tail!r}' if log_tail else ''}" + ) self._container_port = await get_free_port() - startup_token = str(uuid.uuid4()) - + bash = dedent(f"""\ set -euo pipefail - WORKDIR="$(mktemp -d /dev/shm/data_workspace.XXXXXX)" + WORKDIR="{node_workdir}" trap 'rm -rf "$WORKDIR"' EXIT - mkdir -p /dev/shm/$WORKDIR + mkdir -p $WORKDIR cp -a /data_workspace/. $WORKDIR/ cd $WORKDIR @@ -335,27 +351,31 @@ async def initialize(self, work_dir: Path, language: NBLanguage) -> None: self._proc_label(), attempt, work_dir, startup_token[:8], ) - # Create HTTP client - self._http_client = httpx.AsyncClient( - base_url=f"http://localhost:{self._container_port}", - timeout=httpx.Timeout(self.execution_timeout + 10, connect=30.0), - ) + # Create HTTP client (outside semaphore — no need to hold the + # concurrency slot while waiting for the container to come up) + self._http_client = httpx.AsyncClient( + base_url=f"http://localhost:{self._container_port}", + timeout=httpx.Timeout(self.execution_timeout + 10, connect=30.0), + ) - # Wait for health check - try: - await self._wait_for_health(expected_startup_token=startup_token) - launch_ms = (time.perf_counter() - launch_t0) * 1000.0 - logger.log( - _CONTAINER_LOG_LEVEL, - "[%s] Container online after %.1fms (attempt #%d)", - self._proc_label(), launch_ms, attempt, - ) - online = True - except Exception as e: - launch_ms = (time.perf_counter() - launch_t0) * 1000.0 - self._log_container_failure(attempt, launch_ms, e) - await self._cleanup_failed_startup() - await asyncio.sleep(10) + # Wait for health check + try: + await self._wait_for_health(expected_startup_token=startup_token) + launch_ms = (time.perf_counter() - launch_t0) * 1000.0 + logger.log( + _CONTAINER_LOG_LEVEL, + "[%s] Container online after %.1fms (attempt #%d)", + self._proc_label(), launch_ms, attempt, + ) + online = True + except Exception as e: + last_err = e + launch_ms = (time.perf_counter() - launch_t0) * 1000.0 + self._log_container_failure(attempt, launch_ms, e) + await self._cleanup_failed_startup() + if not isinstance(e, _PortCollisionError): + backoff = min(_RETRY_BASE_SLEEP * 2 ** (attempt - 1), _RETRY_MAX_SLEEP) + await asyncio.sleep(backoff) def _log_container_failure(self, attempt: int, launch_ms: float, error: Exception) -> None: """Log detailed diagnostics when a container fails to start.""" @@ -374,12 +394,15 @@ def _log_container_failure(self, attempt: int, launch_ms: float, error: Exceptio if rc is not None: diag_parts.append(f"returncode={rc}") - # Read tail of the container log file for diagnostics + logger.warning("[%s] Container startup FAILED: %s", label, ", ".join(diag_parts)) + + # Log container output separately so tracebacks are readable log_tail = self._read_container_log_tail() if log_tail: - diag_parts.append(f"log_tail={log_tail!r}") - - logger.warning("[%s] Container startup FAILED: %s", label, ", ".join(diag_parts)) + logger.warning( + "[%s] Container log output (last %d chars):\n%s", + label, len(log_tail), log_tail, + ) def _read_container_log_tail(self, max_chars: int = 2000) -> str: """Read the tail of the container log file for diagnostics.""" @@ -430,6 +453,7 @@ async def _wait_for_health(self, expected_startup_token: str | None = None) -> N start_time = time.perf_counter() poll_count = 0 last_status = "no_attempt" + consecutive_token_mismatches = 0 # Use a short per-request timeout for health checks so that a single # poll can never block longer than a few seconds. Without this, the # client's default timeout (execution_timeout + 10 ≈ 190s) means one @@ -471,16 +495,27 @@ async def _wait_for_health(self, expected_startup_token: str | None = None) -> N return else: last_status = f"token_mismatch(got={payload.get('startup_token', '?')[:8]})" + consecutive_token_mismatches += 1 + if consecutive_token_mismatches >= 3: + raise _PortCollisionError( + f"Port {self._container_port} appears to be owned by another server " + f"({consecutive_token_mismatches} consecutive token mismatches)" + ) else: last_status = f"http_{response.status_code}" + consecutive_token_mismatches = 0 except httpx.ConnectError: last_status = "connect_error" + consecutive_token_mismatches = 0 except httpx.ReadError: last_status = "read_error" + consecutive_token_mismatches = 0 except httpx.TimeoutException: last_status = "timeout" + consecutive_token_mismatches = 0 except httpx.RemoteProtocolError: last_status = "protocol_error" + consecutive_token_mismatches = 0 # Log progress every 5s if poll_count % 10 == 0: @@ -494,9 +529,16 @@ async def _wait_for_health(self, expected_startup_token: str | None = None) -> N await asyncio.sleep(0.5) total_elapsed = time.perf_counter() - start_time + if last_status.startswith("token_mismatch"): + raise _PortCollisionError( + f"Port {self._container_port} health-check timed out with token_mismatch " + f"({poll_count} polls, elapsed={total_elapsed:.1f}s)" + ) + log_tail = self._read_container_log_tail(500) raise TimeoutError( f"Kernel server did not become healthy within {cfg.KERNEL_SERVER_STARTUP_TIMEOUT}s " f"({poll_count} polls, last_status={last_status}, elapsed={total_elapsed:.1f}s)" + f"{f' log_tail={log_tail!r}' if log_tail else ''}" ) # @tenacity.retry( @@ -687,8 +729,16 @@ async def _start_enroot_container(self) -> None: online = False attempt = 0 + last_err: Exception | None = None while not online: attempt += 1 + if attempt > MAX_CONTAINER_LAUNCH_RETRIES: + log_tail = self._read_container_log_tail(500) + raise RuntimeError( + f"Container failed to start after {MAX_CONTAINER_LAUNCH_RETRIES} attempts " + f"(last_error={last_err!r})" + f"{f' log_tail={log_tail!r}' if log_tail else ''}" + ) self._container_port = await get_free_port() startup_token = str(uuid.uuid4()) @@ -736,36 +786,44 @@ async def _start_enroot_container(self) -> None: self._enroot_label(), attempt, self.work_dir, startup_token[:8], ) - # Create HTTP client - self._http_client = httpx.AsyncClient( - base_url=f"http://localhost:{self._container_port}", - timeout=httpx.Timeout(self.execution_timeout + 10, connect=30.0), - ) + # Create HTTP client (outside semaphore — no need to hold the + # concurrency slot while waiting for the container to come up) + self._http_client = httpx.AsyncClient( + base_url=f"http://localhost:{self._container_port}", + timeout=httpx.Timeout(self.execution_timeout + 10, connect=30.0), + ) - # Wait for health check - try: - await self._wait_for_health(expected_startup_token=startup_token) - launch_ms = (time.perf_counter() - launch_t0) * 1000.0 - logger.log( - _CONTAINER_LOG_LEVEL, - "[%s] Container online after %.1fms (attempt #%d)", - self._enroot_label(), launch_ms, attempt, + # Wait for health check + try: + await self._wait_for_health(expected_startup_token=startup_token) + launch_ms = (time.perf_counter() - launch_t0) * 1000.0 + logger.log( + _CONTAINER_LOG_LEVEL, + "[%s] Container online after %.1fms (attempt #%d)", + self._enroot_label(), launch_ms, attempt, + ) + online = True + except Exception as e: + last_err = e + launch_ms = (time.perf_counter() - launch_t0) * 1000.0 + label = self._enroot_label() + diag = [f"attempt=#{attempt}", f"elapsed={launch_ms:.0f}ms", f"error={e!r}"] + if self._enroot_proc is not None: + rc = self._enroot_proc.poll() + diag.append(f"process_alive={rc is None}") + if rc is not None: + diag.append(f"returncode={rc}") + logger.warning("[%s] Container startup FAILED: %s", label, ", ".join(diag)) + log_tail = self._read_container_log_tail() + if log_tail: + logger.warning( + "[%s] Container log output (last %d chars):\n%s", + label, len(log_tail), log_tail, ) - online = True - except Exception as e: - launch_ms = (time.perf_counter() - launch_t0) * 1000.0 - diag = [f"attempt=#{attempt}", f"elapsed={launch_ms:.0f}ms", f"error={e!r}"] - if self._enroot_proc is not None: - rc = self._enroot_proc.poll() - diag.append(f"process_alive={rc is None}") - if rc is not None: - diag.append(f"returncode={rc}") - log_tail = self._read_container_log_tail() - if log_tail: - diag.append(f"log_tail={log_tail!r}") - logger.warning("[%s] Container startup FAILED: %s", self._enroot_label(), ", ".join(diag)) - await self._cleanup_failed_startup() - await asyncio.sleep(10) + await self._cleanup_failed_startup() + if not isinstance(e, _PortCollisionError): + backoff = min(_RETRY_BASE_SLEEP * 2 ** (attempt - 1), _RETRY_MAX_SLEEP) + await asyncio.sleep(backoff) def _read_container_log_tail(self, max_chars: int = 2000) -> str: """Read the tail of the container log file for diagnostics.""" @@ -813,7 +871,8 @@ async def _cleanup_failed_startup(self) -> None: async def _start_docker_container(self) -> None: """Start a Docker container with the kernel server.""" self._docker_client = aiodocker.Docker() - self._container_port = get_free_port() + self._container_port = await get_free_port() + startup_token = str(uuid.uuid4()) docker_config = { "Image": cfg.NB_ENVIRONMENT_DOCKER_IMAGE, @@ -824,6 +883,8 @@ async def _start_docker_container(self) -> None: "/data_workspace", "--language", self.language.value, + "--startup-token", + startup_token, ], "HostConfig": { "Binds": [f"{self.work_dir}:/data_workspace"], @@ -844,13 +905,14 @@ async def _start_docker_container(self) -> None: ) # Wait for health check - await self._wait_for_health() + await self._wait_for_health(expected_startup_token=startup_token) async def _wait_for_health(self, expected_startup_token: str | None = None) -> None: """Wait for the kernel server to become healthy.""" start_time = time.perf_counter() poll_count = 0 last_status = "no_attempt" + consecutive_token_mismatches = 0 # Short per-request timeout (see EnrootKernelServer._wait_for_health # for rationale). health_timeout = httpx.Timeout(5.0, connect=3.0) @@ -889,16 +951,27 @@ async def _wait_for_health(self, expected_startup_token: str | None = None) -> N return else: last_status = f"token_mismatch(got={payload.get('startup_token', '?')[:8]})" + consecutive_token_mismatches += 1 + if consecutive_token_mismatches >= 3: + raise _PortCollisionError( + f"Port {self._container_port} appears to be owned by another server " + f"({consecutive_token_mismatches} consecutive token mismatches)" + ) else: last_status = f"http_{response.status_code}" + consecutive_token_mismatches = 0 except httpx.ConnectError: last_status = "connect_error" + consecutive_token_mismatches = 0 except httpx.ReadError: last_status = "read_error" + consecutive_token_mismatches = 0 except httpx.TimeoutException: last_status = "timeout" + consecutive_token_mismatches = 0 except httpx.RemoteProtocolError: last_status = "protocol_error" + consecutive_token_mismatches = 0 # Log progress every 5s if poll_count % 10 == 0: @@ -912,9 +985,16 @@ async def _wait_for_health(self, expected_startup_token: str | None = None) -> N await asyncio.sleep(0.5) total_elapsed = time.perf_counter() - start_time + if last_status.startswith("token_mismatch"): + raise _PortCollisionError( + f"Port {self._container_port} health-check timed out with token_mismatch " + f"({poll_count} polls, elapsed={total_elapsed:.1f}s)" + ) + log_tail = self._read_container_log_tail(500) raise TimeoutError( f"Kernel server did not become healthy within {cfg.KERNEL_SERVER_STARTUP_TIMEOUT}s " f"({poll_count} polls, last_status={last_status}, elapsed={total_elapsed:.1f}s)" + f"{f' log_tail={log_tail!r}' if log_tail else ''}" ) @tenacity.retry( @@ -1116,6 +1196,29 @@ async def execute_and_add_cell( Returns: Tuple of (ExecutionResult, actual_cell_index) """ + # Safety check: block dangerous code before execution + from .code_safety import check_code_safety + block_reason = check_code_safety(code, self.language) + if block_reason is not None: + logger.warning("Blocked code execution in execute_and_add_cell: %s", block_reason) + error_output = nbformat.v4.new_output( + output_type="error", + ename="SecurityError", + evalue=block_reason, + traceback=[f"SecurityError: {block_reason}"], + ) + result = ExecutionResult( + notebook_outputs=[error_output], + error_occurred=True, + execution_time=0.0, + ) + if cell_idx is None or cell_idx >= len(self.nb.cells): + actual_idx = self._add_cell(code, result) + else: + self._update_cell(cell_idx, code, result) + actual_idx = cell_idx + return result, actual_idx + execute_start = time.perf_counter() if self.use_ray and self.use_enroot: rpc_wait_start = time.perf_counter() @@ -1377,7 +1480,6 @@ async def run_cell( Message with multimodal content if images present, otherwise string. The response includes the cell number (e.g., "[Cell #0] output..."). """ - start = time.perf_counter() remaining_seconds = self.get_remaining_time() if remaining_seconds <= self.execution_config.force_submit_threshold: @@ -1424,16 +1526,7 @@ async def run_cell( images=cast(list[np.ndarray | str], image_urls), ) - result = cell_info + result.get_truncated_text() - duration_ms = (time.perf_counter() - start) * 1000.0 - _record_perf_metric("tool_call_ms.run_cell", duration_ms) - if self._perf_sample_step: - self._emit_perf_event( - "tool_call_end", - tool_name="run_cell", - duration_ms=round(duration_ms, 3), - ) - return result + return cell_info + result.get_truncated_text() async def reset_kernel(self) -> str: """Reset the kernel to a clean state. diff --git a/src/hypotest/env/kernel_server.py b/src/hypotest/env/kernel_server.py index b62f4e3..78303e2 100644 --- a/src/hypotest/env/kernel_server.py +++ b/src/hypotest/env/kernel_server.py @@ -16,6 +16,7 @@ import logging import os import random +import re import time from enum import StrEnum, auto from pathlib import Path @@ -145,6 +146,40 @@ def to_notebook_output(self, content: dict[str, Any]) -> NotebookNode | None: # ============================================================================= +# --------------------------------------------------------------------------- +# Lightweight regex safety check (defense-in-depth, standalone — no hypotest imports) +# --------------------------------------------------------------------------- +_KERNEL_SAFETY_PATTERNS: list[tuple[re.Pattern[str], str]] = [ + # Process killing + (re.compile(r"\bos\s*\.\s*kill\s*\("), "restricted function"), + (re.compile(r"\bos\s*\.\s*killpg\s*\("), "restricted function"), + (re.compile(r"\bos\s*\.\s*system\s*\("), "restricted function"), + (re.compile(r"\bos\s*\.\s*popen\s*\("), "restricted function"), + (re.compile(r"\bos\s*\.\s*fork\s*\("), "restricted function"), + (re.compile(r"\bos\s*\.\s*exec\w*\s*\("), "restricted function"), + (re.compile(r"\bsubprocess\s*\.\s*(run|Popen|call|check_call|check_output)\s*\("), "restricted function"), + # Blocked modules + (re.compile(r"\bimport\s+ctypes\b"), "restricted module"), + (re.compile(r"\bimport\s+signal\b"), "restricted module"), + (re.compile(r"\bfrom\s+ctypes\b"), "restricted module"), + (re.compile(r"\bfrom\s+signal\b"), "restricted module"), + # Shell commands + (re.compile(r"\bkillall\b"), "restricted shell command"), + (re.compile(r"\bpkill\b"), "restricted shell command"), +] + + +def _kernel_check_code_safety(code: str) -> str | None: + """Lightweight regex safety check for the kernel server. + + Returns None if safe, or a message if blocked. + """ + for pattern, category in _KERNEL_SAFETY_PATTERNS: + if pattern.search(code): + return f"Code blocked: calls a {category}." + return None + + class ExecuteRequest(BaseModel): """Request model for /execute endpoint.""" @@ -226,6 +261,21 @@ async def execute(self, code: str, timeout: float | None = None) -> ExecuteRespo if not self._client or not self._is_ready: raise RuntimeError("Kernel not ready") + # Defense-in-depth: lightweight regex safety check + block_reason = _kernel_check_code_safety(code) + if block_reason is not None: + logger.warning("Kernel safety block: %s code=%r", block_reason, code[:200]) + error_output = MessageType.ERROR.to_notebook_output({ + "ename": "SecurityError", + "evalue": block_reason, + "traceback": [f"SecurityError: {block_reason}"], + }) + return ExecuteResponse( + notebook_outputs=[dict(error_output)] if error_output else [], + error_occurred=True, + execution_time=0.0, + ) + effective_timeout = timeout if timeout is not None else self.default_timeout start_time = time.perf_counter() should_sample = PERF_METRICS_ENABLED and random.random() < PERF_DETAIL_SAMPLE_RATE @@ -323,6 +373,9 @@ async def _execute_code(self, code: str, deadline: float, log_detail: bool = Fal if msg_type == MessageType.STATUS and content.get("execution_state") == "idle": break + if msg_type == MessageType.ERROR: + logger.debug(f"Error Message:\n{content}") + output = msg_type.to_notebook_output(content) if output: notebook_outputs.append(dict(output)) @@ -417,7 +470,7 @@ async def run_server(work_dir: Path, language: NBLanguage, port: int = 8000, sta logger.setLevel(logging.DEBUG) parser = argparse.ArgumentParser(description="Kernel server for Docker-based execution") - parser.add_argument("--work_dir", type=Path, default=Path("/workspace")) + parser.add_argument("--work_dir", type=Path, default=Path("/")) parser.add_argument("--language", type=str, default="python", choices=["python", "r"]) parser.add_argument("--port", type=int, default=8000) parser.add_argument("--startup-token", type=str, default="") From 51e1d0805634b148425ecf7f55066419c0d5b5d4 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Thu, 12 Mar 2026 20:45:17 +0000 Subject: [PATCH 06/31] added code safety --- src/hypotest/env/code_safety.py | 378 ++++++++++++++++++++++++++++++++ 1 file changed, 378 insertions(+) create mode 100644 src/hypotest/env/code_safety.py diff --git a/src/hypotest/env/code_safety.py b/src/hypotest/env/code_safety.py new file mode 100644 index 0000000..5d939e2 --- /dev/null +++ b/src/hypotest/env/code_safety.py @@ -0,0 +1,378 @@ +"""Code safety checker for model-generated Python/R code. + +Prevents model-generated code from killing processes, destroying filesystems, +or escaping the execution sandbox. Uses AST analysis for Python (with regex +fallback) and regex for R code. +""" + +from __future__ import annotations + +import ast +import logging +import re +from typing import TYPE_CHECKING + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from .kernel_server import NBLanguage + + +# --------------------------------------------------------------------------- +# Category messages — intentionally vague to avoid teaching evasion +# --------------------------------------------------------------------------- +_MSG_BLOCKED_IMPORT = "Code blocked: imports a restricted module." +_MSG_DANGEROUS_CALL = "Code blocked: calls a restricted function." +_MSG_BLOCKED_BUILTIN = "Code blocked: uses a restricted builtin." +_MSG_SHELL_ESCAPE = "Code blocked: shell escape is not allowed." +_MSG_DANGEROUS_SHELL = "Code blocked: contains a restricted shell command." +_MSG_R_RESTRICTED = "Code blocked: calls a restricted R function." + + +# --------------------------------------------------------------------------- +# Blocked modules (importing these is never allowed) +# --------------------------------------------------------------------------- +_BLOCKED_MODULES: frozenset[str] = frozenset({ + "ctypes", + "signal", +}) + +# --------------------------------------------------------------------------- +# Dangerous module.function pairs +# --------------------------------------------------------------------------- +_DANGEROUS_CALLS: frozenset[str] = frozenset({ + # os — process management + "os.kill", + "os.killpg", + "os.system", + "os.popen", + "os.fork", + "os.forkpty", + "os.execl", + "os.execle", + "os.execlp", + "os.execlpe", + "os.execv", + "os.execve", + "os.execvp", + "os.execvpe", + "os._exit", + # subprocess + "subprocess.run", + "subprocess.Popen", + "subprocess.call", + "subprocess.check_call", + "subprocess.check_output", + # shutil — destructive + "shutil.rmtree", + # multiprocessing — process spawning + "multiprocessing.Process", +}) + +# Dangerous attribute names used with getattr() evasion +_DANGEROUS_ATTRS: frozenset[str] = frozenset({ + "kill", "killpg", "system", "popen", "fork", "forkpty", + "execl", "execle", "execlp", "execlpe", + "execv", "execve", "execvp", "execvpe", + "_exit", "rmtree", +}) + +# Modules whose import_module() calls should be checked +_DANGEROUS_IMPORTLIB_TARGETS: frozenset[str] = _BLOCKED_MODULES | frozenset({ + "os", "subprocess", "shutil", "multiprocessing", +}) + + +# --------------------------------------------------------------------------- +# AST Walker for Python +# --------------------------------------------------------------------------- +class _DangerousCodeVisitor(ast.NodeVisitor): + """Walks Python AST to detect dangerous calls, tracking import aliases.""" + + def __init__(self) -> None: + self.aliases: dict[str, str] = {} # alias -> canonical module path + self.blocked_reason: str | None = None + + def _block(self, reason: str) -> None: + if self.blocked_reason is None: + self.blocked_reason = reason + + # -- Import tracking ----------------------------------------------------- + + def visit_Import(self, node: ast.Import) -> None: + for alias in node.names: + module = alias.name + local_name = alias.asname or alias.name + # Check top-level module against blocked list + top_module = module.split(".")[0] + if top_module in _BLOCKED_MODULES: + self._block(_MSG_BLOCKED_IMPORT) + return + self.aliases[local_name] = module + self.generic_visit(node) + + def visit_ImportFrom(self, node: ast.ImportFrom) -> None: + module = node.module or "" + top_module = module.split(".")[0] + if top_module in _BLOCKED_MODULES: + self._block(_MSG_BLOCKED_IMPORT) + return + for alias in node.names: + local_name = alias.asname or alias.name + self.aliases[local_name] = f"{module}.{alias.name}" + self.generic_visit(node) + + # -- Call detection ------------------------------------------------------ + + def _resolve_call_path(self, node: ast.Call) -> str | None: + """Resolve a Call node to a canonical 'module.func' string.""" + func = node.func + if isinstance(func, ast.Attribute): + # e.g. os.kill(...) or sp.run(...) + parts: list[str] = [func.attr] + obj = func.value + while isinstance(obj, ast.Attribute): + parts.append(obj.attr) + obj = obj.value + if isinstance(obj, ast.Name): + parts.append(obj.id) + parts.reverse() + # Resolve the leftmost name via aliases + base = parts[0] + if base in self.aliases: + canonical_base = self.aliases[base] + return f"{canonical_base}.{'.'.join(parts[1:])}" + return ".".join(parts) + elif isinstance(func, ast.Name): + # e.g. kill(...) — might be from `from os import kill` + if func.id in self.aliases: + return self.aliases[func.id] + return func.id + return None + + def visit_Call(self, node: ast.Call) -> None: + call_path = self._resolve_call_path(node) + + if call_path is not None: + # Check exec() builtin + if call_path == "exec" or call_path.endswith(".exec"): + # Only block the builtin exec, not method calls like cursor.exec + func = node.func + if isinstance(func, ast.Name) and func.id == "exec": + self._block(_MSG_BLOCKED_BUILTIN) + self.generic_visit(node) + return + + # Check against dangerous calls + for dangerous in _DANGEROUS_CALLS: + if call_path == dangerous or call_path.endswith(f".{dangerous}"): + self._block(_MSG_DANGEROUS_CALL) + self.generic_visit(node) + return + + # Check __import__('os') style + if call_path == "__import__" and node.args: + arg = node.args[0] + if isinstance(arg, ast.Constant) and isinstance(arg.value, str): + top = arg.value.split(".")[0] + if top in _DANGEROUS_IMPORTLIB_TARGETS or top in _BLOCKED_MODULES: + self._block(_MSG_DANGEROUS_CALL) + self.generic_visit(node) + return + + # Check getattr(os, 'kill') style + if call_path in ("getattr", "builtins.getattr") and len(node.args) >= 2: + attr_arg = node.args[1] + if isinstance(attr_arg, ast.Constant) and isinstance(attr_arg.value, str): + if attr_arg.value in _DANGEROUS_ATTRS: + self._block(_MSG_DANGEROUS_CALL) + self.generic_visit(node) + return + + # Check importlib.import_module('signal') style + if call_path.endswith("import_module") and node.args: + arg = node.args[0] + if isinstance(arg, ast.Constant) and isinstance(arg.value, str): + top = arg.value.split(".")[0] + if top in _DANGEROUS_IMPORTLIB_TARGETS or top in _BLOCKED_MODULES: + self._block(_MSG_BLOCKED_IMPORT) + self.generic_visit(node) + return + + self.generic_visit(node) + + +# --------------------------------------------------------------------------- +# Jupyter Magic Patterns (Python pre-pass) +# --------------------------------------------------------------------------- +# Shell escape line: !cmd +_MAGIC_SHELL_ESCAPE_LINE = re.compile(r"^\s*!(?P.+)$", re.MULTILINE) + +# Dangerous IPython magics +_MAGIC_DANGEROUS = re.compile( + r"^\s*%(run|system)\b|get_ipython\(\)\s*\.\s*system\s*\(", re.MULTILINE +) + + +def _check_jupyter_magics(code: str) -> str | None: + """Check for dangerous Jupyter magics. Returns reason or None.""" + # Check dangerous magics first (always blocked) + if _MAGIC_DANGEROUS.search(code): + return _MSG_SHELL_ESCAPE + + # For !cmd shell escapes, extract the command and check against _BASH_PATTERNS. + # Everything is allowed unless the command matches a dangerous bash pattern. + for match in _MAGIC_SHELL_ESCAPE_LINE.finditer(code): + cmd_text = match.group("cmd").strip() + if _check_bash_patterns(cmd_text) is not None: + return _MSG_DANGEROUS_SHELL + + return None + + +# --------------------------------------------------------------------------- +# Bash-level patterns (both languages) +# --------------------------------------------------------------------------- +_BASH_PATTERNS: list[re.Pattern[str]] = [ + re.compile(r"\bkillall\b"), + re.compile(r"\bpkill\b"), + re.compile(r"\bkill\s+.*\$\("), + re.compile(r"\bkill\s+.*`"), + re.compile(r"\bkill\s+(-\d+\s+|-[A-Z]+\s+|-SIG[A-Z]+\s+)*\$\w+"), + re.compile(r"\bkill\s+(-\d+\s+|-[A-Z]+\s+|-SIG[A-Z]+\s+)*-1\b"), + re.compile(r"\bkill\s+(-\d+\s+|-[A-Z]+\s+|-SIG[A-Z]+\s+)*0\b"), + re.compile(r"\b(shutdown|reboot|poweroff|halt|init\s+[06])\b"), + re.compile(r"\bdd\s+.*of=\s*/dev/"), + re.compile(r"\brm\s+(-\w+\s+)*(/\s*$|/\*)"), + re.compile(r"\brm\s+(-\w+\s+)*(/(bin|usr|etc|var|home|root|opt|lib|lib64|sbin|boot|dev|proc|sys))\b"), +] + + +def _check_bash_patterns(code: str) -> str | None: + """Check for dangerous bash-level patterns. Returns reason or None.""" + for pattern in _BASH_PATTERNS: + if pattern.search(code): + return _MSG_DANGEROUS_SHELL + return None + + +# --------------------------------------------------------------------------- +# R regex patterns +# --------------------------------------------------------------------------- +_R_PATTERNS: list[re.Pattern[str]] = [ + re.compile(r"\bsystem\s*\("), + re.compile(r"\bsystem2\s*\("), + re.compile(r"\bshell\s*\("), + re.compile(r"\bshell\.exec\s*\("), + re.compile(r"\bproc\.kill\b"), + re.compile(r"\bquit\s*\("), + re.compile(r"\bq\s*\(\s*\)"), +] + + +def _check_r_patterns(code: str) -> str | None: + """Check R code for dangerous patterns. Returns reason or None.""" + for pattern in _R_PATTERNS: + if pattern.search(code): + return _MSG_R_RESTRICTED + return None + + +# --------------------------------------------------------------------------- +# Python regex fallback (when AST parse fails) +# --------------------------------------------------------------------------- +_PYTHON_REGEX_FALLBACK: list[tuple[re.Pattern[str], str]] = [ + (re.compile(r"\bos\s*\.\s*kill\s*\("), _MSG_DANGEROUS_CALL), + (re.compile(r"\bos\s*\.\s*killpg\s*\("), _MSG_DANGEROUS_CALL), + (re.compile(r"\bos\s*\.\s*system\s*\("), _MSG_DANGEROUS_CALL), + (re.compile(r"\bos\s*\.\s*popen\s*\("), _MSG_DANGEROUS_CALL), + (re.compile(r"\bos\s*\.\s*fork\s*\("), _MSG_DANGEROUS_CALL), + (re.compile(r"\bos\s*\.\s*exec\w*\s*\("), _MSG_DANGEROUS_CALL), + (re.compile(r"\bos\s*\.\s*_exit\s*\("), _MSG_DANGEROUS_CALL), + (re.compile(r"\bsubprocess\s*\.\s*(run|Popen|call|check_call|check_output)\s*\("), _MSG_DANGEROUS_CALL), + (re.compile(r"\bshutil\s*\.\s*rmtree\s*\("), _MSG_DANGEROUS_CALL), + (re.compile(r"\bmultiprocessing\s*\.\s*Process\s*\("), _MSG_DANGEROUS_CALL), + (re.compile(r"\b__import__\s*\("), _MSG_DANGEROUS_CALL), + (re.compile(r"\bexec\s*\("), _MSG_BLOCKED_BUILTIN), + (re.compile(r"\bimport\s+ctypes\b"), _MSG_BLOCKED_IMPORT), + (re.compile(r"\bimport\s+signal\b"), _MSG_BLOCKED_IMPORT), + (re.compile(r"\bfrom\s+ctypes\b"), _MSG_BLOCKED_IMPORT), + (re.compile(r"\bfrom\s+signal\b"), _MSG_BLOCKED_IMPORT), +] + + +def _check_python_regex_fallback(code: str) -> str | None: + """Regex fallback for Python when AST parsing fails.""" + for pattern, msg in _PYTHON_REGEX_FALLBACK: + if pattern.search(code): + return msg + return None + + +# --------------------------------------------------------------------------- +# Python AST check with fallback +# --------------------------------------------------------------------------- +def _check_python(code: str) -> str | None: + """Three-layer Python check: magic pre-pass, AST, regex fallback.""" + # Layer 1: Jupyter magic pre-pass (not valid Python AST) + magic_result = _check_jupyter_magics(code) + if magic_result is not None: + return magic_result + + # Layer 2: AST analysis + try: + tree = ast.parse(code) + except SyntaxError: + # Try stripping magic lines and re-parsing + stripped_lines = [] + for line in code.splitlines(): + stripped = line.lstrip() + if stripped.startswith("!") or stripped.startswith("%"): + stripped_lines.append("") # blank out magic lines + else: + stripped_lines.append(line) + stripped_code = "\n".join(stripped_lines) + try: + tree = ast.parse(stripped_code) + except SyntaxError: + # Layer 3: Regex fallback + return _check_python_regex_fallback(code) + + visitor = _DangerousCodeVisitor() + visitor.visit(tree) + return visitor.blocked_reason + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- +def check_code_safety(code: str, language: "NBLanguage") -> str | None: + """Check if code is safe to execute. + + Args: + code: The code string to check. + language: The notebook language (PYTHON or R). + + Returns: + None if the code is safe, or a category-level error message if blocked. + """ + from .kernel_server import NBLanguage + + # Bash-level patterns apply to both languages + bash_result = _check_bash_patterns(code) + if bash_result is not None: + code_snippet = code[:200] + ("..." if len(code) > 200 else "") + logger.warning("Code safety block: reason=%r language=%s code=%r", bash_result, language.value, code_snippet) + return bash_result + + result: str | None = None + if language == NBLanguage.PYTHON: + result = _check_python(code) + elif language == NBLanguage.R: + result = _check_r_patterns(code) + + if result is not None: + code_snippet = code[:200] + ("..." if len(code) > 200 else "") + logger.warning("Code safety block: reason=%r language=%s code=%r", result, language.value, code_snippet) + + return result From 80d6660271291cc0840712d12ea0a8d6c98eaec2 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Thu, 12 Mar 2026 20:51:31 +0000 Subject: [PATCH 07/31] fixed resources_servers imports --- src/hypotest/env/interpreter_env.py | 45 ++++++++++++++++++++++-- src/hypotest/env/tools/filesystem.py | 4 +-- src/hypotest/env/utils/__init__.py | 2 +- src/hypotest/env/utils/notebook_utils.py | 2 +- 4 files changed, 46 insertions(+), 7 deletions(-) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index d822298..01cead0 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -1263,7 +1263,6 @@ class InterpreterEnvConfig(BaseModel): max_steps: int = cfg.AGENT_MAX_STEPS use_docker: bool = cfg.USE_DOCKER use_enroot: bool = False - use_ray: bool = False container_sqsh_path: Path | None = None normalize_reward: bool = True @@ -1369,7 +1368,6 @@ async def reset(self) -> tuple[Messages, list[Tool]]: } | self.extra_envs, save_dir=self.save_dir, - use_ray=self.config.use_ray, use_docker=self.config.use_docker, use_enroot=self.config.use_enroot, container_sqsh_path=self.config.container_sqsh_path @@ -1385,7 +1383,8 @@ async def reset(self) -> tuple[Messages, list[Tool]]: self._filesystem_tool = FilesystemTool(self.work_dir) self.tools = [ - Tool.from_function(self.run_cell), + # Tool.from_function(self.run_cell), + Tool.from_function(self.stateful_python_code_exec), Tool.from_function(self.reset_kernel), Tool.from_function(self.submit_answer), Tool.from_function(self.list_dir), @@ -1453,6 +1452,46 @@ async def step(self, action: ToolRequestMessage) -> tuple[Messages, float, bool, # ========== Tools ========== + async def stateful_python_code_exec( + self, + code: str, + idx: int | None = None, + ) -> Message | str | list[dict[str, Any]]: + """Run code in a Python Jupyter notebook cell and return the execution output. + + This method allows running code in a new cell (append) or re-running + an existing cell with updated code. + + Error Recovery: + When a cell fails with an error, you MUST fix it by calling run_cell + with the corrected code and the SAME idx as the failed cell. + + The cell number is shown in the output prefix (e.g., "[Cell #3]"). + Do NOT create a new cell to fix an error - always edit the failed cell. + + Args: + code: Code to execute + idx: Cell index to run. If None or >= len(cells), appends a new cell. + If provided, updates and re-runs the existing cell at that index. + Use this to fix errors in existing cells. + + Returns: + Message with multimodal content if images present, otherwise string. + The response includes the cell number (e.g., "[Cell #0] output..."). + """ + start = time.perf_counter() + result = await self.run_cell(code, idx=idx) + duration_ms = (time.perf_counter() - start) * 1000.0 + _record_perf_metric("tool_call_ms.stateful_python_code_exec", duration_ms) + if self._perf_sample_step: + self._emit_perf_event( + "tool_call_end", + tool_name="stateful_python_code_exec", + duration_ms=round(duration_ms, 3), + ) + return result + + async def run_cell( self, code: str, diff --git a/src/hypotest/env/tools/filesystem.py b/src/hypotest/env/tools/filesystem.py index 642177d..de036cc 100644 --- a/src/hypotest/env/tools/filesystem.py +++ b/src/hypotest/env/tools/filesystem.py @@ -23,8 +23,8 @@ import numpy as np from aviary.core import Message, Tool -from resources_servers.aviary_hypotest.env.utils import img_utils -from resources_servers.aviary_hypotest.env.utils.notebook_utils import view_notebook +from hypotest.env.utils import img_utils +from hypotest.env.utils.notebook_utils import view_notebook # ============================================================================= # CONSTANTS diff --git a/src/hypotest/env/utils/__init__.py b/src/hypotest/env/utils/__init__.py index b7451c1..d1f3c1a 100644 --- a/src/hypotest/env/utils/__init__.py +++ b/src/hypotest/env/utils/__init__.py @@ -19,7 +19,7 @@ "view_notebook", ] -from resources_servers.aviary_hypotest.env.kernel_server import NBLanguage +from hypotest.env.kernel_server import NBLanguage from .core import extract_code_from_markdown, extract_xml_content from .img_utils import ( diff --git a/src/hypotest/env/utils/notebook_utils.py b/src/hypotest/env/utils/notebook_utils.py index d2ec722..e727aab 100644 --- a/src/hypotest/env/utils/notebook_utils.py +++ b/src/hypotest/env/utils/notebook_utils.py @@ -8,7 +8,7 @@ import nbformat from aiodocker.containers import DockerContainer -from resources_servers.aviary_hypotest.env import config as cfg +from hypotest.env import config as cfg from .img_utils import encode_image_to_base64 From 3bb30772038c53aa5075605e27d25afa217de601 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Thu, 12 Mar 2026 21:18:40 +0000 Subject: [PATCH 08/31] cleaned up interpreter env by dropping perf logging --- src/hypotest/env/interpreter_env.py | 200 +--------------------------- 1 file changed, 2 insertions(+), 198 deletions(-) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index 01cead0..1a56102 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -14,10 +14,7 @@ import os import shutil import socket -import statistics -import threading import time -import random from pathlib import Path from tempfile import mkdtemp from typing import TYPE_CHECKING, Any, cast @@ -152,72 +149,6 @@ def _kill_process_group(proc: subprocess.Popen, label: str = "enroot", sigterm_t except subprocess.TimeoutExpired: logger.error("[%s] Process pid=%d still alive after SIGKILL — possible zombie", label, proc.pid) -PERF_METRICS_ENABLED = os.getenv("NEMO_GYM_PERF_METRICS", "0") == "1" -PERF_DETAIL_SAMPLE_RATE = float(os.getenv("NEMO_GYM_PERF_DETAIL_SAMPLE_RATE", "0.2")) -PERF_SLOW_STEP_MS = float(os.getenv("NEMO_GYM_PERF_SLOW_STEP_MS", "5000")) -PERF_AGGREGATE_INTERVAL_S = float(os.getenv("NEMO_GYM_PERF_AGGREGATE_INTERVAL_S", "30")) -PERF_LOG_LEVEL = getattr(logging, os.getenv("NEMO_GYM_PERF_LOG_LEVEL", "WARNING").upper(), logging.WARNING) - -_PERF_METRIC_SAMPLES: dict[str, list[float]] = {} -_PERF_LAST_FLUSH_TS = time.perf_counter() -_PERF_LOCK = threading.Lock() - - -def _log_perf(payload: dict[str, Any]) -> None: - logger.log(PERF_LOG_LEVEL, json.dumps(payload)) - - -def _percentile(sorted_values: list[float], pct: float) -> float: - if not sorted_values: - return 0.0 - if len(sorted_values) == 1: - return sorted_values[0] - idx = (len(sorted_values) - 1) * pct - lo = int(idx) - hi = min(lo + 1, len(sorted_values) - 1) - frac = idx - lo - return sorted_values[lo] * (1.0 - frac) + sorted_values[hi] * frac - - -def _maybe_flush_perf_aggregates() -> None: - global _PERF_LAST_FLUSH_TS, _PERF_METRIC_SAMPLES - now = time.perf_counter() - if now - _PERF_LAST_FLUSH_TS < PERF_AGGREGATE_INTERVAL_S: - return - - with _PERF_LOCK: - now = time.perf_counter() - if now - _PERF_LAST_FLUSH_TS < PERF_AGGREGATE_INTERVAL_S: - return - snapshots = _PERF_METRIC_SAMPLES - _PERF_METRIC_SAMPLES = {} - _PERF_LAST_FLUSH_TS = now - - for metric_name, values in snapshots.items(): - if not values: - continue - ordered = sorted(values) - payload = { - "component": "aviary_hypotest.interpreter_env", - "event": "perf_aggregate", - "metric": metric_name, - "count": len(values), - "mean_ms": round(statistics.fmean(values), 3), - "p50_ms": round(_percentile(ordered, 0.50), 3), - "p90_ms": round(_percentile(ordered, 0.90), 3), - "p95_ms": round(_percentile(ordered, 0.95), 3), - "p99_ms": round(_percentile(ordered, 0.99), 3), - "max_ms": round(ordered[-1], 3), - } - _log_perf(payload) - - -def _record_perf_metric(metric_name: str, value_ms: float) -> None: - if not PERF_METRICS_ENABLED: - return - with _PERF_LOCK: - _PERF_METRIC_SAMPLES.setdefault(metric_name, []).append(value_ms) - _maybe_flush_perf_aggregates() class ProblemInstance(BaseModel): uuid: UUID @@ -1219,30 +1150,14 @@ async def execute_and_add_cell( actual_idx = cell_idx return result, actual_idx - execute_start = time.perf_counter() if self.use_ray and self.use_enroot: - rpc_wait_start = time.perf_counter() result_ref = self.kernel_container._execute_via_http.remote(code, timeout) result = await result_ref - _record_perf_metric( - "execute_path_ms.ray_enroot_rpc_wait", - (time.perf_counter() - rpc_wait_start) * 1000.0, - ) elif self.use_docker or self.use_enroot: - http_exec_start = time.perf_counter() result = await self._execute_via_http(code, timeout) - _record_perf_metric( - "execute_path_ms.http_kernel_execute", - (time.perf_counter() - http_exec_start) * 1000.0, - ) else: assert self.interpreter is not None - local_exec_start = time.perf_counter() result = await self.interpreter.execute_code(code, timeout) - _record_perf_metric( - "execute_path_ms.local_interpreter_execute", - (time.perf_counter() - local_exec_start) * 1000.0, - ) if cell_idx is None or cell_idx >= len(self.nb.cells): actual_idx = self._add_cell(code, result) @@ -1250,7 +1165,6 @@ async def execute_and_add_cell( self._update_cell(cell_idx, code, result) actual_idx = cell_idx - _record_perf_metric("execute_and_add_cell_total_ms", (time.perf_counter() - execute_start) * 1000.0) return result, actual_idx @@ -1302,8 +1216,6 @@ def __init__( self.execution_timeout = self.execution_config.cell_execution_timeout self.max_steps = self.config.max_steps - logger.error(f"EXECUTION TIMEOUT SET TO {self.execution_timeout} SECONDS") - self.input_data = input_data self.output_data: list[dict[str, str | int]] = [] self.logger = logger @@ -1313,19 +1225,6 @@ def __init__( self.state: InterpreterEnvState # prompting_config is set during reset() after language resolution self.prompting_config: PromptingConfig - self._perf_sample_step = False - - def _emit_perf_event(self, event: str, **fields: Any) -> None: - if not PERF_METRICS_ENABLED: - return - payload: dict[str, Any] = { - "component": "aviary_hypotest.interpreter_env", - "event": event, - "env_id": getattr(self, "_nemo_env_id", None), - "step_idx": self.step_count, - } - payload.update(fields) - _log_perf(payload) @property def language(self) -> NBLanguage: @@ -1383,8 +1282,7 @@ async def reset(self) -> tuple[Messages, list[Tool]]: self._filesystem_tool = FilesystemTool(self.work_dir) self.tools = [ - # Tool.from_function(self.run_cell), - Tool.from_function(self.stateful_python_code_exec), + Tool.from_function(self.run_cell), Tool.from_function(self.reset_kernel), Tool.from_function(self.submit_answer), Tool.from_function(self.list_dir), @@ -1411,15 +1309,10 @@ async def reset(self) -> tuple[Messages, list[Tool]]: async def step(self, action: ToolRequestMessage) -> tuple[Messages, float, bool, bool]: """Execute a step in the environment.""" self.step_count += 1 - self._perf_sample_step = PERF_METRICS_ENABLED and random.random() < PERF_DETAIL_SAMPLE_RATE - step_start = time.perf_counter() - exec_tool_calls_start = time.perf_counter() obs = cast( Messages, await self.exec_tool_calls(action, concurrency=False, handle_tool_exc=True), ) - exec_tool_calls_ms = (time.perf_counter() - exec_tool_calls_start) * 1000.0 - _record_perf_metric("exec_tool_calls_ms", exec_tool_calls_ms) obs = [*obs] if self.include_env_state_msg: @@ -1434,64 +1327,10 @@ async def step(self, action: ToolRequestMessage) -> tuple[Messages, float, bool, self.state.actions.append(str(action)) reward = self.state.score if self.state.done else 0.0 - step_total_ms = (time.perf_counter() - step_start) * 1000.0 - _record_perf_metric("env_step_ms", step_total_ms) - should_log_detail = self._perf_sample_step or step_total_ms >= PERF_SLOW_STEP_MS - if should_log_detail: - obs_payload_chars = sum(len(str(getattr(m, "content", ""))) for m in obs) - self._emit_perf_event( - "env_step_end", - duration_ms=round(step_total_ms, 3), - exec_tool_calls_ms=round(exec_tool_calls_ms, 3), - action_tool_calls_count=len(action.tool_calls), - obs_messages_count=len(obs), - obs_payload_chars=obs_payload_chars, - slow=step_total_ms >= PERF_SLOW_STEP_MS, - ) return obs, reward, self.state.done, False # ========== Tools ========== - async def stateful_python_code_exec( - self, - code: str, - idx: int | None = None, - ) -> Message | str | list[dict[str, Any]]: - """Run code in a Python Jupyter notebook cell and return the execution output. - - This method allows running code in a new cell (append) or re-running - an existing cell with updated code. - - Error Recovery: - When a cell fails with an error, you MUST fix it by calling run_cell - with the corrected code and the SAME idx as the failed cell. - - The cell number is shown in the output prefix (e.g., "[Cell #3]"). - Do NOT create a new cell to fix an error - always edit the failed cell. - - Args: - code: Code to execute - idx: Cell index to run. If None or >= len(cells), appends a new cell. - If provided, updates and re-runs the existing cell at that index. - Use this to fix errors in existing cells. - - Returns: - Message with multimodal content if images present, otherwise string. - The response includes the cell number (e.g., "[Cell #0] output..."). - """ - start = time.perf_counter() - result = await self.run_cell(code, idx=idx) - duration_ms = (time.perf_counter() - start) * 1000.0 - _record_perf_metric("tool_call_ms.stateful_python_code_exec", duration_ms) - if self._perf_sample_step: - self._emit_perf_event( - "tool_call_end", - tool_name="stateful_python_code_exec", - duration_ms=round(duration_ms, 3), - ) - return result - - async def run_cell( self, code: str, @@ -1545,12 +1384,9 @@ async def run_cell( cell_idx = None # Execute code and update notebook atomically - exec_start = time.perf_counter() result, actual_cell_idx = await self.state.execute_and_add_cell( code, cell_idx=cell_idx, timeout=effective_timeout ) - execute_and_add_cell_ms = (time.perf_counter() - exec_start) * 1000.0 - _record_perf_metric("execute_and_add_cell_ms", execute_and_add_cell_ms) # Build response with cell number cell_info = f"[Cell #{actual_cell_idx}] " @@ -1572,7 +1408,6 @@ async def reset_kernel(self) -> str: This clears all variables and execution state. """ - start = time.perf_counter() if self.state.use_ray: reset_ref = self.state.kernel_container._reset_via_http.remote() await reset_ref @@ -1588,14 +1423,6 @@ async def reset_kernel(self) -> str: self.state.notebook_runtime_errors = [] self.state._execution_count = 0 - duration_ms = (time.perf_counter() - start) * 1000.0 - _record_perf_metric("tool_call_ms.reset_kernel", duration_ms) - if self._perf_sample_step: - self._emit_perf_event( - "tool_call_end", - tool_name="reset_kernel", - duration_ms=round(duration_ms, 3), - ) return "Kernel reset successfully." @property @@ -1659,10 +1486,7 @@ async def submit_answer(self, answer: str) -> str: Args: answer: Your final response to the research question """ - start = time.perf_counter() if self.state.done: - duration_ms = (time.perf_counter() - start) * 1000.0 - _record_perf_metric("tool_call_ms.submit_answer", duration_ms) return "Episode already finished." self.state.answer = answer @@ -1670,19 +1494,9 @@ async def submit_answer(self, answer: str) -> str: if self.rubric_model is None: self.logger.warning("No rubric_model configured, skipping scoring") - duration_ms = (time.perf_counter() - start) * 1000.0 - _record_perf_metric("tool_call_ms.submit_answer", duration_ms) return answer correct = await self._score_solution(answer) - duration_ms = (time.perf_counter() - start) * 1000.0 - _record_perf_metric("tool_call_ms.submit_answer", duration_ms) - if self._perf_sample_step: - self._emit_perf_event( - "tool_call_end", - tool_name="submit_answer", - duration_ms=round(duration_ms, 3), - ) return CORRECT_MSG if correct else INCORRECT_MSG def list_dir(self) -> str: @@ -1697,17 +1511,7 @@ def list_dir(self) -> str: max_files: Maximum number of files to display (default: 20) show_hidden: Whether to show hidden files starting with '.' (default: False) """ - start = time.perf_counter() - result = self._filesystem_tool.list_dir() - duration_ms = (time.perf_counter() - start) * 1000.0 - _record_perf_metric("tool_call_ms.list_dir", duration_ms) - if self._perf_sample_step: - self._emit_perf_event( - "tool_call_end", - tool_name="list_dir", - duration_ms=round(duration_ms, 3), - ) - return result + return self._filesystem_tool.list_dir() # ========== Time Management ========== From 4f0cddc01bb97cc72511ea1fa9ec795dd0b69a8f Mon Sep 17 00:00:00 2001 From: Siddharth Narayanan Date: Thu, 12 Mar 2026 17:03:02 -0700 Subject: [PATCH 09/31] Add Apache License 2.0 Added the Apache License 2.0 to the project. --- LICENSE | 201 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 201 insertions(+) create mode 100644 LICENSE diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. From 1f5906a79c56aed16f325f8f266022d6ef048d56 Mon Sep 17 00:00:00 2001 From: Siddharth Narayanan Date: Fri, 13 Mar 2026 14:08:12 -0700 Subject: [PATCH 10/31] Reformat eval schema to match training dataset schema (#12) --- src/hypotest/dataset_server.py | 10 ++++++---- src/hypotest/env/interpreter_env.py | 11 ++++++----- src/hypotest/env/prompts.py | 2 +- tests/conftest.py | 4 ++-- tests/test_interpreter_env.py | 4 ++-- 5 files changed, 17 insertions(+), 14 deletions(-) diff --git a/src/hypotest/dataset_server.py b/src/hypotest/dataset_server.py index 8ac9585..e5662be 100644 --- a/src/hypotest/dataset_server.py +++ b/src/hypotest/dataset_server.py @@ -62,11 +62,13 @@ def __init__(self, config: HypotestDatasetConfig): def get_new_env_by_idx(self, idx: int) -> InterpreterEnv: problem = self.problems[idx] - problem_count = self.problem_counter[problem.uuid] - self.problem_counter[problem.uuid] += 1 - run_id = f"{problem.uuid}-iter{problem_count}" + problem_count = self.problem_counter[problem.id] + self.problem_counter[problem.id] += 1 + run_id = f"{problem.id}-iter{problem_count}" - capsule_path = self.config.capsule_dir / f"CapsuleData-{problem.uuid}" + capsule_path = self.config.capsule_dir / problem.input_data_path + if not capsule_path.exists(): + capsule_path = self.config.capsule_dir / f"CapsuleData-{problem.id}" problem_dir = Path(self.config.work_dir) / run_id if self.config.work_dir else Path(mkdtemp()) problem_dir.mkdir(parents=True, exist_ok=True) shutil.copytree(capsule_path, problem_dir, dirs_exist_ok=True) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index 1a56102..174aa95 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -151,12 +151,13 @@ def _kill_process_group(proc: subprocess.Popen, label: str = "enroot", sigterm_t class ProblemInstance(BaseModel): - uuid: UUID + id: UUID hypothesis: str - objective: str + protocol: str accepted: bool = Field(alias="answer") rubric: str max_score: int = Field(alias="max_points") + input_data_path: str = "" metadata: dict[str, JsonValue] = Field(default_factory=dict) nb_primary_language: str = Field(default=str(NBLanguage.PYTHON)) @@ -1293,7 +1294,7 @@ async def reset(self) -> tuple[Messages, list[Tool]]: content=HYPOTHESIS_TASK_DESC.format( language=self.language.value.capitalize(), hypothesis=self.problem.hypothesis, - objective=self.problem.objective, + protocol=self.problem.protocol, ) ) ) @@ -1612,9 +1613,9 @@ async def main() -> None: print(f"Working directory: {work_dir}") problem = ProblemInstance( - uuid="", + id="", hypothesis="", - objective="", + protocol="", answer=False, rubric="", max_points=0, diff --git a/src/hypotest/env/prompts.py b/src/hypotest/env/prompts.py index ce66465..74e1c0a 100644 --- a/src/hypotest/env/prompts.py +++ b/src/hypotest/env/prompts.py @@ -13,7 +13,7 @@ -{objective} +{protocol} """.strip() diff --git a/tests/conftest.py b/tests/conftest.py index 08aedea..e727a29 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -83,9 +83,9 @@ def skip_if_r_unavailable(request): def default_problem() -> ProblemInstance: """Default ProblemInstance for tests.""" return ProblemInstance( - uuid=UUID("00000000-0000-0000-0000-000000000000"), + id=UUID("00000000-0000-0000-0000-000000000000"), hypothesis="Test hypothesis", - objective="Test objective", + protocol="Test objective", answer=True, rubric="Test rubric", max_points=10, diff --git a/tests/test_interpreter_env.py b/tests/test_interpreter_env.py index e91164d..a793caa 100644 --- a/tests/test_interpreter_env.py +++ b/tests/test_interpreter_env.py @@ -855,9 +855,9 @@ class TestRubricGrading: """Tests for rubric-based grading.""" GRADING_PROBLEM = ProblemInstance( - uuid=UUID("12345678-1234-5678-1234-567812345678"), + id=UUID("12345678-1234-5678-1234-567812345678"), hypothesis="Numbers greater than 10 are large", - objective="Load data, compute statistics, determine if hypothesis holds", + protocol="Load data, compute statistics, determine if hypothesis holds", answer=True, rubric="""* 1 point: Data is loaded correctly * 1 point: Statistics are computed From c1d0711eef3a2c615ef16e4a76d672ecaceb60ed Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Tue, 10 Mar 2026 13:34:09 -0700 Subject: [PATCH 11/31] moved perf/enroot fixes from internal Gym to hypotest --- src/hypotest/env/interpreter_env.py | 151 ++++++++++++++++++++++- src/hypotest/env/tools/filesystem.py | 4 +- src/hypotest/env/utils/__init__.py | 2 +- src/hypotest/env/utils/notebook_utils.py | 2 +- 4 files changed, 153 insertions(+), 6 deletions(-) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index 174aa95..c0fe29f 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -14,7 +14,10 @@ import os import shutil import socket +import statistics +import threading import time +import random from pathlib import Path from tempfile import mkdtemp from typing import TYPE_CHECKING, Any, cast @@ -87,7 +90,6 @@ async def get_free_port() -> int: class _PortCollisionError(Exception): """Port already in use by another server — retry with a new port.""" - def _kill_process_group(proc: subprocess.Popen, label: str = "enroot", sigterm_timeout: float = 15) -> None: """Safely terminate a process group, escalating from SIGTERM to SIGKILL. @@ -149,6 +151,72 @@ def _kill_process_group(proc: subprocess.Popen, label: str = "enroot", sigterm_t except subprocess.TimeoutExpired: logger.error("[%s] Process pid=%d still alive after SIGKILL — possible zombie", label, proc.pid) +PERF_METRICS_ENABLED = os.getenv("NEMO_GYM_PERF_METRICS", "0") == "1" +PERF_DETAIL_SAMPLE_RATE = float(os.getenv("NEMO_GYM_PERF_DETAIL_SAMPLE_RATE", "0.2")) +PERF_SLOW_STEP_MS = float(os.getenv("NEMO_GYM_PERF_SLOW_STEP_MS", "5000")) +PERF_AGGREGATE_INTERVAL_S = float(os.getenv("NEMO_GYM_PERF_AGGREGATE_INTERVAL_S", "30")) +PERF_LOG_LEVEL = getattr(logging, os.getenv("NEMO_GYM_PERF_LOG_LEVEL", "WARNING").upper(), logging.WARNING) + +_PERF_METRIC_SAMPLES: dict[str, list[float]] = {} +_PERF_LAST_FLUSH_TS = time.perf_counter() +_PERF_LOCK = threading.Lock() + + +def _log_perf(payload: dict[str, Any]) -> None: + logger.log(PERF_LOG_LEVEL, json.dumps(payload)) + + +def _percentile(sorted_values: list[float], pct: float) -> float: + if not sorted_values: + return 0.0 + if len(sorted_values) == 1: + return sorted_values[0] + idx = (len(sorted_values) - 1) * pct + lo = int(idx) + hi = min(lo + 1, len(sorted_values) - 1) + frac = idx - lo + return sorted_values[lo] * (1.0 - frac) + sorted_values[hi] * frac + + +def _maybe_flush_perf_aggregates() -> None: + global _PERF_LAST_FLUSH_TS, _PERF_METRIC_SAMPLES + now = time.perf_counter() + if now - _PERF_LAST_FLUSH_TS < PERF_AGGREGATE_INTERVAL_S: + return + + with _PERF_LOCK: + now = time.perf_counter() + if now - _PERF_LAST_FLUSH_TS < PERF_AGGREGATE_INTERVAL_S: + return + snapshots = _PERF_METRIC_SAMPLES + _PERF_METRIC_SAMPLES = {} + _PERF_LAST_FLUSH_TS = now + + for metric_name, values in snapshots.items(): + if not values: + continue + ordered = sorted(values) + payload = { + "component": "aviary_hypotest.interpreter_env", + "event": "perf_aggregate", + "metric": metric_name, + "count": len(values), + "mean_ms": round(statistics.fmean(values), 3), + "p50_ms": round(_percentile(ordered, 0.50), 3), + "p90_ms": round(_percentile(ordered, 0.90), 3), + "p95_ms": round(_percentile(ordered, 0.95), 3), + "p99_ms": round(_percentile(ordered, 0.99), 3), + "max_ms": round(ordered[-1], 3), + } + _log_perf(payload) + + +def _record_perf_metric(metric_name: str, value_ms: float) -> None: + if not PERF_METRICS_ENABLED: + return + with _PERF_LOCK: + _PERF_METRIC_SAMPLES.setdefault(metric_name, []).append(value_ms) + _maybe_flush_perf_aggregates() class ProblemInstance(BaseModel): id: UUID @@ -1156,9 +1224,18 @@ async def execute_and_add_cell( result = await result_ref elif self.use_docker or self.use_enroot: result = await self._execute_via_http(code, timeout) + _record_perf_metric( + "execute_path_ms.http_kernel_execute", + (time.perf_counter() - http_exec_start) * 1000.0, + ) else: assert self.interpreter is not None + local_exec_start = time.perf_counter() result = await self.interpreter.execute_code(code, timeout) + _record_perf_metric( + "execute_path_ms.local_interpreter_execute", + (time.perf_counter() - local_exec_start) * 1000.0, + ) if cell_idx is None or cell_idx >= len(self.nb.cells): actual_idx = self._add_cell(code, result) @@ -1166,6 +1243,7 @@ async def execute_and_add_cell( self._update_cell(cell_idx, code, result) actual_idx = cell_idx + _record_perf_metric("execute_and_add_cell_total_ms", (time.perf_counter() - execute_start) * 1000.0) return result, actual_idx @@ -1217,6 +1295,8 @@ def __init__( self.execution_timeout = self.execution_config.cell_execution_timeout self.max_steps = self.config.max_steps + logger.error(f"EXECUTION TIMEOUT SET TO {self.execution_timeout} SECONDS") + self.input_data = input_data self.output_data: list[dict[str, str | int]] = [] self.logger = logger @@ -1226,6 +1306,19 @@ def __init__( self.state: InterpreterEnvState # prompting_config is set during reset() after language resolution self.prompting_config: PromptingConfig + self._perf_sample_step = False + + def _emit_perf_event(self, event: str, **fields: Any) -> None: + if not PERF_METRICS_ENABLED: + return + payload: dict[str, Any] = { + "component": "aviary_hypotest.interpreter_env", + "event": event, + "env_id": getattr(self, "_nemo_env_id", None), + "step_idx": self.step_count, + } + payload.update(fields) + _log_perf(payload) @property def language(self) -> NBLanguage: @@ -1310,10 +1403,15 @@ async def reset(self) -> tuple[Messages, list[Tool]]: async def step(self, action: ToolRequestMessage) -> tuple[Messages, float, bool, bool]: """Execute a step in the environment.""" self.step_count += 1 + self._perf_sample_step = PERF_METRICS_ENABLED and random.random() < PERF_DETAIL_SAMPLE_RATE + step_start = time.perf_counter() + exec_tool_calls_start = time.perf_counter() obs = cast( Messages, await self.exec_tool_calls(action, concurrency=False, handle_tool_exc=True), ) + exec_tool_calls_ms = (time.perf_counter() - exec_tool_calls_start) * 1000.0 + _record_perf_metric("exec_tool_calls_ms", exec_tool_calls_ms) obs = [*obs] if self.include_env_state_msg: @@ -1328,6 +1426,20 @@ async def step(self, action: ToolRequestMessage) -> tuple[Messages, float, bool, self.state.actions.append(str(action)) reward = self.state.score if self.state.done else 0.0 + step_total_ms = (time.perf_counter() - step_start) * 1000.0 + _record_perf_metric("env_step_ms", step_total_ms) + should_log_detail = self._perf_sample_step or step_total_ms >= PERF_SLOW_STEP_MS + if should_log_detail: + obs_payload_chars = sum(len(str(getattr(m, "content", ""))) for m in obs) + self._emit_perf_event( + "env_step_end", + duration_ms=round(step_total_ms, 3), + exec_tool_calls_ms=round(exec_tool_calls_ms, 3), + action_tool_calls_count=len(action.tool_calls), + obs_messages_count=len(obs), + obs_payload_chars=obs_payload_chars, + slow=step_total_ms >= PERF_SLOW_STEP_MS, + ) return obs, reward, self.state.done, False # ========== Tools ========== @@ -1359,6 +1471,7 @@ async def run_cell( Message with multimodal content if images present, otherwise string. The response includes the cell number (e.g., "[Cell #0] output..."). """ + start = time.perf_counter() remaining_seconds = self.get_remaining_time() if remaining_seconds <= self.execution_config.force_submit_threshold: @@ -1385,9 +1498,12 @@ async def run_cell( cell_idx = None # Execute code and update notebook atomically + exec_start = time.perf_counter() result, actual_cell_idx = await self.state.execute_and_add_cell( code, cell_idx=cell_idx, timeout=effective_timeout ) + execute_and_add_cell_ms = (time.perf_counter() - exec_start) * 1000.0 + _record_perf_metric("execute_and_add_cell_ms", execute_and_add_cell_ms) # Build response with cell number cell_info = f"[Cell #{actual_cell_idx}] " @@ -1402,13 +1518,23 @@ async def run_cell( images=cast(list[np.ndarray | str], image_urls), ) - return cell_info + result.get_truncated_text() + result = cell_info + result.get_truncated_text() + duration_ms = (time.perf_counter() - start) * 1000.0 + _record_perf_metric("tool_call_ms.run_cell", duration_ms) + if self._perf_sample_step: + self._emit_perf_event( + "tool_call_end", + tool_name="run_cell", + duration_ms=round(duration_ms, 3), + ) + return result async def reset_kernel(self) -> str: """Reset the kernel to a clean state. This clears all variables and execution state. """ + start = time.perf_counter() if self.state.use_ray: reset_ref = self.state.kernel_container._reset_via_http.remote() await reset_ref @@ -1424,6 +1550,14 @@ async def reset_kernel(self) -> str: self.state.notebook_runtime_errors = [] self.state._execution_count = 0 + duration_ms = (time.perf_counter() - start) * 1000.0 + _record_perf_metric("tool_call_ms.reset_kernel", duration_ms) + if self._perf_sample_step: + self._emit_perf_event( + "tool_call_end", + tool_name="reset_kernel", + duration_ms=round(duration_ms, 3), + ) return "Kernel reset successfully." @property @@ -1487,7 +1621,10 @@ async def submit_answer(self, answer: str) -> str: Args: answer: Your final response to the research question """ + start = time.perf_counter() if self.state.done: + duration_ms = (time.perf_counter() - start) * 1000.0 + _record_perf_metric("tool_call_ms.submit_answer", duration_ms) return "Episode already finished." self.state.answer = answer @@ -1495,9 +1632,19 @@ async def submit_answer(self, answer: str) -> str: if self.rubric_model is None: self.logger.warning("No rubric_model configured, skipping scoring") + duration_ms = (time.perf_counter() - start) * 1000.0 + _record_perf_metric("tool_call_ms.submit_answer", duration_ms) return answer correct = await self._score_solution(answer) + duration_ms = (time.perf_counter() - start) * 1000.0 + _record_perf_metric("tool_call_ms.submit_answer", duration_ms) + if self._perf_sample_step: + self._emit_perf_event( + "tool_call_end", + tool_name="submit_answer", + duration_ms=round(duration_ms, 3), + ) return CORRECT_MSG if correct else INCORRECT_MSG def list_dir(self) -> str: diff --git a/src/hypotest/env/tools/filesystem.py b/src/hypotest/env/tools/filesystem.py index de036cc..642177d 100644 --- a/src/hypotest/env/tools/filesystem.py +++ b/src/hypotest/env/tools/filesystem.py @@ -23,8 +23,8 @@ import numpy as np from aviary.core import Message, Tool -from hypotest.env.utils import img_utils -from hypotest.env.utils.notebook_utils import view_notebook +from resources_servers.aviary_hypotest.env.utils import img_utils +from resources_servers.aviary_hypotest.env.utils.notebook_utils import view_notebook # ============================================================================= # CONSTANTS diff --git a/src/hypotest/env/utils/__init__.py b/src/hypotest/env/utils/__init__.py index d1f3c1a..b7451c1 100644 --- a/src/hypotest/env/utils/__init__.py +++ b/src/hypotest/env/utils/__init__.py @@ -19,7 +19,7 @@ "view_notebook", ] -from hypotest.env.kernel_server import NBLanguage +from resources_servers.aviary_hypotest.env.kernel_server import NBLanguage from .core import extract_code_from_markdown, extract_xml_content from .img_utils import ( diff --git a/src/hypotest/env/utils/notebook_utils.py b/src/hypotest/env/utils/notebook_utils.py index e727aab..d2ec722 100644 --- a/src/hypotest/env/utils/notebook_utils.py +++ b/src/hypotest/env/utils/notebook_utils.py @@ -8,7 +8,7 @@ import nbformat from aiodocker.containers import DockerContainer -from hypotest.env import config as cfg +from resources_servers.aviary_hypotest.env import config as cfg from .img_utils import encode_image_to_base64 From 24363414655fe780d33929e55de81f0e047f1f2b Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Thu, 12 Mar 2026 20:44:50 +0000 Subject: [PATCH 12/31] added code safety, perf fixes on enroot container launch to prevent hangs --- src/hypotest/env/interpreter_env.py | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index c0fe29f..c628b7d 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -1471,7 +1471,6 @@ async def run_cell( Message with multimodal content if images present, otherwise string. The response includes the cell number (e.g., "[Cell #0] output..."). """ - start = time.perf_counter() remaining_seconds = self.get_remaining_time() if remaining_seconds <= self.execution_config.force_submit_threshold: @@ -1518,16 +1517,7 @@ async def run_cell( images=cast(list[np.ndarray | str], image_urls), ) - result = cell_info + result.get_truncated_text() - duration_ms = (time.perf_counter() - start) * 1000.0 - _record_perf_metric("tool_call_ms.run_cell", duration_ms) - if self._perf_sample_step: - self._emit_perf_event( - "tool_call_end", - tool_name="run_cell", - duration_ms=round(duration_ms, 3), - ) - return result + return cell_info + result.get_truncated_text() async def reset_kernel(self) -> str: """Reset the kernel to a clean state. From bdf01f5bc63be47f033943d950b7362130c843d8 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Thu, 12 Mar 2026 21:18:40 +0000 Subject: [PATCH 13/31] cleaned up interpreter env by dropping perf logging --- src/hypotest/env/interpreter_env.py | 138 ---------------------------- 1 file changed, 138 deletions(-) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index c628b7d..b35bf85 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -14,10 +14,7 @@ import os import shutil import socket -import statistics -import threading import time -import random from pathlib import Path from tempfile import mkdtemp from typing import TYPE_CHECKING, Any, cast @@ -151,72 +148,6 @@ def _kill_process_group(proc: subprocess.Popen, label: str = "enroot", sigterm_t except subprocess.TimeoutExpired: logger.error("[%s] Process pid=%d still alive after SIGKILL — possible zombie", label, proc.pid) -PERF_METRICS_ENABLED = os.getenv("NEMO_GYM_PERF_METRICS", "0") == "1" -PERF_DETAIL_SAMPLE_RATE = float(os.getenv("NEMO_GYM_PERF_DETAIL_SAMPLE_RATE", "0.2")) -PERF_SLOW_STEP_MS = float(os.getenv("NEMO_GYM_PERF_SLOW_STEP_MS", "5000")) -PERF_AGGREGATE_INTERVAL_S = float(os.getenv("NEMO_GYM_PERF_AGGREGATE_INTERVAL_S", "30")) -PERF_LOG_LEVEL = getattr(logging, os.getenv("NEMO_GYM_PERF_LOG_LEVEL", "WARNING").upper(), logging.WARNING) - -_PERF_METRIC_SAMPLES: dict[str, list[float]] = {} -_PERF_LAST_FLUSH_TS = time.perf_counter() -_PERF_LOCK = threading.Lock() - - -def _log_perf(payload: dict[str, Any]) -> None: - logger.log(PERF_LOG_LEVEL, json.dumps(payload)) - - -def _percentile(sorted_values: list[float], pct: float) -> float: - if not sorted_values: - return 0.0 - if len(sorted_values) == 1: - return sorted_values[0] - idx = (len(sorted_values) - 1) * pct - lo = int(idx) - hi = min(lo + 1, len(sorted_values) - 1) - frac = idx - lo - return sorted_values[lo] * (1.0 - frac) + sorted_values[hi] * frac - - -def _maybe_flush_perf_aggregates() -> None: - global _PERF_LAST_FLUSH_TS, _PERF_METRIC_SAMPLES - now = time.perf_counter() - if now - _PERF_LAST_FLUSH_TS < PERF_AGGREGATE_INTERVAL_S: - return - - with _PERF_LOCK: - now = time.perf_counter() - if now - _PERF_LAST_FLUSH_TS < PERF_AGGREGATE_INTERVAL_S: - return - snapshots = _PERF_METRIC_SAMPLES - _PERF_METRIC_SAMPLES = {} - _PERF_LAST_FLUSH_TS = now - - for metric_name, values in snapshots.items(): - if not values: - continue - ordered = sorted(values) - payload = { - "component": "aviary_hypotest.interpreter_env", - "event": "perf_aggregate", - "metric": metric_name, - "count": len(values), - "mean_ms": round(statistics.fmean(values), 3), - "p50_ms": round(_percentile(ordered, 0.50), 3), - "p90_ms": round(_percentile(ordered, 0.90), 3), - "p95_ms": round(_percentile(ordered, 0.95), 3), - "p99_ms": round(_percentile(ordered, 0.99), 3), - "max_ms": round(ordered[-1], 3), - } - _log_perf(payload) - - -def _record_perf_metric(metric_name: str, value_ms: float) -> None: - if not PERF_METRICS_ENABLED: - return - with _PERF_LOCK: - _PERF_METRIC_SAMPLES.setdefault(metric_name, []).append(value_ms) - _maybe_flush_perf_aggregates() class ProblemInstance(BaseModel): id: UUID @@ -1224,18 +1155,9 @@ async def execute_and_add_cell( result = await result_ref elif self.use_docker or self.use_enroot: result = await self._execute_via_http(code, timeout) - _record_perf_metric( - "execute_path_ms.http_kernel_execute", - (time.perf_counter() - http_exec_start) * 1000.0, - ) else: assert self.interpreter is not None - local_exec_start = time.perf_counter() result = await self.interpreter.execute_code(code, timeout) - _record_perf_metric( - "execute_path_ms.local_interpreter_execute", - (time.perf_counter() - local_exec_start) * 1000.0, - ) if cell_idx is None or cell_idx >= len(self.nb.cells): actual_idx = self._add_cell(code, result) @@ -1243,7 +1165,6 @@ async def execute_and_add_cell( self._update_cell(cell_idx, code, result) actual_idx = cell_idx - _record_perf_metric("execute_and_add_cell_total_ms", (time.perf_counter() - execute_start) * 1000.0) return result, actual_idx @@ -1295,8 +1216,6 @@ def __init__( self.execution_timeout = self.execution_config.cell_execution_timeout self.max_steps = self.config.max_steps - logger.error(f"EXECUTION TIMEOUT SET TO {self.execution_timeout} SECONDS") - self.input_data = input_data self.output_data: list[dict[str, str | int]] = [] self.logger = logger @@ -1306,19 +1225,6 @@ def __init__( self.state: InterpreterEnvState # prompting_config is set during reset() after language resolution self.prompting_config: PromptingConfig - self._perf_sample_step = False - - def _emit_perf_event(self, event: str, **fields: Any) -> None: - if not PERF_METRICS_ENABLED: - return - payload: dict[str, Any] = { - "component": "aviary_hypotest.interpreter_env", - "event": event, - "env_id": getattr(self, "_nemo_env_id", None), - "step_idx": self.step_count, - } - payload.update(fields) - _log_perf(payload) @property def language(self) -> NBLanguage: @@ -1403,15 +1309,10 @@ async def reset(self) -> tuple[Messages, list[Tool]]: async def step(self, action: ToolRequestMessage) -> tuple[Messages, float, bool, bool]: """Execute a step in the environment.""" self.step_count += 1 - self._perf_sample_step = PERF_METRICS_ENABLED and random.random() < PERF_DETAIL_SAMPLE_RATE - step_start = time.perf_counter() - exec_tool_calls_start = time.perf_counter() obs = cast( Messages, await self.exec_tool_calls(action, concurrency=False, handle_tool_exc=True), ) - exec_tool_calls_ms = (time.perf_counter() - exec_tool_calls_start) * 1000.0 - _record_perf_metric("exec_tool_calls_ms", exec_tool_calls_ms) obs = [*obs] if self.include_env_state_msg: @@ -1426,20 +1327,6 @@ async def step(self, action: ToolRequestMessage) -> tuple[Messages, float, bool, self.state.actions.append(str(action)) reward = self.state.score if self.state.done else 0.0 - step_total_ms = (time.perf_counter() - step_start) * 1000.0 - _record_perf_metric("env_step_ms", step_total_ms) - should_log_detail = self._perf_sample_step or step_total_ms >= PERF_SLOW_STEP_MS - if should_log_detail: - obs_payload_chars = sum(len(str(getattr(m, "content", ""))) for m in obs) - self._emit_perf_event( - "env_step_end", - duration_ms=round(step_total_ms, 3), - exec_tool_calls_ms=round(exec_tool_calls_ms, 3), - action_tool_calls_count=len(action.tool_calls), - obs_messages_count=len(obs), - obs_payload_chars=obs_payload_chars, - slow=step_total_ms >= PERF_SLOW_STEP_MS, - ) return obs, reward, self.state.done, False # ========== Tools ========== @@ -1497,12 +1384,9 @@ async def run_cell( cell_idx = None # Execute code and update notebook atomically - exec_start = time.perf_counter() result, actual_cell_idx = await self.state.execute_and_add_cell( code, cell_idx=cell_idx, timeout=effective_timeout ) - execute_and_add_cell_ms = (time.perf_counter() - exec_start) * 1000.0 - _record_perf_metric("execute_and_add_cell_ms", execute_and_add_cell_ms) # Build response with cell number cell_info = f"[Cell #{actual_cell_idx}] " @@ -1524,7 +1408,6 @@ async def reset_kernel(self) -> str: This clears all variables and execution state. """ - start = time.perf_counter() if self.state.use_ray: reset_ref = self.state.kernel_container._reset_via_http.remote() await reset_ref @@ -1540,14 +1423,6 @@ async def reset_kernel(self) -> str: self.state.notebook_runtime_errors = [] self.state._execution_count = 0 - duration_ms = (time.perf_counter() - start) * 1000.0 - _record_perf_metric("tool_call_ms.reset_kernel", duration_ms) - if self._perf_sample_step: - self._emit_perf_event( - "tool_call_end", - tool_name="reset_kernel", - duration_ms=round(duration_ms, 3), - ) return "Kernel reset successfully." @property @@ -1611,10 +1486,7 @@ async def submit_answer(self, answer: str) -> str: Args: answer: Your final response to the research question """ - start = time.perf_counter() if self.state.done: - duration_ms = (time.perf_counter() - start) * 1000.0 - _record_perf_metric("tool_call_ms.submit_answer", duration_ms) return "Episode already finished." self.state.answer = answer @@ -1622,19 +1494,9 @@ async def submit_answer(self, answer: str) -> str: if self.rubric_model is None: self.logger.warning("No rubric_model configured, skipping scoring") - duration_ms = (time.perf_counter() - start) * 1000.0 - _record_perf_metric("tool_call_ms.submit_answer", duration_ms) return answer correct = await self._score_solution(answer) - duration_ms = (time.perf_counter() - start) * 1000.0 - _record_perf_metric("tool_call_ms.submit_answer", duration_ms) - if self._perf_sample_step: - self._emit_perf_event( - "tool_call_end", - tool_name="submit_answer", - duration_ms=round(duration_ms, 3), - ) return CORRECT_MSG if correct else INCORRECT_MSG def list_dir(self) -> str: From 3cde42bbe208241d6d91c4fd7fa512d172137bcd Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Fri, 13 Mar 2026 18:45:20 -0700 Subject: [PATCH 14/31] tweaked dataset server and problem instance to support new data format --- src/hypotest/dataset_server.py | 2 +- src/hypotest/env/interpreter_env.py | 9 ++++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/src/hypotest/dataset_server.py b/src/hypotest/dataset_server.py index e5662be..c9e4212 100644 --- a/src/hypotest/dataset_server.py +++ b/src/hypotest/dataset_server.py @@ -99,7 +99,7 @@ def __len__(self) -> int: class ServerConfig(BaseModel): model_config = ConfigDict(extra="forbid") - dataset: DatasetConfig + dataset: HypotestDatasetConfig api_key: str = Field( description="API key to access the server; passed either by value or as an environment variable." ) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index b35bf85..2d2dc56 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -38,7 +38,7 @@ from aviary.env import Environment from lmi import LiteLLMModel from nbformat import NotebookNode -from pydantic import BaseModel, Field, JsonValue +from pydantic import BaseModel, Field, JsonValue, model_validator import subprocess from textwrap import dedent @@ -160,6 +160,13 @@ class ProblemInstance(BaseModel): metadata: dict[str, JsonValue] = Field(default_factory=dict) nb_primary_language: str = Field(default=str(NBLanguage.PYTHON)) + @model_validator(mode="before") + @classmethod + def handle_language(cls, data: dict) -> dict: + if data['nb_primary_language'] is None: + data['nb_primary_language'] = str(NBLanguage.PYTHON) + return data + def _prep_workspace_dir(work_dir: str, workspace_path: str = "/data_workspace") -> None: wd = Path(work_dir) (wd / "pydeps").mkdir(parents=True, exist_ok=True) From 75a7d6a30ab421f07c6be1e15db917eb044e8dce Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Sat, 14 Mar 2026 01:56:06 +0000 Subject: [PATCH 15/31] added rw filesystem to enroot container --- src/hypotest/env/interpreter_env.py | 232 +++++++++++++++++++++++++++- 1 file changed, 229 insertions(+), 3 deletions(-) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index 2d2dc56..3b46b2a 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -14,7 +14,10 @@ import os import shutil import socket +import statistics +import threading import time +import random from pathlib import Path from tempfile import mkdtemp from typing import TYPE_CHECKING, Any, cast @@ -148,6 +151,72 @@ def _kill_process_group(proc: subprocess.Popen, label: str = "enroot", sigterm_t except subprocess.TimeoutExpired: logger.error("[%s] Process pid=%d still alive after SIGKILL — possible zombie", label, proc.pid) +PERF_METRICS_ENABLED = os.getenv("NEMO_GYM_PERF_METRICS", "0") == "1" +PERF_DETAIL_SAMPLE_RATE = float(os.getenv("NEMO_GYM_PERF_DETAIL_SAMPLE_RATE", "0.2")) +PERF_SLOW_STEP_MS = float(os.getenv("NEMO_GYM_PERF_SLOW_STEP_MS", "5000")) +PERF_AGGREGATE_INTERVAL_S = float(os.getenv("NEMO_GYM_PERF_AGGREGATE_INTERVAL_S", "30")) +PERF_LOG_LEVEL = getattr(logging, os.getenv("NEMO_GYM_PERF_LOG_LEVEL", "WARNING").upper(), logging.WARNING) + +_PERF_METRIC_SAMPLES: dict[str, list[float]] = {} +_PERF_LAST_FLUSH_TS = time.perf_counter() +_PERF_LOCK = threading.Lock() + + +def _log_perf(payload: dict[str, Any]) -> None: + logger.log(PERF_LOG_LEVEL, json.dumps(payload)) + + +def _percentile(sorted_values: list[float], pct: float) -> float: + if not sorted_values: + return 0.0 + if len(sorted_values) == 1: + return sorted_values[0] + idx = (len(sorted_values) - 1) * pct + lo = int(idx) + hi = min(lo + 1, len(sorted_values) - 1) + frac = idx - lo + return sorted_values[lo] * (1.0 - frac) + sorted_values[hi] * frac + + +def _maybe_flush_perf_aggregates() -> None: + global _PERF_LAST_FLUSH_TS, _PERF_METRIC_SAMPLES + now = time.perf_counter() + if now - _PERF_LAST_FLUSH_TS < PERF_AGGREGATE_INTERVAL_S: + return + + with _PERF_LOCK: + now = time.perf_counter() + if now - _PERF_LAST_FLUSH_TS < PERF_AGGREGATE_INTERVAL_S: + return + snapshots = _PERF_METRIC_SAMPLES + _PERF_METRIC_SAMPLES = {} + _PERF_LAST_FLUSH_TS = now + + for metric_name, values in snapshots.items(): + if not values: + continue + ordered = sorted(values) + payload = { + "component": "aviary_hypotest.interpreter_env", + "event": "perf_aggregate", + "metric": metric_name, + "count": len(values), + "mean_ms": round(statistics.fmean(values), 3), + "p50_ms": round(_percentile(ordered, 0.50), 3), + "p90_ms": round(_percentile(ordered, 0.90), 3), + "p95_ms": round(_percentile(ordered, 0.95), 3), + "p99_ms": round(_percentile(ordered, 0.99), 3), + "max_ms": round(ordered[-1], 3), + } + _log_perf(payload) + + +def _record_perf_metric(metric_name: str, value_ms: float) -> None: + if not PERF_METRICS_ENABLED: + return + with _PERF_LOCK: + _PERF_METRIC_SAMPLES.setdefault(metric_name, []).append(value_ms) + _maybe_flush_perf_aggregates() class ProblemInstance(BaseModel): id: UUID @@ -259,9 +328,39 @@ async def initialize(self, work_dir: Path, language: NBLanguage) -> None: kernel_server_path = env_dir / "kernel_server.py" assert kernel_server_path.is_file(), f"kernel server must be a valid path, found {kernel_server_path}" + base = Path(f"/tmp/enroot_data/{startup_token}") + + enroot_runtime = base / "runtime" + enroot_config = base / "config" + enroot_cache = base / "cache" + enroot_data = base / "data" + enroot_temp = base / "tmp" + + for path in [ + enroot_runtime, + enroot_config, + enroot_cache, + enroot_data, + enroot_temp, + ]: + path.mkdir(parents=True, exist_ok=True) + os.chmod(path, 0o700) + + env = {} + env["ENROOT_RUNTIME_PATH"] = str(enroot_runtime) + env["ENROOT_CONFIG_PATH"] = str(enroot_config) + env["ENROOT_CACHE_PATH"] = str(enroot_cache) + env["ENROOT_DATA_PATH"] = str(enroot_data) + env["ENROOT_TEMP_PATH"] = str(enroot_temp) + cmd = [ "env", "-i", "PATH=/usr/sbin:/usr/bin:/sbin:/bin", 'HOME="$HOME"', 'USER="$USER"', - "enroot", "start", + f"ENROOT_RUNTIME_PATH={enroot_runtime}", + f"ENROOT_CONFIG_PATH={enroot_config}", + f"ENROOT_CACHE_PATH={enroot_cache}", + f"ENROOT_DATA_PATH={enroot_data}", + f"ENROOT_TEMP_PATH={enroot_temp}", + "enroot", "start", "--rw", "--mount", f"{work_dir}:/data_workspace", "--mount", f"{kernel_server_path.resolve()}:/envs/kernel_server.py", str(self.container_sqsh_path.resolve()), @@ -1157,14 +1256,30 @@ async def execute_and_add_cell( actual_idx = cell_idx return result, actual_idx + execute_start = time.perf_counter() if self.use_ray and self.use_enroot: + rpc_wait_start = time.perf_counter() result_ref = self.kernel_container._execute_via_http.remote(code, timeout) result = await result_ref + _record_perf_metric( + "execute_path_ms.ray_enroot_rpc_wait", + (time.perf_counter() - rpc_wait_start) * 1000.0, + ) elif self.use_docker or self.use_enroot: + http_exec_start = time.perf_counter() result = await self._execute_via_http(code, timeout) + _record_perf_metric( + "execute_path_ms.http_kernel_execute", + (time.perf_counter() - http_exec_start) * 1000.0, + ) else: assert self.interpreter is not None + local_exec_start = time.perf_counter() result = await self.interpreter.execute_code(code, timeout) + _record_perf_metric( + "execute_path_ms.local_interpreter_execute", + (time.perf_counter() - local_exec_start) * 1000.0, + ) if cell_idx is None or cell_idx >= len(self.nb.cells): actual_idx = self._add_cell(code, result) @@ -1172,6 +1287,7 @@ async def execute_and_add_cell( self._update_cell(cell_idx, code, result) actual_idx = cell_idx + _record_perf_metric("execute_and_add_cell_total_ms", (time.perf_counter() - execute_start) * 1000.0) return result, actual_idx @@ -1223,6 +1339,8 @@ def __init__( self.execution_timeout = self.execution_config.cell_execution_timeout self.max_steps = self.config.max_steps + logger.error(f"EXECUTION TIMEOUT SET TO {self.execution_timeout} SECONDS") + self.input_data = input_data self.output_data: list[dict[str, str | int]] = [] self.logger = logger @@ -1232,6 +1350,19 @@ def __init__( self.state: InterpreterEnvState # prompting_config is set during reset() after language resolution self.prompting_config: PromptingConfig + self._perf_sample_step = False + + def _emit_perf_event(self, event: str, **fields: Any) -> None: + if not PERF_METRICS_ENABLED: + return + payload: dict[str, Any] = { + "component": "aviary_hypotest.interpreter_env", + "event": event, + "env_id": getattr(self, "_nemo_env_id", None), + "step_idx": self.step_count, + } + payload.update(fields) + _log_perf(payload) @property def language(self) -> NBLanguage: @@ -1289,7 +1420,8 @@ async def reset(self) -> tuple[Messages, list[Tool]]: self._filesystem_tool = FilesystemTool(self.work_dir) self.tools = [ - Tool.from_function(self.run_cell), + # Tool.from_function(self.run_cell), + Tool.from_function(self.stateful_python_code_exec), Tool.from_function(self.reset_kernel), Tool.from_function(self.submit_answer), Tool.from_function(self.list_dir), @@ -1316,10 +1448,15 @@ async def reset(self) -> tuple[Messages, list[Tool]]: async def step(self, action: ToolRequestMessage) -> tuple[Messages, float, bool, bool]: """Execute a step in the environment.""" self.step_count += 1 + self._perf_sample_step = PERF_METRICS_ENABLED and random.random() < PERF_DETAIL_SAMPLE_RATE + step_start = time.perf_counter() + exec_tool_calls_start = time.perf_counter() obs = cast( Messages, await self.exec_tool_calls(action, concurrency=False, handle_tool_exc=True), ) + exec_tool_calls_ms = (time.perf_counter() - exec_tool_calls_start) * 1000.0 + _record_perf_metric("exec_tool_calls_ms", exec_tool_calls_ms) obs = [*obs] if self.include_env_state_msg: @@ -1334,10 +1471,64 @@ async def step(self, action: ToolRequestMessage) -> tuple[Messages, float, bool, self.state.actions.append(str(action)) reward = self.state.score if self.state.done else 0.0 + step_total_ms = (time.perf_counter() - step_start) * 1000.0 + _record_perf_metric("env_step_ms", step_total_ms) + should_log_detail = self._perf_sample_step or step_total_ms >= PERF_SLOW_STEP_MS + if should_log_detail: + obs_payload_chars = sum(len(str(getattr(m, "content", ""))) for m in obs) + self._emit_perf_event( + "env_step_end", + duration_ms=round(step_total_ms, 3), + exec_tool_calls_ms=round(exec_tool_calls_ms, 3), + action_tool_calls_count=len(action.tool_calls), + obs_messages_count=len(obs), + obs_payload_chars=obs_payload_chars, + slow=step_total_ms >= PERF_SLOW_STEP_MS, + ) return obs, reward, self.state.done, False # ========== Tools ========== + async def stateful_python_code_exec( + self, + code: str, + idx: int | None = None, + ) -> Message | str | list[dict[str, Any]]: + """Run code in a Python Jupyter notebook cell and return the execution output. + + This method allows running code in a new cell (append) or re-running + an existing cell with updated code. + + Error Recovery: + When a cell fails with an error, you MUST fix it by calling stateful_python_code_exec + with the corrected code and the SAME idx as the failed cell. + + The cell number is shown in the output prefix (e.g., "[Cell #3]"). + Do NOT create a new cell to fix an error - always edit the failed cell. + + Args: + code: Code to execute + idx: Cell index to run. If None or >= len(cells), appends a new cell. + If provided, updates and re-runs the existing cell at that index. + Use this to fix errors in existing cells. + + Returns: + Message with multimodal content if images present, otherwise string. + The response includes the cell number (e.g., "[Cell #0] output..."). + """ + start = time.perf_counter() + result = await self.run_cell(code, idx=idx) + duration_ms = (time.perf_counter() - start) * 1000.0 + _record_perf_metric("tool_call_ms.stateful_python_code_exec", duration_ms) + if self._perf_sample_step: + self._emit_perf_event( + "tool_call_end", + tool_name="stateful_python_code_exec", + duration_ms=round(duration_ms, 3), + ) + return result + + async def run_cell( self, code: str, @@ -1391,9 +1582,12 @@ async def run_cell( cell_idx = None # Execute code and update notebook atomically + exec_start = time.perf_counter() result, actual_cell_idx = await self.state.execute_and_add_cell( code, cell_idx=cell_idx, timeout=effective_timeout ) + execute_and_add_cell_ms = (time.perf_counter() - exec_start) * 1000.0 + _record_perf_metric("execute_and_add_cell_ms", execute_and_add_cell_ms) # Build response with cell number cell_info = f"[Cell #{actual_cell_idx}] " @@ -1415,6 +1609,7 @@ async def reset_kernel(self) -> str: This clears all variables and execution state. """ + start = time.perf_counter() if self.state.use_ray: reset_ref = self.state.kernel_container._reset_via_http.remote() await reset_ref @@ -1430,6 +1625,14 @@ async def reset_kernel(self) -> str: self.state.notebook_runtime_errors = [] self.state._execution_count = 0 + duration_ms = (time.perf_counter() - start) * 1000.0 + _record_perf_metric("tool_call_ms.reset_kernel", duration_ms) + if self._perf_sample_step: + self._emit_perf_event( + "tool_call_end", + tool_name="reset_kernel", + duration_ms=round(duration_ms, 3), + ) return "Kernel reset successfully." @property @@ -1493,7 +1696,10 @@ async def submit_answer(self, answer: str) -> str: Args: answer: Your final response to the research question """ + start = time.perf_counter() if self.state.done: + duration_ms = (time.perf_counter() - start) * 1000.0 + _record_perf_metric("tool_call_ms.submit_answer", duration_ms) return "Episode already finished." self.state.answer = answer @@ -1501,9 +1707,19 @@ async def submit_answer(self, answer: str) -> str: if self.rubric_model is None: self.logger.warning("No rubric_model configured, skipping scoring") + duration_ms = (time.perf_counter() - start) * 1000.0 + _record_perf_metric("tool_call_ms.submit_answer", duration_ms) return answer correct = await self._score_solution(answer) + duration_ms = (time.perf_counter() - start) * 1000.0 + _record_perf_metric("tool_call_ms.submit_answer", duration_ms) + if self._perf_sample_step: + self._emit_perf_event( + "tool_call_end", + tool_name="submit_answer", + duration_ms=round(duration_ms, 3), + ) return CORRECT_MSG if correct else INCORRECT_MSG def list_dir(self) -> str: @@ -1518,7 +1734,17 @@ def list_dir(self) -> str: max_files: Maximum number of files to display (default: 20) show_hidden: Whether to show hidden files starting with '.' (default: False) """ - return self._filesystem_tool.list_dir() + start = time.perf_counter() + result = self._filesystem_tool.list_dir() + duration_ms = (time.perf_counter() - start) * 1000.0 + _record_perf_metric("tool_call_ms.list_dir", duration_ms) + if self._perf_sample_step: + self._emit_perf_event( + "tool_call_end", + tool_name="list_dir", + duration_ms=round(duration_ms, 3), + ) + return result # ========== Time Management ========== From c8ecc7d070029864748f2c6e6a5c823ae76422fa Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Sat, 14 Mar 2026 06:34:04 +0000 Subject: [PATCH 16/31] fixed isolation bug --- src/hypotest/dataset_server.py | 2 ++ src/hypotest/env/interpreter_env.py | 10 ++++++++++ 2 files changed, 12 insertions(+) diff --git a/src/hypotest/dataset_server.py b/src/hypotest/dataset_server.py index c9e4212..1a471ad 100644 --- a/src/hypotest/dataset_server.py +++ b/src/hypotest/dataset_server.py @@ -70,6 +70,8 @@ def get_new_env_by_idx(self, idx: int) -> InterpreterEnv: if not capsule_path.exists(): capsule_path = self.config.capsule_dir / f"CapsuleData-{problem.id}" problem_dir = Path(self.config.work_dir) / run_id if self.config.work_dir else Path(mkdtemp()) + if problem_dir.exists(): + shutil.rmtree(problem_dir) problem_dir.mkdir(parents=True, exist_ok=True) shutil.copytree(capsule_path, problem_dir, dirs_exist_ok=True) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index 3b46b2a..4c78d9d 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -1144,6 +1144,9 @@ async def close(self): elif self.interpreter is not None: await self.interpreter.close() + if self.save_dir is None and self.work_dir.exists(): + shutil.rmtree(self.work_dir, ignore_errors=True) + def _add_cell(self, code: str, result: "ExecutionResult") -> int: """Add a new code cell to the notebook with execution results. @@ -1352,6 +1355,13 @@ def __init__( self.prompting_config: PromptingConfig self._perf_sample_step = False + if self.score_info_path.exists(): + self.score_info_path.unlink() + + nb_path = self.work_dir / "notebook.ipynb" + if nb_path.exists(): + nb_path.unlink() + def _emit_perf_event(self, event: str, **fields: Any) -> None: if not PERF_METRICS_ENABLED: return From d8334d3298b4f1f26f0adf461078959c2b936184 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Sat, 14 Mar 2026 13:15:21 -0700 Subject: [PATCH 17/31] cleaned up perf logging --- src/hypotest/env/interpreter_env.py | 412 +++++++++++----------------- src/hypotest/env/kernel_server.py | 57 +--- 2 files changed, 165 insertions(+), 304 deletions(-) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index 4c78d9d..06e3241 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -4,30 +4,30 @@ code in Jupyter kernels. It focuses on direct code execution via run_cell(). """ -import sys -import ray -import signal import argparse import asyncio +import contextlib import json import logging import os import shutil +import signal import socket -import statistics -import threading +import subprocess +import sys import time -import random +import uuid from pathlib import Path from tempfile import mkdtemp +from textwrap import dedent from typing import TYPE_CHECKING, Any, cast from uuid import UUID -import uuid import aiodocker import httpx import nbformat import numpy as np +import ray import tenacity from aviary.core import ( EnvStateMessage, @@ -43,9 +43,6 @@ from nbformat import NotebookNode from pydantic import BaseModel, Field, JsonValue, model_validator -import subprocess -from textwrap import dedent - from . import config as cfg from .config import ExecutionConfig from .interpreter import ExecutionResult, Interpreter @@ -67,6 +64,7 @@ _RETRY_BASE_SLEEP = 1.0 _RETRY_MAX_SLEEP = 16.0 + async def get_free_port() -> int: """Get a free port for the kernel server container.""" while True: @@ -90,6 +88,7 @@ async def get_free_port() -> int: class _PortCollisionError(Exception): """Port already in use by another server — retry with a new port.""" + def _kill_process_group(proc: subprocess.Popen, label: str = "enroot", sigterm_timeout: float = 15) -> None: """Safely terminate a process group, escalating from SIGTERM to SIGKILL. @@ -99,7 +98,9 @@ def _kill_process_group(proc: subprocess.Popen, label: str = "enroot", sigterm_t logger.log( _CONTAINER_LOG_LEVEL, "[%s] Process pid=%d already exited with returncode=%d", - label, proc.pid, proc.returncode, + label, + proc.pid, + proc.returncode, ) return @@ -124,13 +125,17 @@ def _kill_process_group(proc: subprocess.Popen, label: str = "enroot", sigterm_t logger.log( _CONTAINER_LOG_LEVEL, "[%s] Process pid=%d exited after SIGTERM with returncode=%d", - label, proc.pid, proc.returncode, + label, + proc.pid, + proc.returncode, ) return except subprocess.TimeoutExpired: logger.warning( "[%s] Process pid=%d did not exit within %.1fs of SIGTERM, sending SIGKILL", - label, proc.pid, sigterm_timeout, + label, + proc.pid, + sigterm_timeout, ) # SIGKILL the whole group @@ -146,77 +151,13 @@ def _kill_process_group(proc: subprocess.Popen, label: str = "enroot", sigterm_t logger.log( _CONTAINER_LOG_LEVEL, "[%s] Process pid=%d exited after SIGKILL with returncode=%d", - label, proc.pid, proc.returncode, + label, + proc.pid, + proc.returncode, ) except subprocess.TimeoutExpired: - logger.error("[%s] Process pid=%d still alive after SIGKILL — possible zombie", label, proc.pid) + logger.exception("[%s] Process pid=%d still alive after SIGKILL — possible zombie", label, proc.pid) -PERF_METRICS_ENABLED = os.getenv("NEMO_GYM_PERF_METRICS", "0") == "1" -PERF_DETAIL_SAMPLE_RATE = float(os.getenv("NEMO_GYM_PERF_DETAIL_SAMPLE_RATE", "0.2")) -PERF_SLOW_STEP_MS = float(os.getenv("NEMO_GYM_PERF_SLOW_STEP_MS", "5000")) -PERF_AGGREGATE_INTERVAL_S = float(os.getenv("NEMO_GYM_PERF_AGGREGATE_INTERVAL_S", "30")) -PERF_LOG_LEVEL = getattr(logging, os.getenv("NEMO_GYM_PERF_LOG_LEVEL", "WARNING").upper(), logging.WARNING) - -_PERF_METRIC_SAMPLES: dict[str, list[float]] = {} -_PERF_LAST_FLUSH_TS = time.perf_counter() -_PERF_LOCK = threading.Lock() - - -def _log_perf(payload: dict[str, Any]) -> None: - logger.log(PERF_LOG_LEVEL, json.dumps(payload)) - - -def _percentile(sorted_values: list[float], pct: float) -> float: - if not sorted_values: - return 0.0 - if len(sorted_values) == 1: - return sorted_values[0] - idx = (len(sorted_values) - 1) * pct - lo = int(idx) - hi = min(lo + 1, len(sorted_values) - 1) - frac = idx - lo - return sorted_values[lo] * (1.0 - frac) + sorted_values[hi] * frac - - -def _maybe_flush_perf_aggregates() -> None: - global _PERF_LAST_FLUSH_TS, _PERF_METRIC_SAMPLES - now = time.perf_counter() - if now - _PERF_LAST_FLUSH_TS < PERF_AGGREGATE_INTERVAL_S: - return - - with _PERF_LOCK: - now = time.perf_counter() - if now - _PERF_LAST_FLUSH_TS < PERF_AGGREGATE_INTERVAL_S: - return - snapshots = _PERF_METRIC_SAMPLES - _PERF_METRIC_SAMPLES = {} - _PERF_LAST_FLUSH_TS = now - - for metric_name, values in snapshots.items(): - if not values: - continue - ordered = sorted(values) - payload = { - "component": "aviary_hypotest.interpreter_env", - "event": "perf_aggregate", - "metric": metric_name, - "count": len(values), - "mean_ms": round(statistics.fmean(values), 3), - "p50_ms": round(_percentile(ordered, 0.50), 3), - "p90_ms": round(_percentile(ordered, 0.90), 3), - "p95_ms": round(_percentile(ordered, 0.95), 3), - "p99_ms": round(_percentile(ordered, 0.99), 3), - "max_ms": round(ordered[-1], 3), - } - _log_perf(payload) - - -def _record_perf_metric(metric_name: str, value_ms: float) -> None: - if not PERF_METRICS_ENABLED: - return - with _PERF_LOCK: - _PERF_METRIC_SAMPLES.setdefault(metric_name, []).append(value_ms) - _maybe_flush_perf_aggregates() class ProblemInstance(BaseModel): id: UUID @@ -232,10 +173,11 @@ class ProblemInstance(BaseModel): @model_validator(mode="before") @classmethod def handle_language(cls, data: dict) -> dict: - if data['nb_primary_language'] is None: - data['nb_primary_language'] = str(NBLanguage.PYTHON) + if data["nb_primary_language"] is None: + data["nb_primary_language"] = str(NBLanguage.PYTHON) return data + def _prep_workspace_dir(work_dir: str, workspace_path: str = "/data_workspace") -> None: wd = Path(work_dir) (wd / "pydeps").mkdir(parents=True, exist_ok=True) @@ -249,6 +191,7 @@ def _prep_workspace_dir(work_dir: str, workspace_path: str = "/data_workspace") f"target = {workspace_path}/pydeps\n" ) + @ray.remote( scheduling_strategy="SPREAD", runtime_env={ @@ -273,7 +216,7 @@ def _proc_label(self) -> str: async def initialize(self, work_dir: Path, language: NBLanguage) -> None: startup_token = str(uuid.uuid4()) - node_workdir = f"/tmp/data_workspace.{startup_token.split('-')[0]}" + node_workdir = f"/tmp/data_workspace.{startup_token.split('-', maxsplit=1)[0]}" _prep_workspace_dir(work_dir, workspace_path=node_workdir) @@ -290,7 +233,7 @@ async def initialize(self, work_dir: Path, language: NBLanguage) -> None: f"{f' log_tail={log_tail!r}' if log_tail else ''}" ) self._container_port = await get_free_port() - + bash = dedent(f"""\ set -euo pipefail @@ -354,17 +297,27 @@ async def initialize(self, work_dir: Path, language: NBLanguage) -> None: env["ENROOT_TEMP_PATH"] = str(enroot_temp) cmd = [ - "env", "-i", "PATH=/usr/sbin:/usr/bin:/sbin:/bin", 'HOME="$HOME"', 'USER="$USER"', + "env", + "-i", + "PATH=/usr/sbin:/usr/bin:/sbin:/bin", + 'HOME="$HOME"', + 'USER="$USER"', f"ENROOT_RUNTIME_PATH={enroot_runtime}", f"ENROOT_CONFIG_PATH={enroot_config}", f"ENROOT_CACHE_PATH={enroot_cache}", f"ENROOT_DATA_PATH={enroot_data}", f"ENROOT_TEMP_PATH={enroot_temp}", - "enroot", "start", "--rw", - "--mount", f"{work_dir}:/data_workspace", - "--mount", f"{kernel_server_path.resolve()}:/envs/kernel_server.py", + "enroot", + "start", + "--rw", + "--mount", + f"{work_dir}:/data_workspace", + "--mount", + f"{kernel_server_path.resolve()}:/envs/kernel_server.py", str(self.container_sqsh_path.resolve()), - "/bin/bash", "-lc", bash, + "/bin/bash", + "-lc", + bash, ] async with CONTAINER_LAUNCH_SEM: @@ -377,15 +330,21 @@ async def initialize(self, work_dir: Path, language: NBLanguage) -> None: log_dir = work_dir / ".container_logs" log_dir.mkdir(exist_ok=True) self._container_log_path = log_dir / "container.log" - self._container_log_file = open(self._container_log_path, "w") # noqa: SIM115 + self._container_log_file = open(self._container_log_path, "w", encoding="utf-8") # noqa: SIM115 self._enroot_proc = subprocess.Popen( - cmd, text=True, start_new_session=True, - stdout=self._container_log_file, stderr=subprocess.STDOUT, + cmd, + text=True, + start_new_session=True, + stdout=self._container_log_file, + stderr=subprocess.STDOUT, ) logger.log( _CONTAINER_LOG_LEVEL, "[%s] Container launch attempt #%d started (work_dir=%s, token=%s)", - self._proc_label(), attempt, work_dir, startup_token[:8], + self._proc_label(), + attempt, + work_dir, + startup_token[:8], ) # Create HTTP client (outside semaphore — no need to hold the @@ -402,7 +361,9 @@ async def initialize(self, work_dir: Path, language: NBLanguage) -> None: logger.log( _CONTAINER_LOG_LEVEL, "[%s] Container online after %.1fms (attempt #%d)", - self._proc_label(), launch_ms, attempt, + self._proc_label(), + launch_ms, + attempt, ) online = True except Exception as e: @@ -438,7 +399,9 @@ def _log_container_failure(self, attempt: int, launch_ms: float, error: Exceptio if log_tail: logger.warning( "[%s] Container log output (last %d chars):\n%s", - label, len(log_tail), log_tail, + label, + len(log_tail), + log_tail, ) def _read_container_log_tail(self, max_chars: int = 2000) -> str: @@ -457,10 +420,8 @@ def _read_container_log_tail(self, max_chars: int = 2000) -> str: def _close_container_log(self) -> None: """Close the container log file handle.""" if self._container_log_file is not None: - try: + with contextlib.suppress(Exception): self._container_log_file.close() - except Exception: - pass self._container_log_file = None async def _cleanup_failed_startup(self) -> None: @@ -473,10 +434,8 @@ async def _cleanup_failed_startup(self) -> None: self._container_port = None if self._http_client is not None: - try: + with contextlib.suppress(Exception): await self._http_client.aclose() - except Exception: - pass self._http_client = None if self._enroot_proc is not None: @@ -507,8 +466,7 @@ async def _wait_for_health(self, expected_startup_token: str | None = None) -> N rc = self._enroot_proc.returncode log_tail = self._read_container_log_tail(1000) raise RuntimeError( - f"Enroot process exited prematurely with returncode={rc} " - f"after {elapsed:.1f}s. log: {log_tail!r}" + f"Enroot process exited prematurely with returncode={rc} after {elapsed:.1f}s. log: {log_tail!r}" ) try: @@ -519,7 +477,9 @@ async def _wait_for_health(self, expected_startup_token: str | None = None) -> N logger.log( _CONTAINER_LOG_LEVEL, "[%s] Kernel server healthy after %.1fs (%d polls)", - self._proc_label(), elapsed, poll_count, + self._proc_label(), + elapsed, + poll_count, ) return payload = response.json() @@ -527,17 +487,18 @@ async def _wait_for_health(self, expected_startup_token: str | None = None) -> N logger.log( _CONTAINER_LOG_LEVEL, "[%s] Kernel server healthy (token matched) after %.1fs (%d polls)", - self._proc_label(), elapsed, poll_count, + self._proc_label(), + elapsed, + poll_count, ) return - else: - last_status = f"token_mismatch(got={payload.get('startup_token', '?')[:8]})" - consecutive_token_mismatches += 1 - if consecutive_token_mismatches >= 3: - raise _PortCollisionError( - f"Port {self._container_port} appears to be owned by another server " - f"({consecutive_token_mismatches} consecutive token mismatches)" - ) + last_status = f"token_mismatch(got={payload.get('startup_token', '?')[:8]})" + consecutive_token_mismatches += 1 + if consecutive_token_mismatches >= 3: + raise _PortCollisionError( + f"Port {self._container_port} appears to be owned by another server " + f"({consecutive_token_mismatches} consecutive token mismatches)" + ) else: last_status = f"http_{response.status_code}" consecutive_token_mismatches = 0 @@ -560,7 +521,11 @@ async def _wait_for_health(self, expected_startup_token: str | None = None) -> N logger.log( _CONTAINER_LOG_LEVEL, "[%s] Health poll #%d at %.1fs: last_status=%s, process_alive=%s", - self._proc_label(), poll_count, elapsed, last_status, proc_alive, + self._proc_label(), + poll_count, + elapsed, + last_status, + proc_alive, ) await asyncio.sleep(0.5) @@ -604,7 +569,10 @@ async def _execute_via_http(self, code: str, timeout: float | None = None) -> Ex effective_timeout = timeout if timeout is not None else self.execution_timeout logger.warning( "[%s] HTTP %s during /execute (requested kernel timeout=%.1fs): %s", - self._proc_label(), type(e).__name__, effective_timeout, e, + self._proc_label(), + type(e).__name__, + effective_timeout, + e, ) timeout_output = nbformat.v4.new_output( output_type="error", @@ -638,7 +606,9 @@ async def _reset_via_http(self) -> None: except httpx.TimeoutException as e: logger.warning( "[%s] HTTP %s during /reset: %s", - self._proc_label(), type(e).__name__, e, + self._proc_label(), + type(e).__name__, + e, ) raise RuntimeError(f"Kernel reset timed out: {e}") from e @@ -655,13 +625,13 @@ async def close(self): response = await self._http_client.post("/close") response.raise_for_status() except (httpx.ConnectError, httpx.ReadError, httpx.RemoteProtocolError, httpx.TimeoutException): - logger.log(_CONTAINER_LOG_LEVEL, "[%s] Graceful /close request failed (container may already be down)", label) + logger.log( + _CONTAINER_LOG_LEVEL, "[%s] Graceful /close request failed (container may already be down)", label + ) except Exception: logger.warning("[%s] Unexpected error on /close request", label, exc_info=True) - try: + with contextlib.suppress(Exception): await self._http_client.aclose() - except Exception: - pass self._http_client = None if self._enroot_proc is not None: @@ -797,12 +767,21 @@ async def _start_enroot_container(self) -> None: assert kernel_server_path.is_file(), f"kernel server must be a valid path, found {kernel_server_path}" cmd = [ - "env", "-i", "PATH=/usr/sbin:/usr/bin:/sbin:/bin", 'HOME="$HOME"', 'USER="$USER"', - "enroot", "start", - "--mount", f"{self.work_dir}:/data_workspace", - "--mount", f"{kernel_server_path.resolve()}:/envs/kernel_server.py", + "env", + "-i", + "PATH=/usr/sbin:/usr/bin:/sbin:/bin", + 'HOME="$HOME"', + 'USER="$USER"', + "enroot", + "start", + "--mount", + f"{self.work_dir}:/data_workspace", + "--mount", + f"{kernel_server_path.resolve()}:/envs/kernel_server.py", str(self.container_sqsh_path.resolve()), - "/bin/bash", "-lc", bash, + "/bin/bash", + "-lc", + bash, ] async with CONTAINER_LAUNCH_SEM: @@ -812,15 +791,21 @@ async def _start_enroot_container(self) -> None: log_dir = self.work_dir / ".container_logs" log_dir.mkdir(exist_ok=True) self._container_log_path = log_dir / "container.log" - self._container_log_file = open(self._container_log_path, "w") # noqa: SIM115 + self._container_log_file = open(self._container_log_path, "w", encoding="utf-8") # noqa: SIM115 self._enroot_proc = subprocess.Popen( - cmd, text=True, start_new_session=True, - stdout=self._container_log_file, stderr=subprocess.STDOUT, + cmd, + text=True, + start_new_session=True, + stdout=self._container_log_file, + stderr=subprocess.STDOUT, ) logger.log( _CONTAINER_LOG_LEVEL, "[%s] Container launch attempt #%d (work_dir=%s, token=%s)", - self._enroot_label(), attempt, self.work_dir, startup_token[:8], + self._enroot_label(), + attempt, + self.work_dir, + startup_token[:8], ) # Create HTTP client (outside semaphore — no need to hold the @@ -837,7 +822,9 @@ async def _start_enroot_container(self) -> None: logger.log( _CONTAINER_LOG_LEVEL, "[%s] Container online after %.1fms (attempt #%d)", - self._enroot_label(), launch_ms, attempt, + self._enroot_label(), + launch_ms, + attempt, ) online = True except Exception as e: @@ -855,7 +842,9 @@ async def _start_enroot_container(self) -> None: if log_tail: logger.warning( "[%s] Container log output (last %d chars):\n%s", - label, len(log_tail), log_tail, + label, + len(log_tail), + log_tail, ) await self._cleanup_failed_startup() if not isinstance(e, _PortCollisionError): @@ -877,10 +866,8 @@ def _read_container_log_tail(self, max_chars: int = 2000) -> str: def _close_container_log(self) -> None: """Close the container log file handle.""" if self._container_log_file is not None: - try: + with contextlib.suppress(Exception): self._container_log_file.close() - except Exception: - pass self._container_log_file = None async def _cleanup_failed_startup(self) -> None: @@ -893,10 +880,8 @@ async def _cleanup_failed_startup(self) -> None: self._container_port = None if self._http_client is not None: - try: + with contextlib.suppress(Exception): await self._http_client.aclose() - except Exception: - pass self._http_client = None if self._enroot_proc is not None: @@ -963,8 +948,7 @@ async def _wait_for_health(self, expected_startup_token: str | None = None) -> N rc = self._enroot_proc.returncode log_tail = self._read_container_log_tail(1000) raise RuntimeError( - f"Enroot process exited prematurely with returncode={rc} " - f"after {elapsed:.1f}s. log: {log_tail!r}" + f"Enroot process exited prematurely with returncode={rc} after {elapsed:.1f}s. log: {log_tail!r}" ) try: @@ -975,7 +959,9 @@ async def _wait_for_health(self, expected_startup_token: str | None = None) -> N logger.log( _CONTAINER_LOG_LEVEL, "[%s] Kernel server healthy after %.1fs (%d polls)", - self._enroot_label(), elapsed, poll_count, + self._enroot_label(), + elapsed, + poll_count, ) return payload = response.json() @@ -983,17 +969,18 @@ async def _wait_for_health(self, expected_startup_token: str | None = None) -> N logger.log( _CONTAINER_LOG_LEVEL, "[%s] Kernel server healthy (token matched) after %.1fs (%d polls)", - self._enroot_label(), elapsed, poll_count, + self._enroot_label(), + elapsed, + poll_count, ) return - else: - last_status = f"token_mismatch(got={payload.get('startup_token', '?')[:8]})" - consecutive_token_mismatches += 1 - if consecutive_token_mismatches >= 3: - raise _PortCollisionError( - f"Port {self._container_port} appears to be owned by another server " - f"({consecutive_token_mismatches} consecutive token mismatches)" - ) + last_status = f"token_mismatch(got={payload.get('startup_token', '?')[:8]})" + consecutive_token_mismatches += 1 + if consecutive_token_mismatches >= 3: + raise _PortCollisionError( + f"Port {self._container_port} appears to be owned by another server " + f"({consecutive_token_mismatches} consecutive token mismatches)" + ) else: last_status = f"http_{response.status_code}" consecutive_token_mismatches = 0 @@ -1016,7 +1003,11 @@ async def _wait_for_health(self, expected_startup_token: str | None = None) -> N logger.log( _CONTAINER_LOG_LEVEL, "[%s] Health poll #%d at %.1fs: last_status=%s, process_alive=%s", - self._enroot_label(), poll_count, elapsed, last_status, proc_alive, + self._enroot_label(), + poll_count, + elapsed, + last_status, + proc_alive, ) await asyncio.sleep(0.5) @@ -1059,7 +1050,10 @@ async def _execute_via_http(self, code: str, timeout: float | None = None) -> Ex effective_timeout = timeout if timeout is not None else self.execution_timeout logger.warning( "[%s] HTTP %s during /execute (requested kernel timeout=%.1fs): %s", - self._enroot_label(), type(e).__name__, effective_timeout, e, + self._enroot_label(), + type(e).__name__, + effective_timeout, + e, ) timeout_output = nbformat.v4.new_output( output_type="error", @@ -1092,7 +1086,9 @@ async def _reset_via_http(self) -> None: response.raise_for_status() except httpx.TimeoutException as e: logger.warning( - "HTTP %s during /reset: %s", type(e).__name__, e, + "HTTP %s during /reset: %s", + type(e).__name__, + e, ) raise RuntimeError(f"Kernel reset timed out: {e}") from e @@ -1117,10 +1113,8 @@ async def close(self): _USED_PORTS.discard(self._container_port) if self._http_client is not None: - try: + with contextlib.suppress(Exception): await self._http_client.aclose() - except Exception: - pass self._http_client = None if self._container is not None: @@ -1238,6 +1232,7 @@ async def execute_and_add_cell( """ # Safety check: block dangerous code before execution from .code_safety import check_code_safety + block_reason = check_code_safety(code, self.language) if block_reason is not None: logger.warning("Blocked code execution in execute_and_add_cell: %s", block_reason) @@ -1259,30 +1254,14 @@ async def execute_and_add_cell( actual_idx = cell_idx return result, actual_idx - execute_start = time.perf_counter() if self.use_ray and self.use_enroot: - rpc_wait_start = time.perf_counter() result_ref = self.kernel_container._execute_via_http.remote(code, timeout) result = await result_ref - _record_perf_metric( - "execute_path_ms.ray_enroot_rpc_wait", - (time.perf_counter() - rpc_wait_start) * 1000.0, - ) elif self.use_docker or self.use_enroot: - http_exec_start = time.perf_counter() result = await self._execute_via_http(code, timeout) - _record_perf_metric( - "execute_path_ms.http_kernel_execute", - (time.perf_counter() - http_exec_start) * 1000.0, - ) else: assert self.interpreter is not None - local_exec_start = time.perf_counter() result = await self.interpreter.execute_code(code, timeout) - _record_perf_metric( - "execute_path_ms.local_interpreter_execute", - (time.perf_counter() - local_exec_start) * 1000.0, - ) if cell_idx is None or cell_idx >= len(self.nb.cells): actual_idx = self._add_cell(code, result) @@ -1290,7 +1269,6 @@ async def execute_and_add_cell( self._update_cell(cell_idx, code, result) actual_idx = cell_idx - _record_perf_metric("execute_and_add_cell_total_ms", (time.perf_counter() - execute_start) * 1000.0) return result, actual_idx @@ -1342,8 +1320,6 @@ def __init__( self.execution_timeout = self.execution_config.cell_execution_timeout self.max_steps = self.config.max_steps - logger.error(f"EXECUTION TIMEOUT SET TO {self.execution_timeout} SECONDS") - self.input_data = input_data self.output_data: list[dict[str, str | int]] = [] self.logger = logger @@ -1353,7 +1329,6 @@ def __init__( self.state: InterpreterEnvState # prompting_config is set during reset() after language resolution self.prompting_config: PromptingConfig - self._perf_sample_step = False if self.score_info_path.exists(): self.score_info_path.unlink() @@ -1362,18 +1337,6 @@ def __init__( if nb_path.exists(): nb_path.unlink() - def _emit_perf_event(self, event: str, **fields: Any) -> None: - if not PERF_METRICS_ENABLED: - return - payload: dict[str, Any] = { - "component": "aviary_hypotest.interpreter_env", - "event": event, - "env_id": getattr(self, "_nemo_env_id", None), - "step_idx": self.step_count, - } - payload.update(fields) - _log_perf(payload) - @property def language(self) -> NBLanguage: return self.config.language @@ -1417,7 +1380,7 @@ async def reset(self) -> tuple[Messages, list[Tool]]: save_dir=self.save_dir, use_docker=self.config.use_docker, use_enroot=self.config.use_enroot, - container_sqsh_path=self.config.container_sqsh_path + container_sqsh_path=self.config.container_sqsh_path, ) await self.state.start() @@ -1458,15 +1421,10 @@ async def reset(self) -> tuple[Messages, list[Tool]]: async def step(self, action: ToolRequestMessage) -> tuple[Messages, float, bool, bool]: """Execute a step in the environment.""" self.step_count += 1 - self._perf_sample_step = PERF_METRICS_ENABLED and random.random() < PERF_DETAIL_SAMPLE_RATE - step_start = time.perf_counter() - exec_tool_calls_start = time.perf_counter() obs = cast( Messages, await self.exec_tool_calls(action, concurrency=False, handle_tool_exc=True), ) - exec_tool_calls_ms = (time.perf_counter() - exec_tool_calls_start) * 1000.0 - _record_perf_metric("exec_tool_calls_ms", exec_tool_calls_ms) obs = [*obs] if self.include_env_state_msg: @@ -1481,20 +1439,6 @@ async def step(self, action: ToolRequestMessage) -> tuple[Messages, float, bool, self.state.actions.append(str(action)) reward = self.state.score if self.state.done else 0.0 - step_total_ms = (time.perf_counter() - step_start) * 1000.0 - _record_perf_metric("env_step_ms", step_total_ms) - should_log_detail = self._perf_sample_step or step_total_ms >= PERF_SLOW_STEP_MS - if should_log_detail: - obs_payload_chars = sum(len(str(getattr(m, "content", ""))) for m in obs) - self._emit_perf_event( - "env_step_end", - duration_ms=round(step_total_ms, 3), - exec_tool_calls_ms=round(exec_tool_calls_ms, 3), - action_tool_calls_count=len(action.tool_calls), - obs_messages_count=len(obs), - obs_payload_chars=obs_payload_chars, - slow=step_total_ms >= PERF_SLOW_STEP_MS, - ) return obs, reward, self.state.done, False # ========== Tools ========== @@ -1526,18 +1470,7 @@ async def stateful_python_code_exec( Message with multimodal content if images present, otherwise string. The response includes the cell number (e.g., "[Cell #0] output..."). """ - start = time.perf_counter() - result = await self.run_cell(code, idx=idx) - duration_ms = (time.perf_counter() - start) * 1000.0 - _record_perf_metric("tool_call_ms.stateful_python_code_exec", duration_ms) - if self._perf_sample_step: - self._emit_perf_event( - "tool_call_end", - tool_name="stateful_python_code_exec", - duration_ms=round(duration_ms, 3), - ) - return result - + return await self.run_cell(code, idx=idx) async def run_cell( self, @@ -1592,12 +1525,9 @@ async def run_cell( cell_idx = None # Execute code and update notebook atomically - exec_start = time.perf_counter() result, actual_cell_idx = await self.state.execute_and_add_cell( code, cell_idx=cell_idx, timeout=effective_timeout ) - execute_and_add_cell_ms = (time.perf_counter() - exec_start) * 1000.0 - _record_perf_metric("execute_and_add_cell_ms", execute_and_add_cell_ms) # Build response with cell number cell_info = f"[Cell #{actual_cell_idx}] " @@ -1619,7 +1549,6 @@ async def reset_kernel(self) -> str: This clears all variables and execution state. """ - start = time.perf_counter() if self.state.use_ray: reset_ref = self.state.kernel_container._reset_via_http.remote() await reset_ref @@ -1635,14 +1564,6 @@ async def reset_kernel(self) -> str: self.state.notebook_runtime_errors = [] self.state._execution_count = 0 - duration_ms = (time.perf_counter() - start) * 1000.0 - _record_perf_metric("tool_call_ms.reset_kernel", duration_ms) - if self._perf_sample_step: - self._emit_perf_event( - "tool_call_end", - tool_name="reset_kernel", - duration_ms=round(duration_ms, 3), - ) return "Kernel reset successfully." @property @@ -1706,10 +1627,7 @@ async def submit_answer(self, answer: str) -> str: Args: answer: Your final response to the research question """ - start = time.perf_counter() if self.state.done: - duration_ms = (time.perf_counter() - start) * 1000.0 - _record_perf_metric("tool_call_ms.submit_answer", duration_ms) return "Episode already finished." self.state.answer = answer @@ -1717,19 +1635,9 @@ async def submit_answer(self, answer: str) -> str: if self.rubric_model is None: self.logger.warning("No rubric_model configured, skipping scoring") - duration_ms = (time.perf_counter() - start) * 1000.0 - _record_perf_metric("tool_call_ms.submit_answer", duration_ms) return answer correct = await self._score_solution(answer) - duration_ms = (time.perf_counter() - start) * 1000.0 - _record_perf_metric("tool_call_ms.submit_answer", duration_ms) - if self._perf_sample_step: - self._emit_perf_event( - "tool_call_end", - tool_name="submit_answer", - duration_ms=round(duration_ms, 3), - ) return CORRECT_MSG if correct else INCORRECT_MSG def list_dir(self) -> str: @@ -1738,23 +1646,13 @@ def list_dir(self) -> str: Recursively lists files in a directory, with built-in protection against overwhelming the context with too many files. Use this tool instead of writing code to list directories to avoid context bloat. - + Args: directory: Directory path to list (default: current working directory) max_files: Maximum number of files to display (default: 20) show_hidden: Whether to show hidden files starting with '.' (default: False) """ - start = time.perf_counter() - result = self._filesystem_tool.list_dir() - duration_ms = (time.perf_counter() - start) * 1000.0 - _record_perf_metric("tool_call_ms.list_dir", duration_ms) - if self._perf_sample_step: - self._emit_perf_event( - "tool_call_end", - tool_name="list_dir", - duration_ms=round(duration_ms, 3), - ) - return result + return self._filesystem_tool.list_dir() # ========== Time Management ========== diff --git a/src/hypotest/env/kernel_server.py b/src/hypotest/env/kernel_server.py index 78303e2..3c9dd2e 100644 --- a/src/hypotest/env/kernel_server.py +++ b/src/hypotest/env/kernel_server.py @@ -9,15 +9,12 @@ from __future__ import annotations -import uuid import argparse import asyncio -import json import logging -import os -import random import re import time +import uuid from enum import StrEnum, auto from pathlib import Path from queue import Empty @@ -35,15 +32,8 @@ class _DeadlineExceeded(Exception): """Raised when a cooperative deadline check expires.""" -logger = logging.getLogger(__name__) -PERF_METRICS_ENABLED = os.getenv("NEMO_GYM_PERF_METRICS", "0") == "1" -PERF_DETAIL_SAMPLE_RATE = float(os.getenv("NEMO_GYM_PERF_DETAIL_SAMPLE_RATE", "0.2")) -PERF_SLOW_KERNEL_EXEC_MS = float(os.getenv("NEMO_GYM_PERF_SLOW_KERNEL_EXEC_MS", "5000")) -PERF_LOG_LEVEL = getattr(logging, os.getenv("NEMO_GYM_PERF_LOG_LEVEL", "WARNING").upper(), logging.WARNING) - -def _log_perf(payload: dict[str, Any]) -> None: - logger.log(PERF_LOG_LEVEL, json.dumps(payload)) +logger = logging.getLogger(__name__) # ============================================================================= @@ -243,7 +233,9 @@ async def start(self) -> None: conn_uuid = uuid.uuid4() kernel_connect_file = (kernel_runtime_dir / f"conn_{conn_uuid}.json").resolve() - self._kernel_manager = AsyncKernelManager(kernel_name=kernel_name, transport='ipc', connection_file=str(kernel_connect_file)) + self._kernel_manager = AsyncKernelManager( + kernel_name=kernel_name, transport="ipc", connection_file=str(kernel_connect_file) + ) await self._kernel_manager.start_kernel(cwd=str(self.work_dir)) self._client = self._kernel_manager.client() @@ -278,15 +270,13 @@ async def execute(self, code: str, timeout: float | None = None) -> ExecuteRespo effective_timeout = timeout if timeout is not None else self.default_timeout start_time = time.perf_counter() - should_sample = PERF_METRICS_ENABLED and random.random() < PERF_DETAIL_SAMPLE_RATE - timed_out = False try: result = await self._execute_code( - code, deadline=start_time + effective_timeout, log_detail=should_sample, + code, + deadline=start_time + effective_timeout, ) except _DeadlineExceeded: - timed_out = True timeout_output = MessageType.ERROR.to_notebook_output({ "ename": "TimeoutError", "evalue": f"Code execution timed out after {effective_timeout} seconds", @@ -309,23 +299,9 @@ async def execute(self, code: str, timeout: float | None = None) -> ExecuteRespo execution_time=time.perf_counter() - start_time, ) - total_ms = (time.perf_counter() - start_time) * 1000.0 - if PERF_METRICS_ENABLED and (should_sample or total_ms >= PERF_SLOW_KERNEL_EXEC_MS): - _log_perf( - { - "component": "aviary_hypotest.kernel_server", - "event": "kernel_execute_end", - "duration_ms": round(total_ms, 3), - "timed_out": timed_out, - "error_occurred": result.error_occurred, - "notebook_outputs_count": len(result.notebook_outputs), - "slow": total_ms >= PERF_SLOW_KERNEL_EXEC_MS, - } - ) - return result - async def _execute_code(self, code: str, deadline: float, log_detail: bool = False) -> ExecuteResponse: + async def _execute_code(self, code: str, deadline: float) -> ExecuteResponse: """Internal method to execute code and collect outputs. Uses cooperative deadline checking instead of asyncio.timeout, because @@ -338,14 +314,13 @@ async def _execute_code(self, code: str, deadline: float, log_detail: bool = Fal # How long each ZMQ poll waits before we re-check the deadline. # Shorter = more responsive timeout, slightly more overhead. - _POLL_INTERVAL_S = 2.0 + POLL_INTERVAL_S = 2.0 start_time = time.perf_counter() msg_id = self._client.execute(code) notebook_outputs: list[dict[str, Any]] = [] error_occurred = False - iopub_msgs_count = 0 while True: # Check deadline before each poll @@ -356,13 +331,12 @@ async def _execute_code(self, code: str, deadline: float, log_detail: bool = Fal # get_iopub_msg(timeout=T) raises queue.Empty if no message arrives # within T seconds. try: - msg = await self._client.get_iopub_msg(timeout=_POLL_INTERVAL_S) + msg = await self._client.get_iopub_msg(timeout=POLL_INTERVAL_S) except Empty: continue if msg["parent_header"].get("msg_id") != msg_id: continue - iopub_msgs_count += 1 msg_type = MessageType.from_string(msg["msg_type"]) if msg_type is None: @@ -383,17 +357,6 @@ async def _execute_code(self, code: str, deadline: float, log_detail: bool = Fal error_occurred = True execution_time = time.perf_counter() - start_time - if PERF_METRICS_ENABLED and log_detail: - _log_perf( - { - "component": "aviary_hypotest.kernel_server", - "event": "kernel_iopub_summary", - "duration_ms": round(execution_time * 1000.0, 3), - "kernel_iopub_msgs_count": iopub_msgs_count, - "notebook_outputs_count": len(notebook_outputs), - "error_occurred": error_occurred, - } - ) return ExecuteResponse( notebook_outputs=notebook_outputs, From 6f86c6e4bfaf48e1189236e3a4c5a925dd8a0f41 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Sat, 14 Mar 2026 13:18:46 -0700 Subject: [PATCH 18/31] dropped stateful_python_code_exec --- src/hypotest/env/interpreter_env.py | 32 +---------------------------- 1 file changed, 1 insertion(+), 31 deletions(-) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index 06e3241..504a664 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -1393,8 +1393,7 @@ async def reset(self) -> tuple[Messages, list[Tool]]: self._filesystem_tool = FilesystemTool(self.work_dir) self.tools = [ - # Tool.from_function(self.run_cell), - Tool.from_function(self.stateful_python_code_exec), + Tool.from_function(self.run_cell), Tool.from_function(self.reset_kernel), Tool.from_function(self.submit_answer), Tool.from_function(self.list_dir), @@ -1443,35 +1442,6 @@ async def step(self, action: ToolRequestMessage) -> tuple[Messages, float, bool, # ========== Tools ========== - async def stateful_python_code_exec( - self, - code: str, - idx: int | None = None, - ) -> Message | str | list[dict[str, Any]]: - """Run code in a Python Jupyter notebook cell and return the execution output. - - This method allows running code in a new cell (append) or re-running - an existing cell with updated code. - - Error Recovery: - When a cell fails with an error, you MUST fix it by calling stateful_python_code_exec - with the corrected code and the SAME idx as the failed cell. - - The cell number is shown in the output prefix (e.g., "[Cell #3]"). - Do NOT create a new cell to fix an error - always edit the failed cell. - - Args: - code: Code to execute - idx: Cell index to run. If None or >= len(cells), appends a new cell. - If provided, updates and re-runs the existing cell at that index. - Use this to fix errors in existing cells. - - Returns: - Message with multimodal content if images present, otherwise string. - The response includes the cell number (e.g., "[Cell #0] output..."). - """ - return await self.run_cell(code, idx=idx) - async def run_cell( self, code: str, From 68fa59c1f57baa52b74b27a3161fe24d42e060e1 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Sat, 14 Mar 2026 17:03:00 -0700 Subject: [PATCH 19/31] fixed formatting issues --- pyproject.toml | 3 + src/hypotest/dataset_server.py | 9 +- src/hypotest/env/code_safety.py | 50 +- src/hypotest/env/config.py | 4 +- src/hypotest/env/interpreter_env.py | 575 ++++++++++------------- src/hypotest/env/kernel_server.py | 6 +- src/hypotest/env/tools/filesystem.py | 4 +- src/hypotest/env/utils/__init__.py | 2 +- src/hypotest/env/utils/notebook_utils.py | 2 +- 9 files changed, 303 insertions(+), 352 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2b267a6..f6984f0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,9 @@ dev = [ "types-tqdm", "vcrpy>=8", # Pin for dropping unused requests support ] +ray = [ + "ray", +] [tool.codespell] check-filenames = true diff --git a/src/hypotest/dataset_server.py b/src/hypotest/dataset_server.py index 1a471ad..806cd72 100644 --- a/src/hypotest/dataset_server.py +++ b/src/hypotest/dataset_server.py @@ -78,9 +78,7 @@ def get_new_env_by_idx(self, idx: int) -> InterpreterEnv: save_dir = Path(self.config.save_dir) / run_id if self.config.save_dir else None language = ( - NBLanguage.PYTHON - if self.config.force_python - else NBLanguage.from_string(cast(str, problem.nb_primary_language).upper()) + NBLanguage.PYTHON if self.config.force_python else NBLanguage.from_string(problem.nb_primary_language) ) return InterpreterEnv( @@ -137,7 +135,7 @@ async def launch_server(): config = ServerConfig.model_validate(yaml.safe_load(args.config.read_text())) else: config = ServerConfig( - dataset=DatasetConfig( + dataset=HypotestDatasetConfig( problem_jsonl=args.problem_jsonl, capsule_dir=args.capsule_dir, rubric_model=args.rubric_model, @@ -156,12 +154,13 @@ async def launch_server(): ], }, use_docker=args.use_docker, + execution_config={"cell_execution_timeout": 600}, ), port=args.port, api_key=args.api_key, ) - dataset = Dataset(config.dataset) + dataset = HypotestDataset(config.dataset) server = TaskDatasetServer(dataset, port=config.port, api_key=config.api_key) ip_address = socket.gethostbyname(socket.gethostname()) diff --git a/src/hypotest/env/code_safety.py b/src/hypotest/env/code_safety.py index 5d939e2..fad9908 100644 --- a/src/hypotest/env/code_safety.py +++ b/src/hypotest/env/code_safety.py @@ -71,15 +71,30 @@ # Dangerous attribute names used with getattr() evasion _DANGEROUS_ATTRS: frozenset[str] = frozenset({ - "kill", "killpg", "system", "popen", "fork", "forkpty", - "execl", "execle", "execlp", "execlpe", - "execv", "execve", "execvp", "execvpe", - "_exit", "rmtree", + "kill", + "killpg", + "system", + "popen", + "fork", + "forkpty", + "execl", + "execle", + "execlp", + "execlpe", + "execv", + "execve", + "execvp", + "execvpe", + "_exit", + "rmtree", }) # Modules whose import_module() calls should be checked _DANGEROUS_IMPORTLIB_TARGETS: frozenset[str] = _BLOCKED_MODULES | frozenset({ - "os", "subprocess", "shutil", "multiprocessing", + "os", + "subprocess", + "shutil", + "multiprocessing", }) @@ -150,7 +165,7 @@ def _resolve_call_path(self, node: ast.Call) -> str | None: return func.id return None - def visit_Call(self, node: ast.Call) -> None: + def visit_Call(self, node: ast.Call) -> None: # noqa: PLR0912 call_path = self._resolve_call_path(node) if call_path is not None: @@ -181,13 +196,16 @@ def visit_Call(self, node: ast.Call) -> None: return # Check getattr(os, 'kill') style - if call_path in ("getattr", "builtins.getattr") and len(node.args) >= 2: + if call_path in {"getattr", "builtins.getattr"} and len(node.args) >= 2: attr_arg = node.args[1] - if isinstance(attr_arg, ast.Constant) and isinstance(attr_arg.value, str): - if attr_arg.value in _DANGEROUS_ATTRS: - self._block(_MSG_DANGEROUS_CALL) - self.generic_visit(node) - return + if ( + isinstance(attr_arg, ast.Constant) + and isinstance(attr_arg.value, str) + and attr_arg.value in _DANGEROUS_ATTRS + ): + self._block(_MSG_DANGEROUS_CALL) + self.generic_visit(node) + return # Check importlib.import_module('signal') style if call_path.endswith("import_module") and node.args: @@ -209,9 +227,7 @@ def visit_Call(self, node: ast.Call) -> None: _MAGIC_SHELL_ESCAPE_LINE = re.compile(r"^\s*!(?P.+)$", re.MULTILINE) # Dangerous IPython magics -_MAGIC_DANGEROUS = re.compile( - r"^\s*%(run|system)\b|get_ipython\(\)\s*\.\s*system\s*\(", re.MULTILINE -) +_MAGIC_DANGEROUS = re.compile(r"^\s*%(run|system)\b|get_ipython\(\)\s*\.\s*system\s*\(", re.MULTILINE) def _check_jupyter_magics(code: str) -> str | None: @@ -327,7 +343,7 @@ def _check_python(code: str) -> str | None: stripped_lines = [] for line in code.splitlines(): stripped = line.lstrip() - if stripped.startswith("!") or stripped.startswith("%"): + if stripped.startswith(("!", "%")): stripped_lines.append("") # blank out magic lines else: stripped_lines.append(line) @@ -346,7 +362,7 @@ def _check_python(code: str) -> str | None: # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- -def check_code_safety(code: str, language: "NBLanguage") -> str | None: +def check_code_safety(code: str, language: NBLanguage) -> str | None: """Check if code is safe to execute. Args: diff --git a/src/hypotest/env/config.py b/src/hypotest/env/config.py index 5d661e8..9fcf862 100644 --- a/src/hypotest/env/config.py +++ b/src/hypotest/env/config.py @@ -15,7 +15,9 @@ # Kernel server settings (for Docker-based execution) KERNEL_SERVER_PORT = 8000 -KERNEL_SERVER_STARTUP_TIMEOUT = float(os.getenv("KERNEL_SERVER_STARTUP_TIMEOUT", "120.0")) # seconds to wait for health check +KERNEL_SERVER_STARTUP_TIMEOUT = float( + os.getenv("KERNEL_SERVER_STARTUP_TIMEOUT", "120.0") +) # seconds to wait for health check MAX_FILES_TO_UPLOAD = int(os.getenv("MAX_FILES_TO_UPLOAD", "100")) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index 504a664..df7d3bf 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -17,6 +17,7 @@ import sys import time import uuid +from collections.abc import Callable from pathlib import Path from tempfile import mkdtemp from textwrap import dedent @@ -27,7 +28,6 @@ import httpx import nbformat import numpy as np -import ray import tenacity from aviary.core import ( EnvStateMessage, @@ -44,12 +44,19 @@ from pydantic import BaseModel, Field, JsonValue, model_validator from . import config as cfg +from .code_safety import check_code_safety from .config import ExecutionConfig from .interpreter import ExecutionResult, Interpreter from .prompts import CORRECT_MSG, HYPOTHESIS_TASK_DESC, INCORRECT_MSG, RUBRIC_SCORE_PROMPT, PromptingConfig from .tools.filesystem import FilesystemTool from .utils import NBLanguage, view_notebook +RAY_INSTALLED = True +try: + import ray +except ImportError: + RAY_INSTALLED = False + if TYPE_CHECKING: from aiodocker.containers import DockerContainer @@ -89,6 +96,112 @@ class _PortCollisionError(Exception): """Port already in use by another server — retry with a new port.""" +async def _poll_kernel_health( # noqa: PLR0912 + http_client: httpx.AsyncClient, + enroot_proc: subprocess.Popen[str] | None, + container_port: int | None, + expected_startup_token: str | None, + read_log_tail: Callable[[int], str], + label: str, +) -> None: + """Poll kernel server /health endpoint until ready, with token validation and timeout.""" + start_time = time.perf_counter() + poll_count = 0 + last_status = "no_attempt" + consecutive_token_mismatches = 0 + # Use a short per-request timeout for health checks so that a single + # poll can never block longer than a few seconds. Without this, the + # client's default timeout (execution_timeout + 10 ≈ 190s) means one + # health poll can exceed the entire KERNEL_SERVER_STARTUP_TIMEOUT, + # e.g. when a port collision causes us to connect to the wrong server. + health_timeout = httpx.Timeout(5.0, connect=3.0) + + while time.perf_counter() - start_time < cfg.KERNEL_SERVER_STARTUP_TIMEOUT: + poll_count += 1 + elapsed = time.perf_counter() - start_time + + # Check if the enroot process has died before we even get an HTTP response + if enroot_proc is not None and enroot_proc.poll() is not None: + rc = enroot_proc.returncode + log_tail = read_log_tail(1000) + raise RuntimeError( + f"Enroot process exited prematurely with returncode={rc} after {elapsed:.1f}s. log: {log_tail!r}" + ) + + try: + response = await http_client.get("/health", timeout=health_timeout) + if response.status_code == 200: + if expected_startup_token is None: + logger.log( + _CONTAINER_LOG_LEVEL, + "[%s] Kernel server healthy after %.1fs (%d polls)", + label, + elapsed, + poll_count, + ) + return + payload = response.json() + if payload.get("startup_token") == expected_startup_token: + logger.log( + _CONTAINER_LOG_LEVEL, + "[%s] Kernel server healthy (token matched) after %.1fs (%d polls)", + label, + elapsed, + poll_count, + ) + return + last_status = f"token_mismatch(got={payload.get('startup_token', '?')[:8]})" + consecutive_token_mismatches += 1 + if consecutive_token_mismatches >= 3: + raise _PortCollisionError( + f"Port {container_port} appears to be owned by another server " + f"({consecutive_token_mismatches} consecutive token mismatches)" + ) + else: + last_status = f"http_{response.status_code}" + consecutive_token_mismatches = 0 + except httpx.ConnectError: + last_status = "connect_error" + consecutive_token_mismatches = 0 + except httpx.ReadError: + last_status = "read_error" + consecutive_token_mismatches = 0 + except httpx.TimeoutException: + last_status = "timeout" + consecutive_token_mismatches = 0 + except httpx.RemoteProtocolError: + last_status = "protocol_error" + consecutive_token_mismatches = 0 + + # Log progress every 5s + if poll_count % 10 == 0: + proc_alive = enroot_proc.poll() is None if enroot_proc else False + logger.log( + _CONTAINER_LOG_LEVEL, + "[%s] Health poll #%d at %.1fs: last_status=%s, process_alive=%s", + label, + poll_count, + elapsed, + last_status, + proc_alive, + ) + + await asyncio.sleep(0.5) + + total_elapsed = time.perf_counter() - start_time + if last_status.startswith("token_mismatch"): + raise _PortCollisionError( + f"Port {container_port} health-check timed out with token_mismatch " + f"({poll_count} polls, elapsed={total_elapsed:.1f}s)" + ) + log_tail = read_log_tail(500) + raise TimeoutError( + f"Kernel server did not become healthy within {cfg.KERNEL_SERVER_STARTUP_TIMEOUT}s " + f"({poll_count} polls, last_status={last_status}, elapsed={total_elapsed:.1f}s)" + f"{f' log_tail={log_tail!r}' if log_tail else ''}" + ) + + def _kill_process_group(proc: subprocess.Popen, label: str = "enroot", sigterm_timeout: float = 15) -> None: """Safely terminate a process group, escalating from SIGTERM to SIGKILL. @@ -214,11 +327,102 @@ def _proc_label(self) -> str: pid = self._enroot_proc.pid if self._enroot_proc else "?" return f"enroot(port={port}, pid={pid})" + @staticmethod + def _build_kernel_bash_script(node_workdir: str, language: NBLanguage, port: int, startup_token: str) -> str: + """Build the bash script that sets up the workspace and launches the kernel server.""" + return dedent(f"""\ + set -euo pipefail + + WORKDIR="{node_workdir}" + trap 'rm -rf "$WORKDIR"' EXIT + + mkdir -p $WORKDIR + cp -a /data_workspace/. $WORKDIR/ + + cd $WORKDIR + + mkdir -p "$WORKDIR/r_libs" + cat >"$WORKDIR/Rprofile" <<'EOF' + .local_lib <- Sys.getenv("R_LIBS_USER") + if (nzchar(.local_lib)) .libPaths(unique(c(.local_lib, .libPaths()))) + EOF + + mkdir -p $WORKDIR/r_libs + export R_LIBS_USER="$WORKDIR/r_libs" + export R_PROFILE_USER="$WORKDIR/Rprofile" + + export PYTHONPATH="$WORKDIR/pydeps:${{PYTHONPATH}}" + export PIP_CONFIG_FILE=$WORKDIR/pip.conf + export target_platform=${{target_platform:-linux-64}} + + source activate /app/kernel_env + exec /app/kernel_env/bin/python /envs/kernel_server.py \\ + --work_dir $WORKDIR \\ + --language {language.value} \\ + --port {port} \\ + --startup-token {startup_token} + """).strip() + + @staticmethod + def _setup_enroot_env(startup_token: str) -> dict[str, str]: + """Create enroot runtime directories and return env dict.""" + base = Path(f"/tmp/enroot_data/{startup_token}") # noqa: S108 + subdirs = ["runtime", "config", "cache", "data", "tmp"] + env_keys = [ + "ENROOT_RUNTIME_PATH", + "ENROOT_CONFIG_PATH", + "ENROOT_CACHE_PATH", + "ENROOT_DATA_PATH", + "ENROOT_TEMP_PATH", + ] + env: dict[str, str] = {} + for subdir, key in zip(subdirs, env_keys, strict=True): + p = base / subdir + p.mkdir(parents=True, exist_ok=True) + os.chmod(p, 0o700) + env[key] = str(p) + return env + + @staticmethod + def _build_enroot_cmd( + work_dir: Path, + kernel_server_path: Path, + bash: str, + enroot_env: dict[str, str], + container_sqsh_path: Path, + ) -> list[str]: + """Assemble the full ``enroot start`` command.""" + env_args = [f"{k}={v}" for k, v in enroot_env.items()] + return [ + "env", + "-i", + "PATH=/usr/sbin:/usr/bin:/sbin:/bin", + 'HOME="$HOME"', + 'USER="$USER"', + *env_args, + "enroot", + "start", + "--rw", + "--mount", + f"{work_dir}:/data_workspace", + "--mount", + f"{kernel_server_path.resolve()}:/envs/kernel_server.py", + str(container_sqsh_path.resolve()), + "/bin/bash", + "-lc", + bash, + ] + async def initialize(self, work_dir: Path, language: NBLanguage) -> None: startup_token = str(uuid.uuid4()) node_workdir = f"/tmp/data_workspace.{startup_token.split('-', maxsplit=1)[0]}" - _prep_workspace_dir(work_dir, workspace_path=node_workdir) + _prep_workspace_dir(str(work_dir), workspace_path=node_workdir) + + kernel_server_path = Path(__file__).parent / "kernel_server.py" + assert kernel_server_path.is_file(), f"kernel server must be a valid path, found {kernel_server_path}" + + enroot_env = self._setup_enroot_env(startup_token) online = False attempt = 0 @@ -234,91 +438,8 @@ async def initialize(self, work_dir: Path, language: NBLanguage) -> None: ) self._container_port = await get_free_port() - bash = dedent(f"""\ - set -euo pipefail - - WORKDIR="{node_workdir}" - trap 'rm -rf "$WORKDIR"' EXIT - - mkdir -p $WORKDIR - cp -a /data_workspace/. $WORKDIR/ - - cd $WORKDIR - - mkdir -p "$WORKDIR/r_libs" - cat >"$WORKDIR/Rprofile" <<'EOF' - .local_lib <- Sys.getenv("R_LIBS_USER") - if (nzchar(.local_lib)) .libPaths(unique(c(.local_lib, .libPaths()))) - EOF - - mkdir -p $WORKDIR/r_libs - export R_LIBS_USER="$WORKDIR/r_libs" - export R_PROFILE_USER="$WORKDIR/Rprofile" - - export PYTHONPATH="$WORKDIR/pydeps:${{PYTHONPATH}}" - export PIP_CONFIG_FILE=$WORKDIR/pip.conf - export target_platform=${{target_platform:-linux-64}} - - source activate /app/kernel_env - exec /app/kernel_env/bin/python /envs/kernel_server.py \\ - --work_dir $WORKDIR \\ - --language {language.value} \\ - --port {self._container_port} \\ - --startup-token {startup_token} - """).strip() - - env_dir = Path(__file__).parent - kernel_server_path = env_dir / "kernel_server.py" - assert kernel_server_path.is_file(), f"kernel server must be a valid path, found {kernel_server_path}" - - base = Path(f"/tmp/enroot_data/{startup_token}") - - enroot_runtime = base / "runtime" - enroot_config = base / "config" - enroot_cache = base / "cache" - enroot_data = base / "data" - enroot_temp = base / "tmp" - - for path in [ - enroot_runtime, - enroot_config, - enroot_cache, - enroot_data, - enroot_temp, - ]: - path.mkdir(parents=True, exist_ok=True) - os.chmod(path, 0o700) - - env = {} - env["ENROOT_RUNTIME_PATH"] = str(enroot_runtime) - env["ENROOT_CONFIG_PATH"] = str(enroot_config) - env["ENROOT_CACHE_PATH"] = str(enroot_cache) - env["ENROOT_DATA_PATH"] = str(enroot_data) - env["ENROOT_TEMP_PATH"] = str(enroot_temp) - - cmd = [ - "env", - "-i", - "PATH=/usr/sbin:/usr/bin:/sbin:/bin", - 'HOME="$HOME"', - 'USER="$USER"', - f"ENROOT_RUNTIME_PATH={enroot_runtime}", - f"ENROOT_CONFIG_PATH={enroot_config}", - f"ENROOT_CACHE_PATH={enroot_cache}", - f"ENROOT_DATA_PATH={enroot_data}", - f"ENROOT_TEMP_PATH={enroot_temp}", - "enroot", - "start", - "--rw", - "--mount", - f"{work_dir}:/data_workspace", - "--mount", - f"{kernel_server_path.resolve()}:/envs/kernel_server.py", - str(self.container_sqsh_path.resolve()), - "/bin/bash", - "-lc", - bash, - ] + bash = self._build_kernel_bash_script(node_workdir, language, self._container_port, startup_token) + cmd = self._build_enroot_cmd(work_dir, kernel_server_path, bash, enroot_env, self.container_sqsh_path) async with CONTAINER_LAUNCH_SEM: launch_t0 = time.perf_counter() @@ -446,109 +567,16 @@ async def _cleanup_failed_startup(self) -> None: async def _wait_for_health(self, expected_startup_token: str | None = None) -> None: """Wait for the kernel server to become healthy.""" - start_time = time.perf_counter() - poll_count = 0 - last_status = "no_attempt" - consecutive_token_mismatches = 0 - # Use a short per-request timeout for health checks so that a single - # poll can never block longer than a few seconds. Without this, the - # client's default timeout (execution_timeout + 10 ≈ 190s) means one - # health poll can exceed the entire KERNEL_SERVER_STARTUP_TIMEOUT, - # e.g. when a port collision causes us to connect to the wrong server. - health_timeout = httpx.Timeout(5.0, connect=3.0) - - while time.perf_counter() - start_time < cfg.KERNEL_SERVER_STARTUP_TIMEOUT: - poll_count += 1 - elapsed = time.perf_counter() - start_time - - # Check if the enroot process has died before we even get an HTTP response - if self._enroot_proc is not None and self._enroot_proc.poll() is not None: - rc = self._enroot_proc.returncode - log_tail = self._read_container_log_tail(1000) - raise RuntimeError( - f"Enroot process exited prematurely with returncode={rc} after {elapsed:.1f}s. log: {log_tail!r}" - ) - - try: - assert self._http_client is not None - response = await self._http_client.get("/health", timeout=health_timeout) - if response.status_code == 200: - if expected_startup_token is None: - logger.log( - _CONTAINER_LOG_LEVEL, - "[%s] Kernel server healthy after %.1fs (%d polls)", - self._proc_label(), - elapsed, - poll_count, - ) - return - payload = response.json() - if payload.get("startup_token") == expected_startup_token: - logger.log( - _CONTAINER_LOG_LEVEL, - "[%s] Kernel server healthy (token matched) after %.1fs (%d polls)", - self._proc_label(), - elapsed, - poll_count, - ) - return - last_status = f"token_mismatch(got={payload.get('startup_token', '?')[:8]})" - consecutive_token_mismatches += 1 - if consecutive_token_mismatches >= 3: - raise _PortCollisionError( - f"Port {self._container_port} appears to be owned by another server " - f"({consecutive_token_mismatches} consecutive token mismatches)" - ) - else: - last_status = f"http_{response.status_code}" - consecutive_token_mismatches = 0 - except httpx.ConnectError: - last_status = "connect_error" - consecutive_token_mismatches = 0 - except httpx.ReadError: - last_status = "read_error" - consecutive_token_mismatches = 0 - except httpx.TimeoutException: - last_status = "timeout" - consecutive_token_mismatches = 0 - except httpx.RemoteProtocolError: - last_status = "protocol_error" - consecutive_token_mismatches = 0 - - # Log progress every 5s - if poll_count % 10 == 0: - proc_alive = self._enroot_proc.poll() is None if self._enroot_proc else False - logger.log( - _CONTAINER_LOG_LEVEL, - "[%s] Health poll #%d at %.1fs: last_status=%s, process_alive=%s", - self._proc_label(), - poll_count, - elapsed, - last_status, - proc_alive, - ) - - await asyncio.sleep(0.5) - - total_elapsed = time.perf_counter() - start_time - if last_status.startswith("token_mismatch"): - raise _PortCollisionError( - f"Port {self._container_port} health-check timed out with token_mismatch " - f"({poll_count} polls, elapsed={total_elapsed:.1f}s)" - ) - log_tail = self._read_container_log_tail(500) - raise TimeoutError( - f"Kernel server did not become healthy within {cfg.KERNEL_SERVER_STARTUP_TIMEOUT}s " - f"({poll_count} polls, last_status={last_status}, elapsed={total_elapsed:.1f}s)" - f"{f' log_tail={log_tail!r}' if log_tail else ''}" + assert self._http_client is not None + await _poll_kernel_health( + http_client=self._http_client, + enroot_proc=self._enroot_proc, + container_port=self._container_port, + expected_startup_token=expected_startup_token, + read_log_tail=self._read_container_log_tail, + label=self._proc_label(), ) - # @tenacity.retry( - # retry=tenacity.retry_if_exception_type((httpx.ConnectError, httpx.ReadError)), - # stop=tenacity.stop_after_attempt(3), - # wait=tenacity.wait_exponential(multiplier=1, min=1, max=10), - # reraise=True, - # ) async def _execute_via_http(self, code: str, timeout: float | None = None) -> ExecutionResult: # noqa: ASYNC109 """Execute code via HTTP to the containerized kernel server. @@ -675,7 +703,7 @@ def __init__( self.use_enroot = use_enroot self.container_sqsh_path = container_sqsh_path self.save_dir = save_dir - self.use_ray = use_ray + self.use_ray = use_ray if RAY_INSTALLED else False # Local interpreter (only used when use_docker=False) self.interpreter: Interpreter | None = None @@ -732,7 +760,7 @@ def _enroot_label(self) -> str: return f"enroot-state(port={port}, pid={pid})" async def _start_enroot_container(self) -> None: - _prep_workspace_dir(self.work_dir) + _prep_workspace_dir(str(self.work_dir)) online = False attempt = 0 @@ -762,9 +790,9 @@ async def _start_enroot_container(self) -> None: --startup-token {startup_token} """).strip() - env_dir = Path(__file__).parent - kernel_server_path = env_dir / "kernel_server.py" + kernel_server_path = Path(__file__).parent / "kernel_server.py" assert kernel_server_path.is_file(), f"kernel server must be a valid path, found {kernel_server_path}" + assert self.container_sqsh_path is not None, "container_sqsh_path must be set when using enroot container" cmd = [ "env", @@ -788,9 +816,8 @@ async def _start_enroot_container(self) -> None: launch_t0 = time.perf_counter() # Redirect container output to a log file (see EnrootKernelServer # for detailed rationale on why we avoid subprocess.PIPE here). - log_dir = self.work_dir / ".container_logs" - log_dir.mkdir(exist_ok=True) - self._container_log_path = log_dir / "container.log" + (self.work_dir / ".container_logs").mkdir(exist_ok=True) + self._container_log_path = self.work_dir / ".container_logs" / "container.log" self._container_log_file = open(self._container_log_path, "w", encoding="utf-8") # noqa: SIM115 self._enroot_proc = subprocess.Popen( cmd, @@ -830,27 +857,31 @@ async def _start_enroot_container(self) -> None: except Exception as e: last_err = e launch_ms = (time.perf_counter() - launch_t0) * 1000.0 - label = self._enroot_label() - diag = [f"attempt=#{attempt}", f"elapsed={launch_ms:.0f}ms", f"error={e!r}"] - if self._enroot_proc is not None: - rc = self._enroot_proc.poll() - diag.append(f"process_alive={rc is None}") - if rc is not None: - diag.append(f"returncode={rc}") - logger.warning("[%s] Container startup FAILED: %s", label, ", ".join(diag)) - log_tail = self._read_container_log_tail() - if log_tail: - logger.warning( - "[%s] Container log output (last %d chars):\n%s", - label, - len(log_tail), - log_tail, - ) + self._log_enroot_container_failure(attempt, launch_ms, e) await self._cleanup_failed_startup() if not isinstance(e, _PortCollisionError): backoff = min(_RETRY_BASE_SLEEP * 2 ** (attempt - 1), _RETRY_MAX_SLEEP) await asyncio.sleep(backoff) + def _log_enroot_container_failure(self, attempt: int, launch_ms: float, error: Exception) -> None: + """Log detailed diagnostics when a container fails to start.""" + label = self._enroot_label() + diag = [f"attempt=#{attempt}", f"elapsed={launch_ms:.0f}ms", f"error={error!r}"] + if self._enroot_proc is not None: + rc = self._enroot_proc.poll() + diag.append(f"process_alive={rc is None}") + if rc is not None: + diag.append(f"returncode={rc}") + logger.warning("[%s] Container startup FAILED: %s", label, ", ".join(diag)) + log_tail = self._read_container_log_tail() + if log_tail: + logger.warning( + "[%s] Container log output (last %d chars):\n%s", + label, + len(log_tail), + log_tail, + ) + def _read_container_log_tail(self, max_chars: int = 2000) -> str: """Read the tail of the container log file for diagnostics.""" if self._container_log_path is None or not self._container_log_path.exists(): @@ -931,98 +962,14 @@ async def _start_docker_container(self) -> None: async def _wait_for_health(self, expected_startup_token: str | None = None) -> None: """Wait for the kernel server to become healthy.""" - start_time = time.perf_counter() - poll_count = 0 - last_status = "no_attempt" - consecutive_token_mismatches = 0 - # Short per-request timeout (see EnrootKernelServer._wait_for_health - # for rationale). - health_timeout = httpx.Timeout(5.0, connect=3.0) - - while time.perf_counter() - start_time < cfg.KERNEL_SERVER_STARTUP_TIMEOUT: - poll_count += 1 - elapsed = time.perf_counter() - start_time - - # Check if the enroot process has died - if self._enroot_proc is not None and self._enroot_proc.poll() is not None: - rc = self._enroot_proc.returncode - log_tail = self._read_container_log_tail(1000) - raise RuntimeError( - f"Enroot process exited prematurely with returncode={rc} after {elapsed:.1f}s. log: {log_tail!r}" - ) - - try: - assert self._http_client is not None - response = await self._http_client.get("/health", timeout=health_timeout) - if response.status_code == 200: - if expected_startup_token is None: - logger.log( - _CONTAINER_LOG_LEVEL, - "[%s] Kernel server healthy after %.1fs (%d polls)", - self._enroot_label(), - elapsed, - poll_count, - ) - return - payload = response.json() - if payload.get("startup_token") == expected_startup_token: - logger.log( - _CONTAINER_LOG_LEVEL, - "[%s] Kernel server healthy (token matched) after %.1fs (%d polls)", - self._enroot_label(), - elapsed, - poll_count, - ) - return - last_status = f"token_mismatch(got={payload.get('startup_token', '?')[:8]})" - consecutive_token_mismatches += 1 - if consecutive_token_mismatches >= 3: - raise _PortCollisionError( - f"Port {self._container_port} appears to be owned by another server " - f"({consecutive_token_mismatches} consecutive token mismatches)" - ) - else: - last_status = f"http_{response.status_code}" - consecutive_token_mismatches = 0 - except httpx.ConnectError: - last_status = "connect_error" - consecutive_token_mismatches = 0 - except httpx.ReadError: - last_status = "read_error" - consecutive_token_mismatches = 0 - except httpx.TimeoutException: - last_status = "timeout" - consecutive_token_mismatches = 0 - except httpx.RemoteProtocolError: - last_status = "protocol_error" - consecutive_token_mismatches = 0 - - # Log progress every 5s - if poll_count % 10 == 0: - proc_alive = self._enroot_proc.poll() is None if self._enroot_proc else False - logger.log( - _CONTAINER_LOG_LEVEL, - "[%s] Health poll #%d at %.1fs: last_status=%s, process_alive=%s", - self._enroot_label(), - poll_count, - elapsed, - last_status, - proc_alive, - ) - - await asyncio.sleep(0.5) - - total_elapsed = time.perf_counter() - start_time - if last_status.startswith("token_mismatch"): - raise _PortCollisionError( - f"Port {self._container_port} health-check timed out with token_mismatch " - f"({poll_count} polls, elapsed={total_elapsed:.1f}s)" - ) - log_tail = self._read_container_log_tail(500) - raise TimeoutError( - f"Kernel server did not become healthy within {cfg.KERNEL_SERVER_STARTUP_TIMEOUT}s " - f"({poll_count} polls, last_status={last_status}, elapsed={total_elapsed:.1f}s)" - f"{f' log_tail={log_tail!r}' if log_tail else ''}" + assert self._http_client is not None + await _poll_kernel_health( + http_client=self._http_client, + enroot_proc=self._enroot_proc, + container_port=self._container_port, + expected_startup_token=expected_startup_token, + read_log_tail=self._read_container_log_tail, + label=self._enroot_label(), ) @tenacity.retry( @@ -1230,9 +1177,6 @@ async def execute_and_add_cell( Returns: Tuple of (ExecutionResult, actual_cell_index) """ - # Safety check: block dangerous code before execution - from .code_safety import check_code_safety - block_reason = check_code_safety(code, self.language) if block_reason is not None: logger.warning("Blocked code execution in execute_and_add_cell: %s", block_reason) @@ -1279,6 +1223,7 @@ class InterpreterEnvConfig(BaseModel): prompting_config: PromptingConfig = Field(default_factory=PromptingConfig) execution_config: ExecutionConfig = Field(default_factory=ExecutionConfig) max_steps: int = cfg.AGENT_MAX_STEPS + use_ray: bool = False use_docker: bool = cfg.USE_DOCKER use_enroot: bool = False container_sqsh_path: Path | None = None @@ -1396,7 +1341,7 @@ async def reset(self) -> tuple[Messages, list[Tool]]: Tool.from_function(self.run_cell), Tool.from_function(self.reset_kernel), Tool.from_function(self.submit_answer), - Tool.from_function(self.list_dir), + Tool.from_function(self._filesystem_tool.list_dir), ] messages.append( @@ -1413,7 +1358,7 @@ async def reset(self) -> tuple[Messages, list[Tool]]: messages.append(self.get_env_state_msg()) # Always show initial directory listing (with truncation protection) - messages.append(Message(content=self.list_dir())) + messages.append(Message(content=self._filesystem_tool.list_dir())) return messages, self.tools @@ -1610,20 +1555,6 @@ async def submit_answer(self, answer: str) -> str: correct = await self._score_solution(answer) return CORRECT_MSG if correct else INCORRECT_MSG - def list_dir(self) -> str: - """List contents of a directory with truncation protection. - - Recursively lists files in a directory, with built-in protection against - overwhelming the context with too many files. Use this tool instead of - writing code to list directories to avoid context bloat. - - Args: - directory: Directory path to list (default: current working directory) - max_files: Maximum number of files to display (default: 20) - show_hidden: Whether to show hidden files starting with '.' (default: False) - """ - return self._filesystem_tool.list_dir() - # ========== Time Management ========== def get_remaining_time(self) -> int: diff --git a/src/hypotest/env/kernel_server.py b/src/hypotest/env/kernel_server.py index 3c9dd2e..5416de8 100644 --- a/src/hypotest/env/kernel_server.py +++ b/src/hypotest/env/kernel_server.py @@ -29,7 +29,7 @@ from pydantic import BaseModel -class _DeadlineExceeded(Exception): +class DeadlineExceededError(Exception): """Raised when a cooperative deadline check expires.""" @@ -276,7 +276,7 @@ async def execute(self, code: str, timeout: float | None = None) -> ExecuteRespo code, deadline=start_time + effective_timeout, ) - except _DeadlineExceeded: + except DeadlineExceededError: timeout_output = MessageType.ERROR.to_notebook_output({ "ename": "TimeoutError", "evalue": f"Code execution timed out after {effective_timeout} seconds", @@ -325,7 +325,7 @@ async def _execute_code(self, code: str, deadline: float) -> ExecuteResponse: while True: # Check deadline before each poll if time.perf_counter() >= deadline: - raise _DeadlineExceeded + raise DeadlineExceededError # Use a bounded poll so we never block longer than _POLL_INTERVAL_S. # get_iopub_msg(timeout=T) raises queue.Empty if no message arrives diff --git a/src/hypotest/env/tools/filesystem.py b/src/hypotest/env/tools/filesystem.py index 642177d..de036cc 100644 --- a/src/hypotest/env/tools/filesystem.py +++ b/src/hypotest/env/tools/filesystem.py @@ -23,8 +23,8 @@ import numpy as np from aviary.core import Message, Tool -from resources_servers.aviary_hypotest.env.utils import img_utils -from resources_servers.aviary_hypotest.env.utils.notebook_utils import view_notebook +from hypotest.env.utils import img_utils +from hypotest.env.utils.notebook_utils import view_notebook # ============================================================================= # CONSTANTS diff --git a/src/hypotest/env/utils/__init__.py b/src/hypotest/env/utils/__init__.py index b7451c1..d1f3c1a 100644 --- a/src/hypotest/env/utils/__init__.py +++ b/src/hypotest/env/utils/__init__.py @@ -19,7 +19,7 @@ "view_notebook", ] -from resources_servers.aviary_hypotest.env.kernel_server import NBLanguage +from hypotest.env.kernel_server import NBLanguage from .core import extract_code_from_markdown, extract_xml_content from .img_utils import ( diff --git a/src/hypotest/env/utils/notebook_utils.py b/src/hypotest/env/utils/notebook_utils.py index d2ec722..e727aab 100644 --- a/src/hypotest/env/utils/notebook_utils.py +++ b/src/hypotest/env/utils/notebook_utils.py @@ -8,7 +8,7 @@ import nbformat from aiodocker.containers import DockerContainer -from resources_servers.aviary_hypotest.env import config as cfg +from hypotest.env import config as cfg from .img_utils import encode_image_to_base64 From 463afcbdde39720d59e4a5657eb322054e8bdcdc Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Sat, 14 Mar 2026 17:08:53 -0700 Subject: [PATCH 20/31] added ray as dep --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index f6984f0..86da5ca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,7 @@ dependencies = [ "pillow", "pydantic", "tenacity>=8.0", + "ray", ] description = "hypotest package" keywords = [] From dc5740e5012a832aa26264be907e60570c020b98 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Sun, 15 Mar 2026 14:22:10 -0700 Subject: [PATCH 21/31] added fallback on nb_primary_language for ProblemInstance Signed-off-by: Arnav Komaragiri --- src/hypotest/env/interpreter_env.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index df7d3bf..b086b99 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -286,7 +286,7 @@ class ProblemInstance(BaseModel): @model_validator(mode="before") @classmethod def handle_language(cls, data: dict) -> dict: - if data["nb_primary_language"] is None: + if data.get("nb_primary_language", None) is None: data["nb_primary_language"] = str(NBLanguage.PYTHON) return data From 4f63f7e955b6641d7f0f655b3b3f18a2be968690 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Sun, 15 Mar 2026 14:34:18 -0700 Subject: [PATCH 22/31] added toggle to safety checks --- src/hypotest/env/config.py | 3 ++ src/hypotest/env/interpreter_env.py | 83 ++++++++++++++++------------- src/hypotest/env/kernel_server.py | 32 ++++++----- 3 files changed, 68 insertions(+), 50 deletions(-) diff --git a/src/hypotest/env/config.py b/src/hypotest/env/config.py index 9fcf862..d8ecf09 100644 --- a/src/hypotest/env/config.py +++ b/src/hypotest/env/config.py @@ -58,6 +58,9 @@ class ExecutionConfig(BaseModel): force_submit_threshold: int = 10 * 60 cell_execution_timeout: int = 15 * 60 + # safety + safe_execute: bool = True + # Capabilities has_gpu: bool = False diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index b086b99..b1b484b 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -312,9 +312,10 @@ def _prep_workspace_dir(work_dir: str, workspace_path: str = "/data_workspace") }, ) class EnrootKernelServer: - def __init__(self, container_sqsh_path: Path, execution_timeout: float): + def __init__(self, container_sqsh_path: Path, execution_timeout: float, safe_execute: bool = True): self.container_sqsh_path = container_sqsh_path self.execution_timeout = execution_timeout + self.safe_execute = safe_execute self._enroot_proc: subprocess.Popen[str] | None = None self._http_client: httpx.AsyncClient | None = None self._container_port: int | None = None @@ -328,7 +329,7 @@ def _proc_label(self) -> str: return f"enroot(port={port}, pid={pid})" @staticmethod - def _build_kernel_bash_script(node_workdir: str, language: NBLanguage, port: int, startup_token: str) -> str: + def _build_kernel_bash_script(node_workdir: str, language: NBLanguage, port: int, startup_token: str, safe_execute: bool = True) -> str: """Build the bash script that sets up the workspace and launches the kernel server.""" return dedent(f"""\ set -euo pipefail @@ -360,7 +361,7 @@ def _build_kernel_bash_script(node_workdir: str, language: NBLanguage, port: int --work_dir $WORKDIR \\ --language {language.value} \\ --port {port} \\ - --startup-token {startup_token} + --startup-token {startup_token} {'--safe-execute' if safe_execute else ''} """).strip() @staticmethod @@ -438,7 +439,7 @@ async def initialize(self, work_dir: Path, language: NBLanguage) -> None: ) self._container_port = await get_free_port() - bash = self._build_kernel_bash_script(node_workdir, language, self._container_port, startup_token) + bash = self._build_kernel_bash_script(node_workdir, language, self._container_port, startup_token, safe_execute=self.safe_execute) cmd = self._build_enroot_cmd(work_dir, kernel_server_path, bash, enroot_env, self.container_sqsh_path) async with CONTAINER_LAUNCH_SEM: @@ -682,6 +683,7 @@ def __init__( work_dir: Path, language: NBLanguage, execution_timeout: int = 600, + safe_execute: bool = True, use_host_env_vars: bool = False, extra_envs: dict[str, str] | None = None, use_docker: bool = cfg.USE_DOCKER, @@ -693,6 +695,7 @@ def __init__( self.work_dir = work_dir self.language = language self.execution_timeout = execution_timeout + self.safe_execute = safe_execute self.total_reward = 0.0 self.use_host_env_vars = use_host_env_vars self.extra_envs = extra_envs or {} @@ -750,7 +753,7 @@ async def start(self): await self.interpreter.start() async def _start_ray_enroot_container(self) -> None: - self.kernel_container = EnrootKernelServer.remote(self.container_sqsh_path, self.execution_timeout) + self.kernel_container = EnrootKernelServer.remote(self.container_sqsh_path, self.execution_timeout, safe_execute=self.safe_execute) init_ref = self.kernel_container.initialize.remote(self.work_dir, self.language) await init_ref @@ -787,7 +790,7 @@ async def _start_enroot_container(self) -> None: --work_dir /data_workspace \\ --language {self.language.value} \\ --port {self._container_port} \\ - --startup-token {startup_token} + --startup-token {startup_token} {'--safe-execute' if self.safe_execute else ''} """).strip() kernel_server_path = Path(__file__).parent / "kernel_server.py" @@ -927,18 +930,22 @@ async def _start_docker_container(self) -> None: self._container_port = await get_free_port() startup_token = str(uuid.uuid4()) + cmd_list = [ + "/app/kernel_env/bin/python", + "/envs/kernel_server.py", + "--work_dir", + "/data_workspace", + "--language", + self.language.value, + "--startup-token", + startup_token, + ] + if self.safe_execute: + cmd_list += ['--safe-execute'] + docker_config = { "Image": cfg.NB_ENVIRONMENT_DOCKER_IMAGE, - "Cmd": [ - "/app/kernel_env/bin/python", - "/envs/kernel_server.py", - "--work_dir", - "/data_workspace", - "--language", - self.language.value, - "--startup-token", - startup_token, - ], + "Cmd": cmd_list, "HostConfig": { "Binds": [f"{self.work_dir}:/data_workspace"], "PortBindings": {f"{cfg.KERNEL_SERVER_PORT}/tcp": [{"HostPort": str(self._container_port)}]}, @@ -1177,26 +1184,29 @@ async def execute_and_add_cell( Returns: Tuple of (ExecutionResult, actual_cell_index) """ - block_reason = check_code_safety(code, self.language) - if block_reason is not None: - logger.warning("Blocked code execution in execute_and_add_cell: %s", block_reason) - error_output = nbformat.v4.new_output( - output_type="error", - ename="SecurityError", - evalue=block_reason, - traceback=[f"SecurityError: {block_reason}"], - ) - result = ExecutionResult( - notebook_outputs=[error_output], - error_occurred=True, - execution_time=0.0, - ) - if cell_idx is None or cell_idx >= len(self.nb.cells): - actual_idx = self._add_cell(code, result) - else: - self._update_cell(cell_idx, code, result) - actual_idx = cell_idx - return result, actual_idx + if self.safe_execute: + block_reason = check_code_safety(code, self.language) + if block_reason is not None: + logger.warning("Blocked code execution in execute_and_add_cell: %s", block_reason) + error_output = nbformat.v4.new_output( + output_type="error", + ename="SecurityError", + evalue=block_reason, + traceback=[f"SecurityError: {block_reason}"], + ) + result = ExecutionResult( + notebook_outputs=[error_output], + error_occurred=True, + execution_time=0.0, + ) + if cell_idx is None or cell_idx >= len(self.nb.cells): + actual_idx = self._add_cell(code, result) + else: + self._update_cell(cell_idx, code, result) + actual_idx = cell_idx + return result, actual_idx + else: + logger.warning("WARNING: You are running code sandbox with unsafe execution, this may result in destructive commands being run on the system") if self.use_ray and self.use_enroot: result_ref = self.kernel_container._execute_via_http.remote(code, timeout) @@ -1312,6 +1322,7 @@ async def reset(self) -> tuple[Messages, list[Tool]]: work_dir=self.work_dir, language=self.language, execution_timeout=self.execution_timeout, + safe_execute=self.execution_config.safe_execute, use_host_env_vars=self.use_host_env_vars, extra_envs={ # Point to kernel environment's site-packages diff --git a/src/hypotest/env/kernel_server.py b/src/hypotest/env/kernel_server.py index 5416de8..6dc3eb5 100644 --- a/src/hypotest/env/kernel_server.py +++ b/src/hypotest/env/kernel_server.py @@ -211,11 +211,13 @@ def __init__( language: NBLanguage, default_timeout: float = 600, startup_token: str = "", + safe_execute: bool = True, ): self.work_dir = work_dir self.language = language self.default_timeout = default_timeout self.startup_token = startup_token + self.safe_execute = safe_execute self._kernel_manager: AsyncKernelManager | None = None self._client: AsyncKernelClient | None = None @@ -253,20 +255,21 @@ async def execute(self, code: str, timeout: float | None = None) -> ExecuteRespo if not self._client or not self._is_ready: raise RuntimeError("Kernel not ready") - # Defense-in-depth: lightweight regex safety check - block_reason = _kernel_check_code_safety(code) - if block_reason is not None: - logger.warning("Kernel safety block: %s code=%r", block_reason, code[:200]) - error_output = MessageType.ERROR.to_notebook_output({ - "ename": "SecurityError", - "evalue": block_reason, - "traceback": [f"SecurityError: {block_reason}"], - }) - return ExecuteResponse( - notebook_outputs=[dict(error_output)] if error_output else [], - error_occurred=True, - execution_time=0.0, - ) + if self.safe_execute: + # Defense-in-depth: lightweight regex safety check + block_reason = _kernel_check_code_safety(code) + if block_reason is not None: + logger.warning("Kernel safety block: %s code=%r", block_reason, code[:200]) + error_output = MessageType.ERROR.to_notebook_output({ + "ename": "SecurityError", + "evalue": block_reason, + "traceback": [f"SecurityError: {block_reason}"], + }) + return ExecuteResponse( + notebook_outputs=[dict(error_output)] if error_output else [], + error_occurred=True, + execution_time=0.0, + ) effective_timeout = timeout if timeout is not None else self.default_timeout start_time = time.perf_counter() @@ -437,6 +440,7 @@ async def run_server(work_dir: Path, language: NBLanguage, port: int = 8000, sta parser.add_argument("--language", type=str, default="python", choices=["python", "r"]) parser.add_argument("--port", type=int, default=8000) parser.add_argument("--startup-token", type=str, default="") + parser.add_argument("--safe-execute", action='store_true') args = parser.parse_args() language = NBLanguage.PYTHON if args.language == "python" else NBLanguage.R From 7648d855ef6603d0fa309dae41b1c4ee175d0db9 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Sun, 15 Mar 2026 14:41:17 -0700 Subject: [PATCH 23/31] added usage examples back to tool schemas --- src/hypotest/env/interpreter_env.py | 5 +++++ src/hypotest/env/tools/filesystem.py | 6 ++++++ 2 files changed, 11 insertions(+) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index b1b484b..f289495 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -727,6 +727,7 @@ def __init__( self._http_client: httpx.AsyncClient | None = None self._container_log_path: Path | None = None self._container_log_file: Any | None = None + self.kernel_container: EnrootKernelServer | None = None # Initialize notebook structure for state tracking self.nb: NotebookNode = nbformat.v4.new_notebook() @@ -1408,6 +1409,10 @@ async def run_cell( This method allows running code in a new cell (append) or re-running an existing cell with updated code. + Usage Examples: + run_cell("print('Hello, world!')") # Run code in new cell + run_cell("print('Hello, world!')", idx=0) # Run code in existing cell at index 0 + Error Recovery: When a cell fails with an error, you MUST fix it by calling run_cell with the corrected code and the SAME idx as the failed cell. diff --git a/src/hypotest/env/tools/filesystem.py b/src/hypotest/env/tools/filesystem.py index de036cc..2d25686 100644 --- a/src/hypotest/env/tools/filesystem.py +++ b/src/hypotest/env/tools/filesystem.py @@ -343,6 +343,12 @@ def list_dir_tool( overwhelming the context with too many files. Use this tool instead of writing code to list directories to avoid context bloat. + Usage Examples: + list_dir() # List working directory + list_dir("data/") # List specific folder + list_dir(max_files=50) # Show more files + list_dir(show_hidden=True) # Include hidden files + Args: directory: Directory path to list (default: current working directory) max_files: Maximum number of files to display (default: 20) From 2d7dcaa557df6062fb24e48cfb4fb1c22cfbebb8 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Sun, 15 Mar 2026 14:46:41 -0700 Subject: [PATCH 24/31] added redundancy check on kernel_container --- src/hypotest/env/interpreter_env.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index f289495..b929232 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -1480,7 +1480,7 @@ async def reset_kernel(self) -> str: This clears all variables and execution state. """ - if self.state.use_ray: + if self.state.use_ray and self.state.kernel_container is not None: reset_ref = self.state.kernel_container._reset_via_http.remote() await reset_ref elif self.state.use_docker or self.state.use_enroot: From f8cae4d5042a712c65e679d9899163882adac6b4 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Sun, 15 Mar 2026 14:51:11 -0700 Subject: [PATCH 25/31] added usage example for error recovery in run_cell --- src/hypotest/env/interpreter_env.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index b929232..e031c2e 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -1415,7 +1415,9 @@ async def run_cell( Error Recovery: When a cell fails with an error, you MUST fix it by calling run_cell - with the corrected code and the SAME idx as the failed cell. + with the corrected code and the SAME idx as the failed cell: + + run_cell("corrected_code", idx=3) # Fix error in Cell #3 The cell number is shown in the output prefix (e.g., "[Cell #3]"). Do NOT create a new cell to fix an error - always edit the failed cell. From 30bf7df1984a45308a422e73de30b846730b432b Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Sun, 15 Mar 2026 15:00:21 -0700 Subject: [PATCH 26/31] shortened safety warning --- src/hypotest/env/interpreter_env.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index e031c2e..e920e8a 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -1207,7 +1207,7 @@ async def execute_and_add_cell( actual_idx = cell_idx return result, actual_idx else: - logger.warning("WARNING: You are running code sandbox with unsafe execution, this may result in destructive commands being run on the system") + logger.warning("Running code sandbox without safety filter, may result in destructive code running on the node") if self.use_ray and self.use_enroot: result_ref = self.kernel_container._execute_via_http.remote(code, timeout) From ba43a6a3e36ac618aee7e4a4b662c0385e8912a6 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Sun, 15 Mar 2026 15:17:52 -0700 Subject: [PATCH 27/31] replaced suprocess.Popen with asyncio.create_subprocess_exec --- src/hypotest/env/interpreter_env.py | 46 +++++++++++++---------------- 1 file changed, 21 insertions(+), 25 deletions(-) diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index e920e8a..5768a72 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -98,7 +98,7 @@ class _PortCollisionError(Exception): async def _poll_kernel_health( # noqa: PLR0912 http_client: httpx.AsyncClient, - enroot_proc: subprocess.Popen[str] | None, + enroot_proc: asyncio.subprocess.Process | None, container_port: int | None, expected_startup_token: str | None, read_log_tail: Callable[[int], str], @@ -121,7 +121,7 @@ async def _poll_kernel_health( # noqa: PLR0912 elapsed = time.perf_counter() - start_time # Check if the enroot process has died before we even get an HTTP response - if enroot_proc is not None and enroot_proc.poll() is not None: + if enroot_proc is not None and enroot_proc.returncode is not None: rc = enroot_proc.returncode log_tail = read_log_tail(1000) raise RuntimeError( @@ -175,7 +175,7 @@ async def _poll_kernel_health( # noqa: PLR0912 # Log progress every 5s if poll_count % 10 == 0: - proc_alive = enroot_proc.poll() is None if enroot_proc else False + proc_alive = enroot_proc.returncode is None if enroot_proc else False logger.log( _CONTAINER_LOG_LEVEL, "[%s] Health poll #%d at %.1fs: last_status=%s, process_alive=%s", @@ -202,12 +202,12 @@ async def _poll_kernel_health( # noqa: PLR0912 ) -def _kill_process_group(proc: subprocess.Popen, label: str = "enroot", sigterm_timeout: float = 15) -> None: +async def _kill_process_group(proc: asyncio.subprocess.Process, label: str = "enroot", sigterm_timeout: float = 15) -> None: """Safely terminate a process group, escalating from SIGTERM to SIGKILL. Handles all edge cases: already-dead process, missing process group, etc. """ - if proc.poll() is not None: + if proc.returncode is not None: logger.log( _CONTAINER_LOG_LEVEL, "[%s] Process pid=%d already exited with returncode=%d", @@ -233,8 +233,7 @@ def _kill_process_group(proc: subprocess.Popen, label: str = "enroot", sigterm_t return try: - # proc.wait(timeout=sigterm_timeout) - proc.communicate(timeout=sigterm_timeout) + await asyncio.wait_for(proc.communicate(), timeout=sigterm_timeout) logger.log( _CONTAINER_LOG_LEVEL, "[%s] Process pid=%d exited after SIGTERM with returncode=%d", @@ -243,7 +242,7 @@ def _kill_process_group(proc: subprocess.Popen, label: str = "enroot", sigterm_t proc.returncode, ) return - except subprocess.TimeoutExpired: + except asyncio.TimeoutError: logger.warning( "[%s] Process pid=%d did not exit within %.1fs of SIGTERM, sending SIGKILL", label, @@ -259,8 +258,7 @@ def _kill_process_group(proc: subprocess.Popen, label: str = "enroot", sigterm_t return try: - # proc.wait(timeout=5) - proc.communicate(timeout=5) + await asyncio.wait_for(proc.communicate(), timeout=5) logger.log( _CONTAINER_LOG_LEVEL, "[%s] Process pid=%d exited after SIGKILL with returncode=%d", @@ -268,7 +266,7 @@ def _kill_process_group(proc: subprocess.Popen, label: str = "enroot", sigterm_t proc.pid, proc.returncode, ) - except subprocess.TimeoutExpired: + except asyncio.TimeoutError: logger.exception("[%s] Process pid=%d still alive after SIGKILL — possible zombie", label, proc.pid) @@ -316,7 +314,7 @@ def __init__(self, container_sqsh_path: Path, execution_timeout: float, safe_exe self.container_sqsh_path = container_sqsh_path self.execution_timeout = execution_timeout self.safe_execute = safe_execute - self._enroot_proc: subprocess.Popen[str] | None = None + self._enroot_proc: asyncio.subprocess.Process | None = None self._http_client: httpx.AsyncClient | None = None self._container_port: int | None = None self._container_log_path: Path | None = None @@ -453,9 +451,8 @@ async def initialize(self, work_dir: Path, language: NBLanguage) -> None: log_dir.mkdir(exist_ok=True) self._container_log_path = log_dir / "container.log" self._container_log_file = open(self._container_log_path, "w", encoding="utf-8") # noqa: SIM115 - self._enroot_proc = subprocess.Popen( - cmd, - text=True, + self._enroot_proc = await asyncio.create_subprocess_exec( + *cmd, start_new_session=True, stdout=self._container_log_file, stderr=subprocess.STDOUT, @@ -509,7 +506,7 @@ def _log_container_failure(self, attempt: int, launch_ms: float, error: Exceptio ] if proc is not None: - rc = proc.poll() + rc = proc.returncode diag_parts.append(f"process_alive={rc is None}") if rc is not None: diag_parts.append(f"returncode={rc}") @@ -561,7 +558,7 @@ async def _cleanup_failed_startup(self) -> None: self._http_client = None if self._enroot_proc is not None: - _kill_process_group(self._enroot_proc, label=label) + await _kill_process_group(self._enroot_proc, label=label) self._enroot_proc = None self._close_container_log() @@ -664,7 +661,7 @@ async def close(self): self._http_client = None if self._enroot_proc is not None: - _kill_process_group(self._enroot_proc, label=label) + await _kill_process_group(self._enroot_proc, label=label) self._enroot_proc = None self._close_container_log() @@ -722,7 +719,7 @@ def __init__( # Docker/Enroot container state (only used when use_docker=True or use_enroot=True) self._docker_client: aiodocker.Docker | None = None self._container: DockerContainer | None = None - self._enroot_proc = None + self._enroot_proc: asyncio.subprocess.Process | None = None self._container_port: int | None = None self._http_client: httpx.AsyncClient | None = None self._container_log_path: Path | None = None @@ -823,9 +820,8 @@ async def _start_enroot_container(self) -> None: (self.work_dir / ".container_logs").mkdir(exist_ok=True) self._container_log_path = self.work_dir / ".container_logs" / "container.log" self._container_log_file = open(self._container_log_path, "w", encoding="utf-8") # noqa: SIM115 - self._enroot_proc = subprocess.Popen( - cmd, - text=True, + self._enroot_proc = await asyncio.create_subprocess_exec( + *cmd, start_new_session=True, stdout=self._container_log_file, stderr=subprocess.STDOUT, @@ -872,7 +868,7 @@ def _log_enroot_container_failure(self, attempt: int, launch_ms: float, error: E label = self._enroot_label() diag = [f"attempt=#{attempt}", f"elapsed={launch_ms:.0f}ms", f"error={error!r}"] if self._enroot_proc is not None: - rc = self._enroot_proc.poll() + rc = self._enroot_proc.returncode diag.append(f"process_alive={rc is None}") if rc is not None: diag.append(f"returncode={rc}") @@ -920,7 +916,7 @@ async def _cleanup_failed_startup(self) -> None: self._http_client = None if self._enroot_proc is not None: - _kill_process_group(self._enroot_proc, label=label) + await _kill_process_group(self._enroot_proc, label=label) self._enroot_proc = None self._close_container_log() @@ -1085,7 +1081,7 @@ async def close(self): self._docker_client = None if self._enroot_proc is not None: - _kill_process_group(self._enroot_proc, label=self._enroot_label()) + await _kill_process_group(self._enroot_proc, label=self._enroot_label()) self._enroot_proc = None self._close_container_log() From 471765d7db39131b970e0544e9952113e79bf397 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Sun, 15 Mar 2026 23:15:38 -0700 Subject: [PATCH 28/31] refactored container logging to be async, addressed comments --- src/hypotest/env/code_safety.py | 1 - src/hypotest/env/config.py | 4 +- src/hypotest/env/interpreter_env.py | 129 ++++++++++++++++------------ 3 files changed, 78 insertions(+), 56 deletions(-) diff --git a/src/hypotest/env/code_safety.py b/src/hypotest/env/code_safety.py index fad9908..9f12279 100644 --- a/src/hypotest/env/code_safety.py +++ b/src/hypotest/env/code_safety.py @@ -91,7 +91,6 @@ # Modules whose import_module() calls should be checked _DANGEROUS_IMPORTLIB_TARGETS: frozenset[str] = _BLOCKED_MODULES | frozenset({ - "os", "subprocess", "shutil", "multiprocessing", diff --git a/src/hypotest/env/config.py b/src/hypotest/env/config.py index d8ecf09..d8ed10a 100644 --- a/src/hypotest/env/config.py +++ b/src/hypotest/env/config.py @@ -16,7 +16,7 @@ # Kernel server settings (for Docker-based execution) KERNEL_SERVER_PORT = 8000 KERNEL_SERVER_STARTUP_TIMEOUT = float( - os.getenv("KERNEL_SERVER_STARTUP_TIMEOUT", "120.0") + os.getenv("KERNEL_SERVER_STARTUP_TIMEOUT", "30.0") ) # seconds to wait for health check MAX_FILES_TO_UPLOAD = int(os.getenv("MAX_FILES_TO_UPLOAD", "100")) @@ -59,7 +59,7 @@ class ExecutionConfig(BaseModel): cell_execution_timeout: int = 15 * 60 # safety - safe_execute: bool = True + safe_execute: bool = bool(os.getenv("SAFE_EXECUTE_SANDBOX", False)) # Capabilities has_gpu: bool = False diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index 5768a72..f190858 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -17,7 +17,7 @@ import sys import time import uuid -from collections.abc import Callable +from collections.abc import Awaitable, Callable from pathlib import Path from tempfile import mkdtemp from textwrap import dedent @@ -71,6 +71,8 @@ _RETRY_BASE_SLEEP = 1.0 _RETRY_MAX_SLEEP = 16.0 +WARNED_UNSAFE_EXECUTION = False + async def get_free_port() -> int: """Get a free port for the kernel server container.""" @@ -101,7 +103,7 @@ async def _poll_kernel_health( # noqa: PLR0912 enroot_proc: asyncio.subprocess.Process | None, container_port: int | None, expected_startup_token: str | None, - read_log_tail: Callable[[int], str], + read_log_tail: Callable[[int], Awaitable[str]], label: str, ) -> None: """Poll kernel server /health endpoint until ready, with token validation and timeout.""" @@ -123,7 +125,7 @@ async def _poll_kernel_health( # noqa: PLR0912 # Check if the enroot process has died before we even get an HTTP response if enroot_proc is not None and enroot_proc.returncode is not None: rc = enroot_proc.returncode - log_tail = read_log_tail(1000) + log_tail = await read_log_tail(1000) raise RuntimeError( f"Enroot process exited prematurely with returncode={rc} after {elapsed:.1f}s. log: {log_tail!r}" ) @@ -194,7 +196,7 @@ async def _poll_kernel_health( # noqa: PLR0912 f"Port {container_port} health-check timed out with token_mismatch " f"({poll_count} polls, elapsed={total_elapsed:.1f}s)" ) - log_tail = read_log_tail(500) + log_tail = await read_log_tail(500) raise TimeoutError( f"Kernel server did not become healthy within {cfg.KERNEL_SERVER_STARTUP_TIMEOUT}s " f"({poll_count} polls, last_status={last_status}, elapsed={total_elapsed:.1f}s)" @@ -429,7 +431,7 @@ async def initialize(self, work_dir: Path, language: NBLanguage) -> None: while not online: attempt += 1 if attempt > MAX_CONTAINER_LAUNCH_RETRIES: - log_tail = self._read_container_log_tail(500) + log_tail = await self._read_container_log_tail(500) raise RuntimeError( f"Container failed to start after {MAX_CONTAINER_LAUNCH_RETRIES} attempts " f"(last_error={last_err!r})" @@ -450,7 +452,9 @@ async def initialize(self, work_dir: Path, language: NBLanguage) -> None: log_dir = work_dir / ".container_logs" log_dir.mkdir(exist_ok=True) self._container_log_path = log_dir / "container.log" - self._container_log_file = open(self._container_log_path, "w", encoding="utf-8") # noqa: SIM115 + self._container_log_file = await asyncio.to_thread( + open, self._container_log_path, "w", encoding="utf-8" + ) self._enroot_proc = await asyncio.create_subprocess_exec( *cmd, start_new_session=True, @@ -488,13 +492,13 @@ async def initialize(self, work_dir: Path, language: NBLanguage) -> None: except Exception as e: last_err = e launch_ms = (time.perf_counter() - launch_t0) * 1000.0 - self._log_container_failure(attempt, launch_ms, e) + await self._log_container_failure(attempt, launch_ms, e) await self._cleanup_failed_startup() if not isinstance(e, _PortCollisionError): backoff = min(_RETRY_BASE_SLEEP * 2 ** (attempt - 1), _RETRY_MAX_SLEEP) await asyncio.sleep(backoff) - def _log_container_failure(self, attempt: int, launch_ms: float, error: Exception) -> None: + async def _log_container_failure(self, attempt: int, launch_ms: float, error: Exception) -> None: """Log detailed diagnostics when a container fails to start.""" label = self._proc_label() proc = self._enroot_proc @@ -514,7 +518,7 @@ def _log_container_failure(self, attempt: int, launch_ms: float, error: Exceptio logger.warning("[%s] Container startup FAILED: %s", label, ", ".join(diag_parts)) # Log container output separately so tracebacks are readable - log_tail = self._read_container_log_tail() + log_tail = await self._read_container_log_tail() if log_tail: logger.warning( "[%s] Container log output (last %d chars):\n%s", @@ -523,26 +527,33 @@ def _log_container_failure(self, attempt: int, launch_ms: float, error: Exceptio log_tail, ) - def _read_container_log_tail(self, max_chars: int = 2000) -> str: + async def _read_container_log_tail(self, max_chars: int = 2000) -> str: """Read the tail of the container log file for diagnostics.""" - if self._container_log_path is None or not self._container_log_path.exists(): - return "" - try: - # Flush so any buffered output is written to disk - if self._container_log_file and not self._container_log_file.closed: - self._container_log_file.flush() - text = self._container_log_path.read_text() - return text[-max_chars:] if len(text) > max_chars else text - except Exception: - return "" - - def _close_container_log(self) -> None: + def _read() -> str: + if self._container_log_path is None or not self._container_log_path.exists(): + return "" + try: + if self._container_log_file and not self._container_log_file.closed: + self._container_log_file.flush() + text = self._container_log_path.read_text() + return text[-max_chars:] if len(text) > max_chars else text + except Exception: + return "" + + return await asyncio.to_thread(_read) + + async def _close_container_log(self) -> None: """Close the container log file handle.""" if self._container_log_file is not None: - with contextlib.suppress(Exception): - self._container_log_file.close() + f = self._container_log_file self._container_log_file = None + def _do_close() -> None: + with contextlib.suppress(Exception): + f.close() + + await asyncio.to_thread(_do_close) + async def _cleanup_failed_startup(self) -> None: """Best-effort cleanup for failed startup attempts before retrying.""" label = self._proc_label() @@ -561,7 +572,7 @@ async def _cleanup_failed_startup(self) -> None: await _kill_process_group(self._enroot_proc, label=label) self._enroot_proc = None - self._close_container_log() + await self._close_container_log() async def _wait_for_health(self, expected_startup_token: str | None = None) -> None: """Wait for the kernel server to become healthy.""" @@ -664,7 +675,7 @@ async def close(self): await _kill_process_group(self._enroot_proc, label=label) self._enroot_proc = None - self._close_container_log() + await self._close_container_log() logger.log(_CONTAINER_LOG_LEVEL, "[%s] EnrootKernelServer closed", label) @@ -751,8 +762,8 @@ async def start(self): await self.interpreter.start() async def _start_ray_enroot_container(self) -> None: - self.kernel_container = EnrootKernelServer.remote(self.container_sqsh_path, self.execution_timeout, safe_execute=self.safe_execute) - init_ref = self.kernel_container.initialize.remote(self.work_dir, self.language) + self.kernel_container = EnrootKernelServer.remote(self.container_sqsh_path, self.execution_timeout, safe_execute=self.safe_execute) # type: ignore[attr-defined] + init_ref = self.kernel_container.initialize.remote(self.work_dir, self.language) # type: ignore[union-attr] await init_ref def _enroot_label(self) -> str: @@ -769,7 +780,7 @@ async def _start_enroot_container(self) -> None: while not online: attempt += 1 if attempt > MAX_CONTAINER_LAUNCH_RETRIES: - log_tail = self._read_container_log_tail(500) + log_tail = await self._read_container_log_tail(500) raise RuntimeError( f"Container failed to start after {MAX_CONTAINER_LAUNCH_RETRIES} attempts " f"(last_error={last_err!r})" @@ -819,7 +830,9 @@ async def _start_enroot_container(self) -> None: # for detailed rationale on why we avoid subprocess.PIPE here). (self.work_dir / ".container_logs").mkdir(exist_ok=True) self._container_log_path = self.work_dir / ".container_logs" / "container.log" - self._container_log_file = open(self._container_log_path, "w", encoding="utf-8") # noqa: SIM115 + self._container_log_file = await asyncio.to_thread( + open, self._container_log_path, "w", encoding="utf-8" + ) self._enroot_proc = await asyncio.create_subprocess_exec( *cmd, start_new_session=True, @@ -857,13 +870,13 @@ async def _start_enroot_container(self) -> None: except Exception as e: last_err = e launch_ms = (time.perf_counter() - launch_t0) * 1000.0 - self._log_enroot_container_failure(attempt, launch_ms, e) + await self._log_enroot_container_failure(attempt, launch_ms, e) await self._cleanup_failed_startup() if not isinstance(e, _PortCollisionError): backoff = min(_RETRY_BASE_SLEEP * 2 ** (attempt - 1), _RETRY_MAX_SLEEP) await asyncio.sleep(backoff) - def _log_enroot_container_failure(self, attempt: int, launch_ms: float, error: Exception) -> None: + async def _log_enroot_container_failure(self, attempt: int, launch_ms: float, error: Exception) -> None: """Log detailed diagnostics when a container fails to start.""" label = self._enroot_label() diag = [f"attempt=#{attempt}", f"elapsed={launch_ms:.0f}ms", f"error={error!r}"] @@ -873,7 +886,7 @@ def _log_enroot_container_failure(self, attempt: int, launch_ms: float, error: E if rc is not None: diag.append(f"returncode={rc}") logger.warning("[%s] Container startup FAILED: %s", label, ", ".join(diag)) - log_tail = self._read_container_log_tail() + log_tail = await self._read_container_log_tail() if log_tail: logger.warning( "[%s] Container log output (last %d chars):\n%s", @@ -882,25 +895,33 @@ def _log_enroot_container_failure(self, attempt: int, launch_ms: float, error: E log_tail, ) - def _read_container_log_tail(self, max_chars: int = 2000) -> str: + async def _read_container_log_tail(self, max_chars: int = 2000) -> str: """Read the tail of the container log file for diagnostics.""" - if self._container_log_path is None or not self._container_log_path.exists(): - return "" - try: - if self._container_log_file and not self._container_log_file.closed: - self._container_log_file.flush() - text = self._container_log_path.read_text() - return text[-max_chars:] if len(text) > max_chars else text - except Exception: - return "" - - def _close_container_log(self) -> None: + def _read() -> str: + if self._container_log_path is None or not self._container_log_path.exists(): + return "" + try: + if self._container_log_file and not self._container_log_file.closed: + self._container_log_file.flush() + text = self._container_log_path.read_text() + return text[-max_chars:] if len(text) > max_chars else text + except Exception: + return "" + + return await asyncio.to_thread(_read) + + async def _close_container_log(self) -> None: """Close the container log file handle.""" if self._container_log_file is not None: - with contextlib.suppress(Exception): - self._container_log_file.close() + f = self._container_log_file self._container_log_file = None + def _do_close() -> None: + with contextlib.suppress(Exception): + f.close() + + await asyncio.to_thread(_do_close) + async def _cleanup_failed_startup(self) -> None: """Best-effort cleanup for failed startup attempts before retrying.""" label = self._enroot_label() @@ -919,7 +940,7 @@ async def _cleanup_failed_startup(self) -> None: await _kill_process_group(self._enroot_proc, label=label) self._enroot_proc = None - self._close_container_log() + await self._close_container_log() async def _start_docker_container(self) -> None: """Start a Docker container with the kernel server.""" @@ -1054,7 +1075,7 @@ async def close(self): if self.use_ray and self.use_enroot: if self.kernel_container is not None: - close_ref = self.kernel_container.close.remote() + close_ref = self.kernel_container.close.remote() # type: ignore[attr-defined] await close_ref self.kernel_container = None @@ -1084,7 +1105,7 @@ async def close(self): await _kill_process_group(self._enroot_proc, label=self._enroot_label()) self._enroot_proc = None - self._close_container_log() + await self._close_container_log() elif self.interpreter is not None: await self.interpreter.close() @@ -1181,6 +1202,7 @@ async def execute_and_add_cell( Returns: Tuple of (ExecutionResult, actual_cell_index) """ + global WARNED_UNSAFE_EXECUTION if self.safe_execute: block_reason = check_code_safety(code, self.language) if block_reason is not None: @@ -1202,11 +1224,12 @@ async def execute_and_add_cell( self._update_cell(cell_idx, code, result) actual_idx = cell_idx return result, actual_idx - else: + elif not WARNED_UNSAFE_EXECUTION: logger.warning("Running code sandbox without safety filter, may result in destructive code running on the node") + WARNED_UNSAFE_EXECUTION = True if self.use_ray and self.use_enroot: - result_ref = self.kernel_container._execute_via_http.remote(code, timeout) + result_ref = self.kernel_container._execute_via_http.remote(code, timeout) # type: ignore[union-attr] result = await result_ref elif self.use_docker or self.use_enroot: result = await self._execute_via_http(code, timeout) @@ -1479,7 +1502,7 @@ async def reset_kernel(self) -> str: This clears all variables and execution state. """ if self.state.use_ray and self.state.kernel_container is not None: - reset_ref = self.state.kernel_container._reset_via_http.remote() + reset_ref = self.state.kernel_container._reset_via_http.remote() # type: ignore[attr-defined] await reset_ref elif self.state.use_docker or self.state.use_enroot: await self.state._reset_via_http() From 94faebae7febb05c0504f8197ec17445fd845588 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Sun, 15 Mar 2026 23:16:43 -0700 Subject: [PATCH 29/31] readded str->int fallback to list_dir() --- src/hypotest/env/tools/filesystem.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/hypotest/env/tools/filesystem.py b/src/hypotest/env/tools/filesystem.py index 2d25686..ded5fbc 100644 --- a/src/hypotest/env/tools/filesystem.py +++ b/src/hypotest/env/tools/filesystem.py @@ -354,6 +354,14 @@ def list_dir_tool( max_files: Maximum number of files to display (default: 20) show_hidden: Whether to show hidden files starting with '.' (default: False) """ + # LLM can sometimes generate string instead of int; fallback to default when typecast fails. + try: + max_files = int(max_files) + except ValueError: + max_files = 20 + + show_hidden = bool(show_hidden) + try: dir_path = Path(directory).resolve() if not dir_path.exists() or not dir_path.is_dir(): From ca8cd4701f348983d17095d7bb09706caa0e0cc5 Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Sun, 15 Mar 2026 23:46:55 -0700 Subject: [PATCH 30/31] addressed linter errors, cleaned up impl --- pyproject.toml | 2 +- src/hypotest/dataset_server.py | 17 ++++++--- src/hypotest/env/code_safety.py | 8 +--- src/hypotest/env/config.py | 3 +- src/hypotest/env/interpreter_env.py | 56 +++++++++++++++++----------- src/hypotest/env/kernel_server.py | 2 +- src/hypotest/env/tools/filesystem.py | 2 +- 7 files changed, 52 insertions(+), 38 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 14cd503..684f0db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,8 +34,8 @@ dependencies = [ "numpy", "pillow", "pydantic", - "tenacity>=8.0", "ray", + "tenacity>=8.0", ] description = "hypotest package" keywords = [] diff --git a/src/hypotest/dataset_server.py b/src/hypotest/dataset_server.py index 76599a1..37b2e12 100644 --- a/src/hypotest/dataset_server.py +++ b/src/hypotest/dataset_server.py @@ -21,7 +21,7 @@ from hypotest.env.kernel_server import NBLanguage -class HypotestDatasetConfig(BaseModel): +class DatasetConfig(BaseModel): model_config = ConfigDict(extra="forbid") problem_jsonl: FilePath | None = None @@ -68,8 +68,8 @@ def _load_from_hf(self) -> list[ProblemInstance]: return [ProblemInstance.model_validate(row) for row in ds] -class HypotestDataset(TaskDataset[InterpreterEnv]): - def __init__(self, config: HypotestDatasetConfig): +class Dataset(TaskDataset[InterpreterEnv]): + def __init__(self, config: DatasetConfig): self.config = config self.problems = self.config.load_problems() @@ -98,6 +98,7 @@ def get_new_env_by_idx(self, idx: int) -> InterpreterEnv: language = ( NBLanguage.PYTHON if self.config.force_python else NBLanguage.from_string(problem.nb_primary_language) ) + language = language if language is not None else NBLanguage.PYTHON # default auto language to python return InterpreterEnv( problem=problem, @@ -111,13 +112,17 @@ def __len__(self) -> int: return len(self.problems) +HypotestDataset = Dataset +HypotestDatasetConfig = DatasetConfig + + DEFAULT_SERVER_PORT = 8405 class ServerConfig(BaseModel): model_config = ConfigDict(extra="forbid") - dataset: HypotestDatasetConfig + dataset: DatasetConfig api_key: str = Field( description="API key to access the server; passed either by value or as an environment variable." ) @@ -154,7 +159,7 @@ async def launch_server(): config = ServerConfig.model_validate(yaml.safe_load(args.config.read_text())) else: config = ServerConfig( - dataset=HypotestDatasetConfig( + dataset=DatasetConfig( problem_jsonl=args.problem_jsonl, hf_dataset=args.hf_dataset, capsule_dir=args.capsule_dir, @@ -180,7 +185,7 @@ async def launch_server(): api_key=args.api_key, ) - dataset = HypotestDataset(config.dataset) + dataset = Dataset(config.dataset) server = TaskDatasetServer(dataset, port=config.port, api_key=config.api_key) ip_address = socket.gethostbyname(socket.gethostname()) diff --git a/src/hypotest/env/code_safety.py b/src/hypotest/env/code_safety.py index 9f12279..26c4a17 100644 --- a/src/hypotest/env/code_safety.py +++ b/src/hypotest/env/code_safety.py @@ -10,12 +10,10 @@ import ast import logging import re -from typing import TYPE_CHECKING -logger = logging.getLogger(__name__) +from .kernel_server import NBLanguage -if TYPE_CHECKING: - from .kernel_server import NBLanguage +logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- @@ -371,8 +369,6 @@ def check_code_safety(code: str, language: NBLanguage) -> str | None: Returns: None if the code is safe, or a category-level error message if blocked. """ - from .kernel_server import NBLanguage - # Bash-level patterns apply to both languages bash_result = _check_bash_patterns(code) if bash_result is not None: diff --git a/src/hypotest/env/config.py b/src/hypotest/env/config.py index d8ed10a..96cbde9 100644 --- a/src/hypotest/env/config.py +++ b/src/hypotest/env/config.py @@ -12,6 +12,7 @@ ) NB_ENVIRONMENT_DOCKER_IMAGE = os.getenv("NB_ENVIRONMENT_DOCKER_IMAGE", "interpreter-env:latest") KERNEL_ENV_PATH = os.getenv("KERNEL_ENV_PATH", "/app/kernel_env") +CONTAINER_WORKSPACE_PREFIX = os.getenv("CONTAINER_WORKSPACE_PREFIX", "/tmp/data_workspace") # noqa: S108 # Kernel server settings (for Docker-based execution) KERNEL_SERVER_PORT = 8000 @@ -59,7 +60,7 @@ class ExecutionConfig(BaseModel): cell_execution_timeout: int = 15 * 60 # safety - safe_execute: bool = bool(os.getenv("SAFE_EXECUTE_SANDBOX", False)) + safe_execute: bool = os.getenv("SAFE_EXECUTE_SANDBOX", "").lower() == "true" # Capabilities has_gpu: bool = False diff --git a/src/hypotest/env/interpreter_env.py b/src/hypotest/env/interpreter_env.py index f190858..3a1ea62 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -71,7 +71,7 @@ _RETRY_BASE_SLEEP = 1.0 _RETRY_MAX_SLEEP = 16.0 -WARNED_UNSAFE_EXECUTION = False +_warned_unsafe_execution: set[str] = set() async def get_free_port() -> int: @@ -204,7 +204,9 @@ async def _poll_kernel_health( # noqa: PLR0912 ) -async def _kill_process_group(proc: asyncio.subprocess.Process, label: str = "enroot", sigterm_timeout: float = 15) -> None: +async def _kill_process_group( + proc: asyncio.subprocess.Process, label: str = "enroot", sigterm_timeout: float = 15 +) -> None: """Safely terminate a process group, escalating from SIGTERM to SIGKILL. Handles all edge cases: already-dead process, missing process group, etc. @@ -236,6 +238,14 @@ async def _kill_process_group(proc: asyncio.subprocess.Process, label: str = "en try: await asyncio.wait_for(proc.communicate(), timeout=sigterm_timeout) + except TimeoutError: + logger.warning( + "[%s] Process pid=%d did not exit within %.1fs of SIGTERM, sending SIGKILL", + label, + proc.pid, + sigterm_timeout, + ) + else: logger.log( _CONTAINER_LOG_LEVEL, "[%s] Process pid=%d exited after SIGTERM with returncode=%d", @@ -244,13 +254,6 @@ async def _kill_process_group(proc: asyncio.subprocess.Process, label: str = "en proc.returncode, ) return - except asyncio.TimeoutError: - logger.warning( - "[%s] Process pid=%d did not exit within %.1fs of SIGTERM, sending SIGKILL", - label, - proc.pid, - sigterm_timeout, - ) # SIGKILL the whole group try: @@ -268,7 +271,7 @@ async def _kill_process_group(proc: asyncio.subprocess.Process, label: str = "en proc.pid, proc.returncode, ) - except asyncio.TimeoutError: + except TimeoutError: logger.exception("[%s] Process pid=%d still alive after SIGKILL — possible zombie", label, proc.pid) @@ -286,7 +289,7 @@ class ProblemInstance(BaseModel): @model_validator(mode="before") @classmethod def handle_language(cls, data: dict) -> dict: - if data.get("nb_primary_language", None) is None: + if data.get("nb_primary_language") is None: data["nb_primary_language"] = str(NBLanguage.PYTHON) return data @@ -329,7 +332,9 @@ def _proc_label(self) -> str: return f"enroot(port={port}, pid={pid})" @staticmethod - def _build_kernel_bash_script(node_workdir: str, language: NBLanguage, port: int, startup_token: str, safe_execute: bool = True) -> str: + def _build_kernel_bash_script( + node_workdir: str, language: NBLanguage, port: int, startup_token: str, safe_execute: bool = True + ) -> str: """Build the bash script that sets up the workspace and launches the kernel server.""" return dedent(f"""\ set -euo pipefail @@ -361,7 +366,7 @@ def _build_kernel_bash_script(node_workdir: str, language: NBLanguage, port: int --work_dir $WORKDIR \\ --language {language.value} \\ --port {port} \\ - --startup-token {startup_token} {'--safe-execute' if safe_execute else ''} + --startup-token {startup_token} {"--safe-execute" if safe_execute else ""} """).strip() @staticmethod @@ -416,7 +421,7 @@ def _build_enroot_cmd( async def initialize(self, work_dir: Path, language: NBLanguage) -> None: startup_token = str(uuid.uuid4()) - node_workdir = f"/tmp/data_workspace.{startup_token.split('-', maxsplit=1)[0]}" + node_workdir = f"{cfg.CONTAINER_WORKSPACE_PREFIX}.{startup_token.split('-', maxsplit=1)[0]}" _prep_workspace_dir(str(work_dir), workspace_path=node_workdir) @@ -439,7 +444,9 @@ async def initialize(self, work_dir: Path, language: NBLanguage) -> None: ) self._container_port = await get_free_port() - bash = self._build_kernel_bash_script(node_workdir, language, self._container_port, startup_token, safe_execute=self.safe_execute) + bash = self._build_kernel_bash_script( + node_workdir, language, self._container_port, startup_token, safe_execute=self.safe_execute + ) cmd = self._build_enroot_cmd(work_dir, kernel_server_path, bash, enroot_env, self.container_sqsh_path) async with CONTAINER_LAUNCH_SEM: @@ -529,6 +536,7 @@ async def _log_container_failure(self, attempt: int, launch_ms: float, error: Ex async def _read_container_log_tail(self, max_chars: int = 2000) -> str: """Read the tail of the container log file for diagnostics.""" + def _read() -> str: if self._container_log_path is None or not self._container_log_path.exists(): return "" @@ -762,7 +770,9 @@ async def start(self): await self.interpreter.start() async def _start_ray_enroot_container(self) -> None: - self.kernel_container = EnrootKernelServer.remote(self.container_sqsh_path, self.execution_timeout, safe_execute=self.safe_execute) # type: ignore[attr-defined] + self.kernel_container = EnrootKernelServer.remote( # type: ignore[attr-defined] + self.container_sqsh_path, self.execution_timeout, safe_execute=self.safe_execute + ) init_ref = self.kernel_container.initialize.remote(self.work_dir, self.language) # type: ignore[union-attr] await init_ref @@ -799,7 +809,7 @@ async def _start_enroot_container(self) -> None: --work_dir /data_workspace \\ --language {self.language.value} \\ --port {self._container_port} \\ - --startup-token {startup_token} {'--safe-execute' if self.safe_execute else ''} + --startup-token {startup_token} {"--safe-execute" if self.safe_execute else ""} """).strip() kernel_server_path = Path(__file__).parent / "kernel_server.py" @@ -897,6 +907,7 @@ async def _log_enroot_container_failure(self, attempt: int, launch_ms: float, er async def _read_container_log_tail(self, max_chars: int = 2000) -> str: """Read the tail of the container log file for diagnostics.""" + def _read() -> str: if self._container_log_path is None or not self._container_log_path.exists(): return "" @@ -959,7 +970,7 @@ async def _start_docker_container(self) -> None: startup_token, ] if self.safe_execute: - cmd_list += ['--safe-execute'] + cmd_list += ["--safe-execute"] docker_config = { "Image": cfg.NB_ENVIRONMENT_DOCKER_IMAGE, @@ -1202,7 +1213,6 @@ async def execute_and_add_cell( Returns: Tuple of (ExecutionResult, actual_cell_index) """ - global WARNED_UNSAFE_EXECUTION if self.safe_execute: block_reason = check_code_safety(code, self.language) if block_reason is not None: @@ -1224,9 +1234,11 @@ async def execute_and_add_cell( self._update_cell(cell_idx, code, result) actual_idx = cell_idx return result, actual_idx - elif not WARNED_UNSAFE_EXECUTION: - logger.warning("Running code sandbox without safety filter, may result in destructive code running on the node") - WARNED_UNSAFE_EXECUTION = True + elif "unsafe_execution" not in _warned_unsafe_execution: + logger.warning( + "Running code sandbox without safety filter, may result in destructive code running on the node" + ) + _warned_unsafe_execution.add("unsafe_execution") if self.use_ray and self.use_enroot: result_ref = self.kernel_container._execute_via_http.remote(code, timeout) # type: ignore[union-attr] diff --git a/src/hypotest/env/kernel_server.py b/src/hypotest/env/kernel_server.py index 6dc3eb5..565d76f 100644 --- a/src/hypotest/env/kernel_server.py +++ b/src/hypotest/env/kernel_server.py @@ -440,7 +440,7 @@ async def run_server(work_dir: Path, language: NBLanguage, port: int = 8000, sta parser.add_argument("--language", type=str, default="python", choices=["python", "r"]) parser.add_argument("--port", type=int, default=8000) parser.add_argument("--startup-token", type=str, default="") - parser.add_argument("--safe-execute", action='store_true') + parser.add_argument("--safe-execute", action="store_true") args = parser.parse_args() language = NBLanguage.PYTHON if args.language == "python" else NBLanguage.R diff --git a/src/hypotest/env/tools/filesystem.py b/src/hypotest/env/tools/filesystem.py index ded5fbc..8e57805 100644 --- a/src/hypotest/env/tools/filesystem.py +++ b/src/hypotest/env/tools/filesystem.py @@ -361,7 +361,7 @@ def list_dir_tool( max_files = 20 show_hidden = bool(show_hidden) - + try: dir_path = Path(directory).resolve() if not dir_path.exists() or not dir_path.is_dir(): From de76f10ea1189dcc6a5cb5489aca32bf4251b6ac Mon Sep 17 00:00:00 2001 From: Arnav Komaragiri Date: Sun, 15 Mar 2026 23:54:11 -0700 Subject: [PATCH 31/31] set execution config to proper type in dataset config --- src/hypotest/dataset_server.py | 3 +- uv.lock | 88 +++++++++++++++++++++++++++++++++- 2 files changed, 88 insertions(+), 3 deletions(-) diff --git a/src/hypotest/dataset_server.py b/src/hypotest/dataset_server.py index 37b2e12..9788ebe 100644 --- a/src/hypotest/dataset_server.py +++ b/src/hypotest/dataset_server.py @@ -17,6 +17,7 @@ from lmi import LiteLLMModel from pydantic import BaseModel, ConfigDict, DirectoryPath, Field, FilePath, field_validator, model_validator +from hypotest.env.config import ExecutionConfig from hypotest.env.interpreter_env import InterpreterEnv, InterpreterEnvConfig, ProblemInstance from hypotest.env.kernel_server import NBLanguage @@ -42,7 +43,7 @@ class DatasetConfig(BaseModel): normalize_reward: bool = True save_dir: Path | None = None - execution_config: dict[str, Any] + execution_config: ExecutionConfig = Field(default_factory=ExecutionConfig) @model_validator(mode="after") def validate_dataset_source(self) -> Self: diff --git a/uv.lock b/uv.lock index 6b2fd4d..4eacb7d 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.11, <3.14" resolution-markers = [ "sys_platform == 'win32'", @@ -993,6 +993,7 @@ dependencies = [ { name = "numpy" }, { name = "pillow" }, { name = "pydantic" }, + { name = "ray" }, { name = "tenacity" }, ] @@ -1016,6 +1017,9 @@ dev = [ { name = "types-tqdm" }, { name = "vcrpy" }, ] +ray = [ + { name = "ray" }, +] [package.dev-dependencies] dev = [ @@ -1049,13 +1053,15 @@ requires-dist = [ { name = "pytest-timeout", marker = "extra == 'dev'" }, { name = "pytest-timer", extras = ["colorama"], marker = "extra == 'dev'" }, { name = "pytest-xdist", marker = "extra == 'dev'" }, + { name = "ray" }, + { name = "ray", marker = "extra == 'ray'" }, { name = "tenacity", specifier = ">=8.0" }, { name = "types-flask-socketio", marker = "extra == 'dev'", specifier = ">=5.3.0" }, { name = "types-pyyaml", marker = "extra == 'dev'" }, { name = "types-tqdm", marker = "extra == 'dev'" }, { name = "vcrpy", marker = "extra == 'dev'", specifier = ">=8" }, ] -provides-extras = ["dev"] +provides-extras = ["dev", "ray"] [package.metadata.requires-dev] dev = [{ name = "hypotest", extras = ["dev"], editable = "." }] @@ -1609,6 +1615,41 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, ] +[[package]] +name = "msgpack" +version = "1.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4d/f2/bfb55a6236ed8725a96b0aa3acbd0ec17588e6a2c3b62a93eb513ed8783f/msgpack-1.1.2.tar.gz", hash = "sha256:3b60763c1373dd60f398488069bcdc703cd08a711477b5d480eecc9f9626f47e", size = 173581, upload-time = "2025-10-08T09:15:56.596Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2c/97/560d11202bcd537abca693fd85d81cebe2107ba17301de42b01ac1677b69/msgpack-1.1.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:2e86a607e558d22985d856948c12a3fa7b42efad264dca8a3ebbcfa2735d786c", size = 82271, upload-time = "2025-10-08T09:14:49.967Z" }, + { url = "https://files.pythonhosted.org/packages/83/04/28a41024ccbd67467380b6fb440ae916c1e4f25e2cd4c63abe6835ac566e/msgpack-1.1.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:283ae72fc89da59aa004ba147e8fc2f766647b1251500182fac0350d8af299c0", size = 84914, upload-time = "2025-10-08T09:14:50.958Z" }, + { url = "https://files.pythonhosted.org/packages/71/46/b817349db6886d79e57a966346cf0902a426375aadc1e8e7a86a75e22f19/msgpack-1.1.2-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:61c8aa3bd513d87c72ed0b37b53dd5c5a0f58f2ff9f26e1555d3bd7948fb7296", size = 416962, upload-time = "2025-10-08T09:14:51.997Z" }, + { url = "https://files.pythonhosted.org/packages/da/e0/6cc2e852837cd6086fe7d8406af4294e66827a60a4cf60b86575a4a65ca8/msgpack-1.1.2-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:454e29e186285d2ebe65be34629fa0e8605202c60fbc7c4c650ccd41870896ef", size = 426183, upload-time = "2025-10-08T09:14:53.477Z" }, + { url = "https://files.pythonhosted.org/packages/25/98/6a19f030b3d2ea906696cedd1eb251708e50a5891d0978b012cb6107234c/msgpack-1.1.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:7bc8813f88417599564fafa59fd6f95be417179f76b40325b500b3c98409757c", size = 411454, upload-time = "2025-10-08T09:14:54.648Z" }, + { url = "https://files.pythonhosted.org/packages/b7/cd/9098fcb6adb32187a70b7ecaabf6339da50553351558f37600e53a4a2a23/msgpack-1.1.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:bafca952dc13907bdfdedfc6a5f579bf4f292bdd506fadb38389afa3ac5b208e", size = 422341, upload-time = "2025-10-08T09:14:56.328Z" }, + { url = "https://files.pythonhosted.org/packages/e6/ae/270cecbcf36c1dc85ec086b33a51a4d7d08fc4f404bdbc15b582255d05ff/msgpack-1.1.2-cp311-cp311-win32.whl", hash = "sha256:602b6740e95ffc55bfb078172d279de3773d7b7db1f703b2f1323566b878b90e", size = 64747, upload-time = "2025-10-08T09:14:57.882Z" }, + { url = "https://files.pythonhosted.org/packages/2a/79/309d0e637f6f37e83c711f547308b91af02b72d2326ddd860b966080ef29/msgpack-1.1.2-cp311-cp311-win_amd64.whl", hash = "sha256:d198d275222dc54244bf3327eb8cbe00307d220241d9cec4d306d49a44e85f68", size = 71633, upload-time = "2025-10-08T09:14:59.177Z" }, + { url = "https://files.pythonhosted.org/packages/73/4d/7c4e2b3d9b1106cd0aa6cb56cc57c6267f59fa8bfab7d91df5adc802c847/msgpack-1.1.2-cp311-cp311-win_arm64.whl", hash = "sha256:86f8136dfa5c116365a8a651a7d7484b65b13339731dd6faebb9a0242151c406", size = 64755, upload-time = "2025-10-08T09:15:00.48Z" }, + { url = "https://files.pythonhosted.org/packages/ad/bd/8b0d01c756203fbab65d265859749860682ccd2a59594609aeec3a144efa/msgpack-1.1.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:70a0dff9d1f8da25179ffcf880e10cf1aad55fdb63cd59c9a49a1b82290062aa", size = 81939, upload-time = "2025-10-08T09:15:01.472Z" }, + { url = "https://files.pythonhosted.org/packages/34/68/ba4f155f793a74c1483d4bdef136e1023f7bcba557f0db4ef3db3c665cf1/msgpack-1.1.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:446abdd8b94b55c800ac34b102dffd2f6aa0ce643c55dfc017ad89347db3dbdb", size = 85064, upload-time = "2025-10-08T09:15:03.764Z" }, + { url = "https://files.pythonhosted.org/packages/f2/60/a064b0345fc36c4c3d2c743c82d9100c40388d77f0b48b2f04d6041dbec1/msgpack-1.1.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c63eea553c69ab05b6747901b97d620bb2a690633c77f23feb0c6a947a8a7b8f", size = 417131, upload-time = "2025-10-08T09:15:05.136Z" }, + { url = "https://files.pythonhosted.org/packages/65/92/a5100f7185a800a5d29f8d14041f61475b9de465ffcc0f3b9fba606e4505/msgpack-1.1.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:372839311ccf6bdaf39b00b61288e0557916c3729529b301c52c2d88842add42", size = 427556, upload-time = "2025-10-08T09:15:06.837Z" }, + { url = "https://files.pythonhosted.org/packages/f5/87/ffe21d1bf7d9991354ad93949286f643b2bb6ddbeab66373922b44c3b8cc/msgpack-1.1.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2929af52106ca73fcb28576218476ffbb531a036c2adbcf54a3664de124303e9", size = 404920, upload-time = "2025-10-08T09:15:08.179Z" }, + { url = "https://files.pythonhosted.org/packages/ff/41/8543ed2b8604f7c0d89ce066f42007faac1eaa7d79a81555f206a5cdb889/msgpack-1.1.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:be52a8fc79e45b0364210eef5234a7cf8d330836d0a64dfbb878efa903d84620", size = 415013, upload-time = "2025-10-08T09:15:09.83Z" }, + { url = "https://files.pythonhosted.org/packages/41/0d/2ddfaa8b7e1cee6c490d46cb0a39742b19e2481600a7a0e96537e9c22f43/msgpack-1.1.2-cp312-cp312-win32.whl", hash = "sha256:1fff3d825d7859ac888b0fbda39a42d59193543920eda9d9bea44d958a878029", size = 65096, upload-time = "2025-10-08T09:15:11.11Z" }, + { url = "https://files.pythonhosted.org/packages/8c/ec/d431eb7941fb55a31dd6ca3404d41fbb52d99172df2e7707754488390910/msgpack-1.1.2-cp312-cp312-win_amd64.whl", hash = "sha256:1de460f0403172cff81169a30b9a92b260cb809c4cb7e2fc79ae8d0510c78b6b", size = 72708, upload-time = "2025-10-08T09:15:12.554Z" }, + { url = "https://files.pythonhosted.org/packages/c5/31/5b1a1f70eb0e87d1678e9624908f86317787b536060641d6798e3cf70ace/msgpack-1.1.2-cp312-cp312-win_arm64.whl", hash = "sha256:be5980f3ee0e6bd44f3a9e9dea01054f175b50c3e6cdb692bc9424c0bbb8bf69", size = 64119, upload-time = "2025-10-08T09:15:13.589Z" }, + { url = "https://files.pythonhosted.org/packages/6b/31/b46518ecc604d7edf3a4f94cb3bf021fc62aa301f0cb849936968164ef23/msgpack-1.1.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:4efd7b5979ccb539c221a4c4e16aac1a533efc97f3b759bb5a5ac9f6d10383bf", size = 81212, upload-time = "2025-10-08T09:15:14.552Z" }, + { url = "https://files.pythonhosted.org/packages/92/dc/c385f38f2c2433333345a82926c6bfa5ecfff3ef787201614317b58dd8be/msgpack-1.1.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:42eefe2c3e2af97ed470eec850facbe1b5ad1d6eacdbadc42ec98e7dcf68b4b7", size = 84315, upload-time = "2025-10-08T09:15:15.543Z" }, + { url = "https://files.pythonhosted.org/packages/d3/68/93180dce57f684a61a88a45ed13047558ded2be46f03acb8dec6d7c513af/msgpack-1.1.2-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1fdf7d83102bf09e7ce3357de96c59b627395352a4024f6e2458501f158bf999", size = 412721, upload-time = "2025-10-08T09:15:16.567Z" }, + { url = "https://files.pythonhosted.org/packages/5d/ba/459f18c16f2b3fc1a1ca871f72f07d70c07bf768ad0a507a698b8052ac58/msgpack-1.1.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fac4be746328f90caa3cd4bc67e6fe36ca2bf61d5c6eb6d895b6527e3f05071e", size = 424657, upload-time = "2025-10-08T09:15:17.825Z" }, + { url = "https://files.pythonhosted.org/packages/38/f8/4398c46863b093252fe67368b44edc6c13b17f4e6b0e4929dbf0bdb13f23/msgpack-1.1.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:fffee09044073e69f2bad787071aeec727183e7580443dfeb8556cbf1978d162", size = 402668, upload-time = "2025-10-08T09:15:19.003Z" }, + { url = "https://files.pythonhosted.org/packages/28/ce/698c1eff75626e4124b4d78e21cca0b4cc90043afb80a507626ea354ab52/msgpack-1.1.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5928604de9b032bc17f5099496417f113c45bc6bc21b5c6920caf34b3c428794", size = 419040, upload-time = "2025-10-08T09:15:20.183Z" }, + { url = "https://files.pythonhosted.org/packages/67/32/f3cd1667028424fa7001d82e10ee35386eea1408b93d399b09fb0aa7875f/msgpack-1.1.2-cp313-cp313-win32.whl", hash = "sha256:a7787d353595c7c7e145e2331abf8b7ff1e6673a6b974ded96e6d4ec09f00c8c", size = 65037, upload-time = "2025-10-08T09:15:21.416Z" }, + { url = "https://files.pythonhosted.org/packages/74/07/1ed8277f8653c40ebc65985180b007879f6a836c525b3885dcc6448ae6cb/msgpack-1.1.2-cp313-cp313-win_amd64.whl", hash = "sha256:a465f0dceb8e13a487e54c07d04ae3ba131c7c5b95e2612596eafde1dccf64a9", size = 72631, upload-time = "2025-10-08T09:15:22.431Z" }, + { url = "https://files.pythonhosted.org/packages/e5/db/0314e4e2db56ebcf450f277904ffd84a7988b9e5da8d0d61ab2d057df2b6/msgpack-1.1.2-cp313-cp313-win_arm64.whl", hash = "sha256:e69b39f8c0aa5ec24b57737ebee40be647035158f14ed4b40e6f150077e21a84", size = 64118, upload-time = "2025-10-08T09:15:23.402Z" }, +] + [[package]] name = "multidict" version = "6.7.0" @@ -2121,6 +2162,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5b/5a/bc7b4a4ef808fa59a816c17b20c4bef6884daebbdf627ff2a161da67da19/propcache-0.4.1-py3-none-any.whl", hash = "sha256:af2a6052aeb6cf17d3e46ee169099044fd8224cbaf75c76a2ef596e8163e2237", size = 13305, upload-time = "2025-10-08T19:49:00.792Z" }, ] +[[package]] +name = "protobuf" +version = "7.34.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f2/00/04a2ab36b70a52d0356852979e08b44edde0435f2115dc66e25f2100f3ab/protobuf-7.34.0.tar.gz", hash = "sha256:3871a3df67c710aaf7bb8d214cc997342e63ceebd940c8c7fc65c9b3d697591a", size = 454726, upload-time = "2026-02-27T00:30:25.421Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/13/c4/6322ab5c8f279c4c358bc14eb8aefc0550b97222a39f04eb3c1af7a830fa/protobuf-7.34.0-cp310-abi3-macosx_10_9_universal2.whl", hash = "sha256:8e329966799f2c271d5e05e236459fe1cbfdb8755aaa3b0914fa60947ddea408", size = 429248, upload-time = "2026-02-27T00:30:14.924Z" }, + { url = "https://files.pythonhosted.org/packages/45/99/b029bbbc61e8937545da5b79aa405ab2d9cf307a728f8c9459ad60d7a481/protobuf-7.34.0-cp310-abi3-manylinux2014_aarch64.whl", hash = "sha256:9d7a5005fb96f3c1e64f397f91500b0eb371b28da81296ae73a6b08a5b76cdd6", size = 325753, upload-time = "2026-02-27T00:30:17.247Z" }, + { url = "https://files.pythonhosted.org/packages/cc/79/09f02671eb75b251c5550a1c48e7b3d4b0623efd7c95a15a50f6f9fc1e2e/protobuf-7.34.0-cp310-abi3-manylinux2014_s390x.whl", hash = "sha256:4a72a8ec94e7a9f7ef7fe818ed26d073305f347f8b3b5ba31e22f81fd85fca02", size = 340200, upload-time = "2026-02-27T00:30:18.672Z" }, + { url = "https://files.pythonhosted.org/packages/b5/57/89727baef7578897af5ed166735ceb315819f1c184da8c3441271dbcfde7/protobuf-7.34.0-cp310-abi3-manylinux2014_x86_64.whl", hash = "sha256:964cf977e07f479c0697964e83deda72bcbc75c3badab506fb061b352d991b01", size = 324268, upload-time = "2026-02-27T00:30:20.088Z" }, + { url = "https://files.pythonhosted.org/packages/1f/3e/38ff2ddee5cc946f575c9d8cc822e34bde205cf61acf8099ad88ef19d7d2/protobuf-7.34.0-cp310-abi3-win32.whl", hash = "sha256:f791ec509707a1d91bd02e07df157e75e4fb9fbdad12a81b7396201ec244e2e3", size = 426628, upload-time = "2026-02-27T00:30:21.555Z" }, + { url = "https://files.pythonhosted.org/packages/cb/71/7c32eaf34a61a1bae1b62a2ac4ffe09b8d1bb0cf93ad505f42040023db89/protobuf-7.34.0-cp310-abi3-win_amd64.whl", hash = "sha256:9f9079f1dde4e32342ecbd1c118d76367090d4aaa19da78230c38101c5b3dd40", size = 437901, upload-time = "2026-02-27T00:30:22.836Z" }, + { url = "https://files.pythonhosted.org/packages/a4/e7/14dc9366696dcb53a413449881743426ed289d687bcf3d5aee4726c32ebb/protobuf-7.34.0-py3-none-any.whl", hash = "sha256:e3b914dd77fa33fa06ab2baa97937746ab25695f389869afdf03e81f34e45dc7", size = 170716, upload-time = "2026-02-27T00:30:23.994Z" }, +] + [[package]] name = "psutil" version = "7.2.1" @@ -2565,6 +2621,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/01/1b/5dbe84eefc86f48473947e2f41711aded97eecef1231f4558f1f02713c12/pyzmq-27.1.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:c9f7f6e13dff2e44a6afeaf2cf54cee5929ad64afaf4d40b50f93c58fc687355", size = 544862, upload-time = "2025-09-08T23:09:56.509Z" }, ] +[[package]] +name = "ray" +version = "2.54.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "filelock" }, + { name = "jsonschema" }, + { name = "msgpack" }, + { name = "packaging" }, + { name = "protobuf" }, + { name = "pyyaml" }, + { name = "requests" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/08/58/6209b2231947f3c8df09ce1436f1c76c4a11fcafd57c8def852dcbb6d8ef/ray-2.54.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:8e39dd56b47a0a1820d5a5a54385bbe54d1d67e1093736d12d8ed4e99d0fa455", size = 70098998, upload-time = "2026-02-18T04:04:58.801Z" }, + { url = "https://files.pythonhosted.org/packages/ac/29/7871f4206e6b00a9bb784c16dad32ccd01e9df5a93545db92de220eb2871/ray-2.54.0-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:491ae56ab80d8822c4eaf4d5bb96dcf32a6231d8d7b76eb8034400eb9be1bb18", size = 72066630, upload-time = "2026-02-18T04:05:04.957Z" }, + { url = "https://files.pythonhosted.org/packages/1d/e8/d2c8ebd9cd945abc817b01ad02a29df78cdb86cd07d764587e16977389d0/ray-2.54.0-cp311-cp311-manylinux2014_x86_64.whl", hash = "sha256:928bb09245a3c6f7c3c113ba8eafc69f948da9602d7f33e8251ecdf97c157615", size = 72895723, upload-time = "2026-02-18T04:05:10.686Z" }, + { url = "https://files.pythonhosted.org/packages/7e/96/a5ea3a149a943475cda1d68fdcdb14c86251826c652c232ae853600ad7e7/ray-2.54.0-cp311-cp311-win_amd64.whl", hash = "sha256:1e786330de55b3ba2228e36ec305381a9b86f0b01a8b6072c5811c3bc4dd9a3d", size = 27448371, upload-time = "2026-02-18T04:05:16.34Z" }, + { url = "https://files.pythonhosted.org/packages/0e/16/45eefb51eb1767342a6dbf41af0b432279e422e56160705fcd1098a7ec53/ray-2.54.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:cf5c33b4b13850ec24a5bd5f9d9e0a8161f8e586bfd297e52913d170dec447fe", size = 70084880, upload-time = "2026-02-18T04:05:22.007Z" }, + { url = "https://files.pythonhosted.org/packages/60/ad/e07aca3637e9c3ec4857ec4366208099cf8488ece8061a9925ba29b66382/ray-2.54.0-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:795ae21d6b764245d3f521bc5833446d58569e7dfde9c5777417eb285d87450f", size = 72107346, upload-time = "2026-02-18T04:05:27.999Z" }, + { url = "https://files.pythonhosted.org/packages/9e/b9/cc5ea8460c3dc602e6b7198277a7c59ba2b8929374ab22efa8df9f3deac8/ray-2.54.0-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:a972afd5aa3dda99d0b2f369b5f62e5dd95865ab7d37bf2e0a0e0d2cfbd9b325", size = 72967230, upload-time = "2026-02-18T04:05:33.771Z" }, + { url = "https://files.pythonhosted.org/packages/de/d7/744de3b1bb881701330ddcbb2f6efaccd65915d564ece899a3838f9fb105/ray-2.54.0-cp312-cp312-win_amd64.whl", hash = "sha256:2ee074ede491d0aacfa339c003f5d7a15826e1e2a72ce873234ccbc0446e19b3", size = 27427353, upload-time = "2026-02-18T04:05:38.853Z" }, + { url = "https://files.pythonhosted.org/packages/7f/f2/5c0161d10445e703b7d01413ab54ec1cc5e27032555279d296df89b9c4ee/ray-2.54.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:5ad77961fea16c697a0fb0e51216dd39c0bec28868cde54ac668edd58d12b8ae", size = 70030991, upload-time = "2026-02-18T04:05:43.966Z" }, + { url = "https://files.pythonhosted.org/packages/fd/8c/4a4a38eaec6e9614076a96967f58540f4f8d4aa0c793f43150c5df23cb9a/ray-2.54.0-cp313-cp313-manylinux2014_aarch64.whl", hash = "sha256:8952c23a8aa94f10728c2d16e0dc3732d09aa0e6254801757ff494984a214f45", size = 72013826, upload-time = "2026-02-18T04:05:49.866Z" }, + { url = "https://files.pythonhosted.org/packages/42/ac/e7ec2a406bd755f61c7090460fa5ab3f09b00c3c2d8db6d0b559f78a30eb/ray-2.54.0-cp313-cp313-manylinux2014_x86_64.whl", hash = "sha256:ab89e6089abb6e46fb98fdd96d399b31a852d79127cd8ac00746c61d93defa2c", size = 72880209, upload-time = "2026-02-18T04:05:55.498Z" }, +] + [[package]] name = "referencing" version = "0.37.0"