diff --git a/project/paperbench/paperbench/computer_utils.py b/project/paperbench/paperbench/computer_utils.py index 2bac6351..e36181c3 100644 --- a/project/paperbench/paperbench/computer_utils.py +++ b/project/paperbench/paperbench/computer_utils.py @@ -78,15 +78,19 @@ def before_sleep(state: RetryCallState) -> None: f"retrying due to '{exception}'" ) - async for attempt in AsyncRetrying( - stop=stop_after_attempt(max_attempts), - retry=retry_if_exception_type(exception_types if exception_types else ()), - before_sleep=before_sleep, - reraise=True, - ): - with attempt: - async with computer_runtime.run(computer_config) as computer: - yield computer + async with AsyncExitStack() as stack: + computer: ComputerInterface | None = None + async for attempt in AsyncRetrying( + stop=stop_after_attempt(max_attempts), + retry=retry_if_exception_type(exception_types if exception_types else ()), + before_sleep=before_sleep, + reraise=True, + ): + with attempt: + computer = await stack.enter_async_context(computer_runtime.run(computer_config)) + + assert computer is not None + yield computer class ReleasableComputer(ComputerInterface): diff --git a/project/paperbench/tests/unit/test_computer_retry.py b/project/paperbench/tests/unit/test_computer_retry.py new file mode 100644 index 00000000..8bd4a04e --- /dev/null +++ b/project/paperbench/tests/unit/test_computer_retry.py @@ -0,0 +1,56 @@ +from contextlib import asynccontextmanager + +import pytest + +from paperbench.computer_utils import start_computer_with_retry + + +class FakeRuntime: + def __init__(self, startup_failures: int = 0) -> None: + self.startup_failures = startup_failures + self.enter_count = 0 + self.exit_count = 0 + self.computer = object() + + @asynccontextmanager + async def run(self, _config): + self.enter_count += 1 + if self.enter_count <= self.startup_failures: + raise ValueError("startup failed") + try: + yield self.computer + finally: + self.exit_count += 1 + + +@pytest.mark.asyncio +async def test_start_computer_retries_only_startup_failures() -> None: + runtime = FakeRuntime(startup_failures=2) + + async with start_computer_with_retry( + runtime, + object(), + exception_types=ValueError, + max_attempts=3, + ) as computer: + assert computer is runtime.computer + + assert runtime.enter_count == 3 + assert runtime.exit_count == 1 + + +@pytest.mark.asyncio +async def test_start_computer_does_not_retry_caller_body() -> None: + runtime = FakeRuntime() + + with pytest.raises(ValueError, match="caller failed"): + async with start_computer_with_retry( + runtime, + object(), + exception_types=ValueError, + max_attempts=3, + ): + raise ValueError("caller failed") + + assert runtime.enter_count == 1 + assert runtime.exit_count == 1