From d9576706cd2c2fea4c62f94cdba50b2a30aa01cf Mon Sep 17 00:00:00 2001
From: Xu Yang <peteryang@vip.qq.com>
Date: Mon, 17 Feb 2025 17:36:05 +0800
Subject: [PATCH] fix: docker error will trigger retry and data science runner
 loop set to 3 (#602)

* docker error will trigger retry
data science runner loop set to 3

* fix CI
---
 rdagent/components/coder/CoSTEER/__init__.py  |  3 ++-
 .../data_science/dev/runner/__init__.py       |  5 +++-
 rdagent/utils/env.py                          | 27 +++++++++++++++++--
 3 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/rdagent/components/coder/CoSTEER/__init__.py b/rdagent/components/coder/CoSTEER/__init__.py
index 9c66ba96..245e7feb 100644
--- a/rdagent/components/coder/CoSTEER/__init__.py
+++ b/rdagent/components/coder/CoSTEER/__init__.py
@@ -30,10 +30,11 @@ class CoSTEER(Developer[Experiment]):
         with_feedback: bool = True,
         knowledge_self_gen: bool = True,
         filter_final_evo: bool = True,
+        max_loop: int | None = None,
         **kwargs,
     ) -> None:
         super().__init__(*args, **kwargs)
-        self.max_loop = settings.max_loop
+        self.max_loop = settings.max_loop if max_loop is None else max_loop
         self.knowledge_base_path = (
             Path(settings.knowledge_base_path) if settings.knowledge_base_path is not None else None
         )
diff --git a/rdagent/scenarios/data_science/dev/runner/__init__.py b/rdagent/scenarios/data_science/dev/runner/__init__.py
index 0feef719..efb88196 100644
--- a/rdagent/scenarios/data_science/dev/runner/__init__.py
+++ b/rdagent/scenarios/data_science/dev/runner/__init__.py
@@ -90,7 +90,10 @@ class DSCoSTEERRunner(CoSTEER):
         )  # Please specify whether you agree running your eva in parallel or not
         es = DSRunnerMultiProcessEvolvingStrategy(scen=scen, settings=CoSTEER_SETTINGS)
 
-        super().__init__(*args, settings=CoSTEER_SETTINGS, eva=eva, es=es, evolving_version=2, scen=scen, **kwargs)
+        # In runner, we don't need very big loops, so we set max_loop to 3
+        super().__init__(
+            *args, settings=CoSTEER_SETTINGS, eva=eva, es=es, evolving_version=2, scen=scen, max_loop=3, **kwargs
+        )
 
     def develop(self, exp):
         bak_sub_tasks = exp.sub_tasks
diff --git a/rdagent/utils/env.py b/rdagent/utils/env.py
index 7f3741cb..3f80d7c9 100644
--- a/rdagent/utils/env.py
+++ b/rdagent/utils/env.py
@@ -150,6 +150,9 @@ class DockerConf(ExtendedBaseSettings):
 
     enable_cache: bool = True  # enable the cache mechanism
 
+    retry_count: int = 5  # retry count for the docker run
+    retry_wait_seconds: int = 10  # retry wait seconds for the docker run
+
 
 class QlibDockerConf(DockerConf):
     model_config = ExtendedSettingsConfigDict(env_prefix="QLIB_DOCKER_")
@@ -404,6 +407,24 @@ class DockerEnv(Env[DockerConf]):
         except docker.errors.APIError as e:
             raise RuntimeError(f"Error while running the container: {e}")
 
+    def __run_with_retry(
+        self,
+        entry: str | None = None,
+        local_path: str = ".",
+        env: dict | None = None,
+        running_extra_volume: dict | None = None,
+        remove_timestamp: bool = True,
+    ) -> str:
+        for retry_index in range(self.conf.retry_count):
+            try:
+                return self.__run(entry, local_path, env, running_extra_volume, remove_timestamp)
+            except Exception as e:
+                logger.warning(
+                    f"Error while running the container: {e}, current try index: {retry_index + 1}, {self.conf.retry_count - retry_index - 1} retries left."
+                )
+                time.sleep(self.conf.retry_wait_seconds)
+        raise RuntimeError("Error while running the container. Retry count exceeded.")
+
     def zip_a_folder_into_a_file(self, folder_path: str, zip_file_path: str) -> None:
         """
         Zip a folder into a file, use zipfile instead of subprocess
@@ -468,7 +489,7 @@ class DockerEnv(Env[DockerConf]):
                 ret: str = pickle.load(f)
             self.unzip_a_file_into_a_folder(str(target_folder / f"{key}.zip"), local_path)
         else:
-            ret = self.__run(entry, local_path, env, running_extra_volume, remove_timestamp)
+            ret = self.__run_with_retry(entry, local_path, env, running_extra_volume, remove_timestamp)
             with open(target_folder / f"{key}.pkl", "wb") as f:
                 pickle.dump(ret, f)
             self.zip_a_folder_into_a_file(local_path, str(target_folder / f"{key}.zip"))
@@ -491,7 +512,9 @@ class DockerEnv(Env[DockerConf]):
         if self.conf.enable_cache:
             out = self.cached_run(entry_add_timeout, local_path, env, running_extra_volume)
         else:
-            out = self.__run(entry_add_timeout, local_path, env, running_extra_volume, remove_timestamp=False)
+            out = self.__run_with_retry(
+                entry_add_timeout, local_path, env, running_extra_volume, remove_timestamp=False
+            )
         end = time.time()
 
         if end - start + 1 >= self.conf.running_timeout_period: