diff --git a/pkg-py/src/commons/_execution/__init__.py b/pkg-py/src/commons/_execution/__init__.py new file mode 100644 index 00000000..48df49b2 --- /dev/null +++ b/pkg-py/src/commons/_execution/__init__.py @@ -0,0 +1 @@ +"""Running model-written code in a worker process.""" diff --git a/pkg-py/src/commons/_execution/_backend.py b/pkg-py/src/commons/_execution/_backend.py new file mode 100644 index 00000000..047b8f6e --- /dev/null +++ b/pkg-py/src/commons/_execution/_backend.py @@ -0,0 +1,218 @@ +"""The seam between the execution driver and whatever runs the worker. + +Everything above this line talks to a single ``exec``-shaped call, so a +container-hosted backend can be added later as another implementation rather +than as an edit to the driver. +""" + +from __future__ import annotations + +import asyncio +import contextlib +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Protocol, runtime_checkable + +__all__ = ["ExecBackend", "ExecResult", "ExecTimeoutError", "LocalBackend"] + +# Enough for a generous amount of printed output without letting a runaway +# loop hold the whole of it in memory. +DEFAULT_OUTPUT_LIMIT = 1024 * 1024 + +# How long to be patient with a process being shut down: first for it to +# honour SIGTERM, then for its exit to be observed after SIGKILL. +TERMINATE_GRACE = 2.0 + + +class ExecTimeoutError(TimeoutError): + """The command ran past its deadline and was killed.""" + + +@dataclass(frozen=True, kw_only=True) +class ExecResult: + returncode: int + stdout: str + stderr: str + stdout_truncated: bool = False + stderr_truncated: bool = False + + +@runtime_checkable +class ExecBackend(Protocol): + """What the driver needs from whatever runs the worker. + + Kept to one call so that hosting the worker somewhere else — a container, + say — is a new implementation of this, not a change to the driver. + """ + + async def exec( + self, + cmd: Sequence[str], + *, + input: str | None = None, + cwd: str | None = None, + env: Mapping[str, str] | None = None, + timeout: float | None = None, + ) -> ExecResult: + """Run ``cmd``, feeding ``input`` on stdin, and collect its output. + + ``env`` replaces the parent's environment outright rather than + extending it, and omitting it gives the child an empty one: the + parent holds credentials the child has no business seeing, so the + default fails closed. + + ``input`` is encoded as UTF-8, and output is decoded as UTF-8 with + invalid bytes replaced. Both are the contract every backend + implements, not a local choice. + + Raises ``ExecTimeoutError`` if ``timeout`` passes before the command + finishes, having first made sure the process is gone. A command that + cannot be started at all raises the underlying ``OSError`` (usually + ``FileNotFoundError``) instead. + """ + ... + + +async def _read_tail( + stream: asyncio.StreamReader | None, limit: int +) -> tuple[bytes, bool]: + """Drain ``stream``, keeping only its last ``limit`` bytes. + + Draining is the point: a process whose output nobody reads blocks forever + on a full pipe. Dropping the head rather than the tail keeps the part of + the output most likely to hold the result. + """ + if stream is None: + return b"", False + kept = bytearray() + truncated = False + while True: + chunk = await stream.read(64 * 1024) + if not chunk: + return bytes(kept), truncated + kept += chunk + if len(kept) > limit: + del kept[: len(kept) - limit] + truncated = True + + +class LocalBackend: + """Runs the worker as a child of this process, with no isolation.""" + + def __init__( + self, + *, + output_limit: int = DEFAULT_OUTPUT_LIMIT, + terminate_grace: float = TERMINATE_GRACE, + ) -> None: + """Configure output retention and shutdown patience. + + ``output_limit`` caps how many bytes are kept from each of stdout + and stderr; past the cap the oldest bytes are dropped, keeping the + tail. ``terminate_grace`` is how long to wait for SIGTERM to be + honoured before escalating to SIGKILL, and again for the exit to be + observed afterwards. + """ + self._output_limit = output_limit + self._terminate_grace = terminate_grace + # Shutdowns outlive the call that started them, so they need an owner + # that keeps them from being garbage-collected mid-escalation. + self._shutdowns: set[asyncio.Task[None]] = set() + + async def _collect( + self, process: asyncio.subprocess.Process + ) -> tuple[tuple[bytes, bool], tuple[bytes, bool]]: + """Drain both streams, then wait for the process to actually exit. + + Reaching end-of-output is not the same as being finished: code can + close its streams and keep running. Both halves sit inside the + caller's deadline so that neither can outlast it. + """ + streams = await asyncio.gather( + _read_tail(process.stdout, self._output_limit), + _read_tail(process.stderr, self._output_limit), + ) + await process.wait() + return streams[0], streams[1] + + async def exec( + self, + cmd: Sequence[str], + *, + input: str | None = None, + cwd: str | None = None, + env: Mapping[str, str] | None = None, + timeout: float | None = None, + ) -> ExecResult: + process = await asyncio.create_subprocess_exec( + *cmd, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=cwd, + env={} if env is None else dict(env), + ) + try: + if process.stdin is not None: + if input is not None: + process.stdin.write(input.encode()) + process.stdin.close() + stdout, stderr = await asyncio.wait_for(self._collect(process), timeout) + except TimeoutError: + await self._shutdown(process) + raise ExecTimeoutError( + f"the command exceeded its {timeout}-second time limit" + ) from None + except BaseException: + # Whoever started the process ends it. Timeout and cancellation + # are not the only ways out of a call: input can fail to encode, + # a pipe can fail mid-read. Any exit that left the worker running + # would keep holding the parent's file descriptors and go on + # burning CPU with nobody waiting on it. + await self._shutdown(process) + raise + # _collect awaited wait(), so the return code is known here; if that + # invariant ever breaks, fail loudly rather than report "succeeded". + assert process.returncode is not None + return ExecResult( + returncode=process.returncode, + stdout=stdout[0].decode(errors="replace"), + stderr=stderr[0].decode(errors="replace"), + stdout_truncated=stdout[1], + stderr_truncated=stderr[1], + ) + + async def _shutdown(self, process: asyncio.subprocess.Process) -> None: + """Terminate ``process``, outliving cancellation of the caller. + + Shutdown runs in its own task so that a cancellation landing while it + is in flight — a caller cancelling twice, or cancelling while the + post-timeout escalation is still waiting — stops us waiting on it + without stopping the escalation itself. It must not be possible to + leave a SIGTERM-ignoring child alive by cancelling at the wrong + moment. + """ + shutdown = asyncio.ensure_future(_terminate(process, self._terminate_grace)) + self._shutdowns.add(shutdown) + shutdown.add_done_callback(self._shutdowns.discard) + with contextlib.suppress(asyncio.CancelledError): + await asyncio.shield(shutdown) + + +async def _terminate(process: asyncio.subprocess.Process, grace: float) -> None: + """Ask the process to exit, then insist.""" + if process.returncode is not None: + return + process.terminate() + try: + await asyncio.wait_for(asyncio.shield(process.wait()), grace) + return + except TimeoutError: + pass + process.kill() + # The exit can go unobserved if the child watcher misses it, and a killed + # process is gone whether or not we see it go. Wait, but not forever. + try: + await asyncio.wait_for(asyncio.shield(process.wait()), grace) + except TimeoutError: + pass diff --git a/pkg-py/tests/test_execution_backend.py b/pkg-py/tests/test_execution_backend.py new file mode 100644 index 00000000..bef9a1c2 --- /dev/null +++ b/pkg-py/tests/test_execution_backend.py @@ -0,0 +1,417 @@ +"""The exec-shaped seam the execution subsystem sits on. + +Tests drive real subprocesses rather than mocks: the behaviours that matter +here (stdin delivery, output caps, kill escalation) are properties of process +handling, and a mock would only restate the implementation. +""" + +from __future__ import annotations + +import asyncio +import os +import sys +from typing import Any, cast + +import pytest + +from commons._execution import _backend as backend_module +from commons._execution._backend import ( + ExecBackend, + ExecTimeoutError, + LocalBackend, + _terminate, +) + + +async def test_runs_a_command_and_returns_its_output() -> None: + backend = LocalBackend() + + result = await backend.exec([sys.executable, "-c", "print('hello')"]) + + assert result.returncode == 0 + assert result.stdout == "hello\n" + assert result.stderr == "" + + +async def test_input_reaches_the_process_on_stdin() -> None: + # Code goes in on stdin rather than as an argument: no escaping to get + # wrong and no command-line length limit. + backend = LocalBackend() + + result = await backend.exec( + [sys.executable, "-c", "import sys; sys.stdout.write(sys.stdin.read())"], + input="model-written code\n", + ) + + assert result.stdout == "model-written code\n" + + +async def test_the_process_starts_in_the_given_working_directory(tmp_path) -> None: + backend = LocalBackend() + + result = await backend.exec( + [sys.executable, "-c", "import os; print(os.getcwd())"], + cwd=str(tmp_path), + ) + + assert result.stdout.strip() == os.path.realpath(tmp_path) + + +async def test_the_given_environment_replaces_the_parents_rather_than_extending_it( + monkeypatch, +) -> None: + # The parent holds credentials a child has no business seeing, and a + # subprocess inherits the whole environment by default. Passing `env` has + # to mean "exactly this", not "this as well". + monkeypatch.setenv("COMMONS_TEST_SECRET", "sk-not-a-real-key") + backend = LocalBackend() + + result = await backend.exec( + [ + sys.executable, + "-c", + "import os; print(os.environ.get('COMMONS_TEST_SECRET'))", + ], + env={"PATH": os.environ["PATH"]}, + ) + + assert result.stdout.strip() == "None" + + +NOISY = "for i in range(200): print(f'line-{i}:' + 'x' * 1000)" + + +async def test_output_past_the_cap_keeps_the_tail_and_the_process_still_finishes() -> ( + None +): + # Killing the process on the cap would lose a result the code had already + # computed, and simply not reading would deadlock it against a full pipe. + # Keep draining, keep the most recent bytes, let it exit. + backend = LocalBackend(output_limit=2000) + + result = await backend.exec([sys.executable, "-c", NOISY]) + + assert result.returncode == 0 + assert len(result.stdout) <= 2000 + assert result.stdout.rstrip().endswith("x" * 100) + assert "line-199:" in result.stdout + assert "line-0:" not in result.stdout + assert result.stdout_truncated + + +def _sleeper(sentinel: object, *, ignore_sigterm: bool = False) -> str: + """Code that outlives its timeout and records the fact if it is allowed to. + + The sleep sits well past the point where a working shutdown has killed + the process, so a slow or loaded machine delays the kill into slack + rather than into a false failure. + """ + guard = ( + "import signal; signal.signal(signal.SIGTERM, signal.SIG_IGN)\n" + if ignore_sigterm + else "" + ) + return f"{guard}import time; time.sleep(1.5); open({str(sentinel)!r}, 'w').close()" + + +async def test_a_call_past_the_timeout_raises_and_the_process_does_not_survive( + tmp_path, +) -> None: + sentinel = tmp_path / "survived" + backend = LocalBackend() + + with pytest.raises(ExecTimeoutError): + await backend.exec([sys.executable, "-c", _sleeper(sentinel)], timeout=0.15) + + await asyncio.sleep(1.8) + assert not sentinel.exists() + + +async def test_a_process_that_handles_sigterm_gets_to_clean_up_first(tmp_path) -> None: + # SIGKILL first would strand whatever the worker was in the middle of. + # Ask politely, then insist. + marker = tmp_path / "cleaned-up" + code = ( + "import signal, sys, time\n" + f"signal.signal(signal.SIGTERM, lambda *a: (open({str(marker)!r}, 'w').close(), sys.exit(0)))\n" + "time.sleep(5)\n" + ) + backend = LocalBackend() + + with pytest.raises(ExecTimeoutError): + await backend.exec([sys.executable, "-c", code], timeout=0.15) + + assert marker.exists() + + +async def test_a_process_that_ignores_sigterm_is_killed_anyway(tmp_path) -> None: + sentinel = tmp_path / "survived" + backend = LocalBackend(terminate_grace=0.1) + + with pytest.raises(ExecTimeoutError): + await backend.exec( + [sys.executable, "-c", _sleeper(sentinel, ignore_sigterm=True)], + timeout=0.15, + ) + + await asyncio.sleep(1.8) + assert not sentinel.exists() + + +class _NeverReaped: + """A process that takes its signals but whose exit is never observed. + + Stands in for the race where the child watcher misses the exit. There is + no way to provoke that on demand, so this is the one place the suite + substitutes a stand-in for a real process. + """ + + returncode: int | None = None + + def __init__(self) -> None: + self.signals: list[str] = [] + + def terminate(self) -> None: + self.signals.append("term") + + def kill(self) -> None: + self.signals.append("kill") + + async def wait(self) -> int: + await asyncio.sleep(3600) + return 0 + + +async def test_terminate_gives_up_when_the_exit_is_never_reaped() -> None: + process = _NeverReaped() + + # The bound is two grace periods (SIGTERM, then SIGKILL), so 0.5s leaves + # generous margin over the real 0.1s while still failing if the waits + # stop honouring the grace they were given. + await asyncio.wait_for(_terminate(cast(Any, process), 0.05), timeout=0.5) + + assert process.signals == ["term", "kill"] + + +async def test_input_to_a_process_that_never_reads_it_is_not_an_error() -> None: + # A worker that dies during startup leaves nobody on the other end of the + # pipe. That is a failed call to report, not an exception from the plumbing. + backend = LocalBackend() + + result = await backend.exec( + [sys.executable, "-c", "raise SystemExit(3)"], + input="x" * (4 * 1024 * 1024), + ) + + assert result.returncode == 3 + + +async def test_input_larger_than_the_pipe_buffer_arrives_in_full() -> None: + # Handles cross this boundary, so delivery cannot quietly stop at whatever + # the operating system's pipe buffer happens to be. + backend = LocalBackend() + payload = "y" * (4 * 1024 * 1024) + + result = await backend.exec( + [sys.executable, "-c", "import sys; print(len(sys.stdin.read()))"], + input=payload, + ) + + assert result.stdout.strip() == str(len(payload)) + + +def test_the_local_backend_satisfies_the_backend_interface() -> None: + # The annotation is the real assertion: pyrefly rejects an implementation + # whose signature has drifted from the interface a container-hosted + # backend would also have to meet. + backend: ExecBackend = LocalBackend() + + assert isinstance(backend, ExecBackend) + + +async def test_the_timeout_still_applies_after_the_output_streams_close( + tmp_path, +) -> None: + # Reaching end-of-output is not the same as being finished. Code that + # closes its streams and keeps running must still hit the deadline. + sentinel = tmp_path / "survived" + code = ( + "import os, time\n" + "os.close(1); os.close(2)\n" + f"time.sleep(1.5); open({str(sentinel)!r}, 'w').close()\n" + ) + backend = LocalBackend() + + with pytest.raises(ExecTimeoutError): + await backend.exec([sys.executable, "-c", code], timeout=0.15) + + await asyncio.sleep(1.8) + assert not sentinel.exists() + + +async def test_cancelling_a_call_does_not_leave_the_process_running(tmp_path) -> None: + # The driver cancels calls when a conversation goes away or the agent + # shuts down. Whoever started the process has to be the one to end it. + sentinel = tmp_path / "survived" + backend = LocalBackend(terminate_grace=0.1) + call = asyncio.create_task(backend.exec([sys.executable, "-c", _sleeper(sentinel)])) + await asyncio.sleep(0.1) + + call.cancel() + with pytest.raises(asyncio.CancelledError): + await call + + await asyncio.sleep(1.8) + assert not sentinel.exists() + + +async def test_a_command_that_cannot_be_started_raises_os_error() -> None: + # Nothing was spawned, so there is nothing to clean up: the failure + # propagates as-is rather than being dressed up as an exec result. + backend = LocalBackend() + + with pytest.raises(FileNotFoundError): + await backend.exec(["/no/such/binary"]) + + +async def test_cancellation_still_escalates_for_a_process_ignoring_sigterm( + tmp_path, +) -> None: + # The cancellation path does its waiting inside an except block, where an + # await can be cut short. SIGKILL still has to land. + sentinel = tmp_path / "survived" + backend = LocalBackend(terminate_grace=0.1) + call = asyncio.create_task( + backend.exec([sys.executable, "-c", _sleeper(sentinel, ignore_sigterm=True)]) + ) + await asyncio.sleep(0.1) + + call.cancel() + with pytest.raises(asyncio.CancelledError): + await call + + await asyncio.sleep(1.8) + assert not sentinel.exists() + + +async def test_a_second_cancellation_cannot_abort_the_shutdown(tmp_path) -> None: + # Shutdown is not the caller's to interrupt. A cancel landing while the + # grace period is being awaited would otherwise skip SIGKILL and leave a + # SIGTERM-ignoring child running. + sentinel = tmp_path / "survived" + backend = LocalBackend(terminate_grace=0.3) + call = asyncio.create_task( + backend.exec([sys.executable, "-c", _sleeper(sentinel, ignore_sigterm=True)]) + ) + await asyncio.sleep(0.1) + + call.cancel() + await asyncio.sleep(0.05) + call.cancel() + with pytest.raises(asyncio.CancelledError): + await call + + await asyncio.sleep(1.8) + assert not sentinel.exists() + + +async def test_omitting_the_environment_gives_the_child_an_empty_one( + monkeypatch, +) -> None: + # The parent holds credentials a child has no business seeing, so with no + # allowlist in hand the default has to fail closed: no environment at + # all, rather than the whole of the parent's. + monkeypatch.setenv("COMMONS_TEST_SECRET", "sk-not-a-real-key") + backend = LocalBackend() + + result = await backend.exec( + [sys.executable, "-c", "import os; print(sorted(os.environ))"] + ) + + assert "COMMONS_TEST_SECRET" not in result.stdout + + +async def test_stderr_past_the_cap_keeps_the_tail_and_sets_its_own_flag() -> None: + # The two streams are capped and flagged independently; a flag wired to + # the wrong stream is invisible if only stdout is ever exercised. + backend = LocalBackend(output_limit=2000) + code = ( + "import sys\n" + "for i in range(200):\n" + " sys.stderr.write(f'line-{i}:' + 'x' * 1000 + '\\n')\n" + ) + + result = await backend.exec([sys.executable, "-c", code]) + + assert result.returncode == 0 + assert result.stdout == "" + assert not result.stdout_truncated + assert result.stderr_truncated + assert len(result.stderr) <= 2000 + assert "line-199:" in result.stderr + assert "line-0:" not in result.stderr + + +async def test_an_unexpected_error_mid_call_still_kills_the_process( + tmp_path, monkeypatch +) -> None: + # Timeout and cancellation are not the only ways out of a call. A pipe + # failing mid-read (or any other surprise) must not leave the worker + # running with nobody waiting on it. + async def fail_read( + stream: asyncio.StreamReader | None, limit: int + ) -> tuple[bytes, bool]: + raise RuntimeError("pipe failed mid-read") + + monkeypatch.setattr(backend_module, "_read_tail", fail_read) + sentinel = tmp_path / "survived" + backend = LocalBackend(terminate_grace=0.1) + + with pytest.raises(RuntimeError, match="pipe failed"): + await backend.exec([sys.executable, "-c", _sleeper(sentinel)]) + + await asyncio.sleep(1.8) + assert not sentinel.exists() + + +async def test_input_that_cannot_be_encoded_still_kills_the_process( + tmp_path, +) -> None: + # Model-written text can contain unpaired surrogates, which fail at + # encode time — after the child has already been spawned. + sentinel = tmp_path / "survived" + backend = LocalBackend(terminate_grace=0.1) + + with pytest.raises(UnicodeEncodeError): + await backend.exec( + [sys.executable, "-c", _sleeper(sentinel)], input="\ud800" + ) + + await asyncio.sleep(1.8) + assert not sentinel.exists() + + +async def test_cancellation_during_the_timeout_shutdown_cannot_abort_it( + tmp_path, +) -> None: + # A cancel landing while the post-timeout shutdown is in flight must not + # skip SIGKILL. The shutdown runs in its own task precisely so that it + # outlives the call that started it. + sentinel = tmp_path / "survived" + backend = LocalBackend(terminate_grace=0.3) + call = asyncio.create_task( + backend.exec( + [sys.executable, "-c", _sleeper(sentinel, ignore_sigterm=True)], + timeout=0.1, + ) + ) + # The timeout fires at ~0.1s and the SIGTERM grace then runs for 0.3s, + # so a cancel at 0.3s lands in the middle of the shutdown. + await asyncio.sleep(0.3) + + call.cancel() + with pytest.raises((asyncio.CancelledError, ExecTimeoutError)): + await call + + await asyncio.sleep(1.8) + assert not sentinel.exists()