diff --git a/docs/source/environments/harbor.md b/docs/source/environments/harbor.md index ef6132939..86ea9236c 100644 --- a/docs/source/environments/harbor.md +++ b/docs/source/environments/harbor.md @@ -631,6 +631,10 @@ A backend counts as usable only if its class imports **and** Harbor's own prefli Credentials alone are not enough: a provider with valid keys but no SDK installed would otherwise report available and fail at rollout time. +Set `OPENENV_HARBOR_MAX_SANDBOX_STARTS` to cap how many sandboxes may be starting at once when a +backend throttles creation. It does not limit how many rollouts run (`MAX_CONCURRENT_ENVS` does). +On SIGINT or SIGTERM the process stops the sandboxes it started before exiting. + ### Rewards The verifier produces a dictionary. OpenEnv wants one number. The dictionary is forwarded unchanged diff --git a/envs/harbor_env/README.md b/envs/harbor_env/README.md index 52da658b2..5473b75a6 100644 --- a/envs/harbor_env/README.md +++ b/envs/harbor_env/README.md @@ -643,6 +643,10 @@ A backend counts as usable only if its class imports **and** Harbor's own prefli Credentials alone are not enough: a provider with valid keys but no SDK installed would otherwise report available and fail at rollout time. +Set `OPENENV_HARBOR_MAX_SANDBOX_STARTS` to cap how many sandboxes may be starting at once when a +backend throttles creation. It does not limit how many rollouts run (`MAX_CONCURRENT_ENVS` does). +On SIGINT or SIGTERM the process stops the sandboxes it started before exiting. + ### Rewards The verifier produces a dictionary. OpenEnv wants one number. The dictionary is forwarded unchanged diff --git a/rfcs/012-harbor-capture-providers.md b/rfcs/012-harbor-capture-providers.md index 253b9351c..e2f9957df 100644 --- a/rfcs/012-harbor-capture-providers.md +++ b/rfcs/012-harbor-capture-providers.md @@ -41,6 +41,16 @@ Upstream responses are buffered to preserve complete capture, then replayed as p A versioned external evidence report supplies stable, experimental and unstable classifications for the recorded provider/model/adapter combination. No report means unqualified; the UI requires an explicit experimental opt-in. Recorded evidence does not qualify a newly entered endpoint or pin installed agent versions. `harness_profile` selects an explicitly supported ACP or NeMo workflow on a local seam copy, without mutating global defaults. Unknown profiles fail. +### Sandbox supervision + +Harbor stops a trial's sandbox in `Trial._finalize`. `openenv.harbor.supervisor` covers what that cannot: a signal that ends the process, provider throttling of creation, and a cancellation that abandons a start mid-flight. `run_rollout` adopts each trial into one process-wide supervisor, which wraps that trial's `agent_environment.start/stop` rather than replacing Harbor's lifecycle. + +`OPENENV_HARBOR_MAX_SANDBOX_STARTS` caps how many sandboxes may be starting at once, across every event loop in the process. Unset or `0` is uncapped. It is separate from `MAX_CONCURRENT_ENVS`: a rollout holds a creation slot only while its sandbox starts, so rollout concurrency can stay above what a provider accepts for creation. + +The supervisor owns exactly the sandboxes whose start it ran and has not yet stopped. Cleanup stops only those, never by listing or labels, so other processes' sandboxes are untouched. It is safe to call repeatedly and concurrently with Harbor's own stop: each sandbox is stopped once, on the loop that started it. A cancelled start is not abandoned; it completes, and stop waits for it (up to 120 s) before deleting. That wait also applies after Harbor's own start timeout, which delays a hung start's teardown rather than racing it. + +On the first sandbox start from the main thread, handlers for SIGINT and SIGTERM are installed over whatever is there, including a server's. A signal stops owned sandboxes (bounded at 180 s: the 120 s wait for an in-flight start, then 60 s to stop), restores the previous handler and re-raises the signal to it, so default termination, `KeyboardInterrupt` and server shutdown behave as before. A second signal skips the wait. A process that starts sandboxes only from worker threads calls `install_signal_handlers()` itself; the Harbor server does so at app startup, since its UI runs each rollout in a worker thread, and a signal then stops those sandboxes on their own loops. Separate verifier environments and sandboxes on an already-closed loop are out of scope. + ### Text-level SFT export The training contract above is token-level and refuses evaluation rollouts. A second export is text-level: `openenv.harbor.sft.export_sft` (and `openenv harbor export-sft`) turns a rollout into one conversational row, the agent's last request (`request_messages`, `request_tools`) plus its reply. Supervised fine-tuning needs that conversation, not engine tokens, so this export accepts `eval` rollouts against hosted providers as well as `train` ones, and never reads token fields. It does not weaken the token contract: an `eval` rollout still cannot produce a `TrainingTrace`. diff --git a/src/openenv/harbor/rollout.py b/src/openenv/harbor/rollout.py index a7c536fd7..2037fd23d 100644 --- a/src/openenv/harbor/rollout.py +++ b/src/openenv/harbor/rollout.py @@ -43,6 +43,7 @@ HarborTurn, turns_from_document, ) +from .supervisor import get_supervisor # Harbor's own retry is deliberately off. A Harbor-level retry re-runs the agent against the SAME # capture session, so two attempts merge into one graph and the trace cross-check compares a @@ -540,6 +541,7 @@ async def run_rollout( if not proc_env: config = build_trial_config(**config_kwargs) trial = await Trial.create(config) + get_supervisor().adopt(trial) trial_result = await _run_trial(trial, on_trial_created) else: # The override has to span `run()`, not just construction. Harbor's goose wrapper reads @@ -572,6 +574,7 @@ async def run_rollout( with ctx_env_overlay(proc_env): config = build_trial_config(**config_kwargs) trial = await Trial.create(config) + get_supervisor().adopt(trial) trial_result = await _run_trial(trial, on_trial_created) else: # Fallback: mutate the real environment under a lock, one rollout at a time. @@ -580,6 +583,7 @@ async def run_rollout( with process_env(seam.name, proc_env): config = build_trial_config(**config_kwargs) trial = await Trial.create(config) + get_supervisor().adopt(trial) trial_result = await _run_trial(trial, on_trial_created) finally: _PROC_ENV_LOCK.release() diff --git a/src/openenv/harbor/serving.py b/src/openenv/harbor/serving.py index 36d44feca..46d3dd967 100644 --- a/src/openenv/harbor/serving.py +++ b/src/openenv/harbor/serving.py @@ -25,6 +25,8 @@ from openenv.core.harness.capture import CaptureServer +from .supervisor import get_supervisor + # Where the capture app is mounted when the env server hosts it directly. CAPTURE_MOUNT = "/capture" @@ -338,6 +340,10 @@ def gradio_builder( # their browser start rollouts on this server's endpoint, or reach a server on their own machine # that the page itself cannot. Nothing a real visitor does is cross-site. app.add_middleware(SameOrigin) + # The UI runs each rollout under `asyncio.run` in a worker thread, where a sandbox start cannot + # install signal handlers. Startup runs on the main thread inside uvicorn's own signal capture, + # so the supervisor wraps uvicorn's handlers and a signal stops those sandboxes on their loops. + app.router.on_startup.append(lambda: get_supervisor().install_signal_handlers()) return app diff --git a/src/openenv/harbor/supervisor.py b/src/openenv/harbor/supervisor.py new file mode 100644 index 000000000..1fad94e42 --- /dev/null +++ b/src/openenv/harbor/supervisor.py @@ -0,0 +1,363 @@ +"""Own the sandboxes this process starts: cap how many start at once, stop them on SIGINT/SIGTERM. + +Harbor already tears a sandbox down at the end of every trial: `Trial._finalize` shields +`agent_environment.stop()`. Three gaps remain once a run is large enough: + +- SIGTERM's default action kills the process before any `finally` runs, so every live sandbox is + stranded. SIGINT cancels, but `asyncio.run` closes its loop right after, which cancels the shielded + stop as well. +- Nothing caps how many sandboxes are being created at once. Providers throttle creation well below + the rollout concurrency a trainer wants, and a throttled create reads as a failed rollout. +- Cancelling a rollout while its sandbox is starting abandons the start mid-flight, and teardown + then races it. Harbor's docker backend leaves `docker compose up` running when `communicate()` is + cancelled (only its timeout path terminates the child); a remote backend can finish a create + request whose handle was dropped. Either way the sandbox can appear after `stop()` has run. + +`SandboxSupervisor.adopt(trial)` wraps the trial's own `agent_environment.start/stop`, so Harbor's +lifecycle is unchanged and the supervisor only sees sandboxes this process created. +""" + +from __future__ import annotations + +import asyncio +import collections +import os +import signal +import threading +from typing import Any + +# How long `stop()` waits for a start that is still in flight. Stopping first is what leaks, so it +# waits; but a start that never returns must not hold teardown forever. +_START_GRACE_S = 120.0 + +# How long a cleanup triggered by a signal may take before the signal is passed on anyway. A wedged +# backend must not turn Ctrl-C into a hang; a second signal skips the wait entirely. It covers a stop +# that first waits out an in-flight start, or a signal during creation would end the process with +# the start still running and nothing deleted. +_CLEANUP_TIMEOUT_S = _START_GRACE_S + 60.0 + +_SIGNALS = (signal.SIGINT, signal.SIGTERM) + + +class _StartLimit: + """At most `limit` concurrent starts across every event loop in the process. + + Not an `asyncio.Semaphore`: rollouts run on several loops (the UI runs each one under its own + `asyncio.run` in a worker thread), and an asyncio primitive binds to the first loop that waits on + it. Waiters queue FIFO and are woken on their own loop. + """ + + def __init__(self, limit: int): + self._limit = limit + self._active = 0 + self._waiters: collections.deque[asyncio.Future] = collections.deque() + self._lock = threading.Lock() + + async def acquire(self) -> None: + with self._lock: + if self._active < self._limit and not self._waiters: + self._active += 1 + return + waiter = asyncio.get_running_loop().create_future() + self._waiters.append(waiter) + try: + await waiter + except asyncio.CancelledError: + with self._lock: + if waiter in self._waiters: + self._waiters.remove(waiter) + raise + # The permit was handed over just as this waiter was cancelled: pass it on. + self.release() + raise + + def release(self) -> None: + with self._lock: + if not self._waiters: + self._active -= 1 + return + waiter = self._waiters.popleft() + waiter.get_loop().call_soon_threadsafe(self._grant, waiter) + + def _grant(self, waiter: asyncio.Future) -> None: + # A waiter cancelled after it was dequeued releases the permit itself, in `acquire`. + if not waiter.done(): + waiter.set_result(None) + + +class _Sandbox: + """One adopted agent environment and the loop its handles belong to.""" + + def __init__(self, env: Any, delete: bool): + self.env = env + self.delete = delete + self.loop: asyncio.AbstractEventLoop | None = None + self.start_task: asyncio.Future | None = None + self.stop_task: asyncio.Future | None = None + + +class _LoopHandler: + """A signal callback that was registered on an event loop with `add_signal_handler`.""" + + def __init__(self, loop: asyncio.AbstractEventLoop, callback: Any, args: tuple): + self.loop = loop + self.callback = callback + self.args = args + + +class SandboxSupervisor: + """Track the sandboxes this process starts and tear them down on request or on a signal. + + Args: + max_starts (`int`, *optional*): + How many sandboxes may be starting at once. `None` or `0` means no limit. This bounds + creation only; how many rollouts run at once is still `MAX_CONCURRENT_ENVS`, so a + rollout that has its sandbox no longer counts against it. + """ + + def __init__(self, max_starts: int | None = None): + self._limit = _StartLimit(max_starts) if max_starts else None + self._owned: set[_Sandbox] = set() + self._lock = threading.Lock() + self._previous: dict[int, Any] = {} + # Signals whose previous handler the event loop delivers by itself; see `install_signal_handlers`. + self._loop_delivers: set[int] = set() + self._handling = False + self._cleanup: asyncio.Future | None = None + # Set by `aclose` under `_lock`. A start that has not registered yet must not open a sandbox + # cleanup will never see, so it fails instead. + self._closing = False + + @property + def owned(self) -> int: + """How many sandboxes this supervisor started and has not yet stopped.""" + with self._lock: + return len(self._owned) + + def adopt(self, trial: Any) -> Any: + """Route `trial.agent_environment`'s start and stop through this supervisor. + + Call after `Trial.create` and before `trial.run()`. Harbor's own calls are unchanged; they + now pass through the creation limit, and `stop()` waits for a pending start. Concurrent + callers share cleanup; successful cleanup is cached and failed cleanup can be retried. + """ + # `getattr`, because callers and tests pass their own trial objects; one without an agent + # environment has no sandbox to supervise. + env = getattr(trial, "agent_environment", None) + if env is None or getattr(env, "_openenv_sandbox", None) is not None: + return trial + box = _Sandbox(env, delete=trial.config.environment.delete) + real_start, real_stop = env.start, env.stop + + async def start(force_build: bool) -> None: + box.loop = asyncio.get_running_loop() + self.install_signal_handlers() + if self._limit is not None: + await self._limit.acquire() + with self._lock: + closing = self._closing + if not closing: + self._owned.add(box) + if closing: + if self._limit is not None: + self._limit.release() + raise RuntimeError( + "sandbox supervisor is shutting down; not starting a new sandbox" + ) + # Shielded: a caller cancelled mid-start gets its CancelledError, but the start itself + # runs to completion so that `stop()` has a whole sandbox to delete, not half of one. + box.start_task = asyncio.ensure_future(real_start(force_build=force_build)) + if self._limit is not None: + box.start_task.add_done_callback(lambda _: self._limit.release()) + await asyncio.shield(box.start_task) + + async def stop(delete: bool) -> None: + # Share an in-flight stop and cache success, but allow failed cleanup to retry. + if box.stop_task is None or ( + box.stop_task.done() + and (box.stop_task.cancelled() or box.stop_task.exception() is not None) + ): + box.stop_task = asyncio.ensure_future(_stop_once(delete)) + await asyncio.shield(box.stop_task) + + async def _stop_once(delete: bool) -> None: + if box.start_task is not None: + await asyncio.wait({box.start_task}, timeout=_START_GRACE_S) + if not box.start_task.done(): + box.start_task.cancel() + # The backend can finish creating the sandbox while the start unwinds. + await asyncio.wait({box.start_task}) + if not box.start_task.cancelled(): + # A start failure is the trial's to report; retrieved so asyncio stays quiet. + box.start_task.exception() + await real_stop(delete=delete) + with self._lock: + self._owned.discard(box) + + env.start, env.stop = start, stop + env._openenv_sandbox = box + return trial + + async def aclose(self, timeout: float | None = None) -> None: + """Stop every sandbox this supervisor owns. Safe to call more than once, and concurrently. + + From the first call on, new starts are refused, so a rollout still queued for a creation + permit cannot open a sandbox while cleanup is running. + + Each stop runs on the loop that started the sandbox, since backends hold loop-bound handles + (docker's are asyncio subprocesses). A sandbox whose loop has already closed cannot be + reached from here and is left alone. `timeout` defaults to `_CLEANUP_TIMEOUT_S`. + """ + with self._lock: + # Same lock as the registration in `start`, so every sandbox is either in this snapshot + # or refused: none can be created after cleanup has looked. + self._closing = True + boxes = list(self._owned) + here = asyncio.get_running_loop() + pending = [] + for box in boxes: + if box.loop is here: + pending.append(asyncio.ensure_future(box.env.stop(delete=box.delete))) + elif box.loop is not None and box.loop.is_running(): + future = asyncio.run_coroutine_threadsafe( + box.env.stop(delete=box.delete), box.loop + ) + pending.append(asyncio.wrap_future(future)) + if not pending: + return + done, _ = await asyncio.wait( + pending, timeout=_CLEANUP_TIMEOUT_S if timeout is None else timeout + ) + for task in done: + task.exception() # stop failures belong to the trial; retrieved so asyncio stays quiet + + def install_signal_handlers(self, signals: tuple[int, ...] = _SIGNALS) -> bool: + """Stop owned sandboxes on these signals, then hand the signal to the previous handler. + + Called automatically when the first sandbox starts. Only the main thread can install signal + handlers, so this returns `False` elsewhere; a process that starts every sandbox from worker + threads should call it once from the main thread. The Harbor server does, at app startup + (`serving.build_app`), because the UI runs each rollout in a worker thread. Installing late + is deliberate: a server such as uvicorn replaces handlers when it starts serving, and this + one has to wrap its. + + A handler registered with `loop.add_signal_handler` (uvicorn before 0.29) is delivered by + the loop through the wakeup fd, whatever `signal.signal` holds, so wrapping that would let + the server start shutting down while cleanup is still running and then see the signal a + second time. On an asyncio loop the supervisor takes the loop's slot instead and puts the + callback back when cleanup is done. A loop that does not expose its callbacks (uvloop) + cannot be wrapped: cleanup then runs alongside the server's shutdown, and the signal is not + delivered again. + """ + if threading.current_thread() is not threading.main_thread(): + return False + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + for sig in signals: + current = signal.getsignal(sig) + # Private, but unchanged since asyncio was added: the public API can replace a loop's + # signal callback, not read it back. + handle = getattr(loop, "_signal_handlers", {}).get(sig) + if current == self._on_signal or ( + handle is not None and handle._callback == self._on_loop_signal + ): + continue + if handle is not None: + self._previous[sig] = _LoopHandler(loop, handle._callback, handle._args) + loop.add_signal_handler(sig, self._on_loop_signal, sig) + continue + if isinstance( + getattr(current, "__self__", None), asyncio.AbstractEventLoop + ): + self._loop_delivers.add(sig) + self._previous[sig] = current + signal.signal(sig, self._on_signal) + return True + + def _on_loop_signal(self, sig: int) -> None: + # Runs as an ordinary loop callback, so unlike `_on_signal` it may start the cleanup itself. + if self._handling: + self._chain(sig) + return + self._handling = True + self._start_cleanup(sig) + + def _on_signal(self, sig: int, frame: Any) -> None: + if self._handling: + # A second signal means "stop waiting". + self._chain(sig) + return + self._handling = True + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + if loop is not None: + # The handler interrupts the loop mid-callback; scheduling is all it may do safely. + loop.call_soon_threadsafe(self._start_cleanup, sig) + return + # No loop on the main thread, so the sandboxes live on other threads' loops and waiting + # for them here cannot deadlock. + worker = threading.Thread( + target=asyncio.run, args=(self.aclose(),), daemon=True + ) + worker.start() + worker.join(_CLEANUP_TIMEOUT_S + 5) + self._chain(sig) + + def _start_cleanup(self, sig: int) -> None: + async def cleanup_then_chain() -> None: + try: + await self.aclose() + finally: + if self._handling: + self._chain(sig) + + self._cleanup = asyncio.ensure_future(cleanup_then_chain()) + + def _chain(self, sig: int) -> None: + """Put the previous handler back and re-deliver the signal to it. + + Re-raising rather than calling it keeps its semantics exactly: SIG_DFL still terminates, + the default SIGINT handler still raises KeyboardInterrupt, uvicorn still shuts down. + """ + self._handling = False + previous = self._previous.pop(sig, signal.SIG_DFL) + if isinstance(previous, _LoopHandler): + if previous.loop.is_closed(): + return + previous.loop.add_signal_handler(sig, previous.callback, *previous.args) + else: + signal.signal(sig, signal.SIG_DFL if previous is None else previous) + if sig in self._loop_delivers: + # The loop already ran its callback for this signal. + self._loop_delivers.discard(sig) + return + signal.raise_signal(sig) + + +_supervisor: SandboxSupervisor | None = None +_supervisor_lock = threading.Lock() + + +def get_supervisor() -> SandboxSupervisor: + """The process-wide supervisor that `run_rollout` adopts every trial into. + + `OPENENV_HARBOR_MAX_SANDBOX_STARTS` caps concurrent sandbox creation; unset or `0` leaves it + uncapped. It is read once, on first use. + """ + global _supervisor + with _supervisor_lock: + if _supervisor is None: + raw = os.environ.get("OPENENV_HARBOR_MAX_SANDBOX_STARTS", "0") or "0" + if not raw.isdigit(): + raise ValueError( + f"OPENENV_HARBOR_MAX_SANDBOX_STARTS must be a non-negative integer, got {raw!r}" + ) + _supervisor = SandboxSupervisor(max_starts=int(raw)) + return _supervisor + + +__all__ = ["SandboxSupervisor", "get_supervisor"] diff --git a/tests/envs/test_harbor_hosted_serving.py b/tests/envs/test_harbor_hosted_serving.py index 482648ffb..1efc12ffb 100644 --- a/tests/envs/test_harbor_hosted_serving.py +++ b/tests/envs/test_harbor_hosted_serving.py @@ -121,6 +121,30 @@ def test_harbor_app_serves_its_own_state_class(monkeypatch): assert state.json()["llm_url"] == UNUSED_LLM +def test_harbor_app_installs_sandbox_signal_handlers_at_startup(monkeypatch): + """UI rollouts start sandboxes on worker threads, so the app installs the handlers itself.""" + from fastapi.testclient import TestClient + from openenv.harbor.environment import HarborEnvironment + + monkeypatch.setenv("ENABLE_WEB_INTERFACE", "false") + monkeypatch.setattr(HarborService, "_instance", None) + for name in ("_datasets", "_llm_url", "_model", "_llm"): + monkeypatch.setattr(HarborEnvironment, name, getattr(HarborEnvironment, name)) + calls: list[str] = [] + + class Spy: + def install_signal_handlers(self) -> bool: + calls.append("install") + return True + + monkeypatch.setattr(serving, "get_supervisor", lambda: Spy()) + + app = serving.build_app(datasets=[], llm_url=UNUSED_LLM) + assert calls == [] + with TestClient(app): + assert calls == ["install"] + + def test_a_failed_forwarder_does_not_leave_the_capture_server_running(monkeypatch): """A half-started service poisons every later attempt. diff --git a/tests/envs/test_harbor_supervisor.py b/tests/envs/test_harbor_supervisor.py new file mode 100644 index 000000000..aa5eb3129 --- /dev/null +++ b/tests/envs/test_harbor_supervisor.py @@ -0,0 +1,564 @@ +"""The sandbox supervisor: ownership, idempotent cleanup, the creation limit and signal chaining. + +Fake environments stand in for Harbor's; the real docker backend is exercised in +`test_harbor_supervisor_docker.py`. +""" + +from __future__ import annotations + +import asyncio +import signal +import threading +import time +from types import SimpleNamespace + +import pytest + +supervisor_mod = pytest.importorskip("openenv.harbor.supervisor") +SandboxSupervisor = supervisor_mod.SandboxSupervisor + + +@pytest.fixture(autouse=True) +def restore_signal_handlers(): + """Starting a sandbox installs handlers in this (main) thread; keep them out of other tests.""" + saved = {sig: signal.getsignal(sig) for sig in (signal.SIGINT, signal.SIGTERM)} + yield + for sig, handler in saved.items(): + signal.signal(sig, handler) + + +class FakeEnv: + def __init__(self, log: list, name: str = "env", start_s: float = 0.0): + self.log = log + self.name = name + self.start_s = start_s + + async def start(self, force_build: bool) -> None: + self.log.append((self.name, "start")) + await asyncio.sleep(self.start_s) + self.log.append((self.name, "started")) + + async def stop(self, delete: bool) -> None: + self.log.append((self.name, "stop", delete)) + + +def make_trial(env, delete: bool = True): + return SimpleNamespace( + agent_environment=env, + config=SimpleNamespace(environment=SimpleNamespace(delete=delete)), + ) + + +def stops(log): + return [entry for entry in log if entry[1] == "stop"] + + +async def test_cleanup_only_touches_sandboxes_this_supervisor_started(): + log: list = [] + mine, theirs = SandboxSupervisor(), SandboxSupervisor() + a, b, never_started = FakeEnv(log, "a"), FakeEnv(log, "b"), FakeEnv(log, "c") + mine.adopt(make_trial(a)) + mine.adopt(make_trial(never_started)) + theirs.adopt(make_trial(b)) + await a.start(force_build=False) + await b.start(force_build=False) + + await mine.aclose() + + assert stops(log) == [("a", "stop", True)] + assert mine.owned == 0 + assert theirs.owned == 1 + + +async def test_cleanup_is_safe_to_call_more_than_once(): + log: list = [] + sup = SandboxSupervisor() + env = FakeEnv(log) + sup.adopt(make_trial(env, delete=False)) + await env.start(force_build=False) + + # Signal cleanup and Harbor's own `_finalize` reaching stop at the same time. + await asyncio.gather(sup.aclose(), sup.aclose(), env.stop(delete=False)) + await sup.aclose() + await env.stop(delete=False) + + assert stops(log) == [("env", "stop", False)] + + +async def test_no_new_sandbox_starts_once_cleanup_begins(): + log: list = [] + sup = SandboxSupervisor(max_starts=1) + first, queued = FakeEnv(log, "first", start_s=0.05), FakeEnv(log, "queued") + sup.adopt(make_trial(first)) + sup.adopt(make_trial(queued)) + + first_start = asyncio.ensure_future(first.start(force_build=False)) + await asyncio.sleep(0) # `first` holds the only permit + queued_start = asyncio.ensure_future(queued.start(force_build=False)) + await asyncio.sleep(0) # `queued` is waiting for it + + await sup.aclose() + await first_start + with pytest.raises(RuntimeError, match="shutting down"): + await queued_start + + assert ("queued", "start") not in log + assert stops(log) == [("first", "stop", True)] + assert sup.owned == 0 + + +async def test_adopting_twice_wraps_once(): + log: list = [] + sup = SandboxSupervisor(max_starts=1) + trial = make_trial(FakeEnv(log)) + sup.adopt(trial) + sup.adopt(trial) + await trial.agent_environment.start(force_build=False) + await trial.agent_environment.stop(delete=True) + assert log == [("env", "start"), ("env", "started"), ("env", "stop", True)] + + +async def test_creation_limit_does_not_limit_rollouts(): + sup = SandboxSupervisor(max_starts=2) + starting = peak_starting = live = peak_live = 0 + + class CountingEnv: + async def start(self, force_build: bool) -> None: + nonlocal starting, peak_starting + starting += 1 + peak_starting = max(peak_starting, starting) + await asyncio.sleep(0.02) + starting -= 1 + + async def stop(self, delete: bool) -> None: + pass + + async def rollout(): + nonlocal live, peak_live + env = CountingEnv() + sup.adopt(make_trial(env)) + await env.start(force_build=False) + live += 1 + peak_live = max(peak_live, live) + await asyncio.sleep( + 0.2 + ) # the agent runs; its sandbox no longer holds a creation slot + live -= 1 + await env.stop(delete=True) + + await asyncio.gather(*(rollout() for _ in range(6))) + + assert peak_starting == 2 + assert peak_live == 6 + + +async def test_cancel_during_start_stops_only_after_the_start_finishes(): + log: list = [] + sup = SandboxSupervisor(max_starts=1) + env = FakeEnv(log, start_s=0.2) + sup.adopt(make_trial(env)) + + starting = asyncio.ensure_future(env.start(force_build=False)) + await asyncio.sleep(0.05) + starting.cancel() + with pytest.raises(asyncio.CancelledError): + await starting + # What Harbor's `Trial.run` does next: `_finalize` -> shielded `stop`. + await env.stop(delete=True) + + assert log == [("env", "start"), ("env", "started"), ("env", "stop", True)] + assert sup.owned == 0 + # The permit came back once the abandoned start actually finished. + nxt = FakeEnv(log, "next") + sup.adopt(make_trial(nxt)) + await asyncio.wait_for(nxt.start(force_build=False), 1) + + +async def test_stop_deletes_a_sandbox_created_while_its_start_is_cancelled(monkeypatch): + # The backend finishes creating the sandbox while its start unwinds. + monkeypatch.setattr(supervisor_mod, "_START_GRACE_S", 0) + release = asyncio.Event() + + class SlowCancelEnv(FakeEnv): + async def start(self, force_build: bool) -> None: + self.log.append((self.name, "start")) + try: + await asyncio.sleep(10) + except asyncio.CancelledError: + await release.wait() + self.log.append((self.name, "created")) + raise + + log: list = [] + sup = supervisor_mod.SandboxSupervisor() + env = SlowCancelEnv(log) + sup.adopt(make_trial(env)) + starting = asyncio.ensure_future(env.start(force_build=False)) + await asyncio.sleep(0.01) + stopping = asyncio.ensure_future(env.stop(delete=True)) + await asyncio.sleep(0.01) + release.set() + await asyncio.wait_for(stopping, 1) + with pytest.raises(asyncio.CancelledError): + await starting + + # The stop must come after the creation, or the sandbox outlives cleanup. + assert log == [("env", "start"), ("env", "created"), ("env", "stop", True)] + assert sup.owned == 0 + + +async def test_cancel_while_waiting_for_a_permit_owns_nothing(): + log: list = [] + sup = SandboxSupervisor(max_starts=1) + first, queued = FakeEnv(log, "first", start_s=0.2), FakeEnv(log, "queued") + sup.adopt(make_trial(first)) + sup.adopt(make_trial(queued)) + + running = asyncio.ensure_future(first.start(force_build=False)) + await asyncio.sleep(0.01) + waiting = asyncio.ensure_future(queued.start(force_build=False)) + await asyncio.sleep(0.01) + waiting.cancel() + with pytest.raises(asyncio.CancelledError): + await waiting + await running + + assert ("queued", "start") not in log + assert sup.owned == 1 + third = FakeEnv(log, "third") + sup.adopt(make_trial(third)) + await asyncio.wait_for(third.start(force_build=False), 1) + + +def test_creation_limit_holds_across_event_loops(): + sup = SandboxSupervisor(max_starts=1) + windows: list[tuple[float, float]] = [] + + class TimedEnv: + async def start(self, force_build: bool) -> None: + began = time.monotonic() + await asyncio.sleep(0.1) + windows.append((began, time.monotonic())) + + async def stop(self, delete: bool) -> None: + pass + + def worker(): + env = TimedEnv() + sup.adopt(make_trial(env)) + asyncio.run(env.start(force_build=False)) + + threads = [threading.Thread(target=worker) for _ in range(3)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(5) + + windows.sort() + assert len(windows) == 3 + assert all(prev[1] <= nxt[0] for prev, nxt in zip(windows, windows[1:])) + + +async def _until(predicate, timeout: float = 2.0) -> None: + deadline = time.monotonic() + timeout + while not predicate(): + assert time.monotonic() < deadline + await asyncio.sleep(0.01) + + +async def test_signal_stops_sandboxes_then_reaches_the_previous_handler(): + log: list = [] + seen: list[int] = [] + + def previous(sig, frame): + log.append(("previous", "called")) + seen.append(sig) + + signal.signal(signal.SIGTERM, previous) + sup = SandboxSupervisor() + env = FakeEnv(log) + sup.adopt(make_trial(env)) + await env.start(force_build=False) # installs the handler on top of `previous` + assert signal.getsignal(signal.SIGTERM) == sup._on_signal + + signal.raise_signal(signal.SIGTERM) + await _until(lambda: seen) + + assert log[-2:] == [("env", "stop", True), ("previous", "called")] + assert seen == [signal.SIGTERM] + assert signal.getsignal(signal.SIGTERM) is previous + + +async def test_second_signal_skips_the_wait(): + seen: list[int] = [] + release = asyncio.Event() + + class HangingEnv: + async def start(self, force_build: bool) -> None: + pass + + async def stop(self, delete: bool) -> None: + await release.wait() + + signal.signal(signal.SIGTERM, lambda sig, frame: seen.append(sig)) + sup = SandboxSupervisor() + env = HangingEnv() + sup.adopt(make_trial(env)) + await env.start(force_build=False) + + signal.raise_signal(signal.SIGTERM) + await asyncio.sleep(0.05) + assert seen == [] # still cleaning up + signal.raise_signal(signal.SIGTERM) + await _until(lambda: seen) + release.set() + await asyncio.sleep(0.05) + + assert seen == [signal.SIGTERM] # chained once, not again when cleanup finished + + +async def test_signal_reaches_a_loop_level_handler_only_after_cleanup(): + # uvicorn before 0.29 registers its handler with `loop.add_signal_handler`. The loop delivers + # that one itself, so a `signal.signal` wrapper would let it run at once, and a second time + # when the signal is chained. + log: list = [] + loop = asyncio.get_running_loop() + + def previous(sig): + log.append(("previous", sig)) + + loop.add_signal_handler(signal.SIGTERM, previous, signal.SIGTERM) + try: + sup = SandboxSupervisor() + env = FakeEnv(log) + sup.adopt(make_trial(env)) + await env.start(force_build=False) + + signal.raise_signal(signal.SIGTERM) + await _until(lambda: ("previous", signal.SIGTERM) in log) + await asyncio.sleep(0.05) + + assert log[-2:] == [("env", "stop", True), ("previous", signal.SIGTERM)] + assert log.count(("previous", signal.SIGTERM)) == 1 + assert loop._signal_handlers[signal.SIGTERM]._callback is previous + finally: + loop.remove_signal_handler(signal.SIGTERM) + + +async def test_signal_a_loop_delivers_by_itself_is_not_delivered_again(): + # uvloop does not expose its signal callbacks, so the supervisor cannot take their place. Its + # Python-level handler is a method of the loop; the loop has already run its callback by the + # time cleanup ends, and re-raising would hand the server a second signal. + class OpaqueLoop(asyncio.SelectorEventLoop): + def sighandler(self, sig, frame): + seen.append(sig) + + seen: list[int] = [] + log: list = [] + opaque = OpaqueLoop() + try: + signal.signal(signal.SIGTERM, opaque.sighandler) + sup = SandboxSupervisor() + env = FakeEnv(log) + sup.adopt(make_trial(env)) + await env.start(force_build=False) + + signal.raise_signal(signal.SIGTERM) + await _until(lambda: stops(log)) + await asyncio.sleep(0.05) + + assert signal.getsignal(signal.SIGTERM) == opaque.sighandler + assert seen == [] + finally: + opaque.close() + + +async def test_signal_stops_a_sandbox_started_on_a_worker_thread(): + # The UI runs each rollout under `asyncio.run` in a worker thread. A start there cannot install + # handlers, so the server installs them on the main thread, and cleanup reaches the worker loop. + log: list = [] + seen: list[int] = [] + started = threading.Event() + stop_threads: list[str] = [] + + class WorkerEnv(FakeEnv): + async def stop(self, delete: bool) -> None: + stop_threads.append(threading.current_thread().name) + await super().stop(delete) + + def previous(sig, frame): + log.append(("previous", "called")) + seen.append(sig) + + signal.signal(signal.SIGTERM, previous) + sup = SandboxSupervisor() + env = WorkerEnv(log) + sup.adopt(make_trial(env)) + + async def rollout() -> None: + await env.start(force_build=False) + started.set() + while not stops(log): # the rollout is still running when the signal arrives + await asyncio.sleep(0.01) + + worker = threading.Thread(target=asyncio.run, args=(rollout(),), name="rollout") + worker.start() + await asyncio.to_thread(started.wait, 2) + # The worker could not install them. + assert signal.getsignal(signal.SIGTERM) is previous + + assert sup.install_signal_handlers() + signal.raise_signal(signal.SIGTERM) + await _until(lambda: seen) + worker.join(2) + + assert log[-2:] == [("env", "stop", True), ("previous", "called")] + assert stop_threads == ["rollout"] + + +def test_signal_cleanup_outlasts_the_wait_for_an_in_flight_start(): + # `stop()` waits up to `_START_GRACE_S` for a start still in flight before deleting anything; a + # shorter signal budget would chain the fatal signal with the sandbox still being created. + assert supervisor_mod._CLEANUP_TIMEOUT_S > supervisor_mod._START_GRACE_S + + +async def test_signal_during_start_deletes_before_reaching_the_previous_handler( + monkeypatch, +): + monkeypatch.setattr(supervisor_mod, "_START_GRACE_S", 2.0) + monkeypatch.setattr(supervisor_mod, "_CLEANUP_TIMEOUT_S", 2.5) + log: list = [] + seen: list[int] = [] + + def previous(sig, frame): + log.append(("previous", "called")) + seen.append(sig) + + signal.signal(signal.SIGTERM, previous) + sup = SandboxSupervisor() + env = FakeEnv(log, start_s=0.3) + sup.adopt(make_trial(env)) + starting = asyncio.ensure_future(env.start(force_build=False)) + await _until(lambda: ("env", "start") in log) + + signal.raise_signal(signal.SIGTERM) + await _until(lambda: seen) + await starting + + assert log == [ + ("env", "start"), + ("env", "started"), + ("env", "stop", True), + ("previous", "called"), + ] + + +def test_max_starts_comes_from_the_environment(monkeypatch): + monkeypatch.setattr(supervisor_mod, "_supervisor", None) + monkeypatch.setenv("OPENENV_HARBOR_MAX_SANDBOX_STARTS", "lots") + with pytest.raises(ValueError, match="OPENENV_HARBOR_MAX_SANDBOX_STARTS"): + supervisor_mod.get_supervisor() + + monkeypatch.setenv("OPENENV_HARBOR_MAX_SANDBOX_STARTS", "3") + assert supervisor_mod.get_supervisor()._limit._limit == 3 + assert supervisor_mod.get_supervisor() is supervisor_mod.get_supervisor() + + +@pytest.mark.parametrize("retry_via_close", [False, True]) +async def test_failed_stop_retains_ownership_and_can_be_retried(retry_via_close): + sup = SandboxSupervisor() + attempts = 0 + + class FlakyEnv(FakeEnv): + async def stop(self, delete: bool) -> None: + nonlocal attempts + attempts += 1 + if attempts == 1: + raise RuntimeError("temporary backend failure") + await super().stop(delete) + + log = [] + env = FlakyEnv(log) + sup.adopt(make_trial(env)) + await env.start(force_build=False) + with pytest.raises(RuntimeError, match="temporary backend failure"): + await env.stop(delete=True) + retained = sup.owned + + if retry_via_close: + await sup.aclose() + else: + await env.stop(delete=True) + assert retained == 1 + assert attempts == 2 + assert sup.owned == 0 + assert stops(log) == [("env", "stop", True)] + await env.stop(delete=True) + assert attempts == 2 + + +async def test_cancelled_stop_attempt_can_be_retried(): + sup = SandboxSupervisor() + attempts = 0 + + class CancelledEnv(FakeEnv): + async def stop(self, delete: bool) -> None: + nonlocal attempts + attempts += 1 + if attempts == 1: + raise asyncio.CancelledError + await super().stop(delete) + + env = CancelledEnv([]) + sup.adopt(make_trial(env)) + await env.start(force_build=False) + with pytest.raises(asyncio.CancelledError): + await env.stop(delete=True) + retained = sup.owned + await env.stop(delete=True) + assert retained == 1 + assert attempts == 2 + assert sup.owned == 0 + + +async def test_retry_shares_cleanup_when_a_waiter_is_cancelled(): + sup = SandboxSupervisor() + entered = asyncio.Event() + release = asyncio.Event() + attempts = 0 + + class BlockingRetryEnv(FakeEnv): + async def stop(self, delete: bool) -> None: + nonlocal attempts + attempts += 1 + if attempts == 1: + raise RuntimeError("temporary backend failure") + entered.set() + await release.wait() + await super().stop(delete) + + log = [] + env = BlockingRetryEnv(log) + sup.adopt(make_trial(env)) + await env.start(force_build=False) + with pytest.raises(RuntimeError, match="temporary backend failure"): + await env.stop(delete=True) + + waiter = asyncio.create_task(env.stop(delete=True)) + try: + await asyncio.wait_for(entered.wait(), 1) + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + # Cancelling a waiter must not cancel or forget the shared cleanup. + assert sup.owned == 1 + await sup.aclose(timeout=0) + assert attempts == 2 + assert sup.owned == 1 + finally: + release.set() + await asyncio.gather(waiter, return_exceptions=True) + await asyncio.gather(sup.aclose(), sup.aclose(), env.stop(delete=True)) + assert attempts == 2 + assert stops(log) == [("env", "stop", True)] + assert sup.owned == 0 diff --git a/tests/envs/test_harbor_supervisor_docker.py b/tests/envs/test_harbor_supervisor_docker.py new file mode 100644 index 000000000..af7d4c627 --- /dev/null +++ b/tests/envs/test_harbor_supervisor_docker.py @@ -0,0 +1,238 @@ +"""The sandbox supervisor against Harbor's real docker backend. + +Opt-in, like the other real-Docker tests: + + OPENENV_DOCKER_INTEGRATION=1 PYTHONPATH=src:envs uv run pytest \ + tests/envs/test_harbor_supervisor_docker.py -v + +Needs Harbor (Python >= 3.12) and a locally cached `alpine:latest`; nothing is pulled. Every +container is found by its compose project, so nothing else on the host is touched. +""" + +from __future__ import annotations + +import asyncio +import os +import shutil +import signal +import subprocess +import sys +import textwrap +import time +import uuid +from pathlib import Path + +import pytest + +pytestmark = [ + pytest.mark.docker, + pytest.mark.integration, + pytest.mark.skipif( + os.environ.get("OPENENV_DOCKER_INTEGRATION") != "1" + or shutil.which("docker") is None, + reason="Set OPENENV_DOCKER_INTEGRATION=1 to run real Docker lifecycle tests", + ), +] + +pytest.importorskip("harbor") +from harbor.environments.docker.docker import ( # noqa: E402 + _sanitize_docker_compose_project_name, +) +from openenv.harbor.supervisor import SandboxSupervisor # noqa: E402 + +IMAGE = os.environ.get("OPENENV_DOCKER_TEST_IMAGE", "alpine:latest") + + +@pytest.fixture(autouse=True) +def restore_signal_handlers(): + """Starting a sandbox installs handlers in this (main) thread; keep them out of other tests.""" + saved = {sig: signal.getsignal(sig) for sig in (signal.SIGINT, signal.SIGTERM)} + yield + for sig, handler in saved.items(): + signal.signal(sig, handler) + + +# Builds a real Harbor trial, starts its sandbox through the supervisor, then waits to be signalled. +CHILD = textwrap.dedent( + """ + import asyncio, sys + from pathlib import Path + from harbor.models.trial.config import ( + AgentConfig, EnvironmentConfig, TaskConfig, TrialConfig, VerifierConfig, + ) + from harbor.trial.trial import Trial + from openenv.harbor.supervisor import get_supervisor + + async def main(): + task, trials, name = Path(sys.argv[1]), Path(sys.argv[2]), sys.argv[3] + trial = await Trial.create(TrialConfig( + task=TaskConfig(path=task), agent=AgentConfig(name="nop"), + environment=EnvironmentConfig(type="docker", delete=True), + verifier=VerifierConfig(), trial_name=name, trials_dir=trials, + )) + get_supervisor().adopt(trial) + await trial.agent_environment.start(force_build=False) + print("READY", trial.agent_environment.session_id, flush=True) + await asyncio.sleep(3600) + + asyncio.run(main()) + """ +) + + +@pytest.fixture +def task_dir(tmp_path: Path) -> Path: + task = tmp_path / "task" + (task / "environment").mkdir(parents=True) + (task / "tests").mkdir() + (task / "instruction.md").write_text("noop\n") + (task / "task.toml").write_text("") + (task / "environment" / "Dockerfile").write_text( + f'FROM {IMAGE}\nCMD ["sleep", "infinity"]\n' + ) + (task / "tests" / "test.sh").write_text("#!/bin/sh\nexit 0\n") + return task + + +def containers(session_id: str) -> list[str]: + project = _sanitize_docker_compose_project_name(session_id) + out = subprocess.run( + [ + "docker", + "ps", + "-aq", + "--filter", + f"label=com.docker.compose.project={project}", + ], + capture_output=True, + text=True, + check=True, + ) + return out.stdout.split() + + +@pytest.fixture +def bystander(): + """A container this process did not start; cleanup must leave it running.""" + name = f"oe-supervisor-bystander-{uuid.uuid4().hex[:8]}" + subprocess.run( + ["docker", "run", "-d", "--name", name, IMAGE, "sleep", "300"], + check=True, + capture_output=True, + ) + yield name + subprocess.run(["docker", "rm", "-f", name], capture_output=True) + + +@pytest.mark.parametrize("sig", [signal.SIGTERM, signal.SIGINT]) +def test_signal_tears_down_owned_sandbox_only(sig, task_dir, tmp_path, bystander): + script = tmp_path / "child.py" + script.write_text(CHILD) + name = f"oe-sup-{uuid.uuid4().hex[:8]}" + child = subprocess.Popen( + [sys.executable, str(script), str(task_dir), str(tmp_path / "trials"), name], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + session_id = None + try: + for line in child.stdout: + if line.startswith("READY "): + session_id = line.split()[1] + break + assert session_id, "child exited before its sandbox was up" + assert containers(session_id), "sandbox should be running before the signal" + + child.send_signal(sig) + child.wait(timeout=120) + + # The signal still reached its original handler: SIG_DFL for SIGTERM, and for SIGINT + # asyncio's KeyboardInterrupt, which CPython turns into death-by-SIGINT on exit. + assert child.returncode == -sig + assert containers(session_id) == [] + running = subprocess.run( + ["docker", "inspect", "-f", "{{.State.Running}}", bystander], + capture_output=True, + text=True, + ) + assert running.stdout.strip() == "true" + finally: + if child.poll() is None: + child.kill() + if session_id: + for container in containers(session_id): + subprocess.run(["docker", "rm", "-f", container], capture_output=True) + + +async def _trial(task_dir: Path, trials: Path): + from harbor.models.trial.config import ( + AgentConfig, + EnvironmentConfig, + TaskConfig, + TrialConfig, + VerifierConfig, + ) + from harbor.trial.trial import Trial + + return await Trial.create( + TrialConfig( + task=TaskConfig(path=task_dir), + agent=AgentConfig(name="nop"), + environment=EnvironmentConfig(type="docker", delete=True), + verifier=VerifierConfig(), + trial_name=f"oe-sup-{uuid.uuid4().hex[:8]}", + trials_dir=trials, + ) + ) + + +async def test_creation_limit_serialises_real_starts(task_dir, tmp_path): + sup = SandboxSupervisor(max_starts=1) + windows: list[tuple[float, float]] = [] + trials = [await _trial(task_dir, tmp_path / "trials") for _ in range(3)] + for trial in trials: + real = trial.agent_environment.start + + async def timed(force_build: bool, real=real) -> None: + began = time.monotonic() + await real(force_build=force_build) + windows.append((began, time.monotonic())) + + trial.agent_environment.start = timed + sup.adopt(trial) + try: + await asyncio.gather( + *(t.agent_environment.start(force_build=False) for t in trials) + ) + assert sup.owned == 3 + finally: + await sup.aclose() + + windows.sort() + assert all(prev[1] <= nxt[0] for prev, nxt in zip(windows, windows[1:])) + assert sup.owned == 0 + for trial in trials: + assert containers(trial.agent_environment.session_id) == [] + + +async def test_cancel_during_creation_leaves_no_container(task_dir, tmp_path): + sup = SandboxSupervisor() + trial = await _trial(task_dir, tmp_path / "trials") + sup.adopt(trial) + env = trial.agent_environment + + starting = asyncio.ensure_future(env.start(force_build=False)) + # Cancel once `docker compose up` is under way, so the start is abandoned mid-flight. + deadline = time.monotonic() + 120 + while not containers(env.session_id): + assert time.monotonic() < deadline and not starting.done() + await asyncio.sleep(0.05) + starting.cancel() + with pytest.raises(asyncio.CancelledError): + await starting + # Harbor's `_finalize` stops the environment after a cancelled start. + await env.stop(delete=True) + + assert sup.owned == 0 + assert containers(env.session_id) == []