From 6ec56a3f6cf8c59da6fb883bc39438e2bcbbe66c Mon Sep 17 00:00:00 2001
From: Xu Yang <peteryang@vip.qq.com>
Date: Fri, 24 Jan 2025 13:34:47 +0800
Subject: [PATCH] feat: add restart and fix unzip (#538)

* add restart function

* update

* fix ci

* update code

* fix restart

---------

Co-authored-by: yuanteli <1957922024@qq.com>
Co-authored-by: Xu Yang <xuyang1@microsoft.com>
---
 rdagent/app/data_science/conf.py           |  2 ++
 rdagent/app/data_science/loop.py           | 13 +++++++++++++
 rdagent/scenarios/kaggle/kaggle_crawler.py | 17 +++++++++++++----
 rdagent/utils/env.py                       |  2 +-
 4 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/rdagent/app/data_science/conf.py b/rdagent/app/data_science/conf.py
index 25a92f3e..1b72c247 100644
--- a/rdagent/app/data_science/conf.py
+++ b/rdagent/app/data_science/conf.py
@@ -45,5 +45,7 @@ class DataScienceBasePropSetting(KaggleBasePropSetting):
     summarizer: str = "rdagent.scenarios.kaggle.developer.feedback.KGExperiment2Feedback"
     """Summarizer class"""
 
+    consecutive_errors: int = 5
+
 
 DS_RD_SETTING = DataScienceBasePropSetting()
diff --git a/rdagent/app/data_science/loop.py b/rdagent/app/data_science/loop.py
index 3d1f8449..a2b5566d 100644
--- a/rdagent/app/data_science/loop.py
+++ b/rdagent/app/data_science/loop.py
@@ -118,6 +118,19 @@ class DataScienceRDLoop(RDLoop):
                     ExperimentFeedback.from_exception(e),
                 )
             )
+            if len(self.trace.hist) > DS_RD_SETTING.consecutive_errors:
+                cons_errors = 0
+                for _, feedback in reversed(self.trace.hist[-DS_RD_SETTING.consecutive_errors :]):
+                    if feedback == ExperimentFeedback.from_exception(e) or feedback.reason.startswith(
+                        "The experiment fails due to"
+                    ):
+                        cons_errors += 1
+                    else:
+                        break
+                if cons_errors == DS_RD_SETTING.consecutive_errors + 1:
+                    logger.error("Consecutive errors reached the limit. Dumping trace.")
+                    logger.log_object(self.trace, tag="trace before restart")
+                    self.trace = DSTrace(scen=self.trace.scen, knowledge_base=self.trace.knowledge_base)
         logger.log_object(self.trace, tag="trace")
         logger.log_object(self.trace.sota_experiment(), tag="SOTA experiment")
 
diff --git a/rdagent/scenarios/kaggle/kaggle_crawler.py b/rdagent/scenarios/kaggle/kaggle_crawler.py
index 265d4d21..0fac2c5a 100644
--- a/rdagent/scenarios/kaggle/kaggle_crawler.py
+++ b/rdagent/scenarios/kaggle/kaggle_crawler.py
@@ -128,10 +128,19 @@ def download_data(competition: str, settings: ExtendedBaseSettings = KAGGLE_IMPL
             (Path(local_path) / competition).mkdir(parents=True, exist_ok=True)
 
             mleb_env.run(f"cp -r ./zip_files/{competition}/prepared/public/* ./{competition}", local_path=local_path)
-            mleb_env.run(
-                f'for zip_file in ./{competition}/*.zip; do dir_name="${{zip_file%.zip}}"; mkdir -p "$dir_name"; unzip -o "$zip_file" -d "$dir_name"; done',
-                local_path=local_path,
-            )
+
+            for zip_path in (Path(local_path) / competition).rglob("*.zip"):
+                with zipfile.ZipFile(zip_path, "r") as zip_ref:
+                    if len(zip_ref.namelist()) == 1:
+                        mleb_env.run(
+                            f"unzip -o ./{zip_path.relative_to(local_path)} -d {zip_path.parent.relative_to(local_path)}",
+                            local_path=local_path,
+                        )
+                    else:
+                        mleb_env.run(
+                            f"mkdir -p ./{zip_path.parent.relative_to(local_path)}/{zip_path.stem}; unzip -o ./{zip_path.relative_to(local_path)} -d ./{zip_path.parent.relative_to(local_path)}/{zip_path.stem}",
+                            local_path=local_path,
+                        )
             # NOTE:
             # Patching:  due to mle has special renaming mechanism for different competition;
             # We have to switch the schema back to a uniform one;
diff --git a/rdagent/utils/env.py b/rdagent/utils/env.py
index 23d89f9d..3de734ff 100644
--- a/rdagent/utils/env.py
+++ b/rdagent/utils/env.py
@@ -464,7 +464,7 @@ 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, local_path, env, running_extra_volume, remove_timestamp=False)
+            out = self.__run(entry_add_timeout, local_path, env, running_extra_volume, remove_timestamp=False)
         end = time.time()
 
         if end - start + 1 >= self.conf.running_timeout_period: