feat: add max time config to costeer in data science (microsoft#645)

* add max time config to costeer

* fix a small bug

---------

Co-authored-by: Xu Yang <xuyang1@microsoft.com>

Author: peteryang1

rdagent/components/coder/CoSTEER/__init__.py

@@ -1,4 +1,5 @@
 import pickle
+from datetime import datetime
 from pathlib import Path
 
 from rdagent.components.coder.CoSTEER.config import CoSTEERSettings
@@ -35,6 +36,7 @@ def __init__(
     ) -> None:
         super().__init__(*args, **kwargs)
         self.max_loop = settings.max_loop if max_loop is None else max_loop
+        self.max_seconds = settings.max_seconds
         self.knowledge_base_path = (
             Path(settings.knowledge_base_path) if settings.knowledge_base_path is not None else None
         )
@@ -96,11 +98,14 @@ def develop(self, exp: Experiment) -> Experiment:
             knowledge_self_gen=self.knowledge_self_gen,
         )
 
+        start_datetime = datetime.now()
         for evo_exp in self.evolve_agent.multistep_evolve(evo_exp, self.evaluator):
             assert isinstance(evo_exp, Experiment)  # multiple inheritance
             logger.log_object(evo_exp.sub_workspace_list, tag="evolving code")
             for sw in evo_exp.sub_workspace_list:
                 logger.info(f"evolving code workspace: {sw}")
+            if (datetime.now() - start_datetime).seconds > self.max_seconds:
+                break
 
         if self.with_feedback and self.filter_final_evo:
             evo_exp = self._exp_postprocess_by_feedback(evo_exp, self.evolve_agent.evolving_trace[-1].feedback)

rdagent/components/coder/CoSTEER/config.py

@@ -35,5 +35,7 @@ class Config:
 
     select_threshold: int = 10
 
+    max_seconds: int = 10**6
+
 
 CoSTEER_SETTINGS = CoSTEERSettings()

rdagent/components/coder/data_science/conf.py

@@ -0,0 +1,10 @@
+from rdagent.components.coder.CoSTEER.config import CoSTEERSettings
+
+
+class DSCoderCoSTEERSettings(CoSTEERSettings):
+    """Data Science CoSTEER settings"""
+
+    class Config:
+        env_prefix = "DS_Coder_CoSTEER_"
+
+    max_seconds: int = 2400

rdagent/components/coder/data_science/ensemble/__init__.py

@@ -14,7 +14,6 @@
 import json
 
 from rdagent.components.coder.CoSTEER import CoSTEER
-from rdagent.components.coder.CoSTEER.config import CoSTEER_SETTINGS
 from rdagent.components.coder.CoSTEER.evaluators import (
     CoSTEERMultiEvaluator,
     CoSTEERSingleFeedback,
@@ -25,6 +24,7 @@
 from rdagent.components.coder.CoSTEER.knowledge_management import (
     CoSTEERQueriedKnowledge,
 )
+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.core.exception import CoderError
@@ -122,7 +122,8 @@ def __init__(
         *args,
         **kwargs,
     ) -> None:
+        settings = DSCoderCoSTEERSettings()
         eva = CoSTEERMultiEvaluator(EnsembleCoSTEEREvaluator(scen=scen), scen=scen)
-        es = EnsembleMultiProcessEvolvingStrategy(scen=scen, settings=CoSTEER_SETTINGS)
+        es = EnsembleMultiProcessEvolvingStrategy(scen=scen, settings=settings)
 
-        super().__init__(*args, settings=CoSTEER_SETTINGS, eva=eva, es=es, evolving_version=2, scen=scen, **kwargs)
+        super().__init__(*args, settings=settings, eva=eva, es=es, evolving_version=2, scen=scen, **kwargs)

rdagent/components/coder/data_science/feature/__init__.py

@@ -1,7 +1,6 @@
 import json
 
 from rdagent.components.coder.CoSTEER import CoSTEER
-from rdagent.components.coder.CoSTEER.config import CoSTEER_SETTINGS
 from rdagent.components.coder.CoSTEER.evaluators import (
     CoSTEERMultiEvaluator,
     CoSTEERSingleFeedback,
@@ -12,6 +11,7 @@
 from rdagent.components.coder.CoSTEER.knowledge_management import (
     CoSTEERQueriedKnowledge,
 )
+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.core.exception import CoderError
@@ -107,9 +107,10 @@ def __init__(
         *args,
         **kwargs,
     ) -> None:
+        settings = DSCoderCoSTEERSettings()
         eva = CoSTEERMultiEvaluator(
             FeatureCoSTEEREvaluator(scen=scen), scen=scen
         )  # Please specify whether you agree running your eva in parallel or not
-        es = FeatureMultiProcessEvolvingStrategy(scen=scen, settings=CoSTEER_SETTINGS)
+        es = FeatureMultiProcessEvolvingStrategy(scen=scen, settings=settings)
 
-        super().__init__(*args, settings=CoSTEER_SETTINGS, eva=eva, es=es, evolving_version=2, scen=scen, **kwargs)
+        super().__init__(*args, settings=settings, eva=eva, es=es, evolving_version=2, scen=scen, **kwargs)

rdagent/components/coder/data_science/model/__init__.py

@@ -1,10 +1,4 @@
-import json
-from pathlib import Path
-
-from jinja2 import Environment, StrictUndefined
-
 from rdagent.components.coder.CoSTEER import CoSTEER
-from rdagent.components.coder.CoSTEER.config import CoSTEER_SETTINGS
 from rdagent.components.coder.CoSTEER.evaluators import (
     CoSTEERMultiEvaluator,
     CoSTEERSingleFeedback,
@@ -15,6 +9,7 @@
 from rdagent.components.coder.CoSTEER.knowledge_management import (
     CoSTEERQueriedKnowledge,
 )
+from rdagent.components.coder.data_science.conf import DSCoderCoSTEERSettings
 from rdagent.components.coder.data_science.model.eval import (
     ModelGeneralCaseSpecEvaluator,
 )
@@ -133,10 +128,11 @@ def __init__(
         *args,
         **kwargs,
     ) -> None:
+        settings = DSCoderCoSTEERSettings()
         eva = CoSTEERMultiEvaluator(
             ModelGeneralCaseSpecEvaluator(scen=scen), scen=scen
         )  # Please specify whether you agree running your eva in parallel or not
         # eva = ModelGeneralCaseSpecEvaluator(scen=scen)
-        es = ModelMultiProcessEvolvingStrategy(scen=scen, settings=CoSTEER_SETTINGS)
+        es = ModelMultiProcessEvolvingStrategy(scen=scen, settings=settings)
 
-        super().__init__(*args, settings=CoSTEER_SETTINGS, eva=eva, es=es, evolving_version=2, scen=scen, **kwargs)
+        super().__init__(*args, settings=settings, eva=eva, es=es, evolving_version=2, scen=scen, **kwargs)

rdagent/components/coder/data_science/raw_data_loader/__init__.py

@@ -25,7 +25,6 @@
 import json
 
 from rdagent.components.coder.CoSTEER import CoSTEER
-from rdagent.components.coder.CoSTEER.config import CoSTEER_SETTINGS
 from rdagent.components.coder.CoSTEER.evaluators import (
     CoSTEERMultiEvaluator,
     CoSTEERSingleFeedback,
@@ -35,8 +34,8 @@
 )
 from rdagent.components.coder.CoSTEER.knowledge_management import (
     CoSTEERQueriedKnowledge,
-    CoSTEERQueriedKnowledgeV2,
 )
+from rdagent.components.coder.data_science.conf import DSCoderCoSTEERSettings
 from rdagent.components.coder.data_science.raw_data_loader.eval import (
     DataLoaderCoSTEEREvaluator,
 )
@@ -187,9 +186,10 @@ def __init__(
         *args,
         **kwargs,
     ) -> None:
+        settings = DSCoderCoSTEERSettings()
         eva = CoSTEERMultiEvaluator(
             DataLoaderCoSTEEREvaluator(scen=scen), scen=scen
         )  # Please specify whether you agree running your eva in parallel or not
-        es = DataLoaderMultiProcessEvolvingStrategy(scen=scen, settings=CoSTEER_SETTINGS)
+        es = DataLoaderMultiProcessEvolvingStrategy(scen=scen, settings=settings)
 
-        super().__init__(*args, settings=CoSTEER_SETTINGS, eva=eva, es=es, evolving_version=2, scen=scen, **kwargs)
+        super().__init__(*args, settings=settings, eva=eva, es=es, evolving_version=2, scen=scen, **kwargs)

rdagent/components/coder/data_science/workflow/__init__.py

@@ -1,7 +1,6 @@
 import json
 
 from rdagent.components.coder.CoSTEER import CoSTEER
-from rdagent.components.coder.CoSTEER.config import CoSTEER_SETTINGS
 from rdagent.components.coder.CoSTEER.evaluators import (
     CoSTEERMultiEvaluator,
     CoSTEERSingleFeedback,
@@ -12,6 +11,7 @@
 from rdagent.components.coder.CoSTEER.knowledge_management import (
     CoSTEERQueriedKnowledge,
 )
+from rdagent.components.coder.data_science.conf import DSCoderCoSTEERSettings
 from rdagent.components.coder.data_science.workflow.eval import (
     WorkflowGeneralCaseSpecEvaluator,
 )
@@ -110,8 +110,9 @@ def __init__(
         *args,
         **kwargs,
     ) -> None:
+        settings = DSCoderCoSTEERSettings()
         eva = CoSTEERMultiEvaluator(
             WorkflowGeneralCaseSpecEvaluator(scen=scen), scen=scen
         )  # Please specify whether you agree running your eva in parallel or not
-        es = WorkflowMultiProcessEvolvingStrategy(scen=scen, settings=CoSTEER_SETTINGS)
-        super().__init__(*args, settings=CoSTEER_SETTINGS, eva=eva, es=es, evolving_version=2, scen=scen, **kwargs)
+        es = WorkflowMultiProcessEvolvingStrategy(scen=scen, settings=settings)
+        super().__init__(*args, settings=settings, eva=eva, es=es, evolving_version=2, scen=scen, **kwargs)

rdagent/core/evolving_agent.py

@@ -85,7 +85,6 @@ def multistep_evolve(
                     evolving_trace=self.evolving_trace,
                     queried_knowledge=queried_knowledge,
                 )
-                yield evo  # yield the control to caller for process control and logging.
 
                 # 4. Pack evolve results
                 es = EvoStep(evo, queried_knowledge)
@@ -100,6 +99,8 @@ def multistep_evolve(
                 # 6. update trace
                 self.evolving_trace.append(es)
 
+                yield evo  # yield the control to caller for process control and logging.
+
                 # 7. check if all tasks are completed
                 if self.with_feedback and es.feedback:
                     logger.info("All tasks in evolving subject have been completed.")