Bound the AutoSetup wait, and stop the task holding it from outliving the run

composer/cli/console_autoprove.py

@@ -1,13 +1,16 @@
 """Entry point for the auto-prove pipeline — console (no TUI) mode."""
 
 import asyncio
+import logging
 
 import composer.bind as _
 
 from composer.diagnostics.timing import RunSummary
 from composer.ui.autoprove_console import AutoProveConsoleHandler
 from composer.spec.source.autoprove_common import _entry_point
 
+_logger = logging.getLogger(__name__)
+
 
 # ---------------------------------------------------------------------------
 # Main
@@ -37,6 +40,30 @@ async def _main() -> int:
         return 0
 
 
+def _name_pending(loop: asyncio.AbstractEventLoop) -> None:
+    """Say what the loop still holds before it is asked to close.
+
+    Closing cancels whatever is left and then waits for all of it, with no bound. Anything named
+    here is something that wait can be spent on.
+    """
+    pending = [t for t in asyncio.all_tasks(loop) if not t.done()]
+    if pending:
+        _logger.info(
+            "%d task(s) still running at shutdown: %s",
+            len(pending),
+            ", ".join(sorted(t.get_name() for t in pending)),
+        )
+
+
 def main() -> int:
-    return asyncio.run(_main())
+    runner = asyncio.Runner()
+    try:
+        try:
+            return runner.run(_main())
+        finally:
+            _logger.info("pipeline returned; shutting the event loop down")
+            _name_pending(runner.get_loop())
+    finally:
+        runner.close()
+        _logger.info("event loop closed")
 

composer/pipeline/cli.py

@@ -1,5 +1,6 @@
 from typing import Protocol, AsyncIterator, TYPE_CHECKING
 import json
+import logging
 import sys
 import pathlib
 import enum
@@ -53,6 +54,8 @@
 from composer.spec.util import fs_forbidden_read
 import hashlib
 
+_logger = logging.getLogger(__name__)
+
 
 def autoprover_version() -> str:
     """The git commit recorded for a ``git+``-installed ``ai-composer`` (its ``direct_url.json``),
@@ -477,3 +480,4 @@ async def cont[FormT: BackendResult, A: ArtifactIdentifier, U: FeatureUnit, Main
                     await at_exit(init_source, data_logger)
                 except Exception:
                     pass
+                _logger.info("run exit handlers done")

composer/pipeline/core.py

@@ -37,7 +37,8 @@
     Protocol, Any, ClassVar, Concatenate, cast, Awaitable, Sequence, Callable, ContextManager, overload
 )
 from abc import ABC, abstractmethod
-from contextlib import nullcontext
+from contextlib import nullcontext, asynccontextmanager
+from collections.abc import AsyncIterator, Coroutine
 
 
 from composer.io.multi_job import TaskInfo
@@ -464,6 +465,24 @@ async def _run_pipeline_inner[P: enum.Enum, FormT: BackendResult, H, A: Artifact
             extra_context=extra_context, max_bug_rounds=max_bug_rounds, ecosystem=ecosystem,
         )
 
+
+@asynccontextmanager
+async def _alongside[T](coro: Coroutine[Any, Any, T]) -> AsyncIterator[asyncio.Task[T]]:
+    """Run ``coro`` beside the body, and end it with the body rather than letting it outlive it.
+
+    A task still running when its scope is left is reachable only by ``asyncio.run``'s teardown,
+    which cancels it after the run's exit handlers have finished and then waits for it with no
+    bound of its own. Whatever the task is holding at that point is held there.
+    """
+    task = asyncio.create_task(coro)
+    try:
+        yield task
+    except BaseException:
+        task.cancel()
+        await asyncio.gather(task, return_exceptions=True)
+        raise
+
+
 # ---- the driver --------------------------------------------------------------
 async def run_pipeline_inner[P: enum.Enum, FormT: BackendResult, H, A: ArtifactIdentifier, U: FeatureUnit, Main, App: BaseApplication, Pre](
     backend: PipelineBackend[P, FormT, H, A, U, Main, App, Pre],
@@ -535,45 +554,44 @@ async def _run_analysis():
     async def _prepare_formalization() -> Formalizer[FormT, U] | StagedFormalizer[FormT, U]:
         with named_budget_or_nop("formalization_preparation"):
             return await prepared.prepare_formalization(run)
-    staged_task = asyncio.create_task(_prepare_formalization())
-
-    batches: list[_Batch[U]] = await _extract_all(
-        backend.analysis_spec.properties_key,
-        prepared.main,
-        backend.backend_guidance,
-        run,
-        phases["extraction"],
-        interactive,
-        threat_model,
-        extra_context,
-        max_bug_rounds,
-        ecosystem,
-        plugin_manager.bind_phase(
-            phases.get("extraction_plugin") or phases["extraction"],
-        )
-    )
-    # 3b. Prioritized runs formalize one property, not every component's. This is the first
-    #     point where every component's candidates exist together — inference never sees more
-    #     than the unit it was given, and the per-component plugin hook cannot compare across
-    #     units either.
-    #
-    #     It sits *before* the ``staged_task`` await rather than after because the ranking
-    #     reads nothing that pre-formalization produces: it needs the extracted batches and
-    #     the guidance documents, both of which are already in hand. Awaiting the setup first
-    #     would queue the run's cheapest decision behind its most expensive step — on a real
-    #     contract that is hours during which the focus is knowable but unknown. It also puts
-    #     the focus in hand earlier than the setup finishes, which is what a pre-formalization
-    #     step would need in order to aim at it.
-    #     The empty-batch check deliberately stays *after* the await below rather than moving
-    #     up here to guard this: a setup that failed is the more useful error than "nothing was
-    #     extracted", and raising early would mask it.
-    deprioritized: list[DeprioritizedProperty] = []
-    if batches and run.run_mode is RunMode.PRIORITIZED:
-        with named_budget_or_nop("property_extraction"):
-            batches, deprioritized = await _prioritize(
-                backend.analysis_spec.properties_key, backend.artifact_store,
-                batches, run, phases["extraction"], threat_model, extra_context,
+    async with _alongside(_prepare_formalization()) as staged_task:
+        batches: list[_Batch[U]] = await _extract_all(
+            backend.analysis_spec.properties_key,
+            prepared.main,
+            backend.backend_guidance,
+            run,
+            phases["extraction"],
+            interactive,
+            threat_model,
+            extra_context,
+            max_bug_rounds,
+            ecosystem,
+            plugin_manager.bind_phase(
+                phases.get("extraction_plugin") or phases["extraction"],
             )
+        )
+        # 3b. Prioritized runs formalize one property, not every component's. This is the first
+        #     point where every component's candidates exist together — inference never sees more
+        #     than the unit it was given, and the per-component plugin hook cannot compare across
+        #     units either.
+        #
+        #     It sits *before* the ``staged_task`` await rather than after because the ranking
+        #     reads nothing that pre-formalization produces: it needs the extracted batches and
+        #     the guidance documents, both of which are already in hand. Awaiting the setup first
+        #     would queue the run's cheapest decision behind its most expensive step — on a real
+        #     contract that is hours during which the focus is knowable but unknown. It also puts
+        #     the focus in hand earlier than the setup finishes, which is what a pre-formalization
+        #     step would need in order to aim at it.
+        #     The empty-batch check deliberately stays *after* the await below rather than moving
+        #     up here to guard this: a setup that failed is the more useful error than "nothing was
+        #     extracted", and raising early would mask it.
+        deprioritized: list[DeprioritizedProperty] = []
+        if batches and run.run_mode is RunMode.PRIORITIZED:
+            with named_budget_or_nop("property_extraction"):
+                batches, deprioritized = await _prioritize(
+                    backend.analysis_spec.properties_key, backend.artifact_store,
+                    batches, run, phases["extraction"], threat_model, extra_context,
+                )
 
     staged = await staged_task
     if not batches:

composer/spec/source/autosetup.py

@@ -10,6 +10,7 @@
 import re
 import sys
 import tempfile
+import time
 from collections.abc import Callable
 from pydantic import BaseModel, Field
 from pathlib import Path
@@ -85,6 +86,46 @@ async def _drain(
         sink(buf)
 
 
+#: How long to give the AutoSetup child to exit before we stop waiting on it. The wait below sits
+#: in a ``finally`` that a cancellation also travels through, so it has to be bounded: an
+#: unbounded wait there parks the whole run behind one subprocess, with no timeout above it that
+#: can help and nothing left running to say why.
+_REAP_TIMEOUT_S = 30.0
+
+
+async def _reap(proc: asyncio.subprocess.Process) -> int | None:
+    """Wait for ``proc`` to exit, escalating to a kill, and stop waiting rather than block forever.
+
+    Returns the exit status, or ``None`` if the child could not be reaped at all — a child that
+    will not die tells us nothing about whether AutoSetup succeeded, so the caller treats that as
+    a failure.
+    """
+    started = time.monotonic()
+    for escalate in (False, True):
+        if escalate and proc.returncode is None:
+            _logger.warning(
+                "AutoSetup child still alive after %.0fs, killing it", _REAP_TIMEOUT_S
+            )
+            proc.kill()
+        try:
+            async with asyncio.timeout(_REAP_TIMEOUT_S):
+                # Shielded: the timeout must cancel our wait, not the child's exit plumbing.
+                returncode = await asyncio.shield(proc.wait())
+        except TimeoutError:
+            continue
+        # Logged where it ends, not only where it begins: a wait that says nothing on the way out
+        # cannot be told apart, afterwards, from one that never returned.
+        _logger.info(
+            "AutoSetup child exited %d after %.1fs", returncode, time.monotonic() - started
+        )
+        return returncode
+    _logger.error(
+        "AutoSetup child outlived a kill after %.1fs, giving up on it and continuing",
+        time.monotonic() - started,
+    )
+    return proc.returncode
+
+
 async def run_autosetup(
     project_root: Path,
     relative_path: str,
@@ -176,7 +217,12 @@ def log_complete(self, returncode: int):
             raise
         finally:
             _logger.debug("AutoSetup process complete, waiting for exit")
-            returncode = await proc.wait()
+            returncode = await _reap(proc)
+        if returncode is None:
+            return SetupFailure(
+                error="AutoSetup did not exit",
+                stderr="\n".join(stderr_lines),
+            )
         cb.log_complete(returncode)
         if returncode != 0:
             return SetupFailure(

tests/test_autosetup_reap.py

@@ -0,0 +1,74 @@
+"""Tests for the bounded reap of the AutoSetup subprocess (``composer/spec/source/autosetup.py``).
+
+The wait for the child sits in a ``finally``, which is also the path a cancelled run takes on its
+way out. Whatever the child is doing, that wait has to end.
+"""
+
+import asyncio
+import sys
+
+import pytest
+
+from composer.spec.source import autosetup
+
+pytestmark = pytest.mark.asyncio
+
+
+async def _spawn(code: str) -> asyncio.subprocess.Process:
+    return await asyncio.create_subprocess_exec(
+        sys.executable, "-c", code,
+        stdout=asyncio.subprocess.DEVNULL,
+        stderr=asyncio.subprocess.DEVNULL,
+        stdin=asyncio.subprocess.DEVNULL,
+    )
+
+
+async def test_a_child_that_exits_reports_its_status():
+    proc = await _spawn("raise SystemExit(3)")
+    assert await autosetup._reap(proc) == 3
+
+
+async def test_a_child_that_will_not_exit_is_killed(monkeypatch):
+    monkeypatch.setattr(autosetup, "_REAP_TIMEOUT_S", 0.2)
+    # Ignores SIGTERM, so only the kill ends it.
+    proc = await _spawn(
+        "import signal, time; signal.signal(signal.SIGTERM, signal.SIG_IGN); time.sleep(300)"
+    )
+    returncode = await asyncio.wait_for(autosetup._reap(proc), timeout=10)
+    assert returncode is not None and returncode < 0
+
+
+async def test_the_reap_ends_inside_a_cancelled_task(monkeypatch):
+    """The shape a run takes on its way out: the wait sits in a ``finally`` of a cancelled task.
+
+    ``asyncio.timeout`` there has to raise ``TimeoutError`` off its own timer rather than pass the
+    task's pending cancellation through, or the escalation never happens and the task never
+    finishes unwinding.
+    """
+    monkeypatch.setattr(autosetup, "_REAP_TIMEOUT_S", 0.2)
+    proc = await _spawn(
+        "import signal, time; signal.signal(signal.SIGTERM, signal.SIG_IGN); time.sleep(300)"
+    )
+    reaped: list[int | None] = []
+
+    async def run_and_reap() -> None:
+        try:
+            await asyncio.Event().wait()
+        finally:
+            reaped.append(await autosetup._reap(proc))
+
+    task = asyncio.create_task(run_and_reap())
+    await asyncio.sleep(0.05)
+    task.cancel()
+
+    _, pending = await asyncio.wait({task}, timeout=10)
+    assert not pending, "the reap never returned inside a cancelled task"
+    assert task.cancelled()
+    assert reaped and reaped[0] is not None and reaped[0] < 0
+
+
+async def test_the_reap_says_how_it_ended(caplog):
+    proc = await _spawn("raise SystemExit(0)")
+    with caplog.at_level("INFO", logger="composer.spec.source.autosetup"):
+        assert await autosetup._reap(proc) == 0
+    assert any("AutoSetup child exited 0" in r.getMessage() for r in caplog.records)