feat: idea pool integrated to exp_gen & add timer to RD-Agent & pause-resume to RD-loops (#795)

* update all code

* update all code

* dump knowledge base

* rename the tag

* add timer to RD-Agent

* fix CI

* fix CI

* use batch embedding

* fix a small bug

* fix prompt bug

* feat: add pause resume to handle K8S cluster pause (#804)

* add resume to cluster running

* fix non-pickle problem

* fix a small bug

* fix a small bug

* avoid shutil move error

* refine the logic

* move knowledge base out of session

* avoid mistake information to pipeline coding

* avoid load and dump in steps

* archive the right folder

* small improvement

* avoid restart when timer is already started

* fix CI

---------

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

---------

Co-authored-by: Xu Yang <peteryang@vip.qq.com>
Co-authored-by: Xu Yang <xuyang1@microsoft.com>
Co-authored-by: Xu <v-xuminrui@microsoft.com>

rdagent/app/data_science/conf.py

@@ -41,5 +41,18 @@ class DataScienceBasePropSetting(KaggleBasePropSetting):
     enable_doc_dev: bool = False
     model_dump_check_level: Literal["medium", "high"] = "medium"
 
+    ### knowledge base
+    enable_knowledge_base: bool = False
+    knowledge_base_version: str = "v1"
+    knowledge_base_path: str | None = None
+    idea_pool_json_path: str | None = None
+
+    ### archive log folder after each loop
+    enable_log_archive: bool = True
+    log_archive_path: str | None = None
+    log_archive_temp_path: str | None = (
+        None  # This is to store the mid tar file since writing the tar file is preferred in local storage then copy to target storage
+    )
+
 
 DS_RD_SETTING = DataScienceBasePropSetting()

rdagent/app/data_science/loop.py

@@ -1,5 +1,8 @@
+import shutil
+import subprocess
+from datetime import datetime
 from pathlib import Path
-from typing import Any
+from typing import Any, Optional, Union
 
 import fire
 
@@ -28,10 +31,8 @@ from rdagent.scenarios.data_science.dev.feedback import DSExperiment2Feedback
 from rdagent.scenarios.data_science.dev.runner import DSCoSTEERRunner
 from rdagent.scenarios.data_science.experiment.experiment import DSExperiment
 from rdagent.scenarios.data_science.proposal.exp_gen import DSExpGen, DSTrace
-from rdagent.scenarios.data_science.proposal.exp_gen.select import (
-    LatestCKPSelector,
-    SOTAJumpCKPSelector,
-)
+from rdagent.scenarios.data_science.proposal.exp_gen.idea_pool import DSKnowledgeBase
+from rdagent.scenarios.data_science.proposal.exp_gen.select import LatestCKPSelector
 from rdagent.scenarios.kaggle.kaggle_crawler import download_data
 
 
@@ -42,13 +43,6 @@ class DataScienceRDLoop(RDLoop):
         logger.log_object(PROP_SETTING.competition, tag="competition")
         scen: Scenario = import_class(PROP_SETTING.scen)(PROP_SETTING.competition)
 
-        ### shared components in the workflow  # TODO: check if
-        knowledge_base = (
-            import_class(PROP_SETTING.knowledge_base)(PROP_SETTING.knowledge_base_path, scen)
-            if PROP_SETTING.knowledge_base != ""
-            else None
-        )
-
         # 1) task generation from scratch
         # self.scratch_gen: tuple[HypothesisGen, Hypothesis2Experiment] = DummyHypothesisGen(scen),
 
@@ -70,8 +64,13 @@ class DataScienceRDLoop(RDLoop):
         # self.summarizer: Experiment2Feedback = import_class(PROP_SETTING.summarizer)(scen)
         # logger.log_object(self.summarizer, tag="summarizer")
 
-        # self.trace = KGTrace(scen=scen, knowledge_base=knowledge_base)
-        self.trace = DSTrace(scen=scen)
+        if DS_RD_SETTING.enable_knowledge_base and DS_RD_SETTING.knowledge_base_version == "v1":
+            knowledge_base = DSKnowledgeBase(
+                path=DS_RD_SETTING.knowledge_base_path, idea_pool_json_path=DS_RD_SETTING.idea_pool_json_path
+            )
+            self.trace = DSTrace(scen=scen, knowledge_base=knowledge_base)
+        else:
+            self.trace = DSTrace(scen=scen)
         self.summarizer = DSExperiment2Feedback(scen)
         super(RDLoop, self).__init__()
 
@@ -166,10 +165,70 @@ class DataScienceRDLoop(RDLoop):
                     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")
+        if DS_RD_SETTING.enable_knowledge_base and DS_RD_SETTING.knowledge_base_version == "v1":
+            logger.log_object(self.trace.knowledge_base, tag="knowledge_base")
+            self.trace.knowledge_base.dump()
+
+        if (
+            DS_RD_SETTING.enable_log_archive
+            and DS_RD_SETTING.log_archive_path is not None
+            and Path(DS_RD_SETTING.log_archive_path).is_dir()
+        ):
+            start_archive_datetime = datetime.now()
+            logger.info(f"Archiving log folder after loop {self.loop_idx}")
+            tar_path = (
+                Path(
+                    DS_RD_SETTING.log_archive_temp_path
+                    if DS_RD_SETTING.log_archive_temp_path
+                    else DS_RD_SETTING.log_archive_path
+                )
+                / "mid_log.tar"
+            )
+            subprocess.run(["tar", "-cf", str(tar_path), "-C", (Path().cwd() / "log"), "."], check=True)
+            if DS_RD_SETTING.log_archive_temp_path is not None:
+                shutil.move(tar_path, Path(DS_RD_SETTING.log_archive_path) / "mid_log.tar")
+                tar_path = Path(DS_RD_SETTING.log_archive_path) / "mid_log.tar"
+            shutil.copy(
+                tar_path, Path(DS_RD_SETTING.log_archive_path) / "mid_log_bak.tar"
+            )  # backup when upper code line is killed when running
+            self.timer.add_duration(datetime.now() - start_archive_datetime)
+
+    @classmethod
+    def load(
+        cls, path: Union[str, Path], output_path: Optional[Union[str, Path]] = None, do_truncate: bool = False
+    ) -> "LoopBase":
+        session = super().load(path, output_path, do_truncate)
+        if (
+            DS_RD_SETTING.enable_knowledge_base
+            and DS_RD_SETTING.knowledge_base_version == "v1"
+            and Path(DS_RD_SETTING.knowledge_base_path).exists()
+        ):
+            knowledge_base = DSKnowledgeBase(path=DS_RD_SETTING.knowledge_base_path)
+            session.trace.knowledge_base = knowledge_base
+        return session
+
+    def dump(self, path: str | Path) -> None:
+        """
+        Since knowledge_base is big and we don't want to dump it every time
+        So we remove it from the trace before dumping and restore it after.
+        """
+        backup_knowledge_base = None
+        if self.trace.knowledge_base is not None:
+            backup_knowledge_base = self.trace.knowledge_base
+            self.trace.knowledge_base = None
+        super().dump(path)
+        if backup_knowledge_base is not None:
+            self.trace.knowledge_base = backup_knowledge_base
 
 
 def main(
-    path=None, output_path=None, step_n=None, loop_n=None, competition="bms-molecular-translation", do_truncate=True
+    path=None,
+    output_path=None,
+    step_n=None,
+    loop_n=None,
+    competition="bms-molecular-translation",
+    do_truncate=True,
+    timeout=None,
 ):
     """
 
@@ -213,7 +272,7 @@ def main(
         kaggle_loop = DataScienceRDLoop(DS_RD_SETTING)
     else:
         kaggle_loop = DataScienceRDLoop.load(path, output_path, do_truncate)
-    kaggle_loop.run(step_n=step_n, loop_n=loop_n)
+    kaggle_loop.run(step_n=step_n, loop_n=loop_n, all_duration=timeout)
 
 
 if __name__ == "__main__":

rdagent/components/knowledge_management/graph.py

@@ -13,15 +13,17 @@ from rdagent.components.knowledge_management.vector_base import (
     cosine,
 )
 from rdagent.core.knowledge_base import KnowledgeBase
+from rdagent.log import rdagent_logger as logger
 from rdagent.oai.llm_utils import APIBackend
 
 Node = KnowledgeMetaData
 
 
 class UndirectedNode(Node):
-    def __init__(self, content: str = "", label: str = "", embedding: Any = None) -> None:
+    def __init__(self, content: str = "", label: str = "", embedding: Any = None, appendix: Any = None) -> None:
         super().__init__(content, label, embedding)
         self.neighbors: set[UndirectedNode] = set()
+        self.appendix = appendix  # appendix stores any additional information
         assert isinstance(content, str), "content must be a string"
 
     def add_neighbor(self, node: UndirectedNode) -> None:
@@ -86,6 +88,10 @@ class Graph(KnowledgeBase):
         size = 16
         embeddings = []
         for i in range(0, len(contents), size):
+            logger.info(
+                f"Creating embedding for index {i} to {i + size} with {len(contents)} contents",
+                tag="batch embedding",
+            )
             embeddings.extend(
                 APIBackend().create_embedding(input_content=contents[i : i + size]),
             )
@@ -270,7 +276,7 @@ class UndirectedGraph(Graph):
         self,
         node: UndirectedNode | str,
         similarity_threshold: float = 0.0,
-        topk_k: int = 5,
+        topk_k: int = None,
         constraint_labels: list[str] | None = None,
     ) -> list[UndirectedNode]:
         """

rdagent/components/knowledge_management/vector_base.py

@@ -87,7 +87,7 @@ class VectorBase(KnowledgeBase):
         """
         pass
 
-    def search(self, content: str, topk_k: int = 5, similarity_threshold: float = 0) -> List[Document]:
+    def search(self, content: str, topk_k: int | None = None, similarity_threshold: float = 0) -> List[Document]:
         """
         search vector_df by node
         Parameters
@@ -156,7 +156,11 @@ class PDVectorBase(VectorBase):
                 self.add(document=doc)
 
     def search(
-        self, content: str, topk_k: int = 5, similarity_threshold: float = 0, constraint_labels: list[str] | None = None
+        self,
+        content: str,
+        topk_k: int | None = None,
+        similarity_threshold: float = 0,
+        constraint_labels: list[str] | None = None,
     ) -> Tuple[List[Document], List]:
         """
         Search vector by node's embedding.
@@ -192,7 +196,9 @@ class PDVectorBase(VectorBase):
             lambda x: 1 - cosine(x, document.embedding)
         )  # cosine is cosine distance, 1-similarity
 
-        searched_similarities = similarities[similarities > similarity_threshold].nlargest(topk_k)
+        searched_similarities = similarities[similarities > similarity_threshold]
+        if topk_k is not None:
+            searched_similarities = searched_similarities.nlargest(topk_k)
         most_similar_docs = filtered_df.loc[searched_similarities.index]
 
         docs = []

rdagent/log/timer.py

@@ -0,0 +1,84 @@
+import re
+from datetime import datetime, timedelta
+
+from rdagent.core.utils import SingletonBaseClass
+from rdagent.log import rdagent_logger as logger
+
+
+class RDAgentTimer:
+    def __init__(self) -> None:
+        self.started: bool = False
+        self.target_time: datetime | None = None
+        self.all_duration: timedelta | None = None
+        self.remain_time_duration: timedelta | None = None
+
+    def reset(self, all_duration: str | timedelta) -> None:
+        if isinstance(all_duration, str):
+            pattern = re.compile(r"^\s*(\d*\.?\d+)\s*([smhd]?)\s*$")
+
+            match = pattern.match(all_duration)
+            if not match:
+                return None
+            value = float(match.group(1))
+            unit = match.group(2)
+            if unit == "s":
+                self.all_duration = timedelta(seconds=value)
+            elif unit == "m":
+                self.all_duration = timedelta(minutes=value)
+            elif unit == "h":
+                self.all_duration = timedelta(hours=value)
+            elif unit == "d":
+                self.all_duration = timedelta(days=value)
+            else:
+                self.all_duration = timedelta(seconds=value)
+        elif isinstance(all_duration, timedelta):
+            self.all_duration = all_duration
+        self.target_time = datetime.now() + self.all_duration
+        logger.info(f"Timer set to {self.all_duration} seconds and counting down.")
+        self.started = True
+        return None
+
+    def restart_by_remain_time(self) -> None:
+        if self.remain_time_duration is not None:
+            self.target_time = datetime.now() + self.remain_time_duration
+            self.started = True
+            logger.info(f"Timer restarted with remaining time: {self.remain_time_duration}")
+        else:
+            logger.warning("No remaining time to restart the timer.")
+        return None
+
+    def add_duration(self, duration: timedelta) -> None:
+        if self.started and self.target_time is not None:
+            logger.info(f"Adding {duration} to the timer. Currently {self.remain_time()} remains.")
+            self.target_time = self.target_time + duration
+            self.update_remain_time()
+
+    def is_timeout(self) -> bool:
+        if self.started and self.target_time is not None:
+            self.update_remain_time()
+            if datetime.now() > self.target_time:
+                return True
+        return False
+
+    def update_remain_time(self) -> None:
+        if self.started and self.target_time is not None:
+            self.remain_time_duration = self.target_time - datetime.now()
+        return None
+
+    def remain_time(self) -> timedelta | None:
+        if self.started:
+            self.update_remain_time()
+            return self.remain_time_duration
+        return None
+
+
+class RDAgentTimerWrapper(SingletonBaseClass):
+    def __init__(self) -> None:
+        self.timer: RDAgentTimer = RDAgentTimer()
+
+    def replace_timer(self, timer: RDAgentTimer) -> None:
+        self.timer = timer
+        logger.info("Timer replaced successfully.")
+
+
+RD_Agent_TIMER_wrapper = RDAgentTimerWrapper()

rdagent/oai/backend/base.py

@@ -7,6 +7,7 @@ import time
 import uuid
 from abc import ABC, abstractmethod
 from copy import deepcopy
+from datetime import datetime
 from pathlib import Path
 from typing import Any, Optional, cast
 
@@ -15,6 +16,7 @@ from pydantic import TypeAdapter
 from rdagent.core.utils import LLM_CACHE_SEED_GEN, SingletonBaseClass
 from rdagent.log import LogColors
 from rdagent.log import rdagent_logger as logger
+from rdagent.log.timer import RD_Agent_TIMER_wrapper
 from rdagent.oai.llm_conf import LLM_SETTINGS
 from rdagent.utils import md5_hash
 
@@ -330,6 +332,7 @@ class APIBackend(ABC):
         max_retry = LLM_SETTINGS.max_retry if LLM_SETTINGS.max_retry is not None else max_retry
         timeout_count = 0
         for i in range(max_retry):
+            API_start_time = datetime.now()
             try:
                 if embedding:
                     return self._create_embedding_with_cache(*args, **kwargs)
@@ -361,6 +364,8 @@ class APIBackend(ABC):
                         raise e
                 else:
                     time.sleep(self.retry_wait_seconds)
+                    if RD_Agent_TIMER_wrapper.timer.started and not isinstance(e, json.decoder.JSONDecodeError):
+                        RD_Agent_TIMER_wrapper.timer.add_duration(datetime.now() - API_start_time)
                 logger.warning(str(e))
                 logger.warning(f"Retrying {i+1}th time...")
         error_message = f"Failed to create chat completion after {max_retry} retries."

rdagent/oai/backend/litellm.py

@@ -51,20 +51,16 @@ class LiteLLMAPIBackend(APIBackend):
         """
         Call the embedding function
         """
-        response_list = []
-        for input_content_iter in input_content_list:
-            model_name = LITELLM_SETTINGS.embedding_model
-            logger.info(f"{LogColors.GREEN}Using emb model{LogColors.END} {model_name}", tag="debug_litellm_emb")
-            logger.info(f"Creating embedding for: {input_content_iter}", tag="debug_litellm_emb")
-            if not isinstance(input_content_iter, str):
-                raise ValueError("Input content must be a string")
-            response = embedding(
-                model=model_name,
-                input=input_content_iter,
-                *args,
-                **kwargs,
-            )
-            response_list.append(response.data[0]["embedding"])
+        model_name = LITELLM_SETTINGS.embedding_model
+        logger.info(f"{LogColors.GREEN}Using emb model{LogColors.END} {model_name}", tag="debug_litellm_emb")
+        logger.info(f"Creating embedding for: {input_content_list}", tag="debug_litellm_emb")
+        response = embedding(
+            model=model_name,
+            input=input_content_list,
+            *args,
+            **kwargs,
+        )
+        response_list = [data["embedding"] for data in response.data]
         return response_list
 
     def _create_chat_completion_inner_function(  # type: ignore[no-untyped-def] # noqa: C901, PLR0912, PLR0915

rdagent/scenarios/data_science/dev/feedback.py

@@ -12,6 +12,7 @@ from rdagent.core.proposal import (
 from rdagent.oai.llm_utils import APIBackend
 from rdagent.scenarios.data_science.experiment.experiment import DSExperiment
 from rdagent.scenarios.data_science.proposal.exp_gen import DSTrace
+from rdagent.scenarios.data_science.proposal.exp_gen.idea_pool import DSIdea
 from rdagent.utils import convert2bool
 from rdagent.utils.agent.tpl import T
 from rdagent.utils.repo.diff import generate_diff_from_dict
@@ -111,10 +112,23 @@ class DSExperiment2Feedback(Experiment2Feedback):
 
         # Currently, we do not use `observations`, `hypothesis_evaluation`, and `new_hypothesis` in the framework.
         # `new_hypothesis` should not exist in the feedback.
-        return HypothesisFeedback(
+        hypothesis_feedback = HypothesisFeedback(
             observations=resp_dict.get("Observations", "No observations provided"),
             hypothesis_evaluation=resp_dict.get("Feedback for Hypothesis", "No feedback provided"),
             new_hypothesis=resp_dict.get("New Hypothesis", "No new hypothesis provided"),
             reason=resp_dict.get("Reasoning", "No reasoning provided"),
             decision=convert2bool(resp_dict.get("Replace Best Result", "no")),
         )
+
+        if hypothesis_feedback and DS_RD_SETTING.enable_knowledge_base:
+            ds_idea = DSIdea(
+                {
+                    "competition": self.scen.get_competition_full_desc(),
+                    "idea": exp.hypothesis.hypothesis,
+                    "method": exp.pending_tasks_list[0][0].get_task_information(),
+                    "hypothesis": {exp.hypothesis.problem_label: exp.hypothesis.problem_desc},
+                }
+            )
+            trace.knowledge_base.add_idea(idea=ds_idea)
+
+        return hypothesis_feedback

rdagent/scenarios/data_science/proposal/exp_gen/base.py

@@ -18,20 +18,25 @@ class DSHypothesis(Hypothesis):
         concise_observation: str = "",
         concise_justification: str = "",
         concise_knowledge: str = "",
-        problem: str = "",
+        problem_name: str = "",
+        problem_desc: str = "",
+        problem_label: Literal["SCENARIO_PROBLEM", "FEEDBACK_PROBLEM"] = "FEEDBACK_PROBLEM",
     ) -> None:
         super().__init__(
             hypothesis, reason, concise_reason, concise_observation, concise_justification, concise_knowledge
         )
         self.component = component
-        self.problem = problem
+        self.problem_name = problem_name
+        self.problem_desc = problem_desc
+        self.problem_label = problem_label
 
     def __str__(self) -> str:
         if self.hypothesis == "":
             return f"No hypothesis available. Trying to construct the first runnable {self.component} component."
         lines = []
-        if self.problem is not None:
-            lines.append(f"Target Problem: {self.problem}")
+        if self.problem_name is not None and self.problem_desc is not None:
+            lines.append(f"Target Problem name: {self.problem_name}")
+            lines.append(f"Target Problem: {self.problem_desc}")
         lines.extend(
             [f"Chosen Component: {self.component}", f"Hypothesis: {self.hypothesis}", f"Reason: {self.reason}"]
         )

rdagent/scenarios/data_science/proposal/exp_gen/idea_pool.py

@@ -0,0 +1,185 @@
+import json
+from pathlib import Path
+from typing import Dict, List
+
+from tqdm import tqdm
+
+from rdagent.components.knowledge_management.graph import (
+    UndirectedNode,  # TODO: add appendix attribute to node
+)
+from rdagent.components.knowledge_management.graph import (
+    UndirectedGraph,
+)
+from rdagent.log import rdagent_logger as logger
+from rdagent.oai.llm_utils import APIBackend
+from rdagent.utils.agent.tpl import T
+
+
+class DSIdea:
+    def __init__(self, raw_knowledge: Dict | str) -> None:
+        """
+        {
+            "idea": "A concise label summarizing the core concept of this idea.",
+            "method": "A specific method used in this idea, described in a general and implementable way (e.g., 'applied a stacking ensemble method to combine predictions from multiple base models'). Avoid mentioning specific models or dataset-specific details to ensure better generalization",
+            "context": "A detailed example of how the notebook implements this idea (e.g., 'the notebook used XGBoost, Random Forest, and LightGBM as base models and logistic regression as the meta-model').",
+            "hypothesis": {
+                "scenario_problem": "The nature of problem the idea addresses, described without referencing the method itself (e.g., 'a classification problem with complex decision boundaries').",
+                "feedback_problem": "The characteristics of the data (e.g., imbalance, high dimensionality, collinearity, outliers, missing data, skewed distribution, time-based pattern, etc.) that justify the use of this method.",
+            }
+        }
+        """
+        # TODO: add competition name -> avoid using self-generated ideas
+        # TODO: align Scenario and Feedback problem (for key and label)
+        if isinstance(raw_knowledge, str):
+            raw_knowledge = json.loads(raw_knowledge)
+        self.competition = raw_knowledge.get("competition", None)
+        self.idea = raw_knowledge["idea"]
+        self.method = raw_knowledge.get("method", None)
+        self.context = raw_knowledge.get("context", None)
+        self.hypothesis = raw_knowledge["hypothesis"].copy()
+
+    def __str__(self) -> str:
+        return json.dumps(
+            {
+                "competition": self.competition,
+                "idea": self.idea,
+                "method": self.method,
+                "context": self.context,
+                "hypothesis": self.hypothesis,
+            }
+        )
+
+    def to_formatted_str(self) -> str:
+        return f"Idea Name: {self.idea}\nIdea Method: {self.method}\nIdea Context: {self.context}"
+
+
+class DSKnowledgeBase(UndirectedGraph):
+    def __init__(self, path: str | Path | None = None, idea_pool_json_path: str | Path | None = None):
+        super().__init__(path)
+        self.used_idea_id_set = set()
+        if idea_pool_json_path is not None:
+            self.build_idea_pool(idea_pool_json_path)
+        self.dump()
+
+    def add_idea(self, idea: List[DSIdea] | DSIdea) -> None:
+        if not isinstance(idea, list):
+            idea_list = [idea]
+        else:
+            idea_list = idea
+
+        node_list = []
+        add_pairs = []
+        for one_idea in idea_list:
+            idea_name = one_idea.idea
+            idea_node = UndirectedNode(content=idea_name, label="IDEA", appendix=str(one_idea))
+            node_list.append(idea_node)
+
+            competition = one_idea.competition
+            if competition is not None:
+                competition_node = UndirectedNode(content=competition, label="competition")
+                node_list.append(competition_node)
+                add_pairs.append((idea_node, [competition_node]))
+
+            data = one_idea.hypothesis.get("SCENARIO_PROBLEM", None)
+            problem = one_idea.hypothesis.get("FEEDBACK_PROBLEM", None)
+            if data is not None:
+                sp_node = UndirectedNode(content=data, label="SCENARIO_PROBLEM")
+                node_list.append(sp_node)
+                add_pairs.append((idea_node, [sp_node]))
+            if problem is not None:
+                fp_node = UndirectedNode(content=problem, label="FEEDBACK_PROBLEM")
+                node_list.append(fp_node)
+                add_pairs.append((idea_node, [fp_node]))
+        self.batch_embedding(node_list)
+        for idea_node, neighbor_list in add_pairs:
+            self.add_nodes(idea_node, neighbor_list)
+
+    def build_idea_pool(self, idea_pool_json_path: str | Path):
+        if len(self.vector_base.vector_df) > 0:
+            logger.warning("Knowledge graph is not empty, please clear it first. Ignore reading from json file.")
+            return
+        else:
+            logger.info(f"Building knowledge graph from idea pool json file: {idea_pool_json_path}")
+        with open(idea_pool_json_path, "r", encoding="utf-8") as f:
+            idea_pool_dict = json.load(f)
+
+        to_add_ideas = []
+        for i, raw_idea in tqdm(enumerate(idea_pool_dict), desc="Building Knowledge Graph from Ideas"):
+            try:
+                idea = DSIdea(raw_idea)
+                to_add_ideas.append(idea)
+            except Exception as e:
+                print(f"The {i}-th idea process failed due to error {e}")
+                continue
+        self.add_idea(to_add_ideas)
+
+    def sample_ideas(
+        self,
+        problems: Dict,
+        scenario_desc: str,
+        exp_feedback_list_desc: str,
+        sota_exp_desc: str,
+        competition_desc: str,
+    ) -> Dict:
+        # sample ideas by cosine similarity
+        text = ""
+        problem_to_sampled_idea_node_id = {}
+        competition_node = self.get_node_by_content(competition_desc)
+
+        for i, (problem_name, problem_dict) in enumerate(problems.items()):
+            sampled_nodes = self.semantic_search(
+                node=problem_dict["problem"], constraint_labels=[problem_dict["label"]]
+            )
+
+            text += f"# Problem Name {i+1}: {problem_name}\n"
+            text += f"- Problem Description: {problem_dict['problem']}\n"
+            problem_to_sampled_idea_node_id[problem_name] = []
+            for node in sampled_nodes:
+                idea_node = self.get_nodes_within_steps(start_node=node, steps=1, constraint_labels="IDEA")[0]
+
+                if idea_node.id not in self.used_idea_id_set and (
+                    competition_node is None or competition_node not in idea_node.neighbors
+                ):
+                    idea = DSIdea(raw_knowledge=idea_node.appendix)
+                    problem_to_sampled_idea_node_id[problem_name].append(idea_node)
+                    text += f"## Idea {len(problem_to_sampled_idea_node_id[problem_name])}\n"
+                    text += f"- Idea Name: {idea.idea}\n"
+                    text += f"- Idea Method: {idea.method}\n"
+                    text += f"- Idea Context: {idea.context}\n\n"
+                if len(problem_to_sampled_idea_node_id[problem_name]) >= 5:
+                    break
+            text += "\n\n"
+
+        # select ideas by LLM
+        sys_prompt = T(".prompts_v2:idea_sample.system").r(
+            idea_spec=T(".prompts_v2:specification.idea").r(),
+            idea_output_format=T(".prompts_v2:output_format.idea").r(),
+        )
+        user_prompt = T(".prompts_v2:idea_sample.user").r(
+            scenario_desc=scenario_desc,
+            exp_feedback_list_desc=exp_feedback_list_desc,
+            sota_exp_desc=sota_exp_desc,
+            problem_ideas=text,
+        )
+        response = APIBackend().build_messages_and_create_chat_completion(
+            user_prompt=user_prompt,
+            system_prompt=sys_prompt,
+            json_mode=True,
+            json_target_type=Dict[str, int],
+        )
+        resp_dict = json.loads(response)
+
+        # update problems with selected ideas
+        for problem_name, picked_id in resp_dict.items():
+            if problem_name in problem_to_sampled_idea_node_id and picked_id < len(
+                problem_to_sampled_idea_node_id[problem_name]
+            ):
+                problems[problem_name]["idea"] = problem_to_sampled_idea_node_id[problem_name][picked_id - 1].appendix
+                problems[problem_name]["idea_node_id"] = problem_to_sampled_idea_node_id[problem_name][picked_id - 1].id
+
+        return problems
+
+    def update_pickled_problem(self, problems: Dict, pickled_problem_name: str) -> None:
+        pickled_id = problems[pickled_problem_name].get("idea_node_id", None)
+        if pickled_id is not None:
+            self.used_idea_id_set.add(pickled_id)

rdagent/scenarios/data_science/proposal/exp_gen/prompts_v2.yaml

@@ -67,6 +67,13 @@ hypothesis_gen:
     1. **Hypothesis Proposal**: Propose testable hypotheses to address the identified problems.
     2. **Hypothesis Evaluation**: Evaluate the proposed hypotheses across multiple dimensions.
 
+    {% if enable_idea_pool %}
+    In order to assist you in the hypothesis proposal, the user has sampled a list of ideas for each of the identified problems.
+    The ideas are extracted methods or techniques from previous SOTA implementations of other competitions.
+    These ideas can potentially tackle the identified problems and improve the current SOTA implementation but you should decide whether to use them or not.
+    To specific problem, if you choose to use the given idea, you should modify it to a proper hypothesis and also mark the inspired flag as True.
+    {% endif %}
+
     # Task 1: Hypothesis Proposal
     For each identified problem, you are required to propose a hypothesis aimed at improving the current SOTA implementation. 
     A hypothesis is a precise, testable, and actionable statement that proposes a specific modification or improvement to address an identified problem in a Kaggle competition implementation.
@@ -91,6 +98,10 @@ hypothesis_gen:
       - If the problem relates to time/memory constraints, suggest smaller model sizes or alternative algorithms with reduced complexity.
       - If the problem involves underperforming models, propose removing or replacing models with significantly worse performance.
       - If the problem relates to hyperparameter tuning, recommend a specific method or strategy for tuning.
+    {% if enable_idea_pool %}
+    4. Idea Reference
+      - Each idea is a method, technique or trick that contributes to high performance from other competition implementation under similar problem. You are free to use them as an inspiration for your hypothesis proposal.
+    {% endif %}
 
     ## Hypothesis Specification
     {{ hypothesis_spec }}
@@ -114,16 +125,15 @@ hypothesis_gen:
     # Scenario Description
     {{ scenario_desc }}
 
-    # Previous Experiments and Feedbacks:
+    # Previous Experiments and Feedbacks
     {{ exp_and_feedback_list_desc }}
 
     # Current SOTA Implementation
     {{ sota_exp_desc }}
 
-    # Identified Problems
+    # Identified Problems{% if enable_idea_pool %} with Sampled Ideas{% endif %}
     {{ problems }}
 
-
 task_gen:
   system: |-
     {% include "scenarios.data_science.share:scen.role" %}
@@ -171,6 +181,35 @@ task_gen:
     # Feedback from Previous Failed Experiments (e.g., experiments that did not pass evaluation, encountered bugs, or failed to surpass SOTA performance):
     {{ failed_exp_and_feedback_list_desc }}
 
+idea_sample:
+  system: |-
+    You are a Kaggle Grandmaster and expert ML engineer with deep expertise in statistics, machine learning, and competition optimization.
+    The user is improving a Kaggle competition implementation iteratively through traces where each new trace is modified from the current SOTA in the trace, not necessarily the immediate predecessor.
+    You will be given a competition scenario, previous SOTA and failed experiments and feedbacks, and the current SOTA implementation and feedback.
+    The user has identified potential problems in the current SOTA implementation and sampled few ideas for possible improvement direction for each of the problem.
+    Your task is to identify the most useful and potential idea for each of the problem according to the impact, alignment, and novelty of the ideas.
+
+    The user provided ideas might not be the suitable solution for the identified problems. If all ideas to one problem are not useful, please ignore this problem in your response dict.
+
+    ### Specification
+    {{ idea_spec }}
+
+    ### Output Format
+    {{ idea_output_format }}
+
+  user: |-
+    # Scenario Description
+    {{ scenario_desc }}
+    
+    # Previous Experiments and Feedbacks
+    {{ exp_feedback_list_desc }}    
+
+    # Current SOTA Implementation
+    {{ sota_exp_desc }}
+
+    # Problem-Ideas Pairs
+    {{ problem_ideas }}
+
 specification:
   problem: |-
     1. The problem should be specific and fine-grained. Avoid general or vague statements. 
@@ -181,6 +220,12 @@ specification:
     2. Each hypothesis should focus on a single direction per experiment. Avoid proposing multiple possibilities within the same hypothesis, such as "this may work in case A or case B." Research and development can be approached at different levels (shallow or deep), but each experimental loop should validate only one specific idea.
     3. The hypothesis should based on current SOTA solution. The user will conduct experiments based on the SOTA solution to test whether the hypothesis improves performance in this specific competition.
 
+  idea: |-
+    1. Alignment: The idea should be aligned with the identified problem. It should be a potential solution to the problem.
+    2. Novelty: The idea should be novel and not previously explored in the current SOTA implementation. Avoid ideas that have already been tried and failed.
+    3. Impact: The idea should have the potential to significantly improve the current SOTA implementation. It should be a promising direction for further exploration.
+    4. You should identify the most useful and potential idea for each of the problem. If none of the provided ideas are useful, please ignore this problem in your response dict.
+
 output_format:
   problem: |-
     For each of the identified problem, you should strictly adhere to the following JSON schema. 
@@ -199,7 +244,8 @@ output_format:
   hypothesis: |-
     For each of the identified problem, you should propose a hypothesis strictly following to the JSON schema. Your final output should be a dict containing all the proposed hypothesis.
     {
-      "problem name 1 (Should be exactly same as the problem name provided)": {
+      "problem name 1 (should be exactly same as the problem name provided)": {
+        {% if enable_idea_pool %}"inspired": "True or False. Set to True if the hypothesis is inspired by the user provided ideas. Otherwise, set it to False.",{% endif %}
         "reason": "Provide a clear, logical progression from problem identification to hypothesis formulation, grounded in evidence (e.g., trace history, domain principles, or competition constraints). Refer to the Hypothesis Guidelines for better understanding. Reason should be short with no more than two sentences.",
         "component": "The component name that the hypothesis {% if pipeline %}mainly {% endif %}focuses on. Must be one of ('DataLoadSpec', 'FeatureEng', 'Model', 'Ensemble', 'Workflow').",
         "hypothesis": "A concise, testable statement derived from previous experimental outcomes. Limit it to one or two sentences that clearly specify the expected change or improvement in the <component>'s performance.",
@@ -212,3 +258,11 @@ output_format:
         }
       },
     }
+  idea: |-
+    For each of the problems, you should identified the most useful and potential idea strictly following to the JSON schema.
+    Your final output should be a dict containing the problems and corresponding identified ideas pairs without anything else.
+    Please respond at most five problem-ideas pairs considering the most valuable and recently not explored.
+    {
+      "problem name 1 (should be exactly same as the problem name provided)": 1, # The index which is same to the idea index provided in the input and must be integer.
+      "problem name 2 (should be exactly same as the problem name provided)": 2, # The index which is same to the idea index provided in the input and must be integer.
+    }

rdagent/scenarios/data_science/proposal/exp_gen/proposal.py

@@ -1,5 +1,5 @@
 import json
-from typing import Dict, List
+from typing import Dict, Tuple
 
 import pandas as pd
 
@@ -14,6 +14,9 @@ from rdagent.core.proposal import ExpGen
 from rdagent.oai.llm_utils import APIBackend, md5_hash
 from rdagent.scenarios.data_science.experiment.experiment import DSExperiment
 from rdagent.scenarios.data_science.proposal.exp_gen.base import DSHypothesis, DSTrace
+from rdagent.scenarios.data_science.proposal.exp_gen.idea_pool import (
+    DSIdea,
+)
 from rdagent.utils.agent.tpl import T
 from rdagent.utils.repo.diff import generate_diff_from_dict
 from rdagent.utils.workflow import wait_retry
@@ -266,20 +269,34 @@ class DSProposalV2ExpGen(ExpGen):
         scenario_desc: str,
         exp_feedback_list_desc: str,
         sota_exp_desc: str,
-        problems: list,
+        problems: dict,
         pipeline: bool,
+        enable_idea_pool: bool,
     ) -> Dict:
+        problem_formatted_str = ""
+        for problem_name, problem_dict in problems.items():
+            problem_formatted_str += f"# Problem Name: {problem_name}\n"
+            problem_formatted_str += f"- Problem Description: {problem_dict['problem']}\n"
+            if "idea" in problem_dict:
+                idea_formatted_str = DSIdea(problem_dict["idea"]).to_formatted_str()
+                problem_formatted_str += f"- Sampled Idea by user: \n{idea_formatted_str}\n"
+            problem_formatted_str += "\n\n"
+
         sys_prompt = T(".prompts_v2:hypothesis_gen.system").r(
             component_desc=component_desc,
             hypothesis_spec=T(".prompts_v2:specification.hypothesis").r(),
-            hypothesis_output_format=T(".prompts_v2:output_format.hypothesis").r(pipeline=pipeline),
+            hypothesis_output_format=T(".prompts_v2:output_format.hypothesis").r(
+                pipeline=pipeline, enable_idea_pool=enable_idea_pool
+            ),
             pipeline=pipeline,
+            enable_idea_pool=enable_idea_pool,
         )
         user_prompt = T(".prompts_v2:hypothesis_gen.user").r(
             scenario_desc=scenario_desc,
             exp_and_feedback_list_desc=exp_feedback_list_desc,
             sota_exp_desc=sota_exp_desc,
-            problems=json.dumps(problems, indent=2),
+            problems=problem_formatted_str,
+            enable_idea_pool=enable_idea_pool,
         )
         response = APIBackend().build_messages_and_create_chat_completion(
             user_prompt=user_prompt,
@@ -288,14 +305,9 @@ class DSProposalV2ExpGen(ExpGen):
             json_target_type=Dict[str, Dict[str, str | Dict[str, str | int]]],
         )
         resp_dict = json.loads(response)
-        for key, value in resp_dict.items():
-            assert "reason" in value, "Reason not provided."
-            assert "component" in value, "Component not provided."
-            assert "hypothesis" in value, "Hypothesis not provided."
-            assert "evaluation" in value, "Evaluation not provided."
         return resp_dict
 
-    def hypothesis_rank(self, hypothesis_dict: dict, problem_dict: dict, pipeline: bool) -> DSHypothesis:
+    def hypothesis_rank(self, hypothesis_dict: dict, problem_dict: dict, pipeline: bool) -> Tuple[str, DSHypothesis]:
         weights = {
             "alignment_score": 0.2,
             "impact_score": 0.4,
@@ -305,6 +317,8 @@ class DSProposalV2ExpGen(ExpGen):
         }
         scores_dict = {}
         for problem_name in hypothesis_dict:
+            if "hypothesis" not in hypothesis_dict[problem_name]:
+                continue
             scores_dict[problem_name] = {}
             for score_key in weights:
                 if score_key not in hypothesis_dict[problem_name]["evaluation"]:
@@ -316,22 +330,35 @@ class DSProposalV2ExpGen(ExpGen):
                         )
                     except (ValueError, TypeError):
                         scores_dict[problem_name][score_key] = 0
+
         scores = pd.DataFrame(scores_dict)
         scores_sorted = scores.sum().sort_values(ascending=False)
         if len(scores_sorted) > 5:
             scores_sorted = scores_sorted[: len(scores_sorted) // 2]
 
-        reproducible_int = int.from_bytes(bytes.fromhex(md5_hash(scores_sorted.to_string())), byteorder="big") % len(
-            scores_sorted
-        )
-        max_score_problem_name = scores_sorted.index[reproducible_int]
-        problem = problem_dict.get(max_score_problem_name, {}).get("problem", "Problem not provided")
+        # Increase the weight of the hypothesis that is inspired by the idea pool
+        index_to_pick_pool_list = []
+        for j, problem_name in enumerate(scores_sorted.index):
+            if hypothesis_dict[problem_name].get("inspired", False):
+                index_to_pick_pool_list.extend([j] * 3)
+            else:
+                index_to_pick_pool_list.append(j)
 
-        return DSHypothesis(
-            component=hypothesis_dict[max_score_problem_name]["component"],
-            hypothesis=hypothesis_dict[max_score_problem_name]["hypothesis"],
-            reason=hypothesis_dict[max_score_problem_name]["reason"],
-            problem=problem,
+        # Create a random but reproducible integer
+        reproducible_int = int.from_bytes(bytes.fromhex(md5_hash(scores_sorted.to_string())), byteorder="big") % len(
+            index_to_pick_pool_list
+        )
+        selected_idx = index_to_pick_pool_list[reproducible_int]
+        max_score_problem_name = scores_sorted.index[selected_idx]
+        problem_dict = problem_dict.get(max_score_problem_name, {})
+
+        return max_score_problem_name, DSHypothesis(
+            component=hypothesis_dict[max_score_problem_name].get("component", "Model"),
+            hypothesis=hypothesis_dict[max_score_problem_name].get("hypothesis", "Hypothesis not provided"),
+            reason=hypothesis_dict[max_score_problem_name].get("reason", "Reason not provided"),
+            problem_name=max_score_problem_name,
+            problem_desc=problem_dict.get("problem", "Problem description not provided"),
+            problem_label=problem_dict.get("label", "FEEDBACK_PROBLEM"),
         )
 
     def task_gen(
@@ -402,12 +429,15 @@ class DSProposalV2ExpGen(ExpGen):
         return exp
 
     def gen(self, trace: DSTrace, pipeline: bool = False) -> DSExperiment:
-        component_desc = "\n".join(
-            [
-                f"[{key}] {value}"
-                for key, value in T("scenarios.data_science.share:component_description").template.items()
-            ]
-        )
+        if pipeline:
+            component_desc = T("scenarios.data_science.share:component_description_in_pipeline").r()
+        else:
+            component_desc = "\n".join(
+                [
+                    f"[{key}] {value}"
+                    for key, value in T("scenarios.data_science.share:component_description").template.items()
+                ]
+            )
 
         sota_exp = trace.sota_experiment()
         if not isinstance(sota_exp, DSExperiment):
@@ -436,14 +466,28 @@ class DSProposalV2ExpGen(ExpGen):
             competition_desc=competition_desc,
             sota_exp_desc=sota_exp_desc,
         )
+        for problem_name in scen_problems:
+            scen_problems[problem_name]["label"] = "SCENARIO_PROBLEM"
         fb_problems = self.identify_feedback_problem(
             scenario_desc=scenario_desc,
             exp_feedback_list_desc=exp_feedback_list_desc,
             sota_exp_desc=sota_exp_desc,
         )
+        for problem_name in fb_problems:
+            fb_problems[problem_name]["label"] = "FEEDBACK_PROBLEM"
         all_problems = {**scen_problems, **fb_problems}
 
-        # Step 2: Propose hypothesis based on the identified problems
+        # Step 1.5: Sample ideas from idea pool
+        if DS_RD_SETTING.enable_knowledge_base:
+            all_problems = trace.knowledge_base.sample_ideas(
+                problems=all_problems,
+                scenario_desc=scenario_desc,
+                exp_feedback_list_desc=exp_feedback_list_desc,
+                sota_exp_desc=sota_exp_desc,
+                competition_desc=self.scen.get_competition_full_desc(),
+            )
+
+        # Step 2: Propose hypothesis based on the identified problems (and sampled ideas)
         hypothesis_dict = self.hypothesis_gen(
             component_desc=component_desc,
             scenario_desc=scenario_desc,
@@ -451,6 +495,7 @@ class DSProposalV2ExpGen(ExpGen):
             sota_exp_desc=sota_exp_desc,
             problems=all_problems,
             pipeline=pipeline,
+            enable_idea_pool=DS_RD_SETTING.enable_knowledge_base,
         )
         if not pipeline:
             sota_exp_model_file_count = len(
@@ -469,11 +514,14 @@ class DSProposalV2ExpGen(ExpGen):
                     hypothesis_dict.pop(name)
 
         # Step 3: Select the best hypothesis
-        new_hypothesis = self.hypothesis_rank(
+        pickled_problem_name, new_hypothesis = self.hypothesis_rank(
             hypothesis_dict=hypothesis_dict,
             problem_dict=all_problems,
             pipeline=pipeline,
         )
+        # Step 3.5: Update knowledge base with the picked problem
+        if DS_RD_SETTING.enable_knowledge_base:
+            trace.knowledge_base.update_pickled_problem(all_problems, pickled_problem_name)
 
         return self.task_gen(
             component_desc=component_desc,

rdagent/scenarios/data_science/share.yaml

@@ -41,14 +41,7 @@ describe: # some template to describe some object
 
   trace: |-
     {% if exp_and_feedback_list|length == 0 %}
-    No previous 
-    {% if type == "success" %}
-    successful
-    {% elif type == "failure" %}
-    failed
-    {% else %}
-    successful or failed
-    {% endif %} trial available.
+    No previous {% if type == "success" %}successful{% elif type == "failure" %}failed{% else %}successful or failed{% endif %} trial available.
     {% else %}
     {% if type == "success" %}
     ## {{ heading | default('Trace of the successful trial') }}
@@ -115,6 +108,13 @@ component_description:
     Integrates all pipeline components, from data loading to ensemble prediction, ensuring efficient execution and correct output formatting.
     - When focusing on this component, the corresponding editing files will be: "main.py"
 
+component_description_in_pipeline: |-
+  [DataLoadSpec]: Focus on the data loading and preprocessing aspects of the pipeline, ensuring that the data is correctly formatted and ready for feature engineering.
+  [FeatureEng]: Concentrate on transforming the raw data into meaningful features while maintaining the integrity of the dataset.
+  [Model]: Focus on the model building, tuning of the pipeline, ensuring that the model is optimized for performance.
+  [Ensemble]: Concentrate on combining predictions from multiple models and evaluating their performance.
+  [Workflow]: Focus on the overall integration of the pipeline or parts not included in the other components, ensuring that all components work together seamlessly.
+
 component_spec:
   general: |-
     {{ spec }}

rdagent/utils/agent/tpl.py

@@ -121,7 +121,6 @@ class RDAT:
         )
         while "\n\n\n" in rendered:
             rendered = rendered.replace("\n\n\n", "\n\n")
-        rendered = "\n".join(line for line in rendered.splitlines() if line.strip())
         logger.log_object(
             obj={
                 "uri": self.uri,

rdagent/utils/workflow.py

@@ -19,6 +19,7 @@ from typing import Any, Callable, Optional, TypeVar, Union, cast
 from tqdm.auto import tqdm
 
 from rdagent.log import rdagent_logger as logger
+from rdagent.log.timer import RD_Agent_TIMER_wrapper, RDAgentTimer
 
 
 class LoopMeta(type):
@@ -36,7 +37,7 @@ class LoopMeta(type):
         steps = []
         for base in bases:
             for step in LoopMeta._get_steps(base.__bases__) + getattr(base, "steps", []):
-                if step not in steps:
+                if step not in steps and step not in ["load", "dump"]:  # incase user override the load/dump method
                     steps.append(step)
         return steps
 
@@ -55,7 +56,7 @@ class LoopMeta(type):
         steps = LoopMeta._get_steps(bases)  # all the base classes of parents
         for name, attr in attrs.items():
             if not name.startswith("_") and callable(attr):
-                if name not in steps:
+                if name not in steps and name not in ["load", "dump"]:  # incase user override the load/dump method
                     # NOTE: if we override the step in the subclass
                     # Then it is not the new step. So we skip it.
                     steps.append(name)
@@ -90,8 +91,9 @@ class LoopBase:
         self.loop_prev_out: dict[str, Any] = {}  # the step results of current loop
         self.loop_trace = defaultdict(list[LoopTrace])  # the key is the number of loop
         self.session_folder = logger.log_trace_path / "__session__"
+        self.timer: RDAgentTimer = RD_Agent_TIMER_wrapper.timer
 
-    def run(self, step_n: int | None = None, loop_n: int | None = None) -> None:
+    def run(self, step_n: int | None = None, loop_n: int | None = None, all_duration: str | None = None) -> None:
         """
 
         Parameters
@@ -103,6 +105,10 @@ class LoopBase:
             How many steps to run; if current loop is incomplete, it will be counted as the first loop for completion
             `None` indicates to run forever until error or KeyboardInterrupt
         """
+
+        if all_duration is not None and not self.timer:
+            self.timer.reset(all_duration=all_duration)
+
         with tqdm(total=len(self.steps), desc="Workflow Progress", unit="step") as pbar:
             while True:
                 if step_n is not None:
@@ -113,6 +119,13 @@ class LoopBase:
                     if loop_n <= 0:
                         break
 
+                if self.timer.started:
+                    if self.timer.is_timeout():
+                        logger.warning("Timeout, exiting the loop.")
+                        break
+                    else:
+                        logger.info(f"Timer remaining time: {self.timer.remain_time()}")
+
                 li, si = self.loop_idx, self.step_idx
                 name = self.steps[si]
                 logger.info(f"Start Loop {li}, Step {si}: {name}")
@@ -155,6 +168,8 @@ class LoopBase:
                 self.dump(self.session_folder / f"{li}" / f"{si}_{name}")  # save a snapshot after the session
 
     def dump(self, path: str | Path) -> None:
+        if RD_Agent_TIMER_wrapper.timer.started:
+            RD_Agent_TIMER_wrapper.timer.update_remain_time()
         path = Path(path)
         path.parent.mkdir(parents=True, exist_ok=True)
         with path.open("wb") as f:
@@ -181,6 +196,10 @@ class LoopBase:
         if do_truncate:
             max_loop = max(session.loop_trace.keys())
             logger.storage.truncate(time=session.loop_trace[max_loop][-1].end)
+
+        if session.timer.started:
+            RD_Agent_TIMER_wrapper.replace_timer(session.timer)
+            RD_Agent_TIMER_wrapper.timer.restart_by_remain_time()
         return session