feat: refactor CoSTEER classes to use DSCoSTEER and update max seconds handling (#1156)

peteryang1 · web-flow · commit 6d01e3e1ca1e · 2025-08-05T18:16:10.000+08:00
* feat: refactor CoSTEER classes to use DSCoSTEER and update max seconds handling

* remove useless line

* enable time_ratio_limit_to_enable_hyperparameter_tuning
diff --git a/rdagent/app/kaggle/conf.py b/rdagent/app/kaggle/conf.py
@@ -75,8 +75,8 @@ class KaggleBasePropSetting(ExtendedBaseSettings):
     mini_case: bool = False
     """Enable mini-case study for experiments"""
 
-    time_ratio_limit_to_enable_hyperparameter_tuning: int | None = None
-    """Time ratio limit to enable hyperparameter tuning, if None, hyperparameter tuning is always enabled in the first evolution."""
+    time_ratio_limit_to_enable_hyperparameter_tuning: float = 1
+    """Time ratio limit to enable hyperparameter tuning, if not change, hyperparameter tuning is always enabled in the first evolution."""
 
 
 KAGGLE_IMPLEMENT_SETTING = KaggleBasePropSetting()
diff --git a/rdagent/components/coder/CoSTEER/__init__.py b/rdagent/components/coder/CoSTEER/__init__.py
@@ -1,4 +1,3 @@
-import pickle
 from copy import deepcopy
 from datetime import datetime
 from pathlib import Path
@@ -26,7 +25,6 @@ def __init__(
         es: EvolvingStrategy,
         *args,
         evolving_version: int = 2,
-        max_seconds: int | None = None,
         with_knowledge: bool = True,
         knowledge_self_gen: bool = True,
         max_loop: int | None = None,
@@ -36,7 +34,6 @@ def __init__(
         self.settings = settings
 
         self.max_loop = settings.max_loop if max_loop is None else max_loop
-        self.max_seconds = max_seconds
         self.knowledge_base_path = (
             Path(settings.knowledge_base_path) if settings.knowledge_base_path is not None else None
         )
@@ -67,6 +64,13 @@ def __init__(
             )
         )
 
+    def get_develop_max_seconds(self) -> int | None:
+        """
+        Get the maximum seconds for the develop task.
+        Sub classes might override this method to provide a different value.
+        """
+        return None
+
     def _get_last_fb(self) -> CoSTEERMultiFeedback:
         fb = self.evolve_agent.evolving_trace[-1].feedback
         assert fb is not None, "feedback is None"
@@ -76,6 +80,7 @@ def _get_last_fb(self) -> CoSTEERMultiFeedback:
     def develop(self, exp: Experiment) -> Experiment:
 
         # init intermediate items
+        max_seconds = self.get_develop_max_seconds()
         evo_exp = EvolvingItem.from_experiment(exp)
 
         self.evolve_agent = RAGEvoAgent[EvolvingItem](
@@ -102,8 +107,8 @@ def develop(self, exp: Experiment) -> Experiment:
             logger.log_object(evo_exp.sub_workspace_list, tag="evolving code")
             for sw in evo_exp.sub_workspace_list:
                 logger.info(f"evolving workspace: {sw}")
-            if self.max_seconds is not None and (datetime.now() - start_datetime).seconds > self.max_seconds:
-                logger.info(f"Reached max time limit {self.max_seconds} seconds, stop evolving")
+            if max_seconds is not None and (datetime.now() - start_datetime).seconds > max_seconds:
+                logger.info(f"Reached max time limit {max_seconds} seconds, stop evolving")
                 reached_max_seconds = True
                 break
             if RD_Agent_TIMER_wrapper.timer.started and RD_Agent_TIMER_wrapper.timer.is_timeout():
@@ -140,7 +145,7 @@ def _exp_postprocess_by_feedback(self, evo: Experiment, feedback: CoSTEERMultiFe
         failed_feedbacks = [
             f"- feedback{index + 1:02d}:\n  - execution: {f.execution}\n  - return_checking: {f.return_checking}\n  - code: {f.code}"
             for index, f in enumerate(feedback)
-            if f is not None and not f.final_decision
+            if f is not None and not f.is_acceptable()
         ]
 
         if len(failed_feedbacks) == len(feedback):
diff --git a/rdagent/components/coder/data_science/ensemble/__init__.py b/rdagent/components/coder/data_science/ensemble/__init__.py
@@ -16,7 +16,6 @@
 from jinja2 import Environment, StrictUndefined
 
 from rdagent.app.data_science.conf import DS_RD_SETTING
-from rdagent.components.coder.CoSTEER import CoSTEER
 from rdagent.components.coder.CoSTEER.evaluators import (
     CoSTEERMultiEvaluator,
     CoSTEERSingleFeedback,
@@ -30,6 +29,7 @@
 from rdagent.components.coder.data_science.conf import DSCoderCoSTEERSettings
 from rdagent.components.coder.data_science.ensemble.eval import EnsembleCoSTEEREvaluator
 from rdagent.components.coder.data_science.ensemble.exp import EnsembleTask
+from rdagent.components.coder.data_science.share.ds_costeer import DSCoSTEER
 from rdagent.core.exception import CoderError
 from rdagent.core.experiment import FBWorkspace
 from rdagent.core.scenario import Scenario
@@ -141,7 +141,7 @@ def assign_code_list_to_evo(self, code_list: list[dict[str, str]], evo):
         return evo
 
 
-class EnsembleCoSTEER(CoSTEER):
+class EnsembleCoSTEER(DSCoSTEER):
     def __init__(
         self,
         scen: Scenario,
@@ -160,6 +160,5 @@ def __init__(
             evolving_version=2,
             scen=scen,
             max_loop=DS_RD_SETTING.coder_max_loop,
-            max_seconds=scen.real_debug_timeout() * settings.max_seconds_multiplier,
             **kwargs,
         )
diff --git a/rdagent/components/coder/data_science/feature/__init__.py b/rdagent/components/coder/data_science/feature/__init__.py
@@ -1,9 +1,6 @@
-import json
 from pathlib import Path
-from typing import Dict
 
 from rdagent.app.data_science.conf import DS_RD_SETTING
-from rdagent.components.coder.CoSTEER import CoSTEER
 from rdagent.components.coder.CoSTEER.evaluators import (
     CoSTEERMultiEvaluator,
     CoSTEERSingleFeedback,
@@ -17,6 +14,7 @@
 from rdagent.components.coder.data_science.conf import DSCoderCoSTEERSettings
 from rdagent.components.coder.data_science.feature.eval import FeatureCoSTEEREvaluator
 from rdagent.components.coder.data_science.feature.exp import FeatureTask
+from rdagent.components.coder.data_science.share.ds_costeer import DSCoSTEER
 from rdagent.core.exception import CoderError
 from rdagent.core.experiment import FBWorkspace
 from rdagent.core.scenario import Scenario
@@ -117,7 +115,7 @@ def assign_code_list_to_evo(self, code_list: list[dict[str, str]], evo):
         return evo
 
 
-class FeatureCoSTEER(CoSTEER):
+class FeatureCoSTEER(DSCoSTEER):
     def __init__(
         self,
         scen: Scenario,
@@ -138,6 +136,5 @@ def __init__(
             evolving_version=2,
             scen=scen,
             max_loop=DS_RD_SETTING.coder_max_loop,
-            max_seconds=scen.real_debug_timeout() * settings.max_seconds_multiplier,
             **kwargs,
         )
diff --git a/rdagent/components/coder/data_science/model/__init__.py b/rdagent/components/coder/data_science/model/__init__.py
@@ -1,8 +1,6 @@
 from pathlib import Path
-from typing import Dict
 
 from rdagent.app.data_science.conf import DS_RD_SETTING
-from rdagent.components.coder.CoSTEER import CoSTEER
 from rdagent.components.coder.CoSTEER.evaluators import (
     CoSTEERMultiEvaluator,
     CoSTEERSingleFeedback,
@@ -18,6 +16,7 @@
     ModelGeneralCaseSpecEvaluator,
 )
 from rdagent.components.coder.data_science.model.exp import ModelTask
+from rdagent.components.coder.data_science.share.ds_costeer import DSCoSTEER
 from rdagent.core.exception import CoderError
 from rdagent.core.experiment import FBWorkspace
 from rdagent.core.scenario import Scenario
@@ -148,7 +147,7 @@ def assign_code_list_to_evo(self, code_list: list[dict[str, str]], evo):
         return evo
 
 
-class ModelCoSTEER(CoSTEER):
+class ModelCoSTEER(DSCoSTEER):
     def __init__(
         self,
         scen: Scenario,
@@ -170,6 +169,5 @@ def __init__(
             evolving_version=2,
             scen=scen,
             max_loop=DS_RD_SETTING.coder_max_loop,
-            max_seconds=scen.real_debug_timeout() * settings.max_seconds_multiplier,
             **kwargs,
         )
diff --git a/rdagent/components/coder/data_science/pipeline/__init__.py b/rdagent/components/coder/data_science/pipeline/__init__.py
@@ -25,7 +25,6 @@
 from pathlib import Path
 
 from rdagent.app.data_science.conf import DS_RD_SETTING
-from rdagent.components.coder.CoSTEER import CoSTEER
 from rdagent.components.coder.CoSTEER.evaluators import (
     CoSTEERMultiEvaluator,
     CoSTEERSingleFeedback,
@@ -39,6 +38,7 @@
 from rdagent.components.coder.data_science.conf import DSCoderCoSTEERSettings
 from rdagent.components.coder.data_science.pipeline.eval import PipelineCoSTEEREvaluator
 from rdagent.components.coder.data_science.pipeline.exp import PipelineTask
+from rdagent.components.coder.data_science.share.ds_costeer import DSCoSTEER
 from rdagent.components.coder.data_science.share.eval import ModelDumpEvaluator
 from rdagent.core.exception import CoderError
 from rdagent.core.experiment import FBWorkspace
@@ -130,7 +130,7 @@ def assign_code_list_to_evo(self, code_list: list[dict[str, str]], evo):
         return evo
 
 
-class PipelineCoSTEER(CoSTEER):
+class PipelineCoSTEER(DSCoSTEER):
     def __init__(
         self,
         scen: Scenario,
@@ -159,6 +159,5 @@ def __init__(
             evolving_version=2,
             scen=scen,
             max_loop=DS_RD_SETTING.coder_max_loop,
-            max_seconds=scen.real_debug_timeout() * settings.max_seconds_multiplier,
             **kwargs,
         )
diff --git a/rdagent/components/coder/data_science/raw_data_loader/__init__.py b/rdagent/components/coder/data_science/raw_data_loader/__init__.py
@@ -22,13 +22,10 @@
     - Each coder could be tested.
 """
 
-import json
 import re
 from pathlib import Path
-from typing import Dict
 
 from rdagent.app.data_science.conf import DS_RD_SETTING
-from rdagent.components.coder.CoSTEER import CoSTEER
 from rdagent.components.coder.CoSTEER.evaluators import (
     CoSTEERMultiEvaluator,
     CoSTEERSingleFeedback,
@@ -47,6 +44,7 @@
     DataLoaderCoSTEEREvaluator,
 )
 from rdagent.components.coder.data_science.raw_data_loader.exp import DataLoaderTask
+from rdagent.components.coder.data_science.share.ds_costeer import DSCoSTEER
 from rdagent.core.exception import CoderError
 from rdagent.core.experiment import FBWorkspace
 from rdagent.core.scenario import Scenario
@@ -197,7 +195,7 @@ def assign_code_list_to_evo(self, code_list: list[dict[str, str]], evo):
         return evo
 
 
-class DataLoaderCoSTEER(CoSTEER):
+class DataLoaderCoSTEER(DSCoSTEER):
     def __init__(
         self,
         scen: Scenario,
@@ -218,7 +216,6 @@ def __init__(
             evolving_version=2,
             scen=scen,
             max_loop=DS_RD_SETTING.coder_max_loop,
-            max_seconds=scen.real_debug_timeout() * settings.max_seconds_multiplier,
             **kwargs,
         )
 
diff --git a/rdagent/components/coder/data_science/share/ds_costeer.py b/rdagent/components/coder/data_science/share/ds_costeer.py
@@ -0,0 +1,9 @@
+from rdagent.components.coder.CoSTEER import CoSTEER
+
+
+class DSCoSTEER(CoSTEER):
+    def get_develop_max_seconds(self) -> int | None:
+        """
+        The coder uses the scenario's real debug timeout as the maximum seconds for development.
+        """
+        return int(self.scen.real_debug_timeout() * self.settings.max_seconds_multiplier)
diff --git a/rdagent/components/coder/data_science/workflow/__init__.py b/rdagent/components/coder/data_science/workflow/__init__.py
@@ -1,8 +1,4 @@
-import json
-from typing import Dict
-
 from rdagent.app.data_science.conf import DS_RD_SETTING
-from rdagent.components.coder.CoSTEER import CoSTEER
 from rdagent.components.coder.CoSTEER.evaluators import (
     CoSTEERMultiEvaluator,
     CoSTEERSingleFeedback,
@@ -14,6 +10,7 @@
     CoSTEERQueriedKnowledge,
 )
 from rdagent.components.coder.data_science.conf import DSCoderCoSTEERSettings
+from rdagent.components.coder.data_science.share.ds_costeer import DSCoSTEER
 from rdagent.components.coder.data_science.workflow.eval import (
     WorkflowGeneralCaseSpecEvaluator,
 )
@@ -111,7 +108,7 @@ def assign_code_list_to_evo(self, code_list: list[dict[str, str]], evo):
         return evo
 
 
-class WorkflowCoSTEER(CoSTEER):
+class WorkflowCoSTEER(DSCoSTEER):
     def __init__(
         self,
         scen: Scenario,
@@ -131,6 +128,5 @@ def __init__(
             evolving_version=2,
             scen=scen,
             max_loop=DS_RD_SETTING.coder_max_loop,
-            max_seconds=scen.real_debug_timeout() * settings.max_seconds_multiplier,
             **kwargs,
         )
diff --git a/rdagent/core/conf.py b/rdagent/core/conf.py
@@ -56,6 +56,7 @@ class RDAgentSettings(ExtendedBaseSettings):
     # workspace conf
     workspace_path: Path = Path.cwd() / "git_ignore_folder" / "RD-Agent_workspace"
     workspace_ckp_size_limit: int = 0
+    workspace_ckp_white_list_names: list[str] | None = None
     """
     the checkpoint for the workspace is a zip file.
     0 (or any value <=0) means *no* size limit for files in workspace checkpoints
diff --git a/rdagent/core/experiment.py b/rdagent/core/experiment.py
@@ -318,7 +318,10 @@ def create_ws_ckp(self) -> None:
                     zf.writestr(zi, str(file_path.readlink()))
                 elif file_path.is_file():
                     size_limit = RD_AGENT_SETTINGS.workspace_ckp_size_limit
-                    if size_limit <= 0 or file_path.stat().st_size <= size_limit:
+                    if (
+                        RD_AGENT_SETTINGS.workspace_ckp_white_list_names is not None
+                        and file_path.name in RD_AGENT_SETTINGS.workspace_ckp_white_list_names
+                    ) or (size_limit <= 0 or file_path.stat().st_size <= size_limit):
                         zf.write(file_path, file_path.relative_to(self.workspace_path))
         self.ws_ckp = buf.getvalue()
 
diff --git a/rdagent/scenarios/data_science/dev/runner/__init__.py b/rdagent/scenarios/data_science/dev/runner/__init__.py
@@ -162,10 +162,15 @@ def __init__(
             evolving_version=2,
             scen=scen,
             max_loop=DS_RD_SETTING.runner_max_loop,
-            max_seconds=scen.real_full_timeout() * settings.max_seconds_multiplier,
             **kwargs,
         )
 
+    def get_develop_max_seconds(self) -> int | None:
+        """
+        The coder uses the scenario's real debug timeout as the maximum seconds for development.
+        """
+        return int(self.scen.real_full_timeout() * self.settings.max_seconds_multiplier)
+
     def develop(self, exp):
         bak_sub_tasks = exp.pending_tasks_list
         exp.sub_tasks = [
diff --git a/rdagent/scenarios/data_science/dev/runner/eval.py b/rdagent/scenarios/data_science/dev/runner/eval.py
@@ -76,13 +76,6 @@ def evaluate(
         queried_knowledge: QueriedKnowledge = None,
         **kwargs,
     ) -> DSRunnerFeedback:
-        # Only enalbe hyperparameter tuning on the first evaluation.
-        # Avoid too much time cunsumming.
-        if len(queried_knowledge.task_to_former_failed_traces[target_task.get_task_information()][0]) == 0:
-            enable_hyperparameter_tuning_check = True
-        else:
-            enable_hyperparameter_tuning_check = False
-
         env = get_ds_env(
             extra_volumes={
                 f"{DS_RD_SETTING.local_data_path}/{self.scen.competition}": T(
@@ -168,6 +161,14 @@ def evaluate(
             stdout += f"\n### Submission check:\n{submission_check_out}\nIf Submission check returns a 'Submission is valid' or similar message, despite some warning messages, you should still consider the submission as valid and give a positive final decision. "
 
         time_spent_ratio = implementation.running_info.running_time / env.conf.running_timeout_period
+        # Only enable hyperparameter tuning on the first evaluation.
+        # Avoid too much time consuming.
+        enable_hyperparameter_tuning_check = False
+        if len(queried_knowledge.task_to_former_failed_traces[target_task.get_task_information()][0]) == 0 and (
+            time_spent_ratio < DS_RD_SETTING.time_ratio_limit_to_enable_hyperparameter_tuning
+        ):
+            enable_hyperparameter_tuning_check = True
+
         if (
             DS_RD_SETTING.time_ratio_limit_to_enable_hyperparameter_tuning is not None
             and time_spent_ratio > DS_RD_SETTING.time_ratio_limit_to_enable_hyperparameter_tuning
