diff --git a/.github/workflows/py-check.yaml b/.github/workflows/py-check.yaml index 89420e2e..8408db5d 100644 --- a/.github/workflows/py-check.yaml +++ b/.github/workflows/py-check.yaml @@ -13,6 +13,8 @@ permissions: read-all jobs: check: runs-on: ${{ matrix.os }} + # A deadlock in the worker tests fails the leg instead of hanging it. + timeout-minutes: 20 strategy: fail-fast: false matrix: @@ -75,6 +77,7 @@ jobs: # kernel runs natively while tagging the syscalls with the 32-bit ABI, # so the live tests check the 32-bit tables against the kernel itself. runs-on: ${{ matrix.os }} + timeout-minutes: 10 strategy: fail-fast: false matrix: diff --git a/pkg-py/pyproject.toml b/pkg-py/pyproject.toml index 3eb8824b..0c8bb203 100644 --- a/pkg-py/pyproject.toml +++ b/pkg-py/pyproject.toml @@ -55,6 +55,8 @@ build-backend = "hatchling.build" [tool.pytest.ini_options] testpaths = ["tests"] asyncio_mode = "auto" +# A test still running after this many seconds prints every thread's stack. +faulthandler_timeout = 300 [tool.ruff] src = ["src", "tests"] diff --git a/pkg-py/src/commons/_execution/_backend.py b/pkg-py/src/commons/_execution/_backend.py index 46198ffc..d373a17a 100644 --- a/pkg-py/src/commons/_execution/_backend.py +++ b/pkg-py/src/commons/_execution/_backend.py @@ -1,8 +1,11 @@ """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. +The driver starts a worker session through ``ExecBackend.start`` and talks to +it over the session's pipes. Everything about where the process runs belongs +to the backend: its scratch directory, its environment, how much protection +it gets, how it is signalled, and how it is shut down. A container-hosted +backend is then another implementation of this protocol, and the driver +stays as it is. """ from __future__ import annotations @@ -10,73 +13,91 @@ import asyncio import contextlib import os +import shutil import signal -from collections.abc import Mapping, Sequence -from dataclasses import dataclass -from typing import Protocol, runtime_checkable +import tempfile +from collections.abc import Coroutine +from pathlib import Path +from typing import Any, Literal, Protocol, runtime_checkable -__all__ = ["ExecBackend", "ExecResult", "ExecTimeoutError", "LocalBackend"] +from ._env import worker_command, worker_env +from ._protocol import STREAM_LIMIT +from ._sandbox import needs_single_thread, protection_mode -# 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 +__all__ = ["ExecBackend", "LocalBackend", "LocalSession", "Network", "WorkerSession"] + +Network = Literal["none", "full"] # 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 +# Kept from the worker's stderr for diagnosing a failed start. The worker +# points fd 2 at a sink once it is up, so only startup output is written. +STDERR_TAIL = 8 * 1024 + +WORKER_SCRIPT = Path(__file__).parent / "_runtime" / "_worker.py" + + +@runtime_checkable +class WorkerSession(Protocol): + """One running worker process, as the driver sees it. + + ``stdin`` and ``stdout`` are the protocol channel. ``close()`` is the + only way the session ends early, and it is safe to call more than once. + """ + + @property + def stdin(self) -> asyncio.StreamWriter: ... + + @property + def stdout(self) -> asyncio.StreamReader: ... + + @property + def returncode(self) -> int | None: + """The exit status, or ``None`` while the worker runs.""" + ... + + def interrupt(self) -> None: + """Send SIGINT to the worker and its children. + + The worker raises ``KeyboardInterrupt`` in a call that is running + model code, and ignores the signal otherwise. + """ + ... -class ExecTimeoutError(TimeoutError): - """The command ran past its deadline and was killed.""" + async def close(self) -> None: + """End the worker and everything it spawned, then remove its files. + The shutdown runs to completion even if the caller is cancelled + while it waits; the caller still receives the ``CancelledError``. + """ + ... -@dataclass(frozen=True, kw_only=True) -class ExecResult: - returncode: int - stdout: str - stderr: str - stdout_truncated: bool = False - stderr_truncated: bool = False + async def stderr(self) -> str: + """The tail of what the worker wrote to stderr before it went quiet. + + Waits for the stream to end, so call it after ``close()``. + """ + ... @runtime_checkable class ExecBackend(Protocol): - """What the driver needs from whatever runs the worker. + """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 start(self, *, network: Network) -> WorkerSession: + """Start a worker with the given network access and return its session. - 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. - - ``input`` is buffered whole in the calling process until the child - reads it, so bounding its size is the caller's responsibility. - - 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. + Raises ``OSError`` if the process cannot be started at all. A + cancelled start leaves no process behind. """ ... + async def aclose(self) -> None: + """Wait for every shutdown still in flight.""" + ... + async def _read_tail( stream: asyncio.StreamReader | None, limit: int @@ -85,7 +106,7 @@ async def _read_tail( 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. + the output most likely to explain a failed start. """ if stream is None: return b"", False @@ -106,126 +127,178 @@ async def _read_tail( class LocalBackend: - """Runs the worker as a child of this process, with no isolation. - - The lifecycle below is POSIX-shaped. On Windows ``terminate()`` and - ``kill()`` are both ``TerminateProcess`` — the grace window does not - exist — and an empty environment can keep the child from spawning at - all (``SystemRoot`` is required). What Windows should do instead — - refuse, or run with weaker guarantees and a warning — is the sandbox - unit's decision, not this class's; ``sandbox_capabilities()`` in - ``pkg-r`` is the existing template for reporting "not sandboxable". + """Runs the worker as a sandboxed child of this process. + + How much protection the worker gets is decided at construction (see + ``protection_mode``), so a host commons cannot protect fails here, ahead + of any model asking to run code. There is deliberately no protection + argument: the only way to accept weaker protection is the environment + opt-in that ``protection_mode()`` consults. + + The lifecycle assumes POSIX signals. Windows is reached only through the + unsafe opt-in, and there ``terminate()`` and ``kill()`` are the same + call, so a shutdown has no grace window. """ def __init__( self, *, - output_limit: int = DEFAULT_OUTPUT_LIMIT, terminate_grace: float = TERMINATE_GRACE, + worker_script: Path = WORKER_SCRIPT, ) -> None: - """Configure output retention and shutdown patience. + """Decide the protection mode and configure 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. + ``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. ``worker_script`` is the entry point the process runs; + leave it at its default outside the tests. """ - self._output_limit = output_limit + self._protection = protection_mode() self._terminate_grace = terminate_grace + self._worker_script = worker_script # 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), - # Session leader, so shutdown signals can take the worker's whole - # process group rather than just the worker (see _signal_tree). - start_new_session=True, + async def start(self, *, network: Network) -> LocalSession: + scratch = tempfile.mkdtemp(prefix="commons-worker-") + spawn = asyncio.ensure_future( + asyncio.create_subprocess_exec( + *worker_command(str(self._worker_script), network, self._protection), + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=scratch, + env=worker_env(scratch, single_thread=needs_single_thread()), + limit=STREAM_LIMIT, + # Session leader, so signals can take the worker's whole + # process group rather than just the worker (see _signal_tree). + start_new_session=True, + ) ) 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 + # Shielded: a process created while this start is being + # cancelled must end up tracked and killed, never leaked. + process = await asyncio.shield(spawn) + except asyncio.CancelledError: + self._track(self._reap_spawn(spawn, scratch)) + raise except BaseException: - # an escape hatch to make sure we properly shutdown the process - # regardless of how it exits - await self._shutdown(process) + shutil.rmtree(scratch, ignore_errors=True) 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], - ) + return LocalSession(self, process, scratch) async def aclose(self) -> None: """Wait for any shutdowns still in flight. - The escalation guarantee — that a cancelled call cannot leave a - SIGTERM-ignoring child alive — holds only while the event loop is + The escalation guarantee, that a cancelled call cannot leave a + SIGTERM-ignoring child alive, applies only while the event loop is running. A driver that is tearing down should call this before the - loop closes, so a last-minute cancellation does not strand a - shutdown mid-escalation. Each shutdown is bounded by two grace - periods, so this returns in bounded time. + loop closes. Each shutdown is bounded by two grace periods, so this + returns in bounded time. """ await asyncio.gather(*list(self._shutdowns), return_exceptions=True) - async def _shutdown(self, process: asyncio.subprocess.Process) -> None: - """Terminate ``process``, outliving cancellation of the caller. + def _track(self, work: Coroutine[Any, Any, None]) -> asyncio.Task[None]: + task = asyncio.ensure_future(work) + self._shutdowns.add(task) + task.add_done_callback(self._shutdowns.discard) + return task + + async def _reap_spawn( + self, spawn: asyncio.Future[asyncio.subprocess.Process], scratch: str + ) -> None: + """Kill a child whose start was cancelled mid-spawn, and remove its files.""" + try: + process = await spawn + except Exception: # noqa: BLE001 - a failed spawn has no child to kill + process = None + if process is not None: + await _terminate(process, self._terminate_grace) + shutil.rmtree(scratch, ignore_errors=True) + + +class LocalSession: + """A worker running as a child process, with its scratch directory.""" + + def __init__( + self, backend: LocalBackend, process: asyncio.subprocess.Process, scratch: str + ) -> None: + assert process.stdin is not None and process.stdout is not None # piped + self.process = process + self.scratch = scratch + self._backend = backend + self._stdin = process.stdin + self._stdout = process.stdout + self._drain = asyncio.ensure_future(_read_tail(process.stderr, STDERR_TAIL)) + self._closing: asyncio.Task[None] | None = None + self._exit = asyncio.ensure_future(process.wait()) + self._exit.add_done_callback(self._kill_group) + + @property + def stdin(self) -> asyncio.StreamWriter: + return self._stdin + + @property + def stdout(self) -> asyncio.StreamReader: + return self._stdout + + @property + def returncode(self) -> int | None: + return self.process.returncode + + def interrupt(self) -> None: + # The whole group, as in the shutdown: a call blocked in + # subprocess.run() must take its children with it, or they outlive + # the timeout the interrupt enforces. An exited leader's pid may + # since belong to another process, so it is not signalled. + if self.process.returncode is not None: + return + with contextlib.suppress(ProcessLookupError): + _signal_tree(self.process, signal.SIGINT) + + async def close(self) -> None: + if self._closing is None: + self._closing = self._backend._track(self._finish()) + await asyncio.shield(self._closing) + + async def stderr(self) -> str: + try: + tail, _ = await asyncio.wait_for( + asyncio.shield(self._drain), self._backend._terminate_grace + ) + except TimeoutError: + return "" + return tail.decode(errors="replace").strip() + + def _kill_group(self, _: asyncio.Future[int]) -> None: + """Kill what is left of the worker's process group once the leader exits. - Shutdown runs in its own task to ensure it is actually performed - (e.g. cancellation arriving while a shutdown is already in process - doesn't stop the shutdown itself). It must not be possible to - leave a SIGTERM-ignoring child alive by cancelling at the wrong - moment. + A background child outlives the crash that took the worker, and + one ignoring SIGTERM outlives a shutdown the leader honoured. The + kill is sent as the exit is observed: a group with members keeps + its id reserved, and an empty one's id is too fresh to be reused. """ - 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) + if os.name == "posix": + with contextlib.suppress(ProcessLookupError, PermissionError): + os.killpg(self.process.pid, signal.SIGKILL) + + async def _finish(self) -> None: + """The bounded kill, then the cleanup.""" + process = self.process + await _terminate(process, self._backend._terminate_grace) + if process.returncode is not None: + # Lets the group kill run before the drain waits on its pipe. + await asyncio.shield(self._exit) + # The drain ends at EOF, which the kill forces; the bound covers a + # descendant that kept the pipe open. + with contextlib.suppress(TimeoutError): + await asyncio.wait_for( + asyncio.shield(self._drain), self._backend._terminate_grace + ) + # The process is gone, so its files can be removed; model code may + # have left some unreadable, which is not a reason to fail. + shutil.rmtree(self.scratch, ignore_errors=True) def _signal_tree(process: asyncio.subprocess.Process, sig: signal.Signals) -> None: @@ -241,10 +314,14 @@ def _signal_tree(process: asyncio.subprocess.Process, sig: signal.Signals) -> No try: os.killpg(process.pid, sig) return - except ProcessLookupError: - pass # The group is gone or never had a session; signal directly. + except (ProcessLookupError, PermissionError): + # The group is gone, never had a session, or (on macOS) has + # only a zombie left in it; signal the child directly. + pass if sig == signal.SIGTERM: process.terminate() + elif sig == signal.SIGINT: + process.send_signal(sig) else: process.kill() diff --git a/pkg-py/src/commons/_execution/_driver.py b/pkg-py/src/commons/_execution/_driver.py new file mode 100644 index 00000000..67e3d253 --- /dev/null +++ b/pkg-py/src/commons/_execution/_driver.py @@ -0,0 +1,460 @@ +"""The worker's lifecycle: lazy spawn, one call at a time, timeouts, and the idle reap. + +A ``Worker`` owns one persistent Python session, started through an +``ExecBackend``. No process exists until the first call; one that sits idle +is closed, and the next call starts a fresh one. Calls are serialized +through an ``asyncio.Lock``, because the single session can only run one +piece of code at a time. A call that outruns its time limit is interrupted +first and restarted only if the interrupt goes unanswered. An interrupted +session keeps its variables; a restarted one will lose them. + +The R package runs the same lifecycle in pkg-r/R/run-r.R +(``worker_ensure()``, ``run_r_tool()``, ``worker_await()``, and +``schedule_worker_reap()``), on a callr session driven through promises. +""" + +from __future__ import annotations + +import asyncio +import itertools +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Self + +from .._handles import HandleStore +from . import _protocol +from ._backend import ExecBackend, LocalBackend, Network, WorkerSession +from ._protocol import Call, Error, Ready, Result + +__all__ = ["Failure", "Worker"] + +CALL_TIMEOUT = 60.0 +IDLE_TIMEOUT = 600.0 + +# How long an interrupted worker gets to answer before it is restarted. +INTERRUPT_GRACE = 5.0 + +# Startup is interpreter boot plus the sandbox engage; a start that +# never finishes should fail the call rather than hang. +SPAWN_TIMEOUT = 30.0 + + +@dataclass(frozen=True, kw_only=True) +class Failure: + """A call that never got an answer from the worker. + + ``message`` is written for the model and says whether the session + survived: an interrupted session kept its variables, a restarted one + lost them, and the difference decides what the model should retry. + """ + + message: str + + +class Worker: + """One persistent Python session in a sandboxed worker. + + Construction defers startup to the first ``run()`` call. + ``measure_sources`` are the trusted measures as source text. + The worker is a fresh interpreter and shares no objects with this process, + so each start execs ``measure_sources`` to install the definitions + into the worker. ``backend`` decides where the worker runs; the default + ``LocalBackend`` decides its protection at construction, so a host + commons cannot protect fails here, ahead of any model asking to run code. + """ + + def __init__( + self, + *, + network: Network = "none", + measure_sources: Sequence[str] = (), + backend: ExecBackend | None = None, + call_timeout: float = CALL_TIMEOUT, + idle_timeout: float = IDLE_TIMEOUT, + interrupt_grace: float = INTERRUPT_GRACE, + spawn_timeout: float = SPAWN_TIMEOUT, + ) -> None: + if network not in ("none", "full"): + raise ValueError(f"unknown network access level {network!r}") + self._backend = backend if backend is not None else LocalBackend() + self._network: Network = network + self._closed = False + self._measure_sources = list(measure_sources) + self._call_timeout = call_timeout + self._idle_timeout = idle_timeout + self._interrupt_grace = interrupt_grace + self._spawn_timeout = spawn_timeout + + self._session: WorkerSession | None = None + # How many of the store's handles the running worker has, and which + # store they came from; a different store is synced from its start. + self._synced = 0 + self._synced_store: HandleStore | None = None + self._pending = 0 + self._reap_task: asyncio.Task[None] | None = None + # A close outlives a cancelled aclose(), so it needs an owner that + # keeps it from being garbage-collected. + self._closing: asyncio.Task[None] | None = None + self._last_used = 0.0 + self._lock = asyncio.Lock() + self._ids = itertools.count(1) + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, *exc: object) -> None: + await self.aclose() + + async def run( + self, code: str, handles: HandleStore | None = None + ) -> Result | Error | Failure: + """Run ``code`` in the session and report what came back. + + A ``Result`` is the code's answer, an ``Error`` is the code's own + exception, and a ``Failure`` is the machinery's: the worker timed + out, crashed, or would not start. ``handles`` is the conversation's + result store; entries the worker has not seen yet are sent ahead + of the call, so a restarted worker is re-synced from the store's + beginning. + + The pending count is incremented before the call queues on the + lock and decremented exactly once when it settles, however it + settles, so the reaper can never close the worker beneath a call + in flight. + """ + if self._closed: + return Failure(message="the Python session is closed.") + self._pending += 1 + try: + async with self._lock: + # A call queued behind aclose() must not start a worker that + # the close would then have to kill. + if self._closed: + return Failure(message="the Python session is closed.") + return await self._call(code, handles) + finally: + self._pending -= 1 + self._schedule_reap() + + async def aclose(self) -> None: + """Close the worker, cancelling the idle reap. + + The close waits to acquire the lock until any running call, and + any start under way, releases it; after the close acquires the + lock, no worker starts. The close runs in its own task, so + cancelling ``aclose()`` while it waits out a running call still + closes the worker, and the caller still receives the + ``CancelledError``. The close waits for the backend's shutdowns, + the reaper's included, last. + """ + self._closed = True + if self._reap_task is not None: + self._reap_task.cancel() + self._reap_task = None + if self._closing is None: + self._closing = asyncio.ensure_future(self._close_locked()) + await asyncio.shield(self._closing) + + async def _close_locked(self) -> None: + """The close proper, serialized against any call still in flight.""" + async with self._lock: + await self._shutdown() + await self._backend.aclose() + + async def _call( + self, code: str, handles: HandleStore | None + ) -> Result | Error | Failure: + """One call, run under the lock with the worker up.""" + try: + session = await self._ensure() + if handles is not None: + failure = await self._sync(session, handles) + if failure is not None: + return failure + call_id = f"c{next(self._ids)}" + await self._send(session, Call(id=call_id, code=code)) + return await self._await_reply(session, call_id) + except _protocol.ProtocolError as exc: + # The call did not fit the channel, and was refused before any + # of it was written, so the worker is still in step. + return Failure( + message=f"the code could not be sent: {exc}. The session and " + "its variables remain available." + ) + except TimeoutError: + # The write blocked: the worker stopped reading, and a caller + # waiting on a pipe has no timeout of its own. This one is the + # call's. + await self._shutdown() + return Failure( + message="the Python session stopped reading its calls, so it " + "was restarted. Session variables were reset." + ) + except ConnectionError: + # The worker exited between calls without its exit being seen + # yet, so the write found no reader. + await self._shutdown() + return Failure( + message="the Python session exited and was restarted. " + "Session variables were reset." + ) + except asyncio.CancelledError: + # The caller gave up mid-call, but the worker has not: it is + # still running the code, or about to answer into a channel the + # next call would misread as its own reply. Shut it down before + # releasing the lock. + await self._shutdown() + raise + except Exception as exc: # noqa: BLE001 - any start or write failure fails the call + # A start or write failure leaves the worker's state unknown; + # start over rather than trust it. + await self._shutdown() + return Failure(message=str(exc)) + + async def _sync( + self, session: WorkerSession, handles: HandleStore + ) -> Failure | None: + """Send the store's handles the worker does not have, one message each. + + A restarted worker is sent the whole store again, and in a single + message a large store would outgrow the channel and be shrunk to + reprs. Each handle counts as synced once the worker acknowledges + it; one the worker could not decode is skipped rather than retried + on every call. + """ + if handles is not self._synced_store: + self._synced_store, self._synced = handles, 0 + for id_ in handles.ids()[self._synced :]: + sync_id = f"s{next(self._ids)}" + try: + await self._send( + session, + Call(id=sync_id, code="None", handles={id_: handles.get(id_)}), + ) + except _protocol.ProtocolError: + # Refused before any of it was written; a handle that cannot + # fit the channel never will, so it is skipped. + self._synced += 1 + continue + try: + reply = await asyncio.wait_for( + _protocol.read_message(session.stdout), self._call_timeout + ) + except TimeoutError: + await self._shutdown() + return Failure( + message="the Python session stopped answering, so it was " + "restarted. Session variables were reset." + ) + except (_protocol.ProtocolError, _protocol.ChannelError) as exc: + await self._shutdown() + return _unparseable(exc) + outcome = await self._classify(sync_id, reply) + if isinstance(outcome, Failure): + return outcome + self._synced += 1 + return None + + async def _send(self, session: WorkerSession, message: Call) -> None: + """Write ``message``, giving up on a worker that stops reading. + + A worker that never reads its stdin would otherwise block the + writer in ``drain()`` forever, and no call timeout would ever + start. + """ + await asyncio.wait_for( + _protocol.write_message(session.stdin, message), self._call_timeout + ) + + async def _await_reply( + self, session: WorkerSession, call_id: str + ) -> Result | Error | Failure: + """Wait for the in-flight call's reply, with the timeout and escalation.""" + try: + reply = await asyncio.wait_for( + _protocol.read_message(session.stdout), self._call_timeout + ) + except TimeoutError: + return await self._escalate(session, call_id) + except (_protocol.ProtocolError, _protocol.ChannelError) as exc: + await self._shutdown() + return _unparseable(exc) + return await self._classify(call_id, reply) + + async def _classify( + self, call_id: str, reply: _protocol.Message | None + ) -> Result | Error | Failure: + """Map a reply, or the channel's end, onto the call's outcome.""" + if reply is None: + await self._shutdown() + return Failure( + message="the Python session crashed and was restarted. " + "Session variables were reset." + ) + if not isinstance(reply, (Result, Error)) or reply.id != call_id: + # The one-in-flight contract is what keys a reply to its call; a + # reply outside it means the two sides disagree, and the worker + # cannot be trusted to stay in step. + await self._shutdown() + return Failure( + message="the Python session answered out of turn, so it was " + "restarted. Session variables were reset." + ) + return reply + + async def _escalate(self, session: WorkerSession, call_id: str) -> Failure: + """Interrupt the in-flight call; restart the worker if it stays silent. + + The interrupt raises KeyboardInterrupt inside the worker's call, + which breaks it out of the computation and leaves the session alive. + A worker that answers within the grace window was interrupted; if + the worker does not respond, it is hard-restarted instead. + + The grace-window reply must meet the same contract as any other: + only this call's own ``Result`` or ``Error`` proves the session + survived. + """ + session.interrupt() + try: + reply = await asyncio.wait_for( + _protocol.read_message(session.stdout), self._interrupt_grace + ) + except TimeoutError: + await self._shutdown() + return Failure( + message=f"the code exceeded the {self._call_timeout:g}-second " + "time limit and the Python session did not respond to an " + "interrupt, so it was restarted. Session variables were reset." + ) + except (_protocol.ProtocolError, _protocol.ChannelError) as exc: + await self._shutdown() + return _unparseable(exc) + outcome = await self._classify(call_id, reply) + if isinstance(outcome, Failure): + return outcome + # The reply itself is discarded: whatever the call was doing when + # it was interrupted, the answer the model needs is that its code + # ran out of time. + return Failure( + message=f"the code was interrupted after exceeding the " + f"{self._call_timeout:g}-second time limit. The session and its " + "variables remain available." + ) + + async def _ensure(self) -> WorkerSession: + """Start the worker if none is running, and define its measure sources.""" + if self._session is not None and self._session.returncode is None: + return self._session + # A worker that died between calls is still attached here; close it + # before replacing it, so its files go with it. + await self._shutdown() + try: + session = await self._backend.start(network=self._network) + except OSError as exc: + raise RuntimeError(f"the Python session failed to start: {exc}") from exc + self._session = session + self._synced = 0 + try: + ready = await asyncio.wait_for( + _protocol.read_message(session.stdout), self._spawn_timeout + ) + if not isinstance(ready, Ready): + # A dead or protocol-breaking worker, not a bad argument. + raise RuntimeError( # noqa: TRY004 + "the worker exited before it was ready" + ) + for index, source in enumerate(self._measure_sources): + await self._define(session, index, source) + except TimeoutError: + await self._shutdown() + detail = await session.stderr() + raise RuntimeError( + "the Python session failed to start: the worker stopped " + "responding during startup" + + (f"\n{detail}" if detail else "") + ) from None + except Exception as exc: + await self._shutdown() + # A worker that wrote its failure just before exiting may have + # bytes still in the pipe; the session's stderr waits for them. + detail = await session.stderr() + raise RuntimeError( + f"the Python session failed to start: {exc}" + + (f"\n{detail}" if detail else "") + ) from exc + return session + + async def _define(self, session: WorkerSession, index: int, source: str) -> None: + """Define one measure source in a freshly started worker. + + A source that raises an error is skipped, so one bad harvest cannot fail + every start. Only a worker that stops answering fails the start. + + ``_commons_define_source`` is seeded into every worker's session at + startup: it compiles with annotations deferred and replaces each + default whose evaluation raises an ``Exception`` with a placeholder, + because a harvested source's own module imports and globals do not exist + in the session. + """ + payload = f"_commons_define_source({source!r})" + await asyncio.wait_for( + _protocol.write_message( + session.stdin, Call(id=f"init-{index}", code=payload) + ), + self._spawn_timeout, + ) + reply = await asyncio.wait_for( + _protocol.read_message(session.stdout), self._spawn_timeout + ) + if isinstance(reply, Error): + # The source's own failure: the session lacks only that + # definition, and a call that uses it gets a NameError. + return + if not isinstance(reply, Result): + raise RuntimeError( # noqa: TRY004 - a failed start, not a bad argument + "the worker exited while defining measure sources" + ) + + def _schedule_reap(self) -> None: + """Close the worker after a quiet stretch; the next call restarts it. + + One timer per worker, re-armed as each call settles: the reaper + fires only when nothing has used the worker for the whole idle + window and no call is in flight. A closed worker gets no timer: a + call settling as ``aclose()`` runs must not leave one running. + """ + if self._closed: + return + self._last_used = asyncio.get_running_loop().time() + if self._reap_task is not None: + self._reap_task.cancel() + self._reap_task = asyncio.ensure_future(self._reap_when_idle()) + + async def _reap_when_idle(self) -> None: + try: + await asyncio.sleep(self._idle_timeout + 1) + except asyncio.CancelledError: + return + quiet = asyncio.get_running_loop().time() - self._last_used + if self._pending == 0 and quiet >= self._idle_timeout: + await self._shutdown() + + async def _shutdown(self) -> None: + """Close the worker, if any, and forget its state. + + The state is cleared before the close is awaited, so a call that + starts while the kill is still in flight starts a fresh worker + rather than queue behind a dying one. The backend finishes the + close even if this caller is cancelled while it waits. + """ + session, self._session = self._session, None + self._synced = 0 + if session is not None: + await session.close() + + +def _unparseable(exc: Exception) -> Failure: + return Failure( + message="the Python session answered with a line that does not " + f"parse ({exc}), so it was restarted. Session variables were reset." + ) diff --git a/pkg-py/src/commons/_execution/_runtime/_guardrails.py b/pkg-py/src/commons/_execution/_runtime/_guardrails.py new file mode 100644 index 00000000..5526bb8e --- /dev/null +++ b/pkg-py/src/commons/_execution/_runtime/_guardrails.py @@ -0,0 +1,215 @@ +"""Best-effort guardrails for a worker that runs without a sandbox. + +Guardrails mode exists so a host commons cannot sandbox can still run, and +it provides no security boundary: model code that looks for a way around +these checks will find one. They stop ordinary code from reading outside +the worker's roots, writing outside its scratch directory, starting +processes, or reaching the network under ``network="none"``. + +The checks are an audit hook (``sys.addaudithook``), consulted only while +model code runs, so the worker's own reads and writes are left alone. A +denied operation raises ``PermissionError``, which libraries that probe +the filesystem already tolerate. pkg-r/inst/worker/worker.R applies the +same policy to R by replacing its file, process, and connection functions +for the length of a call. + +Imported by bare name from the worker; nothing here may import ``commons``. +""" + +from __future__ import annotations + +import os +import sys +import threading +from collections.abc import Callable, Iterable +from typing import Any + +__all__ = ["Guardrails", "engage"] + +# Audit events naming a path, mapped to the access each argument needs. +# https://docs.python.org/3/library/audit_events.html lists them all. +# "entry" is a write to the directory that lists the path, for operations +# that act on a symlink itself rather than on its target. A hard link is a +# second name for its source, through which the source can be written. +_PATH_EVENTS: dict[str, tuple[str, ...]] = { + "os.listdir": ("read",), + "os.scandir": ("read",), + "os.chdir": ("read",), + "os.listxattr": ("read",), + "os.getxattr": ("read",), + "os.mkdir": ("entry",), + "os.remove": ("entry",), + "os.rmdir": ("entry",), + "os.rename": ("entry", "entry"), + "os.link": ("write", "entry"), + "os.symlink": ("read", "entry"), + "os.truncate": ("write",), + "os.chmod": ("write",), + "os.chown": ("write",), + "os.chflags": ("write",), + "os.lchflags": ("entry",), + "os.utime": ("write",), + "os.setxattr": ("write",), + "os.removexattr": ("write",), + "shutil.copyfile": ("read", "write"), + "shutil.copymode": ("read", "write"), + "shutil.copystat": ("read", "write"), + "shutil.copytree": ("read", "write"), + "shutil.move": ("entry", "entry"), + "shutil.rmtree": ("entry",), + "shutil.chown": ("write",), + "sqlite3.connect": ("write",), +} + +_PROCESS_EVENTS = frozenset( + { + "subprocess.Popen", + "os.system", + "os.exec", + "os.spawn", + "os.posix_spawn", + "os.fork", + "os.forkpty", + "os.startfile", + "pty.spawn", + "webbrowser.open", + } +) + +_NETWORK_EVENTS = frozenset( + { + "socket.connect", + "socket.bind", + "socket.sendto", + "socket.sendmsg", + "socket.getaddrinfo", + "socket.gethostbyname", + "socket.gethostbyaddr", + "socket.getnameinfo", + "http.client.connect", + "ftplib.connect", + "smtplib.connect", + "imaplib.open", + "poplib.connect", + "nntplib.connect", + "telnetlib.Telnet.open", + } +) + +# Flags that make an open a write, for an ``open`` event that reports them. +_WRITE_FLAGS = os.O_WRONLY | os.O_RDWR | os.O_APPEND | os.O_CREAT | os.O_TRUNC + + +def _canonical(path: str) -> str: + """``path`` made absolute with every symlink in its existing part resolved. + + A path that does not exist yet resolves through its nearest existing + ancestor, so a link to the outside cannot be written through by naming + a file beneath it that has yet to be created. + """ + return os.path.realpath(os.path.join(os.getcwd(), path)) + + +def _entry(path: str) -> str: + """``path`` with its directory canonicalized and its final name kept. + + This names the directory entry itself, so a symlink is not followed. A + trailing separator makes the operation follow the link, so such a path + is canonicalized whole. + """ + head, tail = os.path.split(os.path.join(os.getcwd(), path)) + if tail in ("", ".", ".."): + return _canonical(path) + return os.path.join(_canonical(head), tail) + + +def _within(path: str, roots: Iterable[str]) -> bool: + return any(path == root or path.startswith(root.rstrip("/") + "/") for root in roots) + + +class Guardrails: + """The audit hook's policy, separated from the hook so it can be tested.""" + + def __init__( + self, + read_roots: Iterable[str], + write_roots: Iterable[str], + *, + network: str, + active: Callable[[], bool], + ) -> None: + canonical = lambda roots: [os.path.realpath(r) for r in roots] + self.write_roots = canonical(write_roots) + [os.devnull] + self.read_roots = canonical(read_roots) + self.write_roots + self.network = network + self.active = active + self._checking = threading.local() + + def __call__(self, event: str, args: tuple[Any, ...]) -> None: + if not self.active() or getattr(self._checking, "busy", False): + return + # The checks resolve paths, which may audit events of their own. + self._checking.busy = True + try: + self._check(event, args) + finally: + self._checking.busy = False + + def _check(self, event: str, args: tuple[Any, ...]) -> None: + if event == "open": + path, mode, flags = args + write = any(c in (mode or "") for c in "wax+") or bool( + (flags or 0) & _WRITE_FLAGS + ) + self._path(path, "write" if write else "read") + elif event in _PATH_EVENTS: + for path, access in zip(args, _PATH_EVENTS[event], strict=False): + self._path(path, access) + elif event in _PROCESS_EVENTS: + _deny("subprocess creation") + elif event == "urllib.Request": + url = str(args[0]) + if url.lower().startswith("file:"): + _deny("read access", url) + self._network(url) + elif event in _NETWORK_EVENTS: + self._network() + + def _path(self, path: Any, access: str) -> None: + if path is None or isinstance(path, int): + # A file descriptor was opened, and checked, by an earlier call. + return + name = os.fsdecode(os.fspath(path)) + if access == "entry": + resolved, roots, access = _entry(name), self.write_roots, "write" + else: + resolved = _canonical(name) + roots = self.write_roots if access == "write" else self.read_roots + if not _within(resolved, roots): + _deny(f"{access} access", name) + + def _network(self, target: str | None = None) -> None: + if self.network == "none": + _deny("network access", target) + + +def _deny(kind: str, target: str | None = None) -> None: + detail = "" if target is None else f" to {target!r}" + raise PermissionError(f"commons run_python guardrails denied {kind}{detail}") + + +def engage( + read_roots: Iterable[str], + write_roots: Iterable[str], + *, + network: str, + active: Callable[[], bool], +) -> None: + """Install the guardrails for the rest of this process's life. + + ``active`` says whether model code is running; the checks apply only + then. An audit hook cannot be removed once added. + """ + sys.addaudithook( + Guardrails(read_roots, write_roots, network=network, active=active) + ) diff --git a/pkg-py/src/commons/_execution/_runtime/_landlock.py b/pkg-py/src/commons/_execution/_runtime/_landlock.py index 93dd703c..0da3072d 100644 --- a/pkg-py/src/commons/_execution/_runtime/_landlock.py +++ b/pkg-py/src/commons/_execution/_runtime/_landlock.py @@ -18,6 +18,7 @@ import errno import os import platform +import stat import sys from collections.abc import Iterable @@ -88,6 +89,10 @@ SCOPE_ABSTRACT_UNIX_SOCKET = 1 << 0 SCOPE_MIN_ABI = 6 +# The only rights a rule on a non-directory may name; the kernel refuses the +# rule with EINVAL if it names any other. +FS_FILE_RIGHTS = FS_EXECUTE | FS_WRITE_FILE | FS_READ_FILE | FS_TRUNCATE | FS_IOCTL_DEV + # A kernel that has no Landlock, or that is behind a policy forbidding it, # is a kernel to fall back from rather than fail on. _UNAVAILABLE = frozenset({errno.ENOSYS, errno.EOPNOTSUPP, errno.EPERM}) @@ -228,6 +233,8 @@ def _grant(ruleset_fd: int, path: str, access: int) -> None: except FileNotFoundError: return try: + if not stat.S_ISDIR(os.fstat(parent).st_mode): + access &= FS_FILE_RIGHTS rule = PathBeneathAttr(allowed_access=access, parent_fd=parent) if _syscall( NR_ADD_RULE, diff --git a/pkg-py/src/commons/_execution/_runtime/_worker.py b/pkg-py/src/commons/_execution/_runtime/_worker.py new file mode 100644 index 00000000..124b01d0 --- /dev/null +++ b/pkg-py/src/commons/_execution/_runtime/_worker.py @@ -0,0 +1,371 @@ +"""The worker process's entry point: sandbox itself, then answer calls. + +Launched by the driver as ``python -I -u _worker.py NETWORK PROTECTION`` with +its working directory set to a scratch directory and its environment reduced +to an allowlist. Nothing here may import ``commons``: this process runs +model-written code and is deliberately kept free of the agent's dependencies. +The sibling runtime modules and ``_protocol`` (one directory up) are imported +by path. + +The protocol channel is claimed at the file-descriptor level before any +other startup step: duplicates of fds 0 and 1 become the channel, and the +model-visible fds are pointed at sinks. Model code that writes straight to +fd 1 (``os.write``, a C extension, a thread still printing after its call +returned) goes to the sink rather than mid-message on the channel, and a +thread reading fd 0 gets EOF rather than the next call's bytes. The +sys-level redirection in ``_repl`` cannot close those holes; these are +closed by construction, for the process's whole life. The sibling modules +are imported before the claim, so none of them may print at import. + +fd 2 stays attached until the sandbox is engaged and the ready announcement +is sent, so a startup failure can still reach the driver's diagnostics. Model +code never runs before then. +""" + +from __future__ import annotations + +import json +import os +import signal +import sys + +# The channel takeover runs before the sandbox, which needs the protocol fds +# preserved across its fd sweep. +_PROTOCOL_IN = -1 +_PROTOCOL_OUT = -1 + + +def _claim_protocol_channel() -> tuple[int, int]: + """Dup fds 0 and 1 for the protocol and point the model-visible ones at sinks.""" + protocol_in = os.dup(0) + protocol_out = os.dup(1) + source = os.open(os.devnull, os.O_RDONLY) + os.dup2(source, 0) + if source > 2: + os.close(source) + sink = os.open(os.devnull, os.O_WRONLY) + os.dup2(sink, 1) + if sink > 2: + os.close(sink) + return protocol_in, protocol_out + + +def _silence_stderr() -> None: + """Point fd 2 at fd 1's sink, once startup diagnostics no longer need it. + + A dup rather than a fresh ``/dev/null`` open: on Linux this runs after + the sandbox engaged, and ``/dev`` is outside its roots. + """ + os.dup2(1, 2) + + +# Isolated mode leaves even this script's own directory off the path, so +# both it (for the sibling runtime modules) and its parent (for the +# protocol, which lives next to the driver) are added by hand. +_HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, os.path.dirname(_HERE)) +sys.path.insert(0, _HERE) + +# Imported by bare name from directories put on sys.path at runtime, which +# static analysis cannot follow. +import _limits # pyrefly: ignore[missing-import] +import _protocol # pyrefly: ignore[missing-import] +import _repl # pyrefly: ignore[missing-import] + + +def _sandbox_roots() -> tuple[list[str], list[str]]: + """The read and write roots the sandbox is engaged with. + + The read roots cover the interpreter, its standard library, and every + directory on the import path, plus the operating-system roots a + process cannot start up without. The write root is the scratch + directory alone, which is the working directory the driver launched + this process into. Sandboxes match symlink-free paths, so each root's + resolved spelling is granted alongside the original. + """ + roots = [sys.prefix, sys.base_prefix, sys.exec_prefix] + roots.extend(path for path in sys.path if path) + if sys.platform == "darwin": + roots.extend( + [ + "/usr", + "/bin", + "/sbin", + "/System", + "/Library", + "/private/etc", + "/private/var/db", + "/opt", + "/dev", + ] + ) + else: + roots.extend(["/usr", "/bin", "/sbin", "/lib", "/lib64", "/etc", "/opt"]) + read_roots = [] + for root in roots: + for spelling in (root, os.path.realpath(root)): + if os.path.isdir(spelling) and spelling not in read_roots: + read_roots.append(spelling) + scratch = os.getcwd() + write_roots = list(dict.fromkeys((scratch, os.path.realpath(scratch)))) + return read_roots, write_roots + + +def _engage_sandbox(network: str, protection: str) -> None: + """Restrict this process, permanently, before any model-written code runs. + + Guardrails mode engages no sandbox: it exists so a host commons cannot + protect can still run, and it provides no security boundary. It + installs best-effort checks on the same roots and network access + instead (see ``_guardrails``), consulted only while model code runs. + The address-space cap applies in both modes; it is a plain rlimit that + protects the host from a runaway allocation. The order + below is fixed by the mechanisms themselves: the user-namespace + fallback must precede the seccomp filter that would screen its mount + and unshare calls, and the filter goes last because nothing may widen + access after it. + """ + _limits.apply_address_space_limit() + read_roots, write_roots = _sandbox_roots() + if protection == "guardrails": + import _guardrails # pyrefly: ignore[missing-import] + + _guardrails.engage( + read_roots, write_roots, network=network, active=lambda: _in_call + ) + return + if sys.platform == "darwin": + import _seatbelt # pyrefly: ignore[missing-import] + + _seatbelt.engage_seatbelt(read_roots, write_roots, network=network) + return + import _landlock # pyrefly: ignore[missing-import] + import _seccomp # pyrefly: ignore[missing-import] + import _userns # pyrefly: ignore[missing-import] + + if _landlock.engage(read_roots, write_roots) is None: + _userns.engage( + read_roots, + write_roots, + preserve_fds=[0, 1, 2, _PROTOCOL_IN, _PROTOCOL_OUT], + ) + _seccomp.engage(network=network) + + +def _send(message: _protocol.Message) -> None: + """Write ``message`` to the protocol channel, whole. + + ``os.write`` may write part of the line, and a half-written line poisons + the channel, so the rest is written until the line is out. + """ + view = memoryview(_protocol.encode_message(message)) + while view: + view = view[os.write(_PROTOCOL_OUT, view) :] + + +# Whether model code is running. The driver's SIGINT is meant for the call +# it timed out; one the worker receives while it encodes a reply, reads the +# next line, or waits between calls has nothing to interrupt, and raising +# there would lose the reply or the session. +_in_call = False + + +def _on_interrupt(signum: int, frame: object) -> None: + if _in_call: + raise KeyboardInterrupt + + +# Seeded into the session namespace at startup. Harvested measure sources +# are reference material, which the model reads back with +# inspect.getsource(). They come without their module's imports and +# globals, and a default argument naming one (``def m(region=DEFAULT)``) +# would otherwise fail the define, where R's lazy defaults define fine. +# Each default expression is therefore evaluated behind a guard, and one +# that fails becomes a placeholder naming the expression's source. A call +# that uses the default gets a NameError from the placeholder, never a +# value computed from it. +_DEFINE_SOURCE = ''' +import ast as _ast +import itertools as _itertools +import linecache as _linecache + + +class _CommonsMissing: + """Stand-in for a default that could not be evaluated in this session.""" + + def __init__(self, what): + self._what = what + + def __repr__(self): + return f"<{self._what}: could not be evaluated in this session>" + + def _refuse(self, *args, **kwargs): + raise NameError( + f"the default `{self._what}` could not be evaluated in this " + "session; pass this argument explicitly" + ) + + def __getattr__(self, attr): + if attr.startswith("__"): + raise AttributeError(attr) + self._refuse() + + __call__ = __getitem__ = __iter__ = __len__ = _refuse + __str__ = __format__ = __bool__ = __eq__ = __ne__ = _refuse + __lt__ = __le__ = __gt__ = __ge__ = __contains__ = _refuse + __add__ = __radd__ = __sub__ = __rsub__ = __mul__ = __rmul__ = _refuse + __truediv__ = __rtruediv__ = __mod__ = __rmod__ = _refuse + __hash__ = object.__hash__ + + +def _commons_default(thunk, what): + try: + return thunk() + except Exception: + return _CommonsMissing(what) + + +def _commons_define_source(source, _files=_itertools.count(1)): + import __future__ + + def guard(default): + segment = _ast.get_source_segment(source, default) or "default" + return _ast.Call( + func=_ast.Name(id="_commons_default", ctx=_ast.Load()), + args=[ + _ast.Lambda( + args=_ast.arguments( + posonlyargs=[], + args=[], + kwonlyargs=[], + kw_defaults=[], + defaults=[], + ), + body=default, + ), + _ast.Constant(value=segment), + ], + keywords=[], + ) + + functions = (_ast.FunctionDef, _ast.AsyncFunctionDef, _ast.Lambda) + tree = _ast.parse(source) + # Every function whose defaults are evaluated as the source is defined: + # a def or lambda reached without entering a function body, which runs + # only when called. A harvested helper can be a lambda assignment. + defined = [] + stack = list(tree.body) + while stack: + node = stack.pop() + if isinstance(node, functions): + defined.append(node) + stack.extend(node.args.defaults) + stack.extend(d for d in node.args.kw_defaults if d is not None) + else: + stack.extend(_ast.iter_child_nodes(node)) + for node in defined: + # Each default keeps its slot; only its evaluation is guarded. + node.args.defaults = [guard(d) for d in node.args.defaults] + node.args.kw_defaults = [ + guard(d) if d is not None else None for d in node.args.kw_defaults + ] + _ast.fix_missing_locations(tree) + # Lets inspect.getsource() read it back; a None mtime is never evicted. + filename = f"" + _linecache.cache[filename] = ( + len(source), None, source.splitlines(keepends=True), filename + ) + code = compile( + tree, + filename, + "exec", + flags=__future__.annotations.compiler_flag, + dont_inherit=True, + ) + exec(code, globals()) +''' + + +def _execute(call: _protocol.Call, namespace: dict) -> _protocol.Message: + """Run one call in the session namespace and render its reply.""" + global _in_call + namespace.update(call.handles) + _in_call = True + try: + evaluation = _repl.run(call.code, namespace) + except KeyboardInterrupt: + # The driver's SIGINT broke the call out of its computation. The + # session and its variables survive; the driver reports the + # interrupt in its own words. + evaluation = None + finally: + _in_call = False + if evaluation is None: + return _protocol.Error(id=call.id, message="KeyboardInterrupt") + if evaluation.error: + return _protocol.Error( + id=call.id, message=evaluation.error, traceback=evaluation.traceback + ) + return _protocol.Result( + id=call.id, + value=evaluation.value, + stdout=evaluation.stdout, + stderr=evaluation.stderr, + ) + + +def _call_id(line: bytes) -> str | None: + """The ``id`` of a call line that failed to decode, if it has one.""" + try: + body = json.loads(line) + except (ValueError, RecursionError): + return None + if isinstance(body, dict) and body.get("type") == "call": + call_id = body.get("id") + if isinstance(call_id, str): + return call_id + return None + + +def main() -> None: + network, protection = sys.argv[1], sys.argv[2] + _engage_sandbox(network, protection) + # Installed explicitly, so a worker whose parent ignores SIGINT, as a + # shell does for a background job, can still be interrupted. + signal.signal(signal.SIGINT, _on_interrupt) + # The namespace is the session: the same mapping on every call, so a + # name one call binds is visible to the next. It is seeded before the + # ready announcement, while a failure can still reach fd 2. + namespace: dict = {"__name__": "__main__"} + exec(compile(_DEFINE_SOURCE, "", "exec"), namespace) # noqa: S102 + _send(_protocol.Ready()) + _silence_stderr() + + calls = os.fdopen(_PROTOCOL_IN, "rb") + while True: + line = calls.readline() + if not line: + return # The driver is gone, and with it the reason to be here. + try: + message = _protocol.decode_message(line) + except _protocol.ChannelError: + return # No later line can be trusted; the driver starts over. + except _protocol.ProtocolError as error: + # One bad line costs the line, not the session. The driver is + # waiting on the call's reply, so a call it can still key gets + # its error now rather than a silent wait for the timeout. + call_id = _call_id(line) + if call_id is not None: + message = f"the call could not be decoded: {error}" + _send(_protocol.Error(id=call_id, message=message)) + continue + if not isinstance(message, _protocol.Call): + # The driver only ever sends calls; anything else means the two + # sides disagree about the protocol, and the channel is done. + return + _send(_execute(message, namespace)) + + +if __name__ == "__main__": + _PROTOCOL_IN, _PROTOCOL_OUT = _claim_protocol_channel() + main() diff --git a/pkg-py/src/commons/_measures.py b/pkg-py/src/commons/_measures.py index 8c485fbf..d5310e3e 100644 --- a/pkg-py/src/commons/_measures.py +++ b/pkg-py/src/commons/_measures.py @@ -11,6 +11,7 @@ from __future__ import annotations +import ast import hashlib import importlib.util import inspect @@ -419,7 +420,8 @@ class SemanticLayer: """The trusted calculations an agent can run. ``source_text`` holds the source of the measures and the module-level - helpers they call, keyed by Python function name. Only text is kept: + helpers they call, keyed by Python function name, dedented and without + decorator lines so the worker can exec it as-is. Only text is kept: the agent's worker session reads measure definitions but never receives a callable. Two functions that share a Python name share one entry, and the first definition collected wins, so a measure whose function shares @@ -838,7 +840,54 @@ def _check_directory_importable(directory: Path, requested: Path) -> None: def _source_text(func: Callable[..., Any]) -> str: + unavailable = f"# source unavailable for {func.__name__}" try: - return inspect.getsource(func) + source = inspect.getsource(func) except (OSError, TypeError): - return f"# source unavailable for {func.__name__}" + return unavailable + try: + return _exec_ready(source) + except SyntaxError: + # The worker defines every harvested source at spawn, and one that + # does not parse would fail every spawn. + return unavailable + + +# Line boundaries as the parser counts them, so AST line numbers index this. +_LINE = re.compile(r"[^\r\n]*(?:\r\n|\r|\n)|[^\r\n]+\Z") + + +def _exec_ready(source: str) -> str: + """Return ``source`` dedented with decorators dropped, ready for ``exec``. + + Lines inside multi-line strings keep their text. Raises ``SyntaxError`` + if the source does not parse. + """ + lines = _LINE.findall(source) + if not lines: + return source + indent = len(lines[0]) - len(lines[0].lstrip(" \t")) + # An indented definition parses as the body of a block; the wrapper line + # shifts every AST line number by one. + prefix, offset = ("if True:\n", 1) if indent else ("", 0) + tree = ast.parse(prefix + source) + body = tree.body + if indent and isinstance(body[0], ast.If): + body = body[0].body + keep: set[int] = set() + for node in ast.walk(tree): + if isinstance(node, (ast.Constant, ast.JoinedStr)) and node.end_lineno: + keep.update(range(node.lineno - offset, node.end_lineno - offset)) + for node in body: + if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + continue + for decorator in node.decorator_list: + end = (decorator.end_lineno or decorator.lineno) - offset + for index in range(decorator.lineno - 1 - offset, end): + lines[index] = "" + for index, line in enumerate(lines): + if index in keep or not line: + continue + leading = len(line) - len(line.lstrip(" \t")) + lines[index] = line[min(indent, leading) :] + return "".join(lines) diff --git a/pkg-py/tests/test_execution_backend.py b/pkg-py/tests/test_execution_backend.py index 704a7916..d6d5cf92 100644 --- a/pkg-py/tests/test_execution_backend.py +++ b/pkg-py/tests/test_execution_backend.py @@ -1,586 +1,361 @@ -"""The exec-shaped seam the execution subsystem sits on. +"""The seam the execution subsystem sits on: starting and ending a worker. 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. +here (the channel, the environment, kill escalation) are properties of +process handling, and a mock would only restate the implementation. Each +test hands the backend a small script of its own in place of the worker +entry point, so the process does exactly what the test needs of it. """ from __future__ import annotations import asyncio import os -import pathlib import signal -import sys +import subprocess import time +from pathlib import Path from typing import Any, cast import pytest from commons._execution import _backend as backend_module from commons._execution._backend import ( - DEFAULT_OUTPUT_LIMIT, + STDERR_TAIL, TERMINATE_GRACE, ExecBackend, - ExecTimeoutError, LocalBackend, + LocalSession, + WorkerSession, + _read_tail, + _signal_tree, _terminate, ) +from commons._execution._sandbox import protection_mode +try: + protection_mode() +except RuntimeError: + pytestmark = pytest.mark.skip(reason="this host cannot sandbox the worker") -async def test_runs_a_command_and_returns_its_output() -> None: - backend = LocalBackend() - result = await backend.exec([sys.executable, "-c", "print('hello')"]) +def backend_running(tmp_path: Path, code: str, **kwargs: Any) -> LocalBackend: + """A backend whose worker runs ``code`` in place of the real entry point.""" + script = tmp_path / "worker.py" + script.write_text(code) + return LocalBackend(worker_script=script, **kwargs) - 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 +async def wait_for_file(path: Path, timeout: float = 10.0) -> None: + """Block until the child creates ``path``; polling, since it is another process.""" + deadline = time.monotonic() + timeout + while not path.exists(): + if time.monotonic() > deadline: + raise AssertionError(f"{path.name} did not appear within {timeout}s") + await asyncio.sleep(0.01) -def _sleeper( - sentinel: object, *, ignore_sigterm: bool = False, ready: object = None -) -> str: - """Code that outlives its timeout and records the fact if it is allowed to. +def sleeper(sentinel: Path, ready: Path, *, ignore_sigterm: bool = False) -> str: + """A worker that announces itself, then records it if it outlives its close. 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. - - ``ready`` names a file the child creates for a test that cancels - mid-call, which waits for it rather than sleeping. The child first reads - stdin to EOF, because exec closes stdin from inside the try block that - installs the shutdown: the file therefore proves the cancel has somewhere - to land. Waiting only for the child to start would not -- the spawn can - still be mid-await, and the cancel would leave nothing to shut down and - the assertions passing for want of a process. + the process, so a slow 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 "" ) - announce = ( - "" - if ready is None - else f"import sys; sys.stdin.read(); open({str(ready)!r}, 'w').close(); " - ) return ( - f"{guard}import time; {announce}" - f"time.sleep(1.5); open({str(sentinel)!r}, 'w').close()" + f"{guard}import time\n" + f"open({str(ready)!r}, 'w').close()\n" + f"time.sleep(1.5); open({str(sentinel)!r}, 'w').close()\n" ) -async def _wait_until_cancellable(ready: pathlib.Path, timeout: float = 10.0) -> None: - """Block until the child announces that exec is inside its try block. - - Polling, because the child is a separate process with nothing to await. - The timeout only has to outlast a spawn. - """ - deadline = time.monotonic() + timeout - while not ready.exists(): - if time.monotonic() > deadline: - raise AssertionError(f"the call did not become cancellable within {timeout}s") - await asyncio.sleep(0.01) +async def test_the_channel_is_the_workers_stdin_and_stdout(tmp_path) -> None: + backend = backend_running( + tmp_path, "import sys\nsys.stdout.write(sys.stdin.readline().upper())\n" + ) + session = await backend.start(network="none") + session.stdin.write(b"model-written code\n") + await session.stdin.drain() + assert await session.stdout.readline() == b"MODEL-WRITTEN CODE\n" + await session.close() -async def test_a_call_past_the_timeout_raises_and_the_process_does_not_survive( +async def test_the_worker_starts_in_a_scratch_directory_that_close_removes( tmp_path, ) -> None: - sentinel = tmp_path / "survived" - backend = LocalBackend() + backend = backend_running(tmp_path, "import os\nprint(os.getcwd())\n") + session = await backend.start(network="none") + assert isinstance(session, LocalSession) + cwd = (await session.stdout.readline()).decode().strip() + assert cwd == os.path.realpath(session.scratch) + assert cwd != os.path.realpath(tmp_path) + await session.close() + assert not os.path.exists(session.scratch) - 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_the_worker_gets_the_allowlisted_environment_not_the_parents( + tmp_path, monkeypatch +) -> None: + # The parent's environment has credentials a child has no business + # seeing, and a subprocess inherits the whole environment by default. + monkeypatch.setenv("COMMONS_TEST_SECRET", "sk-not-a-real-key") + backend = backend_running( + tmp_path, + "import os\n" + "print(os.environ.get('COMMONS_TEST_SECRET'), os.path.realpath(os.environ['HOME']) == os.getcwd())\n", + ) + session = await backend.start(network="none") + assert (await session.stdout.readline()).split() == [b"None", b"True"] + await session.close() + + +async def test_the_worker_is_told_its_network_and_protection(tmp_path) -> None: + backend = backend_running(tmp_path, "import sys\nprint(sys.argv[1:])\n") + session = await backend.start(network="full") + line = (await session.stdout.readline()).decode() + assert line.strip() == repr(["full", protection_mode()]) + await session.close() + +async def test_interrupt_raises_keyboard_interrupt_in_the_worker(tmp_path) -> None: + ready = tmp_path / "ready" + backend = backend_running( + tmp_path, + "import time\n" + "try:\n" + f" open({str(ready)!r}, 'w').close(); time.sleep(30)\n" + "except KeyboardInterrupt:\n" + " print('interrupted', flush=True)\n", + ) + session = await backend.start(network="none") + await wait_for_file(ready) + session.interrupt() + line = await asyncio.wait_for(session.stdout.readline(), 5) + assert line == b"interrupted\n" + await session.close() -async def test_a_process_that_handles_sigterm_gets_to_clean_up_first(tmp_path) -> None: + +async def test_a_worker_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 = ( + ready = tmp_path / "ready" + backend = backend_running( + tmp_path, "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" + f"open({str(ready)!r}, 'w').close()\n" + "time.sleep(5)\n", ) - backend = LocalBackend() - - with pytest.raises(ExecTimeoutError): - await backend.exec([sys.executable, "-c", code], timeout=0.15) - + session = await backend.start(network="none") + await wait_for_file(ready) + await session.close() 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, - ) - +async def test_a_worker_that_ignores_sigterm_is_killed_anyway(tmp_path) -> None: + sentinel, ready = tmp_path / "survived", tmp_path / "ready" + backend = backend_running( + tmp_path, sleeper(sentinel, ready, ignore_sigterm=True), terminate_grace=0.1 + ) + session = await backend.start(network="none") + await wait_for_file(ready) + await session.close() + assert session.returncode is not None 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] = [] - - async def wait(self) -> int: - await asyncio.sleep(3600) - return 0 - - -async def test_terminate_gives_up_when_the_exit_is_never_reaped( - monkeypatch, -) -> None: - process = _NeverReaped() - - def record_signals(process: Any, sig: signal.Signals) -> None: - process.signals.append("term" if sig == signal.SIGTERM else "kill") - - monkeypatch.setattr(backend_module, "_signal_tree", record_signals) - - # 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, +@pytest.mark.skipif(os.name != "posix", reason="os.fork is POSIX-only") +async def test_a_workers_own_children_do_not_survive_the_close(tmp_path) -> None: + # The escalation signals the worker's whole process group: code that + # forks a child of its own cannot strand it. + sentinel, ready = tmp_path / "survived", tmp_path / "ready" + backend = backend_running( + tmp_path, + "import os, time\n" + "if os.fork() == 0:\n" + f" time.sleep(1.5); open({str(sentinel)!r}, 'w').close(); os._exit(0)\n" + f"open({str(ready)!r}, 'w').close()\n" + "time.sleep(5)\n", + terminate_grace=0.1, ) - - 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) + session = await backend.start(network="none") + await wait_for_file(ready) + await session.close() + await asyncio.sleep(1.8) + assert not sentinel.exists() -async def test_the_timeout_still_applies_after_the_output_streams_close( +@pytest.mark.skipif(os.name != "posix", reason="os.fork is POSIX-only") +async def test_a_dead_workers_children_are_killed_when_its_exit_is_seen( 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. + # The leader's exit leaves its process group behind; the group is killed + # as the exit is observed, without waiting for a close that may come an + # idle window later. The child lets go of the pipes, as the real + # worker's fds are sinks, so that the leader's exit is observed while + # the child still runs. sentinel = tmp_path / "survived" - code = ( + backend = backend_running( + tmp_path, "import os, time\n" - "os.close(1); os.close(2)\n" - f"time.sleep(1.5); open({str(sentinel)!r}, 'w').close()\n" + "if os.fork() == 0:\n" + " sink = os.open(os.devnull, os.O_RDWR)\n" + " for fd in (0, 1, 2):\n" + " os.dup2(sink, fd)\n" + f" time.sleep(1.5); open({str(sentinel)!r}, 'w').close(); os._exit(0)\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" - ready = tmp_path / "ready" - backend = LocalBackend(terminate_grace=0.1) - call = asyncio.create_task( - backend.exec([sys.executable, "-c", _sleeper(sentinel, ready=ready)]) + session = await backend.start(network="none") + assert isinstance(session, LocalSession) + try: + await session.process.wait() + await asyncio.sleep(1.8) + assert not sentinel.exists() + finally: + await session.close() + + +async def test_a_cancelled_close_still_kills_the_worker_and_reraises(tmp_path) -> None: + # Shutdown is not the caller's to interrupt: a cancel during the SIGTERM + # grace would otherwise skip SIGKILL. The caller is still told it was + # cancelled. + sentinel, ready = tmp_path / "survived", tmp_path / "ready" + backend = backend_running( + tmp_path, sleeper(sentinel, ready, ignore_sigterm=True), terminate_grace=1.0 ) - await _wait_until_cancellable(ready) - - call.cancel() + session = await backend.start(network="none") + await wait_for_file(ready) + closing = asyncio.ensure_future(session.close()) + # The grace runs for 1.0s, so a cancel at 0.2s is mid-shutdown. + await asyncio.sleep(0.2) + closing.cancel() with pytest.raises(asyncio.CancelledError): - await call - + await closing + assert backend._shutdowns + await backend.aclose() + assert not backend._shutdowns + assert session.returncode is not None 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" - ready = tmp_path / "ready" - backend = LocalBackend(terminate_grace=0.1) - call = asyncio.create_task( - backend.exec( - [ - sys.executable, - "-c", - _sleeper(sentinel, ignore_sigterm=True, ready=ready), - ] - ) - ) - await _wait_until_cancellable(ready) - - call.cancel() - with pytest.raises(asyncio.CancelledError): - await call - - await asyncio.sleep(1.8) - assert not sentinel.exists() +async def test_close_is_safe_to_call_twice(tmp_path) -> None: + backend = backend_running(tmp_path, "import time\ntime.sleep(30)\n") + session = await backend.start(network="none") + await asyncio.gather(session.close(), session.close()) + await session.close() + assert session.returncode is not None -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. +async def test_cancelling_a_start_leaves_no_worker_behind(tmp_path) -> None: sentinel = tmp_path / "survived" - ready = tmp_path / "ready" - backend = LocalBackend(terminate_grace=1.0) - call = asyncio.create_task( - backend.exec( - [ - sys.executable, - "-c", - _sleeper(sentinel, ignore_sigterm=True, ready=ready), - ] - ) + backend = backend_running( + tmp_path, + f"import time\ntime.sleep(1.5); open({str(sentinel)!r}, 'w').close()\n", + terminate_grace=0.1, ) - await _wait_until_cancellable(ready) - - call.cancel() - # The SIGTERM grace runs for 1.0s from the first cancel, so the second - # cancel at ~0.3s lands mid-shutdown with margin on both sides — the - # race this test exists for must actually happen, or it quietly - # degenerates into the plain cancellation test above. - await asyncio.sleep(0.2) - call.cancel() + starting = asyncio.ensure_future(backend.start(network="none")) + await asyncio.sleep(0) + starting.cancel() with pytest.raises(asyncio.CancelledError): - await call - + await starting + await backend.aclose() 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 = ( +async def test_stderr_keeps_the_tail_of_what_the_worker_wrote(tmp_path) -> None: + backend = backend_running( + tmp_path, "import sys\n" "for i in range(200):\n" - " sys.stderr.write(f'line-{i}:' + 'x' * 1000 + '\\n')\n" + " sys.stderr.write(f'line-{i}:' + 'x' * 1000 + '\\n')\n", ) + session = await backend.start(network="none") + assert isinstance(session, LocalSession) + await session.process.wait() + await session.close() + tail = await session.stderr() + assert len(tail) <= STDERR_TAIL + assert "line-199:" in tail + assert "line-0:" not in tail - 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=1.0) - 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 1.0s, - # so a cancel at 0.5s lands in the middle of the shutdown with margin on - # both sides — the race this test exists for must actually happen, or it - # quietly degenerates into the plain timeout test above. - await asyncio.sleep(0.5) - - call.cancel() - with pytest.raises((asyncio.CancelledError, ExecTimeoutError)): - await call - - await asyncio.sleep(1.8) - assert not sentinel.exists() +async def test_reading_a_tail_keeps_the_last_bytes_and_flags_the_cut() -> None: + stream = asyncio.StreamReader() + stream.feed_data(b"a" * 5000 + b"end") + stream.feed_eof() + assert await _read_tail(stream, 10) == (b"aaaaaaaend", True) async def test_reading_the_tail_of_a_missing_stream_is_empty() -> None: - # exec always pipes both streams, so this guard is defensive; pin it so - # it cannot be broken or deleted unnoticed. - assert await backend_module._read_tail(None, 10) == (b"", False) + assert await _read_tail(None, 10) == (b"", False) -def test_the_defaults_are_the_documented_constants() -> None: - # Every other test overrides these; pin the wiring itself. - backend = LocalBackend() +class _NeverReaped: + """A process that takes its signals but whose exit is never observed. - assert backend._output_limit == DEFAULT_OUTPUT_LIMIT - assert backend._terminate_grace == TERMINATE_GRACE + 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 -async def test_concurrent_calls_on_one_backend_do_not_interfere(tmp_path) -> None: - # The backend's one piece of shared state is the set of in-flight - # shutdowns; one call timing out must not disturb another's result. - sentinel = tmp_path / "survived" - backend = LocalBackend(terminate_grace=0.1) - slow = asyncio.create_task( - backend.exec([sys.executable, "-c", _sleeper(sentinel)], timeout=0.15) - ) - fast = asyncio.create_task(backend.exec([sys.executable, "-c", "print('ok')"])) + def __init__(self) -> None: + self.signals: list[str] = [] - with pytest.raises(ExecTimeoutError): - await slow - assert (await fast).stdout == "ok\n" + async def wait(self) -> int: + await asyncio.sleep(3600) + return 0 - await asyncio.sleep(1.8) - assert not sentinel.exists() +async def test_terminate_gives_up_when_the_exit_is_never_reaped(monkeypatch) -> None: + process = _NeverReaped() -async def test_aclose_waits_for_shutdowns_still_in_flight(tmp_path) -> None: - # The escalation guarantee only holds while the event loop is running; - # aclose is how a driver honours it while tearing down. A second cancel - # is what leaves a shutdown running detached after the call has ended. - sentinel = tmp_path / "survived" - ready = tmp_path / "ready" - backend = LocalBackend(terminate_grace=1.0) - call = asyncio.create_task( - backend.exec( - [ - sys.executable, - "-c", - _sleeper(sentinel, ignore_sigterm=True, ready=ready), - ] - ) - ) - await _wait_until_cancellable(ready) + def record_signals(process: Any, sig: signal.Signals) -> None: + process.signals.append("term" if sig == signal.SIGTERM else "kill") - call.cancel() - await asyncio.sleep(0.1) - call.cancel() - with pytest.raises(asyncio.CancelledError): - await call + monkeypatch.setattr(backend_module, "_signal_tree", record_signals) - assert backend._shutdowns - await backend.aclose() - assert not backend._shutdowns + # 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) - await asyncio.sleep(1.8) - assert not sentinel.exists() + assert process.signals == ["term", "kill"] -@pytest.mark.skipif(os.name != "posix", reason="os.fork is POSIX-only") -async def test_a_workers_own_children_do_not_survive_the_shutdown(tmp_path) -> None: - # The escalation signals the worker's whole process group: code that - # forks a child of its own cannot strand it, and cannot hold the output - # pipes open past the worker's own exit. - sentinel = tmp_path / "survived" - code = ( - "import os, time\n" - "if os.fork() == 0:\n" - f" time.sleep(1.5); open({str(sentinel)!r}, 'w').close(); os._exit(0)\n" - "time.sleep(5)\n" - ) - backend = LocalBackend(terminate_grace=0.1) +async def test_the_local_backend_satisfies_the_backend_interface(tmp_path) -> None: + # The annotations are 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 = backend_running(tmp_path, "") + session: WorkerSession = await backend.start(network="none") + assert isinstance(backend, ExecBackend) + assert isinstance(session, WorkerSession) + await session.close() - with pytest.raises(ExecTimeoutError): - await backend.exec([sys.executable, "-c", code], timeout=0.15) - await asyncio.sleep(1.8) - assert not sentinel.exists() +def test_the_default_grace_is_the_documented_constant() -> None: + assert LocalBackend()._terminate_grace == TERMINATE_GRACE -async def test_both_streams_past_the_cap_are_capped_independently() -> None: - # The two streams keep separate buffers and flags even when both blow - # past the cap in the same call. - backend = LocalBackend(output_limit=2000) - code = ( - "import sys\n" - "for i in range(200):\n" - " sys.stdout.write(f'o-{i}:' + 'x' * 1000 + '\\n')\n" - " sys.stderr.write(f'e-{i}:' + 'y' * 1000 + '\\n')\n" - ) - result = await backend.exec([sys.executable, "-c", code]) - - assert result.returncode == 0 - assert result.stdout_truncated - assert result.stderr_truncated - assert len(result.stdout) <= 2000 - assert len(result.stderr) <= 2000 - assert "o-199:" in result.stdout - assert "o-0:" not in result.stdout - assert "e-199:" in result.stderr - assert "e-0:" not in result.stderr +@pytest.mark.skipif( + not hasattr(os, "waitid"), reason="os.waitid is POSIX, and on macOS 3.13+" +) +def test_a_group_left_with_only_a_zombie_can_still_be_signalled() -> None: + # The worker can exit before its exit is observed, and a close then + # signals a group whose only member is a zombie, which macOS refuses + # with EPERM. A Popen is used so that no child watcher reaps it first. + process = subprocess.Popen(["/bin/sh", "-c", "exit 0"], start_new_session=True) + os.waitid(os.P_PID, process.pid, os.WEXITED | os.WNOWAIT) # pyrefly: ignore[missing-attribute] + _signal_tree(cast(Any, process), signal.SIGTERM) + assert process.wait() == 0 diff --git a/pkg-py/tests/test_execution_driver.py b/pkg-py/tests/test_execution_driver.py new file mode 100644 index 00000000..91e6abd7 --- /dev/null +++ b/pkg-py/tests/test_execution_driver.py @@ -0,0 +1,1082 @@ +"""The worker's lifecycle: lazy spawn, serialization, timeouts, respawn, and the reap. + +Every test here runs a real worker process through the driver; none of the +machinery is mocked. Hosts that cannot sandbox the worker (and non-POSIX +hosts, whose interrupt escalation differs) skip the module. +""" + +from __future__ import annotations + +import asyncio +import os +import signal +import sys +import time +from datetime import date +from typing import Annotated + +import pandas as pd +import pytest +from pydantic import Field + +from commons import measure, semantic_layer +from commons._execution._backend import WORKER_SCRIPT, LocalBackend, LocalSession +from commons._execution._driver import Failure, Worker +from commons._execution._protocol import Error, Result +from commons._execution._sandbox import protection_mode +from commons._handles import HandleStore +from commons._measures import Injected + + +@measure(description="Total the x column.") +def total_x(df: Injected[pd.DataFrame]) -> float: + return float(df["x"].sum()) + + +DEFAULT_REGION = "EMEA" + + +@measure(description="Count orders in a region.") +def region_count( + region: Annotated[str, Field(description="The region.")] = DEFAULT_REGION, +) -> int: + return len(region) + + +@measure(description="Orders placed on or after a date.") +def orders_since( + when: Annotated[date, Field(description="The cutoff.")] = date(2020, 1, 1), +) -> int: + return 3 + + +@measure(description="Count with a composed default.") +def composed_default( + region: Annotated[str, Field(description="The region.")] = DEFAULT_REGION + "-west", +) -> int: + return len(region) + + +@measure(description="Double a number.") +def double_n(n: Annotated[int, Field(description="The number.")] = 21) -> int: + return 2 * n + + +try: + _PROTECTION = protection_mode() +except RuntimeError: + _PROTECTION = None + +pytestmark = [ + pytest.mark.skipif( + _PROTECTION is None, reason="this host cannot sandbox the worker" + ), + pytest.mark.skipif( + os.name != "posix", reason="the interrupt escalation is POSIX-shaped" + ), +] + + +def make_worker(**kwargs) -> Worker: + """A worker with test-sized timeouts; kwargs override them.""" + kwargs.setdefault("call_timeout", 5) + kwargs.setdefault("idle_timeout", 600) + kwargs.setdefault("interrupt_grace", 2) + return Worker(**kwargs) + + +def process_of(worker: Worker) -> asyncio.subprocess.Process: + """The running worker's process, which the local backend exposes.""" + session = worker._session + assert isinstance(session, LocalSession) + return session.process + + +async def wait_until(condition, message: str, timeout: float = 5) -> None: + """Poll ``condition`` until it holds, failing with ``message`` at the deadline.""" + deadline = time.monotonic() + timeout + while not condition(): + assert time.monotonic() < deadline, message + await asyncio.sleep(0.01) + + +def group_gone(pid: int): + """Whether no process in the group ``pid`` leads is left.""" + + def check() -> bool: + try: + os.killpg(pid, 0) + except ProcessLookupError: + return True + return False + + return check + + +async def test_no_process_exists_until_the_first_call(): + async with make_worker() as worker: + assert worker._session is None + await worker.run("1") + assert worker._session is not None + + +@pytest.mark.parametrize("network", ["none", "full"]) +async def test_the_worker_is_started_with_its_network_access(network): + # The backend passes the access level to the worker's entry point. + async with make_worker(network=network) as worker: + reply = await worker.run("import sys; sys.argv[1]") + assert isinstance(reply, Result) + assert reply.value == network + + +async def test_a_trailing_expression_is_the_result(): + async with make_worker() as worker: + reply = await worker.run("6 * 7") + assert isinstance(reply, Result) + assert reply.value == 42 + + +async def test_state_persists_across_calls(): + async with make_worker() as worker: + await worker.run("x = 41") + reply = await worker.run("x + 1") + assert isinstance(reply, Result) + assert reply.value == 42 + + +async def test_the_codes_own_exception_is_an_error_not_a_failure(): + async with make_worker() as worker: + reply = await worker.run("1 / 0") + assert isinstance(reply, Error) + assert "ZeroDivisionError" in reply.message + + +async def test_printed_output_comes_back_with_the_result(): + async with make_worker() as worker: + reply = await worker.run("print('hello')") + assert isinstance(reply, Result) + assert reply.stdout == "hello\n" + + +async def test_calls_are_serialized_one_at_a_time(): + async with make_worker() as worker: + first, second = await asyncio.gather( + worker.run("import time; time.sleep(0.5); seen = 1"), + worker.run("seen + 1"), + ) + # Had the second call run alongside the first, `seen` would not + # have been bound yet. + assert isinstance(first, Result) + assert isinstance(second, Result) + assert second.value == 2 + + +async def test_writing_fd_1_directly_cannot_corrupt_the_channel(): + async with make_worker() as worker: + reply = await worker.run("import os; os.write(1, b'junk\\n'); 6 * 7") + assert isinstance(reply, Result) + assert reply.value == 42 + # The channel survived: the next call answers normally. + reply = await worker.run("'still here'") + assert isinstance(reply, Result) + assert reply.value == "still here" + + +async def test_a_thread_printing_after_the_call_cannot_corrupt_the_channel(): + async with make_worker() as worker: + code = ( + "import threading, time\n" + "def late():\n" + " time.sleep(0.2)\n" + " print('junk from a late thread')\n" + "threading.Thread(target=late).start()\n" + "'started'\n" + ) + reply = await worker.run(code) + assert isinstance(reply, Result) + await asyncio.sleep(1) + reply = await worker.run("40 + 2") + assert isinstance(reply, Result) + assert reply.value == 42 + + +async def test_a_thread_reading_stdin_cannot_steal_the_next_call(): + async with make_worker() as worker: + code = ( + "import threading, os\n" + "def thief():\n" + " while os.read(0, 4096):\n" + " pass\n" + "threading.Thread(target=thief, daemon=True).start()\n" + "'started'\n" + ) + reply = await worker.run(code) + assert isinstance(reply, Result) + await asyncio.sleep(0.2) + # The thief is blocked on the sink's EOF, not on the channel: the + # next call's bytes reach the worker's loop. + reply = await worker.run("40 + 2") + assert isinstance(reply, Result) + assert reply.value == 42 + + +async def test_a_timeout_interrupts_the_call_and_keeps_the_session(): + async with make_worker(call_timeout=0.5) as worker: + await worker.run("x = 5") + reply = await worker.run("import time; time.sleep(60)") + assert isinstance(reply, Failure) + assert "interrupted" in reply.message + assert "remain available" in reply.message + # The interrupt left the session and its variables alive. + reply = await worker.run("x") + assert isinstance(reply, Result) + assert reply.value == 5 + + +async def test_an_interrupt_between_calls_leaves_the_session_alone(): + # The driver can interrupt a call just as it finishes, so the worker + # receives the SIGINT once it is back to waiting for the next line. + async with make_worker() as worker: + await worker.run("x = 5") + os.killpg(process_of(worker).pid, signal.SIGINT) + await asyncio.sleep(0.3) + reply = await worker.run("x") + assert isinstance(reply, Result) + assert reply.value == 5 + + +async def test_a_host_that_ignores_sigint_still_gets_its_calls_interrupted(): + # A background job's shell starts it with SIGINT ignored, which a child + # inherits; the worker must not lose its interrupt to that. + previous = signal.signal(signal.SIGINT, signal.SIG_IGN) + worker = make_worker(call_timeout=0.5) + try: + await worker.run("x = 5") + finally: + signal.signal(signal.SIGINT, previous) + async with worker: + reply = await worker.run("import time; time.sleep(60)") + assert isinstance(reply, Failure) + assert "remain available" in reply.message + reply = await worker.run("x") + assert isinstance(reply, Result) + assert reply.value == 5 + + +async def test_a_worker_that_ignores_the_interrupt_is_restarted(): + async with make_worker(call_timeout=0.5, interrupt_grace=1) as worker: + await worker.run("x = 5") + code = ( + "import signal, time\n" + "signal.signal(signal.SIGINT, signal.SIG_IGN)\n" + "time.sleep(60)\n" + ) + reply = await worker.run(code) + assert isinstance(reply, Failure) + assert "did not respond" in reply.message + assert "restarted" in reply.message + # The restart reset the session: x is gone. + reply = await worker.run("x") + assert isinstance(reply, Error) + assert "NameError" in reply.message + + +async def test_a_timed_out_call_takes_its_children_with_it(): + async with make_worker(call_timeout=1, interrupt_grace=2) as worker: + probe = await worker.run( + "import subprocess; subprocess.run(['/bin/echo', 'x']).returncode" + ) + if not isinstance(probe, Result) or probe.value != 0: + # The macOS sandbox aborts child processes outright, which also + # settles the question this test asks. + pytest.skip("the sandbox refuses child processes on this host") + code = ( + "import subprocess\n" + # The interrupt fires while the call waits on the child. If only + # the worker receives it, the child survives to write the file. + # Popen().wait() rather than run(), which kills its own child on + # KeyboardInterrupt and would pass either way. + "subprocess.Popen(['sh', '-c', 'touch started; sleep 2; touch survived'])" + ".wait()\n" + ) + reply = await worker.run(code) + assert isinstance(reply, Failure) + assert "interrupted" in reply.message + # Without a child that started, there is nothing to have outlived. + reply = await worker.run("import os; os.path.exists('started')") + assert isinstance(reply, Result) + assert reply.value is True + await asyncio.sleep(2.5) + reply = await worker.run("import os; os.path.exists('survived')") + assert isinstance(reply, Result) + assert reply.value is False + + +async def test_a_crashed_workers_children_do_not_outlive_it(): + async with make_worker() as worker: + probe = await worker.run( + "import subprocess; subprocess.run(['/bin/echo', 'x']).returncode" + ) + if not isinstance(probe, Result) or probe.value != 0: + # The macOS sandbox aborts child processes outright, so there + # is nothing to outlive the worker on this host. + pytest.skip("the sandbox refuses child processes on this host") + process = process_of(worker) + code = "import subprocess, os\nsubprocess.Popen(['sleep', '30'])\nos._exit(1)\n" + reply = await worker.run(code) + assert isinstance(reply, Failure) + # The shutdown killed the group, not just the leader. The kill is sent + # before run() returns, but a killed child takes a moment to finish + # exiting, and until then the group still exists. + await wait_until( + group_gone(process.pid), "a child outlived the crashed worker", timeout=2 + ) + + +async def test_a_child_ignoring_sigterm_does_not_outlive_the_close(): + worker = make_worker() + probe = await worker.run( + "import subprocess; subprocess.run(['/bin/echo', 'x']).returncode" + ) + if not isinstance(probe, Result) or probe.value != 0: + await worker.aclose() + pytest.skip("the sandbox refuses child processes on this host") + process = process_of(worker) + # The worker honours SIGTERM; its child does not, and before the group + # kill it outlived the shutdown. + await worker.run( + "import subprocess\n" + "subprocess.Popen(['sh', '-c', 'trap \"\" TERM; sleep 30'])\n" + "import time; time.sleep(0.2)\n" + ) + await worker.aclose() + await wait_until( + group_gone(process.pid), "a child outlived the closed worker", timeout=2 + ) + + +@pytest.mark.skipif( + sys.platform == "darwin" and sys.version_info >= (3, 14), + reason="asyncio on 3.14 reads macOS's waitid() report of a stopped child " + "as an exit and blocks the event loop in waitpid() until it really exits", +) +async def test_a_worker_that_stops_reading_fails_the_call_instead_of_hanging(): + async with make_worker(call_timeout=0.5) as worker: + await worker.run("1") + process = process_of(worker) + # A stopped worker cannot drain its stdin; a call bigger than the + # pipe buffer would block the write forever without a bound on it. + process.send_signal(signal.SIGSTOP) + reply = await worker.run("x = " + "1" * (1024 * 1024)) + assert isinstance(reply, Failure) + assert "stopped reading" in reply.message + + +async def test_a_crash_is_reported_and_the_next_call_respawns(): + async with make_worker() as worker: + await worker.run("x = 5") + reply = await worker.run("import os; os._exit(1)") + assert isinstance(reply, Failure) + assert "crashed" in reply.message + reply = await worker.run("x") + assert isinstance(reply, Error) + assert "NameError" in reply.message + + +async def test_an_idle_worker_is_reaped_and_the_next_call_respawns(): + async with make_worker(idle_timeout=0.5) as worker: + await worker.run("x = 5") + process = process_of(worker) + # The reaper closes the worker; nothing is running then. + await wait_until(lambda: worker._session is None, "the worker was not reaped") + reply = await worker.run("x") + # The respawned session has never seen x. + assert isinstance(reply, Error) + assert "NameError" in reply.message + assert process_of(worker).pid != process.pid + + +async def test_the_reaper_leaves_a_worker_with_a_call_in_flight_alone(): + async with make_worker(idle_timeout=0.5, call_timeout=10) as worker: + # The first call arms the idle timer; the second runs straight + # through the moment it fires. + await worker.run("1") + reply = await worker.run("import time; time.sleep(2); 6 * 7") + # The idle window passed mid-call, but the call was in flight, so + # the worker lived to answer it. + assert isinstance(reply, Result) + assert reply.value == 42 + + +async def test_measure_sources_are_defined_at_spawn_and_at_respawn(): + async with make_worker(measure_sources=["ANSWER = 42"], idle_timeout=0.5) as worker: + reply = await worker.run("ANSWER") + assert isinstance(reply, Result) + assert reply.value == 42 + await wait_until(lambda: worker._session is None, "the worker was not reaped") + # The respawned worker was given the same sources at spawn. + reply = await worker.run("ANSWER") + assert isinstance(reply, Result) + assert reply.value == 42 + + +async def test_a_measure_source_that_fails_is_skipped(): + sources = ["BEFORE = 1", "1 / 0", "AFTER = 2"] + async with make_worker(measure_sources=sources) as worker: + reply = await worker.run("BEFORE + AFTER") + assert isinstance(reply, Result) + assert reply.value == 3 + + +async def test_a_harvested_source_that_cannot_run_alone_is_skipped(tmp_path): + # Each is a module-level function the harvest keeps, and each fails + # once lifted out of its module: a lambda whose line names a module + # global, and a closure that rebinds a name its module defined. + module = tmp_path / "unliftable.py" + module.write_text( + "from commons import measure\n" + "\n" + "LOG = True\n" + "to_pct = (lambda x: round(x * 100, 1)) if LOG else None\n" + "\n" + "def cached(f):\n" + " value = None\n" + " def wrapper():\n" + " nonlocal value\n" + " value = value or f()\n" + " return value\n" + " return wrapper\n" + "\n" + '@measure(description="Answer.")\n' + "def answer() -> int:\n" + " return 42\n" + ) + layer = semantic_layer(module) + async with make_worker(measure_sources=list(layer.source_text.values())) as worker: + reply = await worker.run("answer()") + assert isinstance(reply, Result) + assert reply.value == 42 + + +async def test_using_a_default_that_could_not_be_evaluated_raises(): + # A placeholder that formatted into a query would return a wrong + # answer instead of an error. + layer = semantic_layer([region_count]) + async with make_worker(measure_sources=list(layer.source_text.values())) as worker: + for code in ( + "f'{region_count.__defaults__[0]}'", + "str(region_count.__defaults__[0])", + "bool(region_count.__defaults__[0])", + "region_count.__defaults__[0] == 'EMEA'", + "region_count.__defaults__[0].upper()", + ): + reply = await worker.run(code) + assert isinstance(reply, Error), code + assert "NameError" in reply.message, code + assert "DEFAULT_REGION" in reply.message, code + + +async def test_a_measure_source_that_ends_the_worker_stops_the_spawn(): + async with make_worker(measure_sources=["import os; os._exit(1)"]) as worker: + reply = await worker.run("1") + assert isinstance(reply, Failure) + assert "exited while defining" in reply.message + assert worker._session is None + + +async def test_a_session_seed_that_fails_is_reported_at_the_start(tmp_path): + # The real worker with a broken seed: its traceback must reach the + # driver, so the seed has to run while fd 2 is still attached. + runtime = WORKER_SCRIPT.parent + script = tmp_path / "worker.py" + script.write_text( + "import sys\n" + f"sys.path.insert(0, {str(runtime)!r})\n" + "import _worker\n" + "_worker._DEFINE_SOURCE = 'raise RuntimeError(\"the seed is broken\")'\n" + "_worker._PROTOCOL_IN, _worker._PROTOCOL_OUT = " + "_worker._claim_protocol_channel()\n" + "_worker.main()\n" + ) + async with make_worker(backend=LocalBackend(worker_script=script)) as worker: + reply = await worker.run("1") + assert isinstance(reply, Failure) + assert "exited before it was ready" in reply.message + assert "the seed is broken" in reply.message + + +async def test_a_worker_that_never_becomes_ready_fails_the_call(tmp_path): + script = tmp_path / "worker.py" + script.write_text( + "import sys, time\n" + "sys.stderr.write('wedged in startup\\n')\n" + "sys.stderr.flush()\n" + "time.sleep(60)\n" + ) + backend = LocalBackend(worker_script=script, terminate_grace=0.1) + async with make_worker(backend=backend, spawn_timeout=0.5) as worker: + reply = await worker.run("1") + assert isinstance(reply, Failure) + assert "stopped responding during startup" in reply.message + assert "wedged in startup" in reply.message + assert worker._session is None + + +async def test_harvested_measure_sources_define_their_names(): + # The real harvest: decorated and annotated, with the decorator and + # the annotation's import both absent from the worker's session. + layer = semantic_layer([total_x]) + async with make_worker(measure_sources=list(layer.source_text.values())) as worker: + reply = await worker.run("total_x.__name__") + assert isinstance(reply, Result) + assert reply.value == "total_x" + + +async def test_the_model_can_read_each_measure_source_back(): + # The sources are reference material, so each must be readable by name. + layer = semantic_layer([total_x, region_count]) + async with make_worker(measure_sources=list(layer.source_text.values())) as worker: + for name, source in layer.source_text.items(): + reply = await worker.run(f"import inspect; inspect.getsource({name})") + assert isinstance(reply, Result), reply + assert reply.value == source + + +async def test_a_measure_default_naming_an_unharvested_global_still_defines(): + # DEFAULT_REGION is a module global rather than a function, so the + # harvest never includes it; the define must not fail on it. + layer = semantic_layer([region_count]) + async with make_worker(measure_sources=list(layer.source_text.values())) as worker: + reply = await worker.run("region_count.__name__") + assert isinstance(reply, Result) + assert reply.value == "region_count" + # The default stands in as a placeholder that names what is missing. + reply = await worker.run("repr(region_count.__defaults__[0])") + assert isinstance(reply, Result) + assert "DEFAULT_REGION" in reply.value + + +async def test_a_compound_default_naming_an_unharvested_import_still_defines(): + # `date` is an import, so the harvest excludes it, and the default is a + # call on top of the missing name: the placeholder must survive both. + layer = semantic_layer([orders_since]) + async with make_worker(measure_sources=list(layer.source_text.values())) as worker: + reply = await worker.run("orders_since.__name__") + assert isinstance(reply, Result) + assert reply.value == "orders_since" + reply = await worker.run("repr(orders_since.__defaults__[0])") + assert isinstance(reply, Result) + assert "date" in reply.value + + +async def test_a_default_composed_from_an_unharvested_global_still_defines(): + layer = semantic_layer([composed_default]) + async with make_worker(measure_sources=list(layer.source_text.values())) as worker: + reply = await worker.run("composed_default.__name__") + assert isinstance(reply, Result) + reply = await worker.run("repr(composed_default.__defaults__[0])") + assert isinstance(reply, Result) + assert "DEFAULT_REGION" in reply.value + + +async def test_lambda_and_keyword_only_defaults_naming_unharvested_globals_define( + tmp_path, +): + # The harvest keeps module-level lambdas as well as defs, and a default + # in either shape is evaluated as the source is defined. + module = tmp_path / "scaled.py" + module.write_text( + "from typing import Annotated\n" + "from pydantic import Field\n" + "from commons import measure\n" + "\n" + "FACTOR = 2\n" + "scale = lambda x, k=FACTOR: x * k\n" + "\n" + '@measure(description="Scaled orders.")\n' + "def scaled(\n" + ' *, n: Annotated[int, Field(description="Count.")] = FACTOR\n' + ") -> int:\n" + " return scale(n)\n" + ) + layer = semantic_layer(module) + async with make_worker(measure_sources=list(layer.source_text.values())) as worker: + reply = await worker.run("scale(3, k=2)") + assert isinstance(reply, Result) + assert reply.value == 6 + reply = await worker.run( + "repr(scale.__defaults__[0]) + repr(scaled.__kwdefaults__['n'])" + ) + assert isinstance(reply, Result) + assert reply.value.count("FACTOR") == 2 + + +async def test_a_default_that_resolves_keeps_its_real_value(): + layer = semantic_layer([double_n]) + async with make_worker(measure_sources=list(layer.source_text.values())) as worker: + # The guard around a default that evaluates fine is invisible. + reply = await worker.run("double_n()") + assert isinstance(reply, Result) + assert reply.value == 42 + + +async def test_a_failed_start_reports_the_workers_stderr(tmp_path): + script = tmp_path / "worker.py" + script.write_text( + "import sys\nprint('boom: the sandbox refused', file=sys.stderr)\nsys.exit(3)\n" + ) + async with make_worker(backend=LocalBackend(worker_script=script)) as worker: + reply = await worker.run("1") + assert isinstance(reply, Failure) + assert "failed to start" in reply.message + assert "boom: the sandbox refused" in reply.message + + +async def test_handles_are_sent_with_the_call(): + store = HandleStore() + store.register(41) + async with make_worker() as worker: + reply = await worker.run("r1 + 1", handles=store) + assert isinstance(reply, Result) + assert reply.value == 42 + # A handle registered later syncs on the next call; the first one + # is still there. + store.register(100) + reply = await worker.run("r1 + r2", handles=store) + assert isinstance(reply, Result) + assert reply.value == 141 + + +async def test_handles_are_resent_to_a_respawned_worker(): + store = HandleStore() + store.register(41) + async with make_worker(idle_timeout=0.5) as worker: + reply = await worker.run("r1", handles=store) + assert isinstance(reply, Result) + await wait_until(lambda: worker._session is None, "the worker was not reaped") + reply = await worker.run("r1 + 1", handles=store) + assert isinstance(reply, Result) + assert reply.value == 42 + + +async def test_closing_the_worker_removes_its_scratch_directory(): + worker = make_worker() + reply = await worker.run("import os; os.getcwd()") + assert isinstance(reply, Result) + scratch = reply.value + assert os.path.isdir(scratch) + await worker.aclose() + assert not os.path.exists(scratch) + + +async def test_a_reaped_workers_scratch_directory_is_removed(): + async with make_worker(idle_timeout=0.5) as worker: + reply = await worker.run("import os; os.getcwd()") + assert isinstance(reply, Result) + scratch = reply.value + # The reap forgets the session first and removes its files last. + await wait_until( + lambda: not os.path.exists(scratch), "the scratch directory remained" + ) + + +async def test_aclose_during_a_reap_still_kills_a_sigterm_ignoring_worker(): + worker = make_worker(idle_timeout=0.2) + await worker.run("import signal; signal.signal(signal.SIGTERM, signal.SIG_IGN)") + process = process_of(worker) + # Wait for the reap to begin (it clears the process synchronously), + # then close while its SIGTERM grace window is still open. Cancelling + # the reaper must not strand the worker before its SIGKILL. + await wait_until(lambda: worker._session is None, "the reap never began") + await worker.aclose() + assert process.returncode is not None + + +async def test_cancelling_a_call_shuts_the_worker_down(): + async with make_worker() as worker: + await worker.run("x = 5") + call = asyncio.ensure_future(worker.run("import time; time.sleep(30)")) + await asyncio.sleep(0.2) + call.cancel() + with pytest.raises(asyncio.CancelledError): + await call + # The cancelled call's worker is gone rather than still computing, + # so its late reply cannot poison the next call's channel. + assert worker._session is None + reply = await worker.run("x") + assert isinstance(reply, Error) + assert "NameError" in reply.message + + +async def test_a_closed_worker_stays_closed(): + worker = make_worker() + await worker.run("1") + await worker.aclose() + reply = await worker.run("1") + assert isinstance(reply, Failure) + assert "closed" in reply.message + + +async def test_cancelling_aclose_still_closes_the_worker(): + worker = make_worker() + call = asyncio.ensure_future(worker.run("import time; time.sleep(1); 1")) + await asyncio.sleep(0.2) + closing = asyncio.ensure_future(worker.aclose()) + await asyncio.sleep(0.2) + # The close is waiting out the in-flight call when it is cancelled. The + # cancellation reaches aclose()'s caller, and the close still happens + # once the call finishes. + process = process_of(worker) + closing.cancel() + with pytest.raises(asyncio.CancelledError): + await closing + reply = await call + assert isinstance(reply, Result) + # The cancelled close finishes on its own; nothing calls aclose() again + # until it has. + await asyncio.wait_for(process.wait(), 5) + assert worker._session is None + await worker.aclose() + + +async def test_a_call_queued_behind_aclose_does_not_start_a_worker(): + worker = make_worker() + first = asyncio.ensure_future(worker.run("import time; time.sleep(0.5); 1")) + await asyncio.sleep(0.2) + queued = asyncio.ensure_future(worker.run("2")) + await asyncio.sleep(0) + await worker.aclose() + assert isinstance(await first, Result) + reply = await queued + assert isinstance(reply, Failure) + assert "closed" in reply.message + assert worker._session is None + + +async def test_aclose_waits_for_a_spawn_in_flight(): + worker = make_worker() + call = asyncio.ensure_future(worker.run("1")) + await asyncio.sleep(0) + await worker.aclose() + # However the two interleaved, the close waited the call out and no + # worker outlived it. + await call + assert worker._session is None + + +async def test_aclose_during_a_failing_spawn_does_not_deadlock(tmp_path): + script = tmp_path / "worker.py" + script.write_text("import sys\nsys.exit(3)\n") + worker = make_worker(backend=LocalBackend(worker_script=script)) + call = asyncio.ensure_future(worker.run("1")) + await asyncio.sleep(0) + closing = asyncio.ensure_future(worker.aclose()) + # asyncio.wait bounds the two together without cancelling either, so + # a deadlock shows up as a task left pending. + done, _ = await asyncio.wait({call, closing}, timeout=10) + assert done == {call, closing} + assert isinstance(call.result(), Failure) + assert worker._session is None + + +async def test_cancelling_during_the_spawn_leaves_nothing_tracked(): + async with make_worker() as worker: + call = asyncio.ensure_future(worker.run("1")) + await asyncio.sleep(0) + call.cancel() + with pytest.raises(asyncio.CancelledError): + await call + # Whether the call was cancelled before or during the spawn, + # nothing half-started survives it. + assert worker._session is None + reply = await worker.run("6 * 7") + assert isinstance(reply, Result) + assert reply.value == 42 + + +async def test_a_worker_that_died_between_calls_loses_its_scratch_directory(): + async with make_worker() as worker: + reply = await worker.run("import os; os.getcwd()") + assert isinstance(reply, Result) + scratch = reply.value + process = process_of(worker) + process.kill() + await process.wait() + # The next call spawns a replacement, and the dead worker's + # scratch directory goes with it. + reply = await worker.run("1 + 1") + assert isinstance(reply, Result) + assert not os.path.exists(scratch) + + +# Model code can reach the worker's own module, and with it the protocol +# channel; these tests use that to put a bad line on the channel on demand. +WORKER_MODULE = "import sys, os; w = sys.modules['__main__']\n" +BOGUS_REPLY = "w._send(w._protocol.Result(id='bogus', value=1, stdout='', stderr=''))\n" + + +async def test_a_reply_that_does_not_parse_restarts_the_worker(): + async with make_worker() as worker: + await worker.run("x = 5") + reply = await worker.run( + WORKER_MODULE + "os.write(w._PROTOCOL_OUT, b'garbage\\n')\n1" + ) + assert isinstance(reply, Failure) + assert "does not parse" in reply.message + reply = await worker.run("x") + assert isinstance(reply, Error) + assert "NameError" in reply.message + + +async def test_a_reply_out_of_turn_restarts_the_worker(): + # Were the stale reply trusted, the next call would read this call's + # real answer as its own. + async with make_worker() as worker: + await worker.run("x = 5") + reply = await worker.run(WORKER_MODULE + BOGUS_REPLY + "1") + assert isinstance(reply, Failure) + assert "out of turn" in reply.message + reply = await worker.run("x") + assert isinstance(reply, Error) + assert "NameError" in reply.message + + +async def test_a_worker_that_dies_on_the_interrupt_is_reported_as_crashed(): + async with make_worker(call_timeout=0.5) as worker: + code = ( + "import os, signal, time\n" + "signal.signal(signal.SIGINT, lambda *a: os._exit(1))\n" + "time.sleep(60)\n" + ) + reply = await worker.run(code) + assert isinstance(reply, Failure) + assert "crashed" in reply.message + assert worker._session is None + + +async def test_an_out_of_turn_answer_to_the_interrupt_restarts_the_worker(): + async with make_worker(call_timeout=0.5) as worker: + code = ( + WORKER_MODULE + + "import signal, time\n" + + "signal.signal(signal.SIGINT, lambda *a: " + + BOGUS_REPLY.strip() + + ")\n" + + "time.sleep(60)\n" + ) + reply = await worker.run(code) + assert isinstance(reply, Failure) + assert "out of turn" in reply.message + assert worker._session is None + + +async def test_garbage_in_answer_to_the_interrupt_restarts_the_worker(): + async with make_worker(call_timeout=0.5) as worker: + code = ( + WORKER_MODULE + + "import signal, time\n" + + "signal.signal(signal.SIGINT, " + + "lambda *a: os.write(w._PROTOCOL_OUT, b'garbage\\n'))\n" + + "time.sleep(60)\n" + ) + reply = await worker.run(code) + assert isinstance(reply, Failure) + assert "does not parse" in reply.message + assert worker._session is None + + +# Installed by a call ahead of the store's first handle, so the patched +# worker is the one the handles are synced to. Sync messages have ids +# starting with "s"; calls start with "c". +def patch_sync(replacement: str) -> str: + return ( + WORKER_MODULE + + "import time\n" + + "_real = w._execute\n" + + "def _patched(call, ns):\n" + + " if call.id.startswith('s'):\n" + + f" {replacement}\n" + + " return _real(call, ns)\n" + + "w._execute = _patched\n" + ) + + +async def test_a_sync_that_gets_no_answer_restarts_the_worker(): + store = HandleStore() + async with make_worker(call_timeout=0.5) as worker: + await worker.run(patch_sync("time.sleep(60)")) + store.register(41) + reply = await worker.run("r1", handles=store) + assert isinstance(reply, Failure) + assert "stopped answering" in reply.message + assert worker._session is None + + +async def test_a_sync_answered_with_garbage_restarts_the_worker(): + store = HandleStore() + async with make_worker() as worker: + await worker.run(patch_sync("os.write(w._PROTOCOL_OUT, b'garbage\\n')")) + store.register(41) + reply = await worker.run("r1", handles=store) + assert isinstance(reply, Failure) + assert "does not parse" in reply.message + assert worker._session is None + + +async def test_a_sync_answered_out_of_turn_restarts_the_worker(): + store = HandleStore() + async with make_worker() as worker: + bogus = "return w._protocol.Result(id='bogus', value=1, stdout='', stderr='')" + await worker.run(patch_sync(bogus)) + store.register(41) + reply = await worker.run("r1", handles=store) + assert isinstance(reply, Failure) + assert "out of turn" in reply.message + assert worker._session is None + + +async def test_a_handle_the_worker_cannot_decode_is_skipped(): + store = HandleStore() + async with make_worker() as worker: + code = ( + WORKER_MODULE + + "import json\n" + + "_real = w._protocol.decode_message\n" + + "def _refuse(line):\n" + + " if json.loads(line).get('id', '').startswith('s'):\n" + + " raise w._protocol.ProtocolError('refused')\n" + + " return _real(line)\n" + + "w._protocol.decode_message = _refuse\n" + + "x = 5\n" + ) + await worker.run(code) + store.register(41) + reply = await worker.run("'r1' in dir(), x", handles=store) + assert isinstance(reply, Result) + assert list(reply.value) == [False, 5] + # Counted as synced, so the next call does not send it again. + assert worker._synced == 1 + + +async def test_a_call_too_large_for_the_channel_keeps_the_session(monkeypatch): + from commons._execution import _protocol + + async with make_worker() as worker: + await worker.run("x = 5") + monkeypatch.setattr(_protocol, "STREAM_LIMIT", 10_000) + reply = await worker.run("y = '" + "a" * 20_000 + "'") + assert isinstance(reply, Failure) + assert "could not be sent" in reply.message + assert "remain available" in reply.message + monkeypatch.undo() + reply = await worker.run("x") + assert isinstance(reply, Result) + assert reply.value == 5 + + +async def test_a_handle_too_large_for_the_channel_is_skipped(monkeypatch): + from commons._execution import _protocol + + store = HandleStore() + async with make_worker() as worker: + await worker.run("x = 5") + # Even shrunk to its repr, this handle overruns the lowered limit; + # the call itself still fits. + monkeypatch.setattr(_protocol, "STREAM_LIMIT", 400) + store.register("a" * 100_000) + reply = await worker.run("'r1' in dir(), x", handles=store) + assert isinstance(reply, Result) + assert list(reply.value) == [False, 5] + assert worker._synced == 1 + + +class _Unstartable: + """A backend whose process cannot be created at all.""" + + async def start(self, *, network): + raise FileNotFoundError(2, "No such file or directory", "python") + + async def aclose(self) -> None: + pass + + +async def test_a_worker_that_cannot_be_created_says_it_failed_to_start(): + async with make_worker(backend=_Unstartable()) as worker: + reply = await worker.run("1") + assert isinstance(reply, Failure) + assert reply.message.startswith("the Python session failed to start:") + assert "No such file" in reply.message + + +async def test_a_line_the_worker_cannot_decode_costs_only_that_line(): + async with make_worker() as worker: + await worker.run("x = 5") + session = worker._session + assert session is not None + session.stdin.write(b"garbage\n") + reply = await worker.run("x") + assert isinstance(reply, Result) + assert reply.value == 5 + + +async def test_a_call_the_worker_cannot_decode_is_answered_with_an_error(): + # The driver is waiting on this id; without an answer it would wait out + # the whole call timeout and then the interrupt grace. + backend = LocalBackend() + session = await backend.start(network="none") + try: + ready = await asyncio.wait_for(session.stdout.readline(), 10) + assert b'"ready"' in ready + session.stdin.write(b'{"type": "call", "id": "c1", "code": 5}\n') + reply = await asyncio.wait_for(session.stdout.readline(), 5) + assert b'"error"' in reply + assert b'"c1"' in reply + assert b"could not be decoded" in reply + finally: + await session.close() + + +async def test_a_call_without_handles_leaves_the_synced_handles_alone(): + # A call that passes no store must not make the next one resend the + # store, which would overwrite a handle the model reassigned. + store = HandleStore() + store.register(41) + async with make_worker() as worker: + await worker.run("r1 = 0", handles=store) + await worker.run("1") + reply = await worker.run("r1", handles=store) + assert isinstance(reply, Result) + assert reply.value == 0 + + +async def test_a_different_store_is_synced_from_its_start(): + first, second = HandleStore(), HandleStore() + first.register(41) + second.register(7) + async with make_worker() as worker: + await worker.run("r1", handles=first) + reply = await worker.run("r1", handles=second) + assert isinstance(reply, Result) + assert reply.value == 7 + + +async def test_handles_too_large_for_one_message_all_reach_the_worker(monkeypatch): + # Each handle fits the channel on its own, but together they do not; sent + # in one message they would be shrunk to reprs. The limit is lowered in + # the driver alone, which is the side that shrinks. + import numpy as np + + from commons._execution import _protocol + + monkeypatch.setattr(_protocol, "STREAM_LIMIT", 150_000) + # About 107 KB each once encoded: one fits the lowered limit, two do not. + frame = pd.DataFrame({"x": np.random.default_rng(0).integers(0, 2**62, 10000)}) + store = HandleStore() + store.register(frame) + store.register(frame.copy()) + async with make_worker() as worker: + # A shrunk handle keeps its type's name, so a value is compared. + reply = await worker.run("int(r1['x'].iloc[-1]), int(r2['x'].iloc[-1])", handles=store) + assert isinstance(reply, Result) + last = int(frame["x"].iloc[-1]) + assert list(reply.value) == [last, last] diff --git a/pkg-py/tests/test_execution_guardrails.py b/pkg-py/tests/test_execution_guardrails.py new file mode 100644 index 00000000..1df33133 --- /dev/null +++ b/pkg-py/tests/test_execution_guardrails.py @@ -0,0 +1,179 @@ +"""Guardrails: the best-effort checks a worker gets when it cannot be sandboxed. + +Each test runs a real worker in guardrails mode. The mode is chosen by +``protection_mode()``, which on a host that can sandbox never picks it, so +the backend is built with that one decision overridden. +""" + +from __future__ import annotations + +import os + +import pytest + +from commons._execution import _backend +from commons._execution._backend import LocalBackend, Network +from commons._execution._driver import Worker +from commons._execution._protocol import Error, Result +from commons._handles import HandleStore + +pytestmark = pytest.mark.skipif( + os.name != "posix", reason="the worker's lifecycle is POSIX-shaped" +) + + +@pytest.fixture +def guarded(monkeypatch): + def make(network: Network = "none") -> Worker: + monkeypatch.setattr(_backend, "protection_mode", lambda: "guardrails") + return Worker(network=network, backend=LocalBackend(), call_timeout=10) + + return make + + +async def denied(worker: Worker, code: str) -> str: + reply = await worker.run(code) + assert isinstance(reply, Error), f"{code!r} was allowed" + assert "PermissionError" in reply.message + return reply.message + + +async def test_the_worker_runs_in_guardrails_mode(guarded): + async with guarded() as worker: + reply = await worker.run("import sys; sys.argv[2]") + assert isinstance(reply, Result) + assert reply.value == "guardrails" + + +async def test_computation_and_scratch_files_are_allowed(guarded): + async with guarded() as worker: + code = ( + "import os, decimal\n" + "with open(os.__file__) as f: first = f.readline()\n" + "with open('notes.txt', 'w') as f: f.write(first)\n" + "os.mkdir('sub'); os.rename('notes.txt', 'sub/notes.txt')\n" + "with open('sub/notes.txt') as f: copied = f.readline()\n" + "copied == first, sorted(os.listdir('.')), sum(range(11))\n" + ) + reply = await worker.run(code) + assert isinstance(reply, Result), reply + assert list(reply.value) == [True, ["sub"], 55] + + +async def test_files_outside_the_roots_are_denied(guarded, tmp_path): + outside = tmp_path / "secret.txt" + outside.write_text("secret") + async with guarded() as worker: + read = await denied(worker, f"open({str(outside)!r}).read()") + assert "denied read access" in read + write = await denied(worker, f"open({str(outside)!r}, 'w').write('x')") + assert "denied write access" in write + await denied(worker, f"import os; os.listdir({str(tmp_path)!r})") + await denied(worker, f"import os; os.remove({str(outside)!r})") + assert outside.read_text() == "secret" + + +async def test_the_interpreter_is_readable_but_not_writable(guarded): + async with guarded() as worker: + await denied(worker, "import os; open(os.__file__, 'a')") + + +async def test_a_symlink_out_of_the_scratch_directory_is_followed(guarded, tmp_path): + outside = tmp_path / "outside" + outside.mkdir() + (outside / "secret.txt").write_text("secret") + async with guarded() as worker: + # Model code cannot link to a target it may not read, so the host + # makes the link, as a file it left in the scratch directory would. + await denied(worker, f"import os; os.symlink({str(outside)!r}, 'mine')") + reply = await worker.run("import os; os.getcwd()") + assert isinstance(reply, Result) + os.symlink(outside, os.path.join(reply.value, "link")) + read = await denied(worker, "open('link/secret.txt').read()") + assert "denied read access" in read + write = await denied(worker, "open('link/new.txt', 'w')") + assert "denied write access" in write + assert not (outside / "new.txt").exists() + + +async def test_a_symlink_in_scratch_can_be_renamed_and_removed(guarded): + async with guarded() as worker: + code = ( + "import os\n" + "os.symlink(os.__file__, 'link')\n" + "os.rename('link', 'moved')\n" + "os.remove('moved')\n" + "os.listdir('.')\n" + ) + reply = await worker.run(code) + assert isinstance(reply, Result), reply + assert list(reply.value) == [] + + +async def test_a_symlink_named_with_a_trailing_slash_is_followed(guarded, tmp_path): + outside = tmp_path / "outside" + outside.mkdir() + async with guarded() as worker: + reply = await worker.run("import os; os.getcwd()") + assert isinstance(reply, Result) + os.symlink(outside, os.path.join(reply.value, "link")) + for code in ("import os; os.rmdir('link/')", "import shutil; shutil.rmtree('link/')"): + message = await denied(worker, code) + assert "denied write access" in message + assert outside.is_dir() + + +async def test_a_hard_link_to_a_file_that_is_only_readable_is_denied(guarded): + async with guarded() as worker: + message = await denied(worker, "import os; os.link(os.__file__, 'alias')") + assert "denied write access" in message + + +async def test_processes_are_denied(guarded): + async with guarded() as worker: + for code in ( + "import subprocess; subprocess.run(['true'])", + "import os; os.system('true')", + ): + message = await denied(worker, code) + assert "denied subprocess creation" in message + + +async def test_the_network_follows_the_configured_access(guarded): + code = ( + "import socket\n" + "with socket.socket() as s: s.bind(('127.0.0.1', 0)); bound = True\n" + "bound\n" + ) + async with guarded("none") as worker: + message = await denied(worker, code) + assert "denied network access" in message + await denied(worker, "import urllib.request; urllib.request.urlopen('https://example.com', timeout=1)") + async with guarded("full") as worker: + reply = await worker.run(code) + assert isinstance(reply, Result), reply + assert reply.value is True + + +async def test_a_file_url_is_denied_even_with_full_network(guarded, tmp_path): + outside = tmp_path / "secret.txt" + outside.write_text("secret") + async with guarded("full") as worker: + message = await denied( + worker, + f"import urllib.request; urllib.request.urlopen('file://{outside}').read()", + ) + assert "denied read access" in message + + +async def test_the_workers_own_work_between_calls_is_not_checked(guarded): + # A frame handle is decoded, and its reply encoded, outside model code; + # both import libraries and read files the checks would otherwise see. + pd = pytest.importorskip("pandas") + store = HandleStore() + store.register(pd.DataFrame({"x": [1, 2, 3]})) + async with guarded() as worker: + reply = await worker.run("r1['x'].sum()", handles=store) + assert isinstance(reply, Result), reply + assert reply.value == 6 + diff --git a/pkg-py/tests/test_execution_landlock.py b/pkg-py/tests/test_execution_landlock.py index 47cbabd0..fa6bc7fb 100644 --- a/pkg-py/tests/test_execution_landlock.py +++ b/pkg-py/tests/test_execution_landlock.py @@ -249,6 +249,39 @@ def test_a_package_symlinked_into_a_store_is_readable_through_its_library( assert reported["at_the_store"] == "allowed" +FILE_SYMLINK_PROBE = ( + """ +import json, os, sys, tempfile +import _landlock + +# The shape /bin has: an entry that is a symlink to a regular file, as pkill +# is to pgrep. Resolving it yields a file, which cannot take directory rights. +base = tempfile.mkdtemp() +tools = os.path.join(base, "tools") +scratch = os.path.join(base, "scratch") +os.makedirs(tools) +os.makedirs(scratch) +target = os.path.join(base, "target.txt") +with open(target, "w") as handle: + handle.write("contents") +os.symlink(target, os.path.join(tools, "link")) +_landlock.engage(["/usr", tools], [scratch]) +""" + + ATTEMPT + + """ +json.dump( + {"through_the_link": attempt(lambda: open(os.path.join(tools, "link")).close())}, + sys.stdout, +) +""" +) + + +def test_a_root_holding_a_symlink_to_a_file_engages(landlock_kernel: None) -> None: + reported = run_on_a_landlock_kernel(FILE_SYMLINK_PROBE) + assert reported["through_the_link"] == "allowed" + + MISSING_ROOT_PROBE = ( """ import json, os, sys, tempfile diff --git a/pkg-py/tests/test_measures.py b/pkg-py/tests/test_measures.py index 351b107c..d0a7558c 100644 --- a/pkg-py/tests/test_measures.py +++ b/pkg-py/tests/test_measures.py @@ -1,5 +1,7 @@ """The semantic layer: measures, their schemas, and injected arguments.""" +import __future__ + import enum import importlib import os @@ -269,6 +271,87 @@ def order_count() -> int: assert m.description == "Count of orders." +def test_source_text_execs_as_a_standalone_definition() -> None: + @measure(description="Count of orders.") + def order_count( + region: Annotated[Literal["EMEA", "AMER"], Field(description="The region.")], + ) -> int: + return 1 + + layer = semantic_layer([order_count]) + source = layer.source_text["order_count"] + # Harvested inside a function, so the raw source was indented and + # decorated; the kept text is neither, and it execs with neither the + # measure decorator nor the annotations' imports existing. + assert "@measure" not in source + namespace: dict[str, Any] = {} + exec( # noqa: S102 - the worker's exec of this text is what is tested + compile(source, "", "exec", flags=__future__.annotations.compiler_flag), + namespace, + ) + assert namespace["order_count"]("EMEA") == 1 + + +def test_source_text_keeps_a_flush_left_string_in_a_nested_measure() -> None: + @measure( + description="Count of orders.", + ) + def order_count() -> int: + query = """ +SELECT count(*) +FROM orders +""" + return len(query) + + # The string's lines sit at column 0, so there is no indent common to + # every line; the dedent still applies to the code lines around it. + source = semantic_layer([order_count]).source_text["order_count"] + assert source.startswith("def order_count(") + namespace: dict[str, Any] = {} + exec(source, namespace) # noqa: S102 - the worker's exec of this text is what is tested + assert namespace["order_count"]() == order_count() + + +def test_source_text_keeps_the_indent_inside_a_string_in_a_nested_measure() -> None: + @measure(description="Count of orders.") + def order_count() -> int: + query = """ + SELECT count(*) + FROM orders + """ + return len(query) + + # Every line shares the nested measure's indent, string lines included; + # dedenting those would change the string's value. + source = semantic_layer([order_count]).source_text["order_count"] + assert source.startswith("def order_count(") + namespace: dict[str, Any] = {} + exec(source, namespace) # noqa: S102 - the worker's exec of this text is what is tested + assert namespace["order_count"]() == order_count() + + +def test_a_helper_whose_source_does_not_parse_is_kept_as_a_comment( + tmp_path: Path, +) -> None: + # A lambda sharing a line with the tail of a bracketed expression: its + # source is that line alone, which does not parse. The worker defines a + # comment harmlessly, where the fragment would fail every spawn. + module = tmp_path / "odd_helpers.py" + module.write_text( + "from commons import measure\n" + "\n" + "LIMITS = [\n" + " 1] ; halve = lambda x: x / 2\n" + "\n" + '@measure(description="Count of orders.")\n' + "def order_count() -> int:\n" + " return 1\n" + ) + layer = semantic_layer(module) + assert layer.source_text[""] == "# source unavailable for " + assert layer.source_text["order_count"].startswith("def order_count(") + + def test_measure_leaves_the_function_callable() -> None: @measure(description="Count of orders.") def order_count() -> int: @@ -391,6 +474,7 @@ def test_measure_is_frozen() -> None: "measure_schema_text" ]["cases"] + def test_schema_fixture_is_not_empty() -> None: assert SCHEMA_CASES @@ -655,7 +739,10 @@ def test_semantic_layer_harvests_helper_source_alongside_measures() -> None: assert set(layer.source_text) >= {"double", "order_count"} assert "x * 2" in layer.source_text["double"] - assert "@measure(" in layer.source_text["order_count"] + # The harvested text execs as-is in the worker, where the decorator + # does not exist, so the decorator line is not part of it. + assert "@measure(" not in layer.source_text["order_count"] + assert layer.source_text["order_count"].startswith("def order_count(") def test_harvested_source_excludes_imported_names() -> None: