diff --git a/verifiers/v1/rollout.py b/verifiers/v1/rollout.py index b3aca523a..dfde47dc1 100644 --- a/verifiers/v1/rollout.py +++ b/verifiers/v1/rollout.py @@ -81,6 +81,7 @@ def __init__( self._interception = interception self.runtime = runtime self._borrowed_runtime = runtime + self._borrow_lock: asyncio.Lock | None = None self.trace: Trace = Trace( task=TraceTask( type=type(task).__name__, @@ -204,6 +205,9 @@ async def open(self) -> bool: if self._borrowed_runtime is None: runtime.env = runtime_env else: + if runtime.network_restricted: + await runtime.borrow_lock.acquire() + self._borrow_lock = runtime.borrow_lock runtime = runtime.with_env(runtime_env) self.runtime = runtime if self.task.data.prompt is None and not self._has_user: @@ -423,17 +427,22 @@ async def abort(self) -> None: (a cancellation mid-setup, a lifetime bug raised to the caller) means the driver will never reach `close()`. Safe after a partial `close()`.""" self._closed = True - if self._harness_session is not None: - with contextlib.suppress(Exception): - await self._harness_session.close() - with contextlib.suppress(Exception): - await self._stack.aclose() - if self.runtime is not None: - with contextlib.suppress(Exception): - await self.harness.cleanup(self.trace, self.runtime) - if self._borrowed_runtime is None and self.runtime is not None: + try: + if self._harness_session is not None: + with contextlib.suppress(Exception): + await self._harness_session.close() with contextlib.suppress(Exception): - await self.runtime.stop() + await self._stack.aclose() + if self.runtime is not None: + with contextlib.suppress(Exception): + await self.harness.cleanup(self.trace, self.runtime) + if self._borrowed_runtime is None and self.runtime is not None: + with contextlib.suppress(Exception): + await self.runtime.stop() + finally: + if self._borrow_lock is not None: + self._borrow_lock.release() + self._borrow_lock = None async def close(self) -> Trace: """Finish the rollout: tool servers and interception down, task `finalize` @@ -522,6 +531,9 @@ async def close(self) -> Trace: logger.warning( "runtime teardown failed (rollout %s)", trace.id, exc_info=True ) + if self._borrow_lock is not None: + self._borrow_lock.release() + self._borrow_lock = None logger.info( "rollout done: id=%s task=%s reward=%.3f turns=%d stop=%s", trace.id, diff --git a/verifiers/v1/runtimes/base.py b/verifiers/v1/runtimes/base.py index 0460749bc..eb38639c2 100644 --- a/verifiers/v1/runtimes/base.py +++ b/verifiers/v1/runtimes/base.py @@ -153,6 +153,9 @@ def __init__(self, name: str | None = None) -> None: # Per-run task values live on the runtime rather than its serializable config/info. # Explicit process values (model credentials, proxy settings, etc.) override these. self.env: dict[str, str] = {} + # Restricted runtimes have one sandbox-wide network policy. Borrowers hold + # this lock for their full rollout so setup cannot widen another agent's policy. + self.borrow_lock = asyncio.Lock() self._uv_interpreters: dict[str, str] = {} self._uv_script_locks: dict[str, asyncio.Lock] = {} self._setup_claimed = False