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 <[email protected]>

---------

Co-authored-by: Xu Yang <[email protected]>
Co-authored-by: Xu Yang <[email protected]>
Co-authored-by: Xu <[email protected]>

Author: 4 people

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.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 @@ def __init__(self, PROP_SETTING: BasePropSetting):
         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 @@ def __init__(self, PROP_SETTING: BasePropSetting):
         # 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 @@ def record(self, prev_out: dict[str, Any]):
                     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 @@
     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 @@ def batch_embedding(nodes: list[Node]) -> list[Node]:
         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 @@ def semantic_search(
         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 @@ def add(self, document: Union[Document, List[Document]]):
         """
         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 @@ def add(self, document: Union[Document, List[Document]]):
                 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 @@ def search(
             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 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 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 @@ def _try_create_chat_completion_or_embedding(  # type: ignore[no-untyped-def]
         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 @@ def _try_create_chat_completion_or_embedding(  # type: ignore[no-untyped-def]
                         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 @@ def _create_embedding_inner_function(
         """
         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