add retries when starting alcatraz computer

Author: thesofakillers

project/paperbench/paperbench/nano/eval.py

@@ -11,7 +11,6 @@
 from dotenv import load_dotenv
 
 load_dotenv()
-from nanoeval_alcatraz.alcatraz_computer_interface import AlcatrazComputerInterface
 from nanoeval_alcatraz.task_to_alcatraz_config import task_to_alcatraz_config
 from typing_extensions import override
 
@@ -43,6 +42,7 @@
 from paperbench.nano.task import PBTask
 from paperbench.nano.utils import SPLIT_TO_EXPECTED_PAPERS, gather_eval_runs
 from paperbench.paper_registry import paper_registry
+from paperbench.scripts.alcatraz_services import start_alcatraz_computer
 from paperbench.utils import (
     create_run_dir,
     create_run_id,
@@ -57,6 +57,7 @@
 
 MIN_UPLOAD_INTERVAL_MESSAGES = 5
 MIN_UPLOAD_INTERVAL_SECONDS = 1800
+MAX_CLUSTER_START_ATTEMPTS = 5
 
 
 GRADER_OPENAI_API_KEY = os.getenv("GRADER_OPENAI_API_KEY") or os.getenv("OPENAI_API_KEY")
@@ -166,8 +167,11 @@ async def _start_computer(self, task: PBTask) -> AsyncGenerator[ComputerInterfac
             is_nvidia_gpu_env=self.is_nvidia_gpu_env,
         )
 
-        async with alcatraz_config.build() as cluster:
-            yield AlcatrazComputerInterface(cluster_value=cluster)
+        async with start_alcatraz_computer(
+            cluster_config=alcatraz_config,
+            max_attempts=MAX_CLUSTER_START_ATTEMPTS,
+        ) as computer:
+            yield computer
 
     @override
     async def run(self, task: ComputerTask) -> AsyncGenerator[Step | FinalResult, None]:

project/paperbench/paperbench/nano/task.py

@@ -383,7 +383,6 @@ async def _run_reproduce(self, submission: str) -> ReproductionMetadata | None:
                 metadata = replace(metadata, executed_submission=reproduce_output_path)
                 return metadata
 
-        # Reproduce on alcatraz and collect metadata
         try:
             metadata = await reproduce_on_computer_with_salvaging(
                 cluster_config=self.reproduction.cluster_config,

project/paperbench/paperbench/scripts/alcatraz_services.py

@@ -4,7 +4,9 @@
 import structlog
 from dotenv import load_dotenv
 from nanoeval_alcatraz.alcatraz_computer_interface import AlcatrazComputerInterface
+from tenacity import RetryCallState, Retrying, retry_if_exception_type, stop_after_attempt
 
+from alcatraz.clusters.interface import AlcatrazException
 from alcatraz.clusters.local import ClusterConfig
 from nanoeval.solvers.computer_tasks.code_execution_interface import ComputerInterface
 from paperbench.infra.alcatraz import put_file_in_computer
@@ -58,7 +60,24 @@ async def put_submission_in_computer(
 @asynccontextmanager
 async def start_alcatraz_computer(
     cluster_config: ClusterConfig,
+    max_attempts: int = 5,
 ) -> AsyncGenerator[ComputerInterface, None]:
     """Helper method for starting an AlcatrazComputerInterface given a ClusterConfig."""
-    async with cluster_config.build() as cluster:
-        yield AlcatrazComputerInterface(cluster_value=cluster)
+
+    def before_sleep(state: RetryCallState) -> None:
+        exception = state.outcome.exception() if state.outcome else None
+        logger.warning(
+            f"Cluster start failed on attempt {state.attempt_number} out of {max_attempts}; "
+            f"retrying due to '{exception}'"
+        )
+
+    retrying = Retrying(
+        stop=stop_after_attempt(max_attempts),
+        retry=retry_if_exception_type(AlcatrazException),
+        before_sleep=before_sleep,
+    )
+
+    for attempt in retrying:
+        with attempt:
+            async with cluster_config.build() as cluster:
+                yield AlcatrazComputerInterface(cluster_value=cluster)