diff --git a/composer/prover/core.py b/composer/prover/core.py index 76b1876d..d620b483 100644 --- a/composer/prover/core.py +++ b/composer/prover/core.py @@ -27,11 +27,12 @@ from contextlib import asynccontextmanager from dataclasses import dataclass, field from pathlib import Path -from typing import AsyncIterator, Callable, Protocol, cast, override, Awaitable +from typing import Any, AsyncIterator, Callable, Protocol, cast, override, Awaitable from abc import ABC, abstractmethod import json import logging import os +import signal import uuid @@ -409,35 +410,93 @@ async def _report_to_todo_list( res = await ainvoke(llm, fresh_messages) return res.text +class ProverSubprocessTimeout(Exception): + """A prover subprocess outlived its bound and was killed.""" + + +# Bound for the compile-and-typecheck subprocesses that support rule listing. +# They build the project and run the CVL typechecker; neither waits on the cloud. +_BUILD_TIMEOUT_S = 1800 + + +def _kill_tree(proc: asyncio.subprocess.Process) -> None: + """Kill the subprocess and everything it started. + + ``certoraRun`` is a python front end that shells out to a JVM, so killing + only the direct child leaves that JVM running: it holds the pipe the parent + was reading, and on a container it survives until the container does. + Children are spawned in their own process group precisely so one signal can + reach the whole tree. + """ + try: + os.killpg(proc.pid, signal.SIGKILL) + except (ProcessLookupError, PermissionError): + proc.kill() + except Exception: + _logger.exception("Could not kill prover subprocess group %d", proc.pid) + + +@asynccontextmanager +async def _bounded_subprocess( + *argv: str, cwd: str, timeout: float, **kwargs: Any +) -> AsyncIterator[asyncio.subprocess.Process]: + """Run a subprocess whose whole tree is killed if the body outlives ``timeout``. + + The prover's local phase (solc, the CVL typechecker) can wedge without + exiting or writing anything further, and every await on it is unbounded: + the caller waits on the child, the child waits on the JVM, and the JVM waits + on a lock it will never get. One observed instance sat that way for three + days, having consumed a second of CPU. Nothing recovers a run from that + state, so bound it here and let the failure surface as an ordinary error the + agent can act on. + + ``process_group=0`` makes the child a group leader, which is what lets + ``_kill_tree`` reach a grandchild JVM. + """ + proc = await asyncio.subprocess.create_subprocess_exec( + *argv, cwd=cwd, process_group=0, **kwargs + ) + try: + async with asyncio.timeout(timeout): + yield proc + except TimeoutError: + _kill_tree(proc) + await proc.wait() + raise ProverSubprocessTimeout( + f"{argv[0]} exceeded {timeout:.0f}s and was killed" + ) from None + + async def run_prover_inner( folder: Path, args: list[str], on_err: Callable[[int | None, str, str], None], - on_stdout: Callable[[str], Awaitable[None]] + on_stdout: Callable[[str], Awaitable[None]], + timeout: float, ) -> tuple[ProverResult | str, str]: # 3-5. Spawn async subprocess, stream stdout, collect stderr wrapper_script = Path(__file__).parent / "certoraRunWrapper.py" with tempfile.NamedTemporaryFile("rb", suffix=".json") as output_file: - proc = await asyncio.subprocess.create_subprocess_exec( + async with _bounded_subprocess( sys.executable, str(wrapper_script), str(output_file.name), *args, cwd=str(folder), + timeout=timeout, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, - ) - - stdout_lines: list[str] = [] - assert proc.stdout is not None - while True: - raw = await proc.stdout.readline() - if not raw: - break - line = raw.decode() - stdout_lines.append(line) - await on_stdout(line.rstrip("\n")) - - stderr_raw = await proc.stderr.read() if proc.stderr else b"" + ) as proc: + stdout_lines: list[str] = [] + assert proc.stdout is not None + while True: + raw = await proc.stdout.readline() + if not raw: + break + line = raw.decode() + stdout_lines.append(line) + await on_stdout(line.rstrip("\n")) + + stderr_raw = await proc.stderr.read() if proc.stderr else b"" await proc.wait() stdout = "".join(stdout_lines) @@ -465,14 +524,14 @@ async def declared_rules_list( if any(m == "--msg" for m in args): raise ValueError("This unholy black magic only works if you don't pass msg") tc_key = uuid.uuid4().hex - proc = await asyncio.subprocess.create_subprocess_exec( + async with _bounded_subprocess( "certoraRun", *args, "--msg", tc_key, "--compilation_steps_only", cwd=str(folder), + timeout=_BUILD_TIMEOUT_S, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE - ) - - rc = await proc.wait() + ) as proc: + rc = await proc.wait() if rc != 0: raise ValueError("Type check failed?") from importlib.resources import files @@ -502,13 +561,14 @@ async def declared_rules_list( if found is None: raise ValueError("Couldn't find build dir") with tempfile.NamedTemporaryFile("r") as f: - proc = await asyncio.subprocess.create_subprocess_exec( + async with _bounded_subprocess( "java", "-jar", str(tc_jar), "-buildDirectory", str(found), "-typeCheck", "true", "-listRules", f.name, cwd=str(folder), + timeout=_BUILD_TIMEOUT_S, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE - ) - tc_rc = await proc.wait() + ) as proc: + tc_rc = await proc.wait() if tc_rc != 0: raise ValueError("Nope, no dice") all_rules = f.read() @@ -544,12 +604,28 @@ async def run_prover( # only submits and returns, so this isn't used — cloud runtime comes from the job's # execution window (see cloud_results / _job_runtime_ms). _t0 = time.perf_counter() - run_result, stdout = await run_prover_inner( - folder, - effective_args, - lambda ret_code, stdout, stderr: _logger.error("Process failed %d\nstdout:%s\nstderr:%s", ret_code, stdout, stderr), - callbacks.on_stdout_line - ) + # The same bound the result poller uses (step 7): whatever the run is allowed + # to take, plus slack. It is a backstop against a wedged local phase, not a + # budget -- a healthy cloud submit finishes in minutes. + subprocess_timeout = prover_opts.global_timeout + 5 * 60 + try: + run_result, stdout = await run_prover_inner( + folder, + effective_args, + lambda ret_code, stdout, stderr: _logger.error("Process failed %d\nstdout:%s\nstderr:%s", ret_code, stdout, stderr), + callbacks.on_stdout_line, + subprocess_timeout, + ) + except ProverSubprocessTimeout as e: + # Returned rather than raised: the agent reads this as a tool result and + # can retry or change approach, which beats stalling the run. + _logger.error("Prover subprocess timed out: %s", e) + return ( + f"The prover did not finish within {subprocess_timeout:.0f}s and was " + f"terminated. This is an infrastructure failure, not a problem with " + f"the specification: the local build or type-check phase stopped " + f"responding. Retrying may succeed." + ) local_runtime_ms = int((time.perf_counter() - _t0) * 1000) if isinstance(run_result, str): return run_result diff --git a/tests/test_prover_subprocess_timeout.py b/tests/test_prover_subprocess_timeout.py new file mode 100644 index 00000000..a758907a --- /dev/null +++ b/tests/test_prover_subprocess_timeout.py @@ -0,0 +1,89 @@ +"""A wedged prover subprocess must not hang the run, and must not leak its JVM. + +`certoraRun` is a python front end that shells out to a JVM for the build and +type-check phase. Every await on that chain used to be unbounded, so when the JVM +stopped responding the whole run stopped with it: the caller waited on +certoraRun, certoraRun waited on the JVM, and the JVM waited on a lock it never +got. One observed run sat that way for three days on a second of CPU. + +Two properties matter, and the second is the one that is easy to get wrong: +the wait has to end, and the kill has to reach the grandchild. Signalling only +the direct child leaves the JVM holding the pipe its parent was reading. +""" + +import asyncio +import os +import signal + +import pytest + +from composer.prover.core import ProverSubprocessTimeout, _bounded_subprocess + +pytestmark = pytest.mark.asyncio + + +# Every test here bounds itself. A regression in the group kill makes the helper +# hang rather than fail -- the orphaned grandchild holds the pipe its parent is +# read-waiting on -- and a test guarding against a hang must not hang CI. +_TEST_DEADLINE_S = 30 + + +def _alive(pid: int) -> bool: + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + return True + + +async def test_a_hung_subprocess_raises_instead_of_hanging(tmp_path) -> None: + async with asyncio.timeout(_TEST_DEADLINE_S): + with pytest.raises(ProverSubprocessTimeout, match="exceeded"): + async with _bounded_subprocess( + "sleep", "300", cwd=str(tmp_path), timeout=0.2 + ) as proc: + await proc.wait() + + +async def test_the_grandchild_is_killed_too(tmp_path) -> None: + """Stands in for the JVM that certoraRun spawns.""" + # The shell prints its child's pid, then outlives the timeout. + script = "sleep 300 & echo $! ; wait" + async with asyncio.timeout(_TEST_DEADLINE_S): + with pytest.raises(ProverSubprocessTimeout): + async with _bounded_subprocess( + "sh", "-c", script, + cwd=str(tmp_path), + timeout=1.0, + stdout=asyncio.subprocess.PIPE, + ) as proc: + assert proc.stdout is not None + grandchild = int((await proc.stdout.readline()).decode().strip()) + assert _alive(grandchild), "grandchild should be running before the timeout" + await proc.wait() + + # The group signal has to have reached it. It exits asynchronously, so allow + # a moment rather than asserting on the instant the exception surfaces. + for _ in range(50): + if not _alive(grandchild): + break + await asyncio.sleep(0.1) + assert not _alive(grandchild), "grandchild survived the timeout — killpg did not reach it" + + +async def test_a_subprocess_that_finishes_is_left_alone(tmp_path) -> None: + async with _bounded_subprocess("true", cwd=str(tmp_path), timeout=30) as proc: + assert await proc.wait() == 0 + + +async def test_the_child_leads_its_own_process_group(tmp_path) -> None: + """`process_group=0` is what makes the group kill reach a grandchild; if a + refactor drops it, the child shares our group and killpg would signal the + test runner instead.""" + async with _bounded_subprocess( + "sh", "-c", "sleep 30", cwd=str(tmp_path), timeout=30 + ) as proc: + assert os.getpgid(proc.pid) == proc.pid + assert os.getpgid(proc.pid) != os.getpgid(0) + os.killpg(proc.pid, signal.SIGKILL) + await proc.wait() diff --git a/tests/test_rules_striping.py b/tests/test_rules_striping.py index bea6cc8b..73f83cf4 100644 --- a/tests/test_rules_striping.py +++ b/tests/test_rules_striping.py @@ -213,8 +213,14 @@ def _install_fake_prover_procs( ``conf_msg``, to simulate a foreign run), java writes ``rules_text`` to its ``-listRules`` target.""" - async def fake_exec(*argv, cwd=None, stdout=None, stderr=None): + # ``**kwargs`` mirrors the real signature: the production call also passes + # ``process_group`` so a timeout can kill the whole tree. + async def fake_exec(*argv, cwd=None, stdout=None, stderr=None, **kwargs): assert cwd is not None + assert kwargs.get("process_group") == 0, ( + "children must lead their own process group so a timeout kill " + "reaches a grandchild JVM" + ) argv = [str(a) for a in argv] if argv[0] == "certoraRun": assert "--compilation_steps_only" in argv