diff --git a/tests/v1/test_e2e.py b/tests/v1/test_e2e.py index a78dcee7f1..a71ae9202c 100644 --- a/tests/v1/test_e2e.py +++ b/tests/v1/test_e2e.py @@ -234,3 +234,111 @@ async def test_agentic(run_v1, harness, harness_runtime, tmp_path): assert trace.errors == [] assert trace.num_turns >= 1 # ran a command, then finished assert trace.reward == 1.0 + + +@pytest.mark.parametrize( + "runtime_type", + [ + pytest.param("subprocess", id="cancel-guard-subprocess"), + pytest.param("docker", marks=pytest.mark.docker, id="cancel-guard-docker"), + pytest.param("prime", marks=pytest.mark.prime, id="cancel-guard-prime"), + ], +) +async def test_cancelled_stop_still_frees_runtime(runtime_type, caplog, monkeypatch): + """The cancellation guard, end-to-end on a real resource: a Ctrl-C landing while + `stop()` is in flight must not truncate teardown (leaking the workdir / container / + paid sandbox), must report the drain, and must still propagate the cancellation + afterwards. This is the trip-wire for the `run_shielded` shield inside + `Runtime.stop` — the gate below holds teardown open so the cancel deterministically + lands before the real cleanup starts, so an unshielded `stop()` genuinely leaks and + fails the leak check. subprocess/docker run in CI; prime is the local full-realism + variant (a real API DELETE observed server-side).""" + import asyncio + import contextlib + import logging + import os + import subprocess + import time + from pathlib import Path + from uuid import uuid4 + + from verifiers.v1.runtimes import make_runtime + from verifiers.v1.runtimes.docker import DockerConfig + from verifiers.v1.runtimes.prime import PrimeConfig + from verifiers.v1.runtimes.subprocess import SubprocessConfig + + if runtime_type == "prime" and not os.environ.get("PRIME_API_KEY"): + pytest.skip("needs PRIME_API_KEY") + + config = { + "subprocess": lambda: SubprocessConfig(), + "docker": lambda: DockerConfig(), + "prime": lambda: PrimeConfig(labels=["vf-ci"]), + }[runtime_type]() + name = f"cancel-guard-{uuid4().hex[:8]}" + runtime = make_runtime(config, name=name) + + # Gate the real teardown: signal entry, hold the window open, then run it. The + # `finished` flag is the provider-agnostic trip — without the shield the cancel + # kills the sleep and the real teardown never runs. + entered, finished = asyncio.Event(), asyncio.Event() + real_teardown = runtime.teardown + + async def gated_teardown() -> None: + entered.set() + await asyncio.sleep(0.2) + await real_teardown() + finished.set() + + monkeypatch.setattr(runtime, "teardown", gated_teardown) + monkeypatch.setattr(logging.getLogger("verifiers"), "propagate", True) + + await runtime.start() + descriptor = runtime.descriptor + try: + + async def owner() -> None: + try: + pass # the rollout body finished; teardown is on the happy path + finally: + await runtime.stop() + + task = asyncio.create_task(owner()) + await entered.wait() + with caplog.at_level(logging.WARNING, logger="verifiers.v1.runtimes.base"): + task.cancel() # Ctrl-C while stop() is in flight + with pytest.raises(asyncio.CancelledError): + await task + assert finished.is_set() # teardown ran to completion despite the cancel + assert task.cancelled() # and the cancellation still propagated after it + assert any("Ctrl-C again" in r.getMessage() for r in caplog.records) + + # Provider-side truth, checked in-process (at exit the atexit backstop would + # mask the result): the resource is actually gone. + if runtime_type == "subprocess": + assert not Path(f"/tmp/{name}").exists(), "workdir leaked" + elif runtime_type == "docker": + leaked = subprocess.run( + ["docker", "ps", "-aq", "--filter", f"name={name}"], + capture_output=True, + text=True, + timeout=30, + ).stdout.strip() + assert not leaked, f"container {name} leaked after a cancelled teardown" + else: + from prime_sandboxes import SandboxClient + from prime_sandboxes.core import APIClient + + client = SandboxClient(APIClient()) + deadline = time.time() + 30 + while ( + status := await asyncio.to_thread(lambda: client.get(descriptor).status) + ) != "TERMINATED" and time.time() < deadline: + await asyncio.sleep(2) + assert status == "TERMINATED", ( + f"sandbox {descriptor} leaked after a cancelled teardown ({status})" + ) + finally: + # A failing run must not keep the resource around (cleanup is idempotent). + with contextlib.suppress(Exception): + runtime.cleanup() diff --git a/verifiers/v1/runtimes/base.py b/verifiers/v1/runtimes/base.py index 3becd19658..a1d3d46b3f 100644 --- a/verifiers/v1/runtimes/base.py +++ b/verifiers/v1/runtimes/base.py @@ -11,6 +11,7 @@ import hashlib import logging import shlex +import time import uuid import weakref from abc import ABC, abstractmethod @@ -86,6 +87,17 @@ def parse_gpu(gpu: str | None) -> tuple[str | None, int]: _LIVE: "weakref.WeakSet[Runtime]" = weakref.WeakSet() _atexit_armed = False +_STOPPING: "dict[Runtime, float]" = {} +"""In-flight teardowns -> start time. When an interrupt lands mid-`stop()`, one aggregate +report tells the user how much work is draining, so they can judge whether to wait or +press Ctrl-C again; also a debugger's view of what a shutdown is stuck on. Entries are +popped when their `stop()` returns.""" + +_drain_reported = False +"""Whether the current drain has been reported — one aggregate line per drain, however +many teardowns absorb the interrupt (512 in-flight must not print 512 lines). Reset when +`_STOPPING` empties, so a later drain in a long-lived process reports again.""" + def register(runtime: "Runtime") -> None: """Track a runtime so the atexit hook can free it if a signal cuts its `finally` short. @@ -146,9 +158,47 @@ async def stop(self) -> None: """Free the provisioned resource on the normal path (the owner's `finally`), shielded from cancellation: a Ctrl-C / SIGTERM cancels that `finally` mid-await, and an interrupted teardown leaks the container / paid sandbox. Runs `teardown` - to completion, then re-raises the cancellation. Framework method — override - `teardown`, not this.""" - await run_shielded(self.teardown()) + to completion, then re-raises the cancellation; when the interrupt lands + mid-teardown, reports the drain (name, in-flight count, age) so the wait reads + as intentional and the user can decide whether to Ctrl-C again. Framework + method — override `teardown`, not this.""" + + def interrupted() -> None: + global _drain_reported + logger.debug( + "finishing teardown of %s (%s)", self.name, self.descriptor or "-" + ) + if _drain_reported: + return + _drain_reported = True + logger.warning( + "finishing %d in-flight teardown(s) (oldest %.0fs) — Ctrl-C again to " + "abort; leftover resources are freed best-effort at exit", + len(_STOPPING), + time.time() - min(_STOPPING.values()), + ) + + global _drain_reported + _STOPPING[self] = time.time() + aborted = False + try: + await run_shielded(self.teardown(), interrupted=interrupted) + except (KeyboardInterrupt, SystemExit): + aborted = ( + True # a user abort cut the drain short (run_shielded propagates it) + ) + raise + finally: + _STOPPING.pop(self, None) + if _drain_reported and not _STOPPING: + _drain_reported = False + if aborted: + logger.warning( + "teardown drain aborted — leftover resources are freed " + "best-effort at exit" + ) + else: + logger.warning("all in-flight teardowns finished") async def teardown(self) -> None: """Free the provisioned resource, off the event loop. Override only for teardown diff --git a/verifiers/v1/utils/aio.py b/verifiers/v1/utils/aio.py index a62887f509..7705d780da 100644 --- a/verifiers/v1/utils/aio.py +++ b/verifiers/v1/utils/aio.py @@ -1,30 +1,38 @@ """Asyncio helpers shared across the framework.""" import asyncio -from collections.abc import Awaitable +from collections.abc import Awaitable, Callable from typing import TypeVar T = TypeVar("T") -async def run_shielded(coro: Awaitable[T]) -> T: +async def run_shielded( + coro: Awaitable[T], interrupted: Callable[[], None] | None = None +) -> T: """Run `coro` to completion even if the surrounding task is cancelled, then re-raise the cancellation. A bare `asyncio.shield` is not enough: on cancellation it re-raises immediately while the inner task runs on orphaned — and the loop's shutdown then cancels that orphan mid-await anyway. This keeps awaiting (absorbing repeated task cancellations) until `coro` finishes; the first CancelledError is re-raised after. + `interrupted`, if given, is invoked once on the first absorbed cancellation — so the + caller can tell the user the wait is intentional, not a hang. If `coro` raises without a cancellation, the error propagates unchanged; with one, the cancellation wins and the error is chained under it (`from`), so it is never - silently lost. A second Ctrl-C that raises KeyboardInterrupt out of the event loop - itself is beyond any task-level shield — that path is the atexit backstop's job.""" + silently lost — except a `BaseException` (KeyboardInterrupt, SystemExit), which + propagates immediately: a user abort must never lose to a pending cancellation. + A second Ctrl-C that raises KeyboardInterrupt out of the event loop itself is beyond + any task-level shield — that path is the atexit backstop's job.""" task = asyncio.ensure_future(coro) cancelled: asyncio.CancelledError | None = None while not task.done(): try: await asyncio.shield(task) except asyncio.CancelledError as e: + if cancelled is None and interrupted is not None: + interrupted() cancelled = e - except BaseException: + except Exception: pass # `coro` itself failed → task is done; re-raised (or chained) below if cancelled is not None: raise cancelled from (None if task.cancelled() else task.exception()) diff --git a/verifiers/v1/utils/logging.py b/verifiers/v1/utils/logging.py index 9ee900a17e..8cf487792f 100644 --- a/verifiers/v1/utils/logging.py +++ b/verifiers/v1/utils/logging.py @@ -12,6 +12,7 @@ from pathlib import Path from loguru import logger +from rich import get_console LIBRARY_LOGGER = "verifiers.v1" FORMAT = ( @@ -41,11 +42,21 @@ def setup_logging( ) -> None: """Route the library's stdlib logs through loguru, to stderr (when `console`) and/or a `log_file`. `console=False` (used under the `--rich` dashboard, which owns the screen) - keeps logs off the terminal while still writing them to `log_file`.""" + keeps routine logs off the terminal while still writing them to `log_file` — but + warnings still surface, printed through the dashboard's own rich Console so an active + `Live` renders them above the live region instead of tearing the frame (a Ctrl-C + teardown drain must not look like a hang).""" lvl = level.upper() logger.remove() if console: logger.add(sys.stderr, level=lvl, format=FORMAT) + else: + rich_console = get_console() + logger.add( + lambda m: rich_console.print(m, end="", markup=False, highlight=False), + level="WARNING", + format=FORMAT, + ) if log_file is not None: Path(log_file).parent.mkdir(parents=True, exist_ok=True) logger.add(log_file, level=lvl, format=FORMAT)