feat: Add runtime measurement for each step and loop in RDLoop. (#281)

* Add runtime measurement for each step and loop in RDLoop.

* refine some codes

* refine the code (#276)

* show variables only when it exists (#277)

* fix: support seed and fix absolute path (#278)

* fix: support seed and fix absolute path

* Absolute path

* lint

* fix: improve_execution_time_in_kaggle_loop (#279)

* improve_execution_time_in_kaggle_loop

* fix CI

* fix CI

* fix CI

* fix: Update runner.py to fix a small bug (#282)

* fix: Update runner.py to fix a small bug

* fix CI

* refine the code

* Update loop.py

* Update rd_loop.py

* Update model_xgb.py

---------

Co-authored-by: XianBW <36835909+XianBW@users.noreply.github.com>
Co-authored-by: you-n-g <you-n-g@users.noreply.github.com>
Co-authored-by: Xu Yang <peteryang@vip.qq.com>

rdagent/app/kaggle/loop.py

@@ -17,16 +17,15 @@ from rdagent.core.proposal import (
 from rdagent.core.scenario import Scenario
 from rdagent.core.utils import import_class
 from rdagent.log import rdagent_logger as logger
-from rdagent.scenarios.kaggle.knowledge_management.vector_base import (
-    KaggleExperienceBase,
-)
+from rdagent.log.time import measure_time
 from rdagent.scenarios.kaggle.proposal.proposal import (
     KG_ACTION_FEATURE_ENGINEERING,
     KG_ACTION_FEATURE_PROCESSING,
 )
 
 
-class ModelRDLoop(RDLoop):
+class KaggleRDLoop(RDLoop):
+    @measure_time
     def __init__(self, PROP_SETTING: BasePropSetting):
         with logger.tag("init"):
             scen: Scenario = import_class(PROP_SETTING.scen)(PROP_SETTING.competition)
@@ -53,6 +52,7 @@ class ModelRDLoop(RDLoop):
             self.trace = Trace(scen=scen)
             super(RDLoop, self).__init__()
 
+    @measure_time
     def coding(self, prev_out: dict[str, Any]):
         with logger.tag("d"):  # develop
             if prev_out["propose"].action in [KG_ACTION_FEATURE_ENGINEERING, KG_ACTION_FEATURE_PROCESSING]:
@@ -62,6 +62,7 @@ class ModelRDLoop(RDLoop):
             logger.log_object(exp.sub_workspace_list, tag="coder result")
         return exp
 
+    @measure_time
     def running(self, prev_out: dict[str, Any]):
         with logger.tag("ef"):  # evaluate and feedback
             if prev_out["propose"].action in [KG_ACTION_FEATURE_ENGINEERING, KG_ACTION_FEATURE_PROCESSING]:
@@ -89,10 +90,10 @@ def main(path=None, step_n=None, competition=None):
     if competition:
         KAGGLE_IMPLEMENT_SETTING.competition = competition
     if path is None:
-        model_loop = ModelRDLoop(KAGGLE_IMPLEMENT_SETTING)
+        kaggle_loop = KaggleRDLoop(KAGGLE_IMPLEMENT_SETTING)
     else:
-        model_loop = ModelRDLoop.load(path)
-    model_loop.run(step_n=step_n)
+        kaggle_loop = KaggleRDLoop.load(path)
+    kaggle_loop.run(step_n=step_n)
 
 
 if __name__ == "__main__":

rdagent/app/qlib_rd_loop/factor.py

@@ -10,11 +10,13 @@ from rdagent.app.qlib_rd_loop.conf import FACTOR_PROP_SETTING
 from rdagent.components.workflow.rd_loop import RDLoop
 from rdagent.core.exception import FactorEmptyError
 from rdagent.log import rdagent_logger as logger
+from rdagent.log.time import measure_time
 
 
 class FactorRDLoop(RDLoop):
     skip_loop_error = (FactorEmptyError,)
 
+    @measure_time
     def running(self, prev_out: dict[str, Any]):
         with logger.tag("ef"):  # evaluate and feedback
             exp = self.runner.develop(prev_out["coding"])

rdagent/app/qlib_rd_loop/factor_from_report.py

@@ -14,6 +14,7 @@ from rdagent.components.document_reader.document_reader import (
 from rdagent.core.prompts import Prompts
 from rdagent.core.proposal import Hypothesis
 from rdagent.log import rdagent_logger as logger
+from rdagent.log.time import measure_time
 from rdagent.oai.llm_utils import APIBackend
 from rdagent.scenarios.qlib.experiment.factor_experiment import QlibFactorExperiment
 from rdagent.scenarios.qlib.factor_experiment_loader.pdf_loader import (
@@ -101,6 +102,7 @@ def extract_hypothesis_and_exp_from_reports(report_file_path: str) -> Tuple[Qlib
 
 
 class FactorReportLoop(FactorRDLoop, metaclass=LoopMeta):
+    @measure_time
     def __init__(self, report_folder: str = None):
         super().__init__(PROP_SETTING=FACTOR_FROM_REPORT_PROP_SETTING)
         if report_folder is None:
@@ -116,6 +118,7 @@ class FactorReportLoop(FactorRDLoop, metaclass=LoopMeta):
         self.current_loop_exp = None
         self.steps = ["propose_hypo_exp", "propose", "exp_gen", "coding", "running", "feedback"]
 
+    @measure_time
     def propose_hypo_exp(self, prev_out: dict[str, Any]):
         with logger.tag("r"):
             while True:
@@ -137,9 +140,11 @@ class FactorReportLoop(FactorRDLoop, metaclass=LoopMeta):
                 self.current_loop_exp = exp
                 return None
 
+    @measure_time
     def propose(self, prev_out: dict[str, Any]):
         return self.current_loop_hypothesis
 
+    @measure_time
     def exp_gen(self, prev_out: dict[str, Any]):
         return self.current_loop_exp
 

rdagent/components/workflow/rd_loop.py

@@ -3,6 +3,7 @@ Model workflow with session control
 It is from `rdagent/app/qlib_rd_loop/model.py` and try to replace `rdagent/app/qlib_rd_loop/RDAgent.py`
 """
 
+import time
 from typing import Any
 
 from rdagent.components.workflow.conf import BasePropSetting
@@ -16,10 +17,12 @@ from rdagent.core.proposal import (
 from rdagent.core.scenario import Scenario
 from rdagent.core.utils import import_class
 from rdagent.log import rdagent_logger as logger
+from rdagent.log.time import measure_time
 from rdagent.utils.workflow import LoopBase, LoopMeta
 
 
 class RDLoop(LoopBase, metaclass=LoopMeta):
+    @measure_time
     def __init__(self, PROP_SETTING: BasePropSetting):
         with logger.tag("init"):
             scen: Scenario = import_class(PROP_SETTING.scen)()
@@ -41,30 +44,35 @@ class RDLoop(LoopBase, metaclass=LoopMeta):
             self.trace = Trace(scen=scen)
             super().__init__()
 
+    @measure_time
     def propose(self, prev_out: dict[str, Any]):
         with logger.tag("r"):  # research
             hypothesis = self.hypothesis_gen.gen(self.trace)
             logger.log_object(hypothesis, tag="hypothesis generation")
         return hypothesis
 
+    @measure_time
     def exp_gen(self, prev_out: dict[str, Any]):
         with logger.tag("r"):  # research
             exp = self.hypothesis2experiment.convert(prev_out["propose"], self.trace)
             logger.log_object(exp.sub_tasks, tag="experiment generation")
         return exp
 
+    @measure_time
     def coding(self, prev_out: dict[str, Any]):
         with logger.tag("d"):  # develop
             exp = self.coder.develop(prev_out["exp_gen"])
             logger.log_object(exp.sub_workspace_list, tag="coder result")
         return exp
 
+    @measure_time
     def running(self, prev_out: dict[str, Any]):
         with logger.tag("ef"):  # evaluate and feedback
             exp = self.runner.develop(prev_out["coding"])
             logger.log_object(exp, tag="runner result")
         return exp
 
+    @measure_time
     def feedback(self, prev_out: dict[str, Any]):
         feedback = self.summarizer.generate_feedback(prev_out["running"], prev_out["propose"], self.trace)
         with logger.tag("ef"):  # evaluate and feedback

rdagent/log/time.py

@@ -0,0 +1,19 @@
+import time
+from functools import wraps
+
+from rdagent.log import rdagent_logger as logger
+
+
+def measure_time(method):
+    @wraps(method)
+    def timed(*args, **kwargs):
+        start_time = time.time()
+        result = method(*args, **kwargs)
+        end_time = time.time()
+        duration = end_time - start_time
+        method_name = method.__name__
+        # logger.log_object(f"{method_name} took {duration:.2f} sec")
+        logger.info(f"{method_name} took {duration:.2f} sec")
+        return result
+
+    return timed

rdagent/scenarios/kaggle/experiment/workspace.py

@@ -75,8 +75,6 @@ class KGFBWorkspace(FBWorkspace):
 
         csv_path = self.workspace_path / "submission_score.csv"
 
-        logger.info(self.workspace_path)
-
         if not csv_path.exists():
             logger.error(f"File {csv_path} does not exist.")
             return None