diff --git a/autohands/build_util.py b/autohands/build_util.py index b89e25e..b8ea25c 100644 --- a/autohands/build_util.py +++ b/autohands/build_util.py @@ -2,6 +2,7 @@ import logging import os import re +import signal import subprocess import sys import time @@ -28,7 +29,7 @@ def timeout_for(env=None) -> int: so on its own it can only ever express ONE cap for a whole run. The per-script environment built by ``env_config.build_env_for_script`` is handed to the child, and a profile may set ``BUILD_SCRIPT_TIMEOUT`` on it - for a matching pattern — but the ``subprocess.run(timeout=...)`` kill timer + for a matching pattern — but the ``run_capped(timeout=...)`` kill timer lives in the PARENT, so that value has no effect unless the parent reads it back out. This resolves it. @@ -87,6 +88,72 @@ def tail(stream) -> str: return "\n".join(parts) + +def kill_group(proc: subprocess.Popen) -> None: + """SIGKILL the child's whole process group, tolerating an already-dead one. + + Public (unprefixed) because the workspace `run_smoke.py` runners import it + rather than each growing a copy. + """ + try: + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + except (ProcessLookupError, PermissionError): # pragma: no cover - race + proc.kill() + + +def run_capped(args, timeout, check=False, stdout=None, stderr=None, + text=False, env=None, cwd=None) -> subprocess.CompletedProcess: + """`subprocess.run`, but a timeout kills the child's whole process GROUP. + + A drop-in for the `subprocess.run(..., timeout=...)` calls this module used + to make: same `TimeoutExpired` and `CalledProcessError`, carrying the same + captured `output`/`stderr`, so every caller's handling is unchanged. + + The group is the point. `subprocess.run` kills only the direct child, so a + grandchild -- a Popen'd helper, a multiprocessing worker, a compile server + -- outlives the cap and keeps running, holding whatever memory and devices + it had. Over a run of hundreds of scripts those accumulate against every + script that follows. `start_new_session=True` puts the child in its own + group, and killing the group takes its descendants with it. + + The same mechanism matters more for a runner that captures output without + a cap at all: there the parent waits for the stdout pipe to reach EOF, and + a grandchild holding that pipe open keeps the read blocked after the child + itself has exited -- a script whose work has finished hangs the runner + indefinitely. That shape is a workspace-side bug, not this module's, but + the fix is this same group kill. + """ + proc = subprocess.Popen( + args, + stdout=stdout, + stderr=stderr, + text=text, + env=env, + cwd=cwd, + start_new_session=True, + ) + try: + output, errs = proc.communicate(timeout=timeout) + except subprocess.TimeoutExpired: + kill_group(proc) + # The group is gone, so this drains whatever was buffered and returns. + output, errs = proc.communicate() + raise subprocess.TimeoutExpired( + args, timeout, output=output, stderr=errs + ) from None + except BaseException: + # Matches subprocess.run's context manager: never leave the child (or + # its group) running when the caller is unwinding, e.g. on Ctrl-C. + kill_group(proc) + proc.wait() + raise + if check and proc.returncode != 0: + raise subprocess.CalledProcessError( + proc.returncode, args, output=output, stderr=errs + ) + return subprocess.CompletedProcess(args, proc.returncode, output, errs) + + def py_to_notebook(filename: Path): subprocess.run( ["python3", f"{BUILD_PATH}/add_notebook_quotes.py", filename, "temp.py"], @@ -305,7 +372,7 @@ def execute_notebook(f, report=None, env=None): # from the repo root. nbconvert has no CLI flag for the kernel cwd, so # the runner sets resources['metadata']['path'] via the Python API. # Still a subprocess, so isolation/timeout/env are unchanged. - subprocess.run( + run_capped( [ sys.executable, str(Path(__file__).parent / "run_notebook.py"), @@ -456,16 +523,17 @@ def execute_script(f, report=None, env=None, extra_args=None): start = time.time() try: if report is not None: - result = subprocess.run( + result = run_capped( args, check=True, timeout=timeout_secs, - capture_output=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, text=True, env=env, ) else: - subprocess.run( + run_capped( args, check=True, timeout=timeout_secs, diff --git a/tests/test_script_timeout.py b/tests/test_script_timeout.py index fd1504e..e82bb79 100644 --- a/tests/test_script_timeout.py +++ b/tests/test_script_timeout.py @@ -18,11 +18,19 @@ script cannot report its own progress, so without that tail a TIMEOUT cannot say which block was executing — the reason the three jax_grad timeouts could not be diagnosed from CI artefacts at all (PyAutoHands#226). + +A third property, covered by ``TestTimeoutKillsProcessGroup``: the kill must +reach the child's whole process GROUP. ``subprocess.run(timeout=...)`` kills +only the direct child, so a grandchild outlives the cap and keeps running -- +over a long mega-run those leak, each holding whatever memory and GPU it had. +``build_util.run_capped`` puts the child in its own session and SIGKILLs the +group instead. """ import os import subprocess import sys +import time from pathlib import Path import pytest @@ -207,3 +215,59 @@ def test_timeout_preserves_child_stdout(self, tmp_path, monkeypatch, real_interp message = report.results[0].error_message assert report.results[0].status == Status.TIMEOUT assert "=== variant 3 ===" in message + + +class TestTimeoutKillsProcessGroup: + """The cap must reap the child's descendants, not just the child.""" + + # A script that finishes its own work immediately but leaves a grandchild + # running and holding the inherited stdout pipe. Under a plain + # `subprocess.run(timeout=...)` the grandchild survives the cap. + _SPAWNS_GRANDCHILD = ( + "import subprocess, sys\n" + "subprocess.Popen([sys.executable, '-c',\n" + " 'import time; MARKER_{marker}=1; time.sleep(120)'])\n" + "print('work done', flush=True)\n" + "import time; time.sleep(120)\n" + ) + + @staticmethod + def _alive(marker: str) -> int: + found = subprocess.run( + ["pgrep", "-f", f"MARKER_{marker}"], capture_output=True, text=True + ).stdout + return len([line for line in found.split() if line.strip()]) + + def test_grandchild_is_reaped_at_the_cap(self, tmp_path, monkeypatch, real_interpreter): + marker = "pyautohands_groupkill" + monkeypatch.setattr(build_util, "TIMEOUT_SECS", 600) + script = _write_script(tmp_path, self._SPAWNS_GRANDCHILD.format(marker=marker)) + + report = RunReport(project="t", directory="d", run_type="script") + try: + execute_script( + str(script), + report=report, + env={**dict(PATH=os.environ.get("PATH", "")), "BUILD_SCRIPT_TIMEOUT": "2"}, + ) + assert report.results[0].status == Status.TIMEOUT + # The point of the group kill. Without it this is 1: the direct + # child is dead but its descendant runs on for its full 120s. + time.sleep(1) + assert self._alive(marker) == 0 + finally: + subprocess.run(["pkill", "-f", f"MARKER_{marker}"], capture_output=True) + + def test_run_capped_reports_the_output_captured_before_the_kill(self, tmp_path): + # The drain after the group kill still has to yield what the child + # printed -- killing the group is what lets that read reach EOF at all. + script = _write_script(tmp_path, "print('before the hang', flush=True)\nimport time\ntime.sleep(60)\n") + with pytest.raises(subprocess.TimeoutExpired) as excinfo: + build_util.run_capped( + [sys.executable, str(script)], + timeout=2, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + assert "before the hang" in (excinfo.value.output or "")