diff --git a/scripts/native_eval/runtime.py b/scripts/native_eval/runtime.py index f26f693..22608cb 100644 --- a/scripts/native_eval/runtime.py +++ b/scripts/native_eval/runtime.py @@ -58,6 +58,7 @@ class RewardFileNotFoundError(NativeEvalError): ) CODEX_STREAM_READ_ATTEMPTS = 20 CODEX_STREAM_READ_DELAY_SECONDS = 0.1 +DOCKER_CLEANUP_TIMEOUT_SEC = 30 def build_judge_env(proxy_url: str, proxy_key: str) -> dict[str, str]: @@ -418,27 +419,30 @@ async def exec( stderr_path=stderr_path, ) except TimeoutError: - await run_process( - ["docker", "kill", self.container_id], - stdout_path=self.trial_dir / "trial.log", - stderr_path=self.trial_dir / "trial.log", - ) + async with asyncio.timeout(DOCKER_CLEANUP_TIMEOUT_SEC): + await run_process( + ["docker", "kill", self.container_id], + stdout_path=self.trial_dir / "trial.log", + stderr_path=self.trial_dir / "trial.log", + ) raise async def stop(self) -> None: if self.task.compose_file: - await run_process( - self._compose_prefix() - + ["down", "--volumes", "--remove-orphans", "--timeout", "10"], - stdout_path=self.trial_dir / "environment-stop.log", - stderr_path=self.trial_dir / "environment-stop.log", - ) + async with asyncio.timeout(DOCKER_CLEANUP_TIMEOUT_SEC): + await run_process( + self._compose_prefix() + + ["down", "--volumes", "--remove-orphans", "--timeout", "10"], + stdout_path=self.trial_dir / "environment-stop.log", + stderr_path=self.trial_dir / "environment-stop.log", + ) else: - await run_process( - ["docker", "rm", "-f", self.container_name], - stdout_path=self.trial_dir / "environment-stop.log", - stderr_path=self.trial_dir / "environment-stop.log", - ) + async with asyncio.timeout(DOCKER_CLEANUP_TIMEOUT_SEC): + await run_process( + ["docker", "rm", "-f", self.container_name], + stdout_path=self.trial_dir / "environment-stop.log", + stderr_path=self.trial_dir / "environment-stop.log", + ) async def run_trial( @@ -1294,6 +1298,31 @@ def _timing(result: CommandResult) -> dict[str, str]: } +REAP_WAIT_SEC = 2 + + +async def _reap_process(process: asyncio.subprocess.Process) -> None: + if process.returncode is not None: + return + try: + process.terminate() + except ProcessLookupError: + return + try: + async with asyncio.timeout(REAP_WAIT_SEC): + await process.wait() + except TimeoutError: + try: + process.kill() + except ProcessLookupError: + return + try: + async with asyncio.timeout(REAP_WAIT_SEC): + await process.wait() + except TimeoutError: + return + + async def run_process( args: list[str], *, @@ -1306,12 +1335,16 @@ async def run_process( stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - await asyncio.gather( - _drain(process.stdout, stdout_path), - _drain(process.stderr, stderr_path), - ) - returncode = await process.wait() - return CommandResult(returncode, started_at, utc_now()) + try: + await asyncio.gather( + _drain(process.stdout, stdout_path), + _drain(process.stderr, stderr_path), + ) + returncode = await process.wait() + return CommandResult(returncode, started_at, utc_now()) + except (TimeoutError, asyncio.CancelledError): + await _reap_process(process) + raise async def capture_process(args: list[str]) -> str: diff --git a/tests/test_native_eval_runtime.py b/tests/test_native_eval_runtime.py new file mode 100644 index 0000000..b387d4e --- /dev/null +++ b/tests/test_native_eval_runtime.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import asyncio +import os +import sys +import time +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from scripts.native_eval.runtime import DockerTaskEnvironment +from scripts.native_eval import runtime as native_runtime + + +def _environment(tmp_path: Path, *, compose: bool = False) -> DockerTaskEnvironment: + compose_file = tmp_path / "docker-compose.yaml" if compose else None + if compose_file: + compose_file.write_text("services: {}\n", encoding="utf-8") + return DockerTaskEnvironment( + task=SimpleNamespace(compose_file=compose_file), # type: ignore[arg-type] + trial_dir=tmp_path, + container_name="trial", + project_name="trial", + toolchain_root=tmp_path, + container_id="container-id", + ) + + +def _hanging_run_process() -> object: + async def hang(_args: list[str], **_kwargs: object) -> None: + await asyncio.Event().wait() + + return hang + + +def _cleanup_deadline(monkeypatch) -> None: + # Production uses 30s. Tests use a short bound so a hang is visible + # as a wait_for miss instead of a minutes-long stall. + monkeypatch.setattr(native_runtime, "DOCKER_CLEANUP_TIMEOUT_SEC", 0.05, raising=False) + + +async def _await_with_deadline(coro): + outcome = {"kind": "hung"} + try: + + async def run() -> None: + try: + await coro + outcome["kind"] = "returned" + except TimeoutError: + outcome["kind"] = "timeout" + + await asyncio.wait_for(run(), timeout=1.0) + except TimeoutError: + outcome["kind"] = "hung" + return outcome["kind"] + + +def test_exec_timeout_ends_when_docker_kill_hangs(tmp_path: Path, monkeypatch) -> None: + _cleanup_deadline(monkeypatch) + monkeypatch.setattr(native_runtime, "run_process", _hanging_run_process()) + environment = _environment(tmp_path) + + kind = asyncio.run(_await_with_deadline(environment.exec("true", timeout=0.01))) + + assert kind == "timeout" + + +def test_stop_ends_when_docker_rm_hangs(tmp_path: Path, monkeypatch) -> None: + _cleanup_deadline(monkeypatch) + monkeypatch.setattr(native_runtime, "run_process", _hanging_run_process()) + environment = _environment(tmp_path) + + kind = asyncio.run(_await_with_deadline(environment.stop())) + + assert kind == "timeout" + + +def test_stop_ends_when_compose_down_hangs(tmp_path: Path, monkeypatch) -> None: + _cleanup_deadline(monkeypatch) + monkeypatch.setattr(native_runtime, "run_process", _hanging_run_process()) + environment = _environment(tmp_path, compose=True) + + kind = asyncio.run(_await_with_deadline(environment.stop())) + + assert kind == "timeout" + + +def _pid_alive(pid: int) -> bool: + try: + os.kill(pid, 0) + except OSError: + return False + return True + + +def test_run_process_kills_and_reaps_hung_child(tmp_path: Path) -> None: + stdout_path = tmp_path / "child.out" + child = [ + sys.executable, + "-c", + "import os, sys, time; print(os.getpid(), flush=True); time.sleep(30)", + ] + + async def run() -> None: + async with asyncio.timeout(0.4): + await native_runtime.run_process( + child, + stdout_path=stdout_path, + stderr_path=tmp_path / "child.err", + ) + + started = time.monotonic() + with pytest.raises(TimeoutError): + asyncio.run(run()) + assert time.monotonic() - started < 3 + + pid_text = stdout_path.read_text(encoding="utf-8").strip() + assert pid_text.isdigit(), pid_text + pid = int(pid_text) + deadline = time.monotonic() + 2 + while time.monotonic() < deadline and _pid_alive(pid): + time.sleep(0.05) + assert not _pid_alive(pid), f"child {pid} still running after cleanup timeout" + + +def test_reap_process_returns_when_wait_hangs_after_kill() -> None: + class _StuckProcess: + returncode = None + + def terminate(self) -> None: + return None + + def kill(self) -> None: + return None + + async def wait(self) -> int: + await asyncio.Event().wait() + return 0 + + async def run() -> float: + started = time.monotonic() + await native_runtime._reap_process(_StuckProcess()) + return time.monotonic() - started + + elapsed = asyncio.run(run()) + assert elapsed < 6