build_util: the per-script cap kills the process group, not just the child

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,

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 "")