diff --git a/pyproject.toml b/pyproject.toml index 6f91882..684f0db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,7 @@ dependencies = [ "numpy", "pillow", "pydantic", + "ray", "tenacity>=8.0", ] description = "hypotest package" @@ -66,6 +67,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 1eecfdf..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 @@ -34,11 +35,16 @@ class DatasetConfig(BaseModel): ) work_dir: Path | None = None - use_docker: bool = True + use_ray: 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: ExecutionConfig = Field(default_factory=ExecutionConfig) + @model_validator(mode="after") def validate_dataset_source(self) -> Self: if not self.problem_jsonl and not self.hf_dataset: @@ -83,16 +89,17 @@ 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) 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.metadata.get("nb_primary_language", "python")).upper()) + 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, @@ -106,6 +113,10 @@ def __len__(self) -> int: return len(self.problems) +HypotestDataset = Dataset +HypotestDatasetConfig = DatasetConfig + + DEFAULT_SERVER_PORT = 8405 @@ -169,6 +180,7 @@ async def launch_server(): ], }, use_docker=args.use_docker, + execution_config={"cell_execution_timeout": 600}, ), port=args.port, api_key=args.api_key, diff --git a/src/hypotest/env/code_safety.py b/src/hypotest/env/code_safety.py new file mode 100644 index 0000000..26c4a17 --- /dev/null +++ b/src/hypotest/env/code_safety.py @@ -0,0 +1,389 @@ +"""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 .kernel_server import NBLanguage + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# 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({ + "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: # noqa: PLR0912 + 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) + 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: + 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(("!", "%")): + 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. + """ + # 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 diff --git a/src/hypotest/env/config.py b/src/hypotest/env/config.py index e43cc8e..96cbde9 100644 --- a/src/hypotest/env/config.py +++ b/src/hypotest/env/config.py @@ -12,10 +12,13 @@ ) 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 -KERNEL_SERVER_STARTUP_TIMEOUT = 30.0 # seconds to wait for health check +KERNEL_SERVER_STARTUP_TIMEOUT = float( + 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")) @@ -56,6 +59,9 @@ class ExecutionConfig(BaseModel): force_submit_threshold: int = 10 * 60 cell_execution_timeout: int = 15 * 60 + # safety + 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 d9dff15..3a1ea62 100644 --- a/src/hypotest/env/interpreter_env.py +++ b/src/hypotest/env/interpreter_env.py @@ -6,14 +6,21 @@ import argparse import asyncio +import contextlib import json import logging import os import shutil +import signal import socket +import subprocess +import sys import time +import uuid +from collections.abc import Awaitable, Callable from pathlib import Path from tempfile import mkdtemp +from textwrap import dedent from typing import TYPE_CHECKING, Any, cast from uuid import UUID @@ -34,36 +41,239 @@ 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 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 # 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) +MAX_CONTAINER_LAUNCH_RETRIES = int(os.getenv("MAX_CONTAINER_LAUNCH_RETRIES", "5")) +_RETRY_BASE_SLEEP = 1.0 +_RETRY_MAX_SLEEP = 16.0 + +_warned_unsafe_execution: set[str] = set() -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 + + +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: asyncio.subprocess.Process | None, + container_port: int | None, + expected_startup_token: str | None, + read_log_tail: Callable[[int], Awaitable[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.returncode is not None: + rc = enroot_proc.returncode + 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}" + ) + + 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.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", + 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 = 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)" + f"{f' log_tail={log_tail!r}' if log_tail else ''}" + ) + + +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.returncode 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: + 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", + label, + proc.pid, + proc.returncode, + ) + return + + # 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: + await asyncio.wait_for(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 TimeoutError: + logger.exception("[%s] Process pid=%d still alive after SIGKILL — possible zombie", label, proc.pid) + class ProblemInstance(BaseModel): id: UUID @@ -74,6 +284,407 @@ class ProblemInstance(BaseModel): 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)) + + @model_validator(mode="before") + @classmethod + def handle_language(cls, data: dict) -> dict: + if data.get("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) + (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, safe_execute: bool = True): + self.container_sqsh_path = container_sqsh_path + self.execution_timeout = execution_timeout + self.safe_execute = safe_execute + 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 + 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})" + + @staticmethod + 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 + + 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} {"--safe-execute" if safe_execute else ""} + """).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"{cfg.CONTAINER_WORKSPACE_PREFIX}.{startup_token.split('-', maxsplit=1)[0]}" + + _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 + last_err: Exception | None = None + while not online: + attempt += 1 + if attempt > MAX_CONTAINER_LAUNCH_RETRIES: + 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})" + f"{f' log_tail={log_tail!r}' if log_tail else ''}" + ) + 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 + ) + 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() + # 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 = 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, + 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 (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: + last_err = e + launch_ms = (time.perf_counter() - launch_t0) * 1000.0 + 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) + + 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 + + diag_parts = [ + f"attempt=#{attempt}", + f"elapsed={launch_ms:.0f}ms", + f"error={error!r}", + ] + + if proc is not None: + rc = proc.returncode + diag_parts.append(f"process_alive={rc is None}") + if rc is not None: + diag_parts.append(f"returncode={rc}") + + logger.warning("[%s] Container startup FAILED: %s", label, ", ".join(diag_parts)) + + # Log container output separately so tracebacks are readable + log_tail = await 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, + ) + + 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 "" + 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: + 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() + + 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: + with contextlib.suppress(Exception): + await self._http_client.aclose() + self._http_client = None + + if self._enroot_proc is not None: + await _kill_process_group(self._enroot_proc, label=label) + self._enroot_proc = None + + 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.""" + 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(), + ) + + 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) + with contextlib.suppress(Exception): + await self._http_client.aclose() + self._http_client = None + + if self._enroot_proc is not None: + await _kill_process_group(self._enroot_proc, label=label) + self._enroot_proc = None + + await self._close_container_log() + logger.log(_CONTAINER_LOG_LEVEL, "[%s] EnrootKernelServer closed", label) class InterpreterEnvState: @@ -88,14 +699,19 @@ 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, + use_enroot: bool = False, + use_ray: bool = True, + container_sqsh_path: Path | None = None, save_dir: Path | None = None, ): 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 {} @@ -103,11 +719,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 if RAY_INSTALLED else False # 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, @@ -116,11 +735,15 @@ 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: 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 + 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() @@ -134,38 +757,235 @@ 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( # 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 + + 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(str(self.work_dir)) + + online = False + attempt = 0 + last_err: Exception | None = None + while not online: + attempt += 1 + if attempt > MAX_CONTAINER_LAUNCH_RETRIES: + 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})" + 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 + 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} {"--safe-execute" if self.safe_execute else ""} + """).strip() + + 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", + "-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). + (self.work_dir / ".container_logs").mkdir(exist_ok=True) + self._container_log_path = self.work_dir / ".container_logs" / "container.log" + 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, + 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 (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, + ) + online = True + except Exception as e: + last_err = e + launch_ms = (time.perf_counter() - launch_t0) * 1000.0 + 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) + + 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}"] + if self._enroot_proc is not None: + rc = self._enroot_proc.returncode + 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 = await 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, + ) + + 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 "" + 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: + 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() + + 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: + with contextlib.suppress(Exception): + await self._http_client.aclose() + self._http_client = None + + if self._enroot_proc is not None: + await _kill_process_group(self._enroot_proc, label=label) + self._enroot_proc = None + + await 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() + 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", - "/workspace", - "--language", - self.language.value, - ], + "Cmd": cmd_list, "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( @@ -174,22 +994,19 @@ async def _start_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) -> 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() - while time.perf_counter() - start_time < cfg.KERNEL_SERVER_STARTUP_TIMEOUT: - try: - assert self._http_client is not None - response = await self._http_client.get("/health") - if response.status_code == 200: - logger.info("Kernel server is healthy") - return - except (httpx.ConnectError, httpx.ReadError, httpx.RemoteProtocolError): - pass - await asyncio.sleep(0.5) - raise TimeoutError(f"Kernel server did not become healthy within {cfg.KERNEL_SERVER_STARTUP_TIMEOUT}s") + 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( retry=tenacity.retry_if_exception_type((httpx.ConnectError, httpx.ReadError)), @@ -198,14 +1015,41 @@ 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 @@ -220,8 +1064,16 @@ 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.""" @@ -232,12 +1084,20 @@ 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() # type: ignore[attr-defined] + 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() + with contextlib.suppress(Exception): + await self._http_client.aclose() self._http_client = None if self._container is not None: @@ -251,9 +1111,19 @@ 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: + await _kill_process_group(self._enroot_proc, label=self._enroot_label()) + self._enroot_proc = None + + await self._close_container_log() + 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. @@ -343,7 +1213,37 @@ async def execute_and_add_cell( Returns: Tuple of (ExecutionResult, actual_cell_index) """ - if self.use_docker: + 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 + 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] + result = await result_ref + elif self.use_docker or self.use_enroot: result = await self._execute_via_http(code, timeout) else: assert self.interpreter is not None @@ -365,7 +1265,10 @@ 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 normalize_reward: bool = True @@ -414,6 +1317,13 @@ def __init__( # prompting_config is set during reset() after language resolution self.prompting_config: PromptingConfig + 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() + @property def language(self) -> NBLanguage: return self.config.language @@ -444,6 +1354,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 @@ -456,6 +1367,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() @@ -600,7 +1513,10 @@ async def reset_kernel(self) -> str: This clears all variables and execution state. """ - if self.state.use_docker: + if self.state.use_ray and self.state.kernel_container is not None: + 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() else: assert self.state.interpreter is not None diff --git a/src/hypotest/env/kernel_server.py b/src/hypotest/env/kernel_server.py index c2c1f2b..565d76f 100644 --- a/src/hypotest/env/kernel_server.py +++ b/src/hypotest/env/kernel_server.py @@ -12,9 +12,12 @@ import argparse import asyncio import logging +import re import time +import uuid from enum import StrEnum, auto from pathlib import Path +from queue import Empty from typing import Any, assert_never import nbformat @@ -25,6 +28,11 @@ from nbformat import NotebookNode from pydantic import BaseModel + +class DeadlineExceededError(Exception): + """Raised when a cooperative deadline check expires.""" + + logger = logging.getLogger(__name__) @@ -128,6 +136,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.""" @@ -152,6 +194,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 +210,14 @@ def __init__( work_dir: Path, 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 @@ -175,7 +229,15 @@ 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() @@ -193,13 +255,31 @@ 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") + 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() 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, + ) + except DeadlineExceededError: timeout_output = MessageType.ERROR.to_notebook_output({ "ename": "TimeoutError", "evalue": f"Code execution timed out after {effective_timeout} seconds", @@ -224,21 +304,39 @@ async def execute(self, code: str, timeout: float | None = None) -> ExecuteRespo 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) -> 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 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 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 + # 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 @@ -252,6 +350,9 @@ async def _execute_code(self, code: str) -> ExecuteResponse: 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)) @@ -287,6 +388,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 +406,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 +417,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) @@ -334,10 +436,12 @@ async def run_server(work_dir: Path, language: NBLanguage, port: int = 8000) -> 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="") + parser.add_argument("--safe-execute", action="store_true") 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/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"