diff --git a/src/agent_env/env/envs/multi_env.py b/src/agent_env/env/envs/multi_env.py index 25720963..24fd0fa5 100644 --- a/src/agent_env/env/envs/multi_env.py +++ b/src/agent_env/env/envs/multi_env.py @@ -17,6 +17,9 @@ # behaviour, or hold it fixed while bisecting a contention problem. _LOAD_CONCURRENCY_ENV_VAR = "AGENT_ENV_LOAD_CONCURRENCY" +# `nproc` is GNU coreutils; macOS has only `getconf`. +_CORE_COUNT_PROBE = "nproc 2>/dev/null || getconf _NPROCESSORS_ONLN" + def _snapshot_after_load_default() -> bool: return os.environ.get(_SNAPSHOT_AFTER_LOAD_ENV_VAR, "").strip().lower() in ("1", "true", "yes") @@ -475,7 +478,7 @@ async def _load_concurrency_limit(self, total: int) -> int: cores = None if self._sandbox is not None: try: - cores = int((await self._sandbox.exec_script("nproc")).strip()) + cores = int((await self._sandbox.exec_script(_CORE_COUNT_PROBE)).strip()) except Exception as e: # noqa: BLE001 — a probe, not a dependency logger.warning("Could not read sandbox core count (%s); loading all %d services at once", e, total) if not cores or cores < 1: diff --git a/src/agent_env/task_step/task_steps/run_code.py b/src/agent_env/task_step/task_steps/run_code.py index 63a44cd2..32bd4bef 100644 --- a/src/agent_env/task_step/task_steps/run_code.py +++ b/src/agent_env/task_step/task_steps/run_code.py @@ -38,7 +38,7 @@ _RUN_DIR = "/tmp/agent_env_run_code" -# The bootstrap we run in the sandbox: `python3 runner.py `. +# The bootstrap we run in the sandbox: `python3 runner.py `. _RUNNER_SOURCE = (Path(__file__).parent / "run_code_runner.py").read_text(encoding="utf-8") # A universe member with one of these names would shadow run_code's own staging files. @@ -68,8 +68,9 @@ class RunCodeTaskStep(TaskStep): # Not an __init__ default (script_file defaults to None) — the universe entry # chosen when several .py members exist. ENTRY_WHEN_AMBIGUOUS: ClassVar[str] = "run.py" - _TIMEOUT_EXIT_CODE: ClassVar[int] = 124 # `timeout` coreutil convention - _STDERR_TRUNCATE: ClassVar[int] = 1000 # max stderr chars surfaced in errors + _TIMEOUT_EXIT_CODE: ClassVar[int] = 124 # the runner's status at the limit, after GNU timeout's + # The end of a failed script's stderr that its error keeps: the error is stored and raised whole. + _STDERR_TAIL_BYTES: ClassVar[int] = 1024 * 1024 # A collect_artifacts universe can hold hundreds of non-.py members; name a few. _SKIPPED_MEMBERS_LOGGED: ClassVar[int] = 10 @@ -325,17 +326,16 @@ async def _stage( await target.write_file(f"{run_dir}/runner.py", _RUNNER_SOURCE) async def _run(self, target: _ExecTarget, run_dir: str, script_file: str) -> None: - """Run the script under a timeout.""" + """Run the script, which the runner stops at the timeout.""" exit_code, _stdout, stderr = await target.run( - "timeout", str(self.timeout_seconds), - "python3", f"{run_dir}/runner.py", run_dir, self.entrypoint, script_file, + "python3", f"{run_dir}/runner.py", run_dir, self.entrypoint, script_file, str(self.timeout_seconds), ) if exit_code == 0: return reason = ( f"timed out after {self.timeout_seconds}s" if exit_code == self._TIMEOUT_EXIT_CODE - else f"exited {exit_code}: {stderr[:self._STDERR_TRUNCATE]}" + else f"exited {exit_code}: {self._stderr_tail(stderr)}" ) raise self._exec_error(reason) @@ -348,14 +348,14 @@ async def _read_output(self, target: _ExecTarget, run_dir: str) -> Any: path = f"{run_dir}/output.json" exit_code, size_out, stderr = await target.run("wc", "-c", path) if exit_code != 0: - raise self._exec_error(f"wrote no result file: {stderr[:self._STDERR_TRUNCATE]}") + raise self._exec_error(f"wrote no result file: {self._stderr_tail(stderr)}") if int(size_out.split()[0]) > self.MAX_OUTPUT_BYTES: raise self._validation_error( f"result exceeds {self.MAX_OUTPUT_BYTES}-byte cap" ) exit_code, stdout, stderr = await target.run("cat", path) if exit_code != 0: - raise self._exec_error(f"could not read result file: {stderr[:self._STDERR_TRUNCATE]}") + raise self._exec_error(f"could not read result file: {self._stderr_tail(stderr)}") try: return json.loads(stdout) except json.JSONDecodeError as exc: @@ -448,6 +448,14 @@ async def _agent_target(self, context: TaskStepContext) -> _ExecTarget: # --- errors --- + def _stderr_tail(self, stderr: str) -> str: + """``stderr`` within ``_STDERR_TAIL_BYTES``, kept from the end, where the exception is.""" + data = stderr.encode("utf-8", "replace") + if len(data) <= self._STDERR_TAIL_BYTES: + return stderr + tail = data[-self._STDERR_TAIL_BYTES:].decode("utf-8", "ignore") # a character cut in two goes + return f"[first {len(data) - len(tail.encode('utf-8')):,} bytes of stderr cut]\n{tail}" + def _exec_error(self, msg: str) -> RunCodeExecutionError: return RunCodeExecutionError(f"run_code '{self.id}': {msg}") diff --git a/src/agent_env/task_step/task_steps/run_code_runner.py b/src/agent_env/task_step/task_steps/run_code_runner.py index 9c8c4da8..cf54bdb9 100644 --- a/src/agent_env/task_step/task_steps/run_code_runner.py +++ b/src/agent_env/task_step/task_steps/run_code_runner.py @@ -3,16 +3,96 @@ Imports the entry script, calls its entrypoint with the JSON input, and writes the JSON result. Reads `run_dir`, `entrypoint` and `script_file` from argv. +Given a fourth argument, `timeout_seconds`, it first supervises: it reruns itself without +it in a new session and stops that whole process group at the limit, so no target needs +GNU `timeout`. + Both the entry script's own directory and `run_dir` go on sys.path, so a multi-file artifact can import its siblings whether it is laid out flat or under a subdirectory. -Stdlib only — it runs in the sandbox, which may not have agent_env installed. +Stdlib only, and Python 3.9 compatible: it runs in the sandbox, which may not have +agent_env installed, or anything newer than macOS's /usr/bin/python3. """ +import sys + +_TIMEOUT_EXIT_CODE = 124 # GNU timeout's, which run_code reports as a timeout +_GRACE_SECONDS = 5 # between SIGTERM and SIGKILL +_POLL_SECONDS = 0.1 + + +def _signal_group(pgid, signum): + try: + os.killpg(pgid, signum) + except (ProcessLookupError, PermissionError): # gone, or (macOS) only zombies are left + pass + + +def _supervise(command, timeout_seconds): + """Run `command` in a new session; return its exit status, or 124 if it hit the limit. + + At the limit, or on a signal forwarded to it, the group gets that signal (SIGTERM at the + limit), then SIGKILL once the child exits or the grace period runs out. + """ + child = None + early = [] # signals that came before the child did + deadline = time.monotonic() + timeout_seconds + kill_at = None + timed_out = False + + def forward(signum, _frame): + nonlocal kill_at + if child is None: + early.append(signum) + else: + _signal_group(child.pid, signum) + kill_at = kill_at or time.monotonic() + _GRACE_SECONDS + + # Caught before the child starts, as GNU timeout does, so it gets their defaults: one ignored here (a background + # job's SIGINT) would stay ignored in it. + for signum in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP): + signal.signal(signum, forward) + child = subprocess.Popen(command, start_new_session=True) + for signum in early: + _signal_group(child.pid, signum) + while child.poll() is None: + now = time.monotonic() + if kill_at is not None and now >= kill_at: + break + if not timed_out and now >= deadline: + timed_out = True + _signal_group(child.pid, signal.SIGTERM) + kill_at = kill_at or now + _GRACE_SECONDS + try: + child.wait(_POLL_SECONDS) + except subprocess.TimeoutExpired: + pass + if kill_at is not None: + _signal_group(child.pid, signal.SIGKILL) + status = child.wait() + if timed_out: + return _TIMEOUT_EXIT_CODE + return status if status >= 0 else 128 - status + + +if len(sys.argv) == 5: + # Only the supervisor imports these, and not from the run dir (sys.path[0]): an artifact member named after one + # would otherwise run here, outside the limit. The script's own process imports just what it always has. + import os # loaded at startup, so never from the run dir + + # Python puts the realpath there, and the run dir is often reached through a symlink (macOS's /tmp). + here = os.path.dirname(os.path.realpath(__file__)) + sys.path[:] = [entry for entry in sys.path if os.path.realpath(entry or ".") != here] + import signal + import subprocess + import time + + sys.exit(_supervise([sys.executable] + sys.argv[:4], float(sys.argv[4]))) + +# Below the supervisor, which must not import these. import importlib.machinery import importlib.util import json import os -import sys run_dir, entrypoint, script_file = sys.argv[1], sys.argv[2], sys.argv[3] diff --git a/tst/unit/env/envs/test_multi_env_load_concurrency.py b/tst/unit/env/envs/test_multi_env_load_concurrency.py index 0f3f7b44..9c5ea033 100644 --- a/tst/unit/env/envs/test_multi_env_load_concurrency.py +++ b/tst/unit/env/envs/test_multi_env_load_concurrency.py @@ -10,11 +10,13 @@ from __future__ import annotations import asyncio +import shutil +import subprocess from unittest.mock import AsyncMock, MagicMock, patch import pytest -from agent_env.env.envs.multi_env import MultiEnv +from agent_env.env.envs.multi_env import _CORE_COUNT_PROBE, MultiEnv from agent_env.providers.env_providers import EnvironmentGatewayProvider from agent_env.providers.env_state import LocalPostgresStateProvider @@ -48,6 +50,22 @@ async def test_limit_scales_with_cores(cores, total, expected): assert await _env(cores)._load_concurrency_limit(total) == expected +def _probe(path): + out = subprocess.run([shutil.which("bash"), "-c", _CORE_COUNT_PROBE], env={"PATH": path}, + capture_output=True, text=True, check=True).stdout + return int(out.strip()) + + +def test_the_core_count_probe_reads_this_hosts_cores(): + assert _probe("/usr/bin:/bin") >= 1 + + +def test_the_core_count_probe_falls_back_to_getconf_without_nproc(tmp_path): + # macOS has no `nproc`; this stands in for it on any host. + (tmp_path / "getconf").symlink_to(shutil.which("getconf")) + assert _probe(str(tmp_path)) >= 1 + + @pytest.mark.asyncio async def test_unreadable_core_count_keeps_todays_behaviour(): """A probe failure must not quietly change how a load runs.""" diff --git a/tst/unit/task_step/test_run_code.py b/tst/unit/task_step/test_run_code.py index 99ffdf5c..aa7de17c 100644 --- a/tst/unit/task_step/test_run_code.py +++ b/tst/unit/task_step/test_run_code.py @@ -5,9 +5,12 @@ from __future__ import annotations +import ast import json +import signal import subprocess import sys +import time from contextlib import contextmanager from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch @@ -171,8 +174,6 @@ async def exec_script(self, script, **_kw): async def exec_with_output(self, *args): args = [a for a in args if a != "sudo"] - if args[0] == "timeout": # macOS has no coreutils `timeout` - args = args[2:] proc = subprocess.run(args, capture_output=True) return proc.returncode, proc.stdout.decode(), proc.stderr.decode() @@ -515,7 +516,7 @@ async def test_host_installed_agent_on_a_vm_runs_on_the_host(): @pytest.mark.asyncio async def test_timeout_is_reported(): - with _sandbox([(124, "", "")]): # `timeout` coreutil exit code + with _sandbox([(124, "", "")]): # the runner's status at the limit with pytest.raises(RunCodeExecutionError, match="timed out"): await _step(timeout_seconds=5).execute(_make_context()) @@ -527,6 +528,44 @@ async def test_script_crash_surfaces_stderr(): await _step().execute(_make_context()) +@pytest.mark.asyncio +async def test_a_long_stderr_keeps_its_end_and_says_how_much_was_cut(): + stderr = "x" * 500 + "\nValueError: the cause\n" + step = _step() + step._STDERR_TAIL_BYTES = 100 + with _sandbox([(1, "", stderr)]): + with pytest.raises(RunCodeExecutionError) as exc: + await step.execute(_make_context()) + assert str(exc.value).endswith(f"[first {len(stderr) - 100:,} bytes of stderr cut]\n{stderr[-100:]}") + + +def test_a_cut_inside_a_character_drops_the_rest_of_it_and_counts_it(): + step = _step() + step._STDERR_TAIL_BYTES = 5 # the last 5 bytes of "éééé" start inside its second é + assert step._stderr_tail("éééé") == "[first 4 bytes of stderr cut]\néé" + assert step._stderr_tail("éé") == "éé" + + +@pytest.mark.asyncio +async def test_a_script_past_its_timeout_is_stopped_and_reported(tmp_path): + artifact = _file_artifact(body=b"import time\ndef run(inp):\n time.sleep(300)\n") + started = time.monotonic() + with _real_run(tmp_path, artifact): + with pytest.raises(RunCodeExecutionError, match="timed out after 1s"): + await _step(timeout_seconds=1).execute(_make_context()) + assert time.monotonic() - started < 30 + + +@pytest.mark.asyncio +async def test_a_crash_reports_all_of_stderr_through_to_the_exception(tmp_path): + body = b"import sys\ndef run(inp):\n print('x' * 5000, file=sys.stderr)\n raise ValueError('the cause')\n" + with _real_run(tmp_path, _file_artifact(body=body)): + with pytest.raises(RunCodeExecutionError) as exc: + await _step().execute(_make_context()) + assert "x" * 5000 in str(exc.value) + assert str(exc.value).rstrip().endswith("ValueError: the cause") + + class _LocalExec: """A sandbox stand-in that runs commands against the real filesystem, decoding strictly like the production sandbox (providers/sandbox.py). Lets the size/read @@ -626,6 +665,144 @@ def test_runner_imports_sibling_modules_from_the_run_dir(tmp_path): assert json.loads((tmp_path / "output.json").read_text()) == {"from": "helper"} +def test_runner_parses_as_python_3_9(): + # macOS's /usr/bin/python3 is 3.9, and the runner runs on whatever python3 a host has. + ast.parse(_RUNNER.read_text(), feature_version=(3, 9)) + + +def _runner_args(run_dir, body, timeout): + """Stage the runner, script and input as run_code does, the runner inside the run dir.""" + (run_dir / "runner.py").write_text(_RUNNER.read_text()) + (run_dir / "script.py").write_text(body) + (run_dir / "input.json").write_text(json.dumps({"args": {}, "results": {}})) + return [sys.executable, str(run_dir / "runner.py"), str(run_dir), "run", "script.py", str(timeout)] + + +def _supervised(run_dir, body, timeout=60): + return subprocess.run(_runner_args(run_dir, body, timeout), capture_output=True, text=True, timeout=60) + + +def _running(pid): + """True while `pid` is a live process; a zombie awaiting its reaper counts as gone.""" + stat = subprocess.run(["ps", "-o", "stat=", "-p", str(pid)], capture_output=True, text=True).stdout.strip() + return bool(stat) and not stat.startswith("Z") + + +def _gone_within(pid, seconds=5.0): + deadline = time.monotonic() + seconds + while _running(pid) and time.monotonic() < deadline: + time.sleep(0.05) + return not _running(pid) + + +# Writes the pid of a `sleep` it starts, so a test can check the whole group was stopped. +_STARTS_A_CHILD = """\ +import os, subprocess, time +def run(inp): + child = subprocess.Popen(["sleep", "300"]) + with open(os.path.join(os.path.dirname(__file__), "pid.tmp"), "w") as f: + f.write(str(child.pid)) + os.rename(f.name, f.name[:-len(".tmp")]) + time.sleep(300) +""" + + +def _wait_for_pid(run_dir): + deadline = time.monotonic() + 30 + while not (run_dir / "pid").exists(): + assert time.monotonic() < deadline, "the script never started its child" + time.sleep(0.05) + return int((run_dir / "pid").read_text()) + + +def test_a_supervised_script_writes_its_result(tmp_path): + result = _supervised(tmp_path, "def run(inp):\n return {'ok': True}\n") + assert result.returncode == 0, result.stderr + assert json.loads((tmp_path / "output.json").read_text()) == {"ok": True} + + +def test_a_supervised_scripts_exit_status_and_stderr_pass_through(tmp_path): + result = _supervised(tmp_path, "import sys\ndef run(inp):\n print('boom', file=sys.stderr)\n sys.exit(3)\n") + assert (result.returncode, result.stderr) == (3, "boom\n") + + +def test_a_script_killed_by_a_signal_exits_128_plus_its_number(tmp_path): + result = _supervised(tmp_path, "import os, signal\ndef run(inp):\n os.kill(os.getpid(), signal.SIGKILL)\n") + assert result.returncode == 128 + signal.SIGKILL + + +def test_the_limit_stops_the_script_and_everything_it_started(tmp_path): + started = time.monotonic() + proc = subprocess.Popen(_runner_args(tmp_path, _STARTS_A_CHILD, 1)) + child = _wait_for_pid(tmp_path) + assert proc.wait(timeout=30) == 124 + assert time.monotonic() - started < 30 + assert _gone_within(child) + + +def test_a_script_that_ignores_sigterm_is_killed_after_the_grace_period(tmp_path): + body = ( + "import signal, time\ndef run(inp):\n signal.signal(signal.SIGTERM, signal.SIG_IGN)\n time.sleep(300)\n" + ) + started = time.monotonic() + result = _supervised(tmp_path, body, timeout=1) + elapsed = time.monotonic() - started + assert result.returncode == 124 + assert 5 <= elapsed < 30 + + +# Starts the runner with SIGINT and SIGHUP ignored, as a background job is (CI runs the tier as one). The script must +# still get their defaults, as it does under GNU timeout. +_WITH_SIGNALS_IGNORED = ( + "import os, signal, sys\n" + "for s in (signal.SIGINT, signal.SIGHUP):\n signal.signal(s, signal.SIG_IGN)\n" + "os.execv(sys.executable, [sys.executable] + sys.argv[1:])\n" +) + + +def _start_runner(run_dir, body, timeout): + return subprocess.Popen([sys.executable, "-c", _WITH_SIGNALS_IGNORED, *_runner_args(run_dir, body, timeout)[1:]]) + + +@pytest.mark.parametrize("signum", [signal.SIGTERM, signal.SIGINT, signal.SIGHUP]) +def test_a_signal_to_the_runner_stops_the_script_and_everything_it_started(tmp_path, signum): + proc = _start_runner(tmp_path, _STARTS_A_CHILD, 60) + child = _wait_for_pid(tmp_path) + proc.send_signal(signum) + assert proc.wait(timeout=30) == 128 + signum + assert _gone_within(child) + + +def test_a_forwarded_signal_the_script_ignores_is_followed_by_sigkill(tmp_path): + ignore_term = "def run(inp):\n signal.signal(signal.SIGTERM, signal.SIG_IGN)\n" + body = _STARTS_A_CHILD.replace("def run(inp):\n", ignore_term) + proc = _start_runner(tmp_path, "import signal\n" + body, 60) + child = _wait_for_pid(tmp_path) + started = time.monotonic() + proc.send_signal(signal.SIGTERM) + assert proc.wait(timeout=30) == 128 + signal.SIGKILL + assert 5 <= time.monotonic() - started < 30 + assert _gone_within(child) + + +def test_artifact_files_named_after_the_supervisors_imports_never_run_in_it(tmp_path): + run_dir = tmp_path / "run" + run_dir.mkdir() + (tmp_path / "link").symlink_to(run_dir) # reached through a symlink, as macOS's /tmp is + for name in ("signal", "subprocess", "time"): + (run_dir / f"{name}.py").write_text("raise SystemExit('ran in the supervisor')\n") + result = _supervised(tmp_path / "link", "def run(inp):\n return {'ok': True}\n") + assert result.returncode == 0, result.stderr + assert json.loads((run_dir / "output.json").read_text()) == {"ok": True} + + +def test_the_script_still_imports_its_own_module_named_like_one_the_supervisor_uses(tmp_path): + (tmp_path / "signal.py").write_text("VALUE = 'mine'\n") + result = _supervised(tmp_path, "import signal\ndef run(inp):\n return {'from': signal.VALUE}\n") + assert result.returncode == 0, result.stderr + assert json.loads((tmp_path / "output.json").read_text()) == {"from": "mine"} + + def test_config_round_trips_through_to_dict_from_dict(): step = _step( entrypoint="filter_prs",