chore: release 0.6.1

Release-As: 0.6.1

CHANGELOG.md

@@ -1,5 +1,31 @@
 # Changelog
 
+## [0.6.1](https://github.com/microsoft/RD-Agent/compare/v0.6.0...v0.6.1) (2025-06-28)
+
+
+### Bug Fixes
+
+* fix mount ([#1001](https://github.com/microsoft/RD-Agent/issues/1001)) ([4ae2f13](https://github.com/microsoft/RD-Agent/commit/4ae2f1303dfcbaea53d459be7c8e85bf85ce5f4f))
+* handle the bug of wrong dag_parant index ([#996](https://github.com/microsoft/RD-Agent/issues/996)) ([bda12ff](https://github.com/microsoft/RD-Agent/commit/bda12ffecf9ae116e0d04eece0c6a1b61413d916))
+* improve log folder sorting and selection UX ([#993](https://github.com/microsoft/RD-Agent/issues/993)) ([b116807](https://github.com/microsoft/RD-Agent/commit/b11680777f116b6c40f9e535e0da10c186c95050))
+
+## [0.6.0](https://github.com/microsoft/RD-Agent/compare/v0.5.0...v0.6.0) (2025-06-26)
+
+
+### Features
+
+* async mechanism for multi-trace ([#981](https://github.com/microsoft/RD-Agent/issues/981)) ([9e60c32](https://github.com/microsoft/RD-Agent/commit/9e60c32cf348481eb55617809c059c359d7603b8))
+
+
+### Bug Fixes
+
+* add async to direct_exp_gen avoid infinite loop ([#992](https://github.com/microsoft/RD-Agent/issues/992)) ([78c203d](https://github.com/microsoft/RD-Agent/commit/78c203d8eefbba67fc120b35cb25e85b2200ac49))
+* docker container cleanup to prevent accumulation and system slowdown ([#975](https://github.com/microsoft/RD-Agent/issues/975)) ([05cf094](https://github.com/microsoft/RD-Agent/commit/05cf094913e48c903c8a4476d6c609d8bfa10681))
+* fix a bug and update the docs ([#978](https://github.com/microsoft/RD-Agent/issues/978)) ([d1ae9e1](https://github.com/microsoft/RD-Agent/commit/d1ae9e1dcc2ccd1ffe05cb1c6db3e905fa70425c))
+* merge datascience v3 and v2 ([#974](https://github.com/microsoft/RD-Agent/issues/974)) ([1ba7548](https://github.com/microsoft/RD-Agent/commit/1ba754853ce2010ce1cb0bbd217b67689fa1ebdf))
+* refine details ([#979](https://github.com/microsoft/RD-Agent/issues/979)) ([25caa3d](https://github.com/microsoft/RD-Agent/commit/25caa3d00c255286dce27915b9355987b87ed2e8))
+* refine prompt ([#987](https://github.com/microsoft/RD-Agent/issues/987)) ([76df96e](https://github.com/microsoft/RD-Agent/commit/76df96ee88212a8aee7f518b9cacf80591dc2939))
+
 ## [0.5.0](https://github.com/microsoft/RD-Agent/compare/v0.4.0...v0.5.0) (2025-06-18)
 
 

docs/scens/data_science.rst

@@ -69,6 +69,32 @@ The Data Science Agent is an agent that can automatically perform feature engine
       dotenv set DS_CODER_ON_WHOLE_PIPELINE True
       dotenv set DS_CODER_COSTEER_ENV_TYPE docker
 
+- 🚀 **Run the Application**
+
+  - You can directly run the application by using the following command:
+    
+    .. code-block:: sh
+
+        rdagent data_science --competition <Competition ID>
+
+  - Then, you can run the test set score corresponding to each round of the loop.
+
+    .. code-block:: sh
+
+        dotenv run -- python rdagent/log/mle_summary.py grade <url_to_log>
+
+    Here, <url_to_log> refers to the parent directory of the log folder generated during the run.
+
+- 📥 **Visualize the R&D Process**
+
+  - We provide a web UI to visualize the log. You just need to run:
+
+    .. code-block:: sh
+
+        streamlit run rdagent/log/ui/dsapp.py
+
+  - Then you can input the log path and visualize the R&D process.
+
 🔍 MLE-bench Guide: Running ML Engineering via MLE-bench
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

rdagent/app/qlib_rd_loop/factor_from_report.py

@@ -11,6 +11,7 @@ from rdagent.components.document_reader.document_reader import (
     extract_first_page_screenshot_from_pdf,
     load_and_process_pdfs_by_langchain,
 )
+from rdagent.core.conf import RD_AGENT_SETTINGS
 from rdagent.core.proposal import Hypothesis
 from rdagent.log import rdagent_logger as logger
 from rdagent.oai.llm_utils import APIBackend
@@ -104,22 +105,31 @@ class FactorReportLoop(FactorRDLoop, metaclass=LoopMeta):
             self.judge_pdf_data_items = [i for i in Path(report_folder).rglob("*.pdf")]
 
         self.loop_n = min(len(self.judge_pdf_data_items), FACTOR_FROM_REPORT_PROP_SETTING.report_limit)
+        self.shift_report = (
+            0  # some reports does not contain viable factor, so we ship some of them to avoid infinite loop
+        )
 
-    def direct_exp_gen(self, prev_out: dict[str, Any]):
+    async def direct_exp_gen(self, prev_out: dict[str, Any]):
         while True:
-            report_file_path = self.judge_pdf_data_items[self.loop_idx]
-            logger.info(f"Processing number {self.loop_idx} report: {report_file_path}")
-            exp = extract_hypothesis_and_exp_from_reports(str(report_file_path))
-            if exp is None:
-                continue
-            exp.based_experiments = [QlibFactorExperiment(sub_tasks=[], hypothesis=exp.hypothesis)] + [
-                t[0] for t in self.trace.hist if t[1]
-            ]
-            exp.sub_workspace_list = exp.sub_workspace_list[: FACTOR_FROM_REPORT_PROP_SETTING.max_factors_per_exp]
-            exp.sub_tasks = exp.sub_tasks[: FACTOR_FROM_REPORT_PROP_SETTING.max_factors_per_exp]
-            logger.log_object(exp.hypothesis, tag="hypothesis generation")
-            logger.log_object(exp.sub_tasks, tag="experiment generation")
-            return exp
+            if self.get_unfinished_loop_cnt(self.loop_idx) < RD_AGENT_SETTINGS.get_max_parallel():
+                report_file_path = self.judge_pdf_data_items[self.loop_idx + self.shift_report]
+                logger.info(f"Processing number {self.loop_idx} report: {report_file_path}")
+                exp = extract_hypothesis_and_exp_from_reports(str(report_file_path))
+                if exp is None:
+                    self.shift_report += 1
+                    self.loop_n -= 1
+                    if self.loop_n < 0:  # NOTE: on every step, we self.loop_n -= 1 at first.
+                        raise self.LoopTerminationError("Reach stop criterion and stop loop")
+                    continue
+                exp.based_experiments = [QlibFactorExperiment(sub_tasks=[], hypothesis=exp.hypothesis)] + [
+                    t[0] for t in self.trace.hist if t[1]
+                ]
+                exp.sub_workspace_list = exp.sub_workspace_list[: FACTOR_FROM_REPORT_PROP_SETTING.max_factors_per_exp]
+                exp.sub_tasks = exp.sub_tasks[: FACTOR_FROM_REPORT_PROP_SETTING.max_factors_per_exp]
+                logger.log_object(exp.hypothesis, tag="hypothesis generation")
+                logger.log_object(exp.sub_tasks, tag="experiment generation")
+                return exp
+            await asyncio.sleep(1)
 
     def coding(self, prev_out: dict[str, Any]):
         exp = self.coder.develop(prev_out["direct_exp_gen"])

rdagent/app/qlib_rd_loop/quant.py

@@ -10,6 +10,7 @@ import fire
 from rdagent.app.qlib_rd_loop.conf import QUANT_PROP_SETTING
 from rdagent.components.workflow.conf import BasePropSetting
 from rdagent.components.workflow.rd_loop import RDLoop
+from rdagent.core.conf import RD_AGENT_SETTINGS
 from rdagent.core.developer import Developer
 from rdagent.core.exception import FactorEmptyError, ModelEmptyError
 from rdagent.core.proposal import (
@@ -64,15 +65,18 @@ class QuantRDLoop(RDLoop):
         self.trace = QuantTrace(scen=scen)
         super(RDLoop, self).__init__()
 
-    def direct_exp_gen(self, prev_out: dict[str, Any]):
-        hypo = self._propose()
-        assert hypo.action in ["factor", "model"]
-        if hypo.action == "factor":
-            exp = self.factor_hypothesis2experiment.convert(hypo, self.trace)
-        else:
-            exp = self.model_hypothesis2experiment.convert(hypo, self.trace)
-        logger.log_object(exp.sub_tasks, tag="experiment generation")
-        return {"propose": hypo, "exp_gen": exp}
+    async def direct_exp_gen(self, prev_out: dict[str, Any]):
+        while True:
+            if self.get_unfinished_loop_cnt(self.loop_idx) < RD_AGENT_SETTINGS.get_max_parallel():
+                hypo = self._propose()
+                assert hypo.action in ["factor", "model"]
+                if hypo.action == "factor":
+                    exp = self.factor_hypothesis2experiment.convert(hypo, self.trace)
+                else:
+                    exp = self.model_hypothesis2experiment.convert(hypo, self.trace)
+                logger.log_object(exp.sub_tasks, tag="experiment generation")
+                return {"propose": hypo, "exp_gen": exp}
+            await asyncio.sleep(1)
 
     def coding(self, prev_out: dict[str, Any]):
         if prev_out["direct_exp_gen"]["propose"].action == "factor":

rdagent/app/utils/health_check.py

@@ -3,22 +3,25 @@ import socket
 import docker
 
 from rdagent.log import rdagent_logger as logger
+from rdagent.utils.env import cleanup_container
 
 
 def check_docker() -> None:
+    container = None
     try:
         client = docker.from_env()
         client.images.pull("hello-world")
         container = client.containers.run("hello-world", detach=True)
         logs = container.logs().decode("utf-8")
         print(logs)
-        container.remove()
         logger.info(f"The docker status is normal")
     except docker.errors.DockerException as e:
         logger.error(f"An error occurred: {e}")
         logger.warning(
             f"Docker status is exception, please check the docker configuration or reinstall it. Refs: https://docs.docker.com/engine/install/ubuntu/."
         )
+    finally:
+        cleanup_container(container, "health check")
 
 
 def is_port_in_use(port):

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

@@ -98,21 +98,12 @@ class PipelineMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
             spec=T("scenarios.data_science.share:component_spec.Pipeline").r(),
             enable_model_dump=DS_RD_SETTING.enable_model_dump,
         )
-        if DS_RD_SETTING.proposal_version == "v3":
-            # FIXME: A temporary patch for BUILD
-            user_prompt = T(".prompts:pipeline_coder.user_v3").r(
-                competition_info=competition_info,
-                folder_spec=data_folder_info,
-                latest_code=workspace.file_dict.get("main.py"),
-                latest_code_feedback=prev_task_feedback,
-            )
-        else:
-            user_prompt = T(".prompts:pipeline_coder.user").r(
-                competition_info=competition_info,
-                folder_spec=data_folder_info,
-                latest_code=workspace.file_dict.get("main.py"),
-                latest_code_feedback=prev_task_feedback,
-            )
+        user_prompt = T(".prompts:pipeline_coder.user").r(
+            competition_info=competition_info,
+            folder_spec=data_folder_info,
+            latest_code=workspace.file_dict.get("main.py"),
+            latest_code_feedback=prev_task_feedback,
+        )
 
         for _ in range(5):
             pipeline_code = PythonAgentOut.extract_output(

rdagent/components/coder/data_science/pipeline/eval.py

@@ -74,13 +74,13 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator):
 
                 # Check model names (index)
                 if not score_df.index.is_unique:
-                    score_check_text += "\n[Error] The score dataframe contains duplicate model names."
+                    score_check_text += "\n[Error] The file 'scores.csv' contains duplicate model names."
                     score_ret_code = 1
                 if "ensemble" not in model_set_in_scores:
-                    score_check_text += "\n[Error] The score dataframe doesn't contain the ensemble model."
+                    score_check_text += "\n[Error] The file 'scores.csv' doesn't contain the ensemble model."
                     score_ret_code = 1
                 if score_ret_code != 0:
-                    score_check_text += f"The score_df is:\n{score_df}"
+                    score_check_text += f"The dataframe in file 'scores.csv' is:\n{score_df}"
 
                 # Check metric name (columns)
                 if score_df.columns.tolist() != [self.scen.metric_name]:
@@ -125,10 +125,6 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator):
                     )
             stdout += "\n" + submission_check_out
 
-        eda_output = implementation.file_dict.get("EDA.md", None)
-
-        eda_output = implementation.file_dict.get("EDA.md", None)
-
         if not isinstance(implementation, FBWorkspace):
             eda_output = None
         else:

rdagent/components/coder/data_science/pipeline/prompts.yaml

@@ -83,27 +83,6 @@ pipeline_coder:
     --------- Data Folder Description (All path are relative to the data folder) ---------
     {{ folder_spec }}
     
-    {% if latest_code %}
-    --------- Former code ---------
-    {{ latest_code }}
-    {% if latest_code_feedback is not none %}
-    --------- Feedback to former code ---------
-    {{ latest_code_feedback }}
-    The former code contains errors. You should correct the code based on the provided information, ensuring you do not repeat the same mistakes.
-    {% else %}
-    The former code is correct. You should try to improve the code based on the provided task while not changing the irrelevant parts.
-    {% endif %}
-    {% endif %} 
-
-    You should strictly follow the code specifications provided by the specification to implement the function.
-
-  user_v3: |-
-    --------- Competition Information ---------
-    {{ competition_info }}
-
-    --------- Data Folder Description (All path are relative to the data folder) ---------
-    {{ folder_spec }}
-    
     {% if latest_code %}
     --------- Former code ---------
     {{ latest_code }}

rdagent/components/workflow/rd_loop.py

@@ -3,9 +3,11 @@ 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 asyncio
 from typing import Any
 
 from rdagent.components.workflow.conf import BasePropSetting
+from rdagent.core.conf import RD_AGENT_SETTINGS
 from rdagent.core.developer import Developer
 from rdagent.core.proposal import (
     Experiment2Feedback,
@@ -55,10 +57,13 @@ class RDLoop(LoopBase, metaclass=LoopMeta):
         return exp
 
     # included steps
-    def direct_exp_gen(self, prev_out: dict[str, Any]):
-        hypo = self._propose()
-        exp = self._exp_gen(hypo)
-        return {"propose": hypo, "exp_gen": exp}
+    async def direct_exp_gen(self, prev_out: dict[str, Any]):
+        while True:
+            if self.get_unfinished_loop_cnt(self.loop_idx) < RD_AGENT_SETTINGS.get_max_parallel():
+                hypo = self._propose()
+                exp = self._exp_gen(hypo)
+                return {"propose": hypo, "exp_gen": exp}
+            await asyncio.sleep(1)
 
     def coding(self, prev_out: dict[str, Any]):
         exp = self.coder.develop(prev_out["direct_exp_gen"]["exp_gen"])

rdagent/core/conf.py

@@ -87,8 +87,7 @@ class RDAgentSettings(ExtendedBaseSettings):
         """Based on the setting of semaphore, return the maximum number of parallel loops"""
         if isinstance(self.step_semaphore, int):
             return self.step_semaphore
-        else:
-            return max(self.step_semaphore.values())
+        return max(self.step_semaphore.values())
 
     # NOTE: for debug
     # the following function only serves as debugging and is necessary in main logic.

rdagent/core/experiment.py

@@ -324,6 +324,9 @@ class Experiment(
             {}
         )  # TODO: in Kaggle, now sub results are all saved in self.result, remove this in the future.
 
+        # For parallel multi-trace support
+        self.local_selection: tuple[int, ...] | None = None
+
 
 ASpecificExp = TypeVar("ASpecificExp", bound=Experiment)
 

rdagent/core/proposal.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import asyncio
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Generic, List, Tuple, TypeVar
+from typing import TYPE_CHECKING, Generic, TypeVar
 
 from rdagent.core.conf import RD_AGENT_SETTINGS
 from rdagent.core.evaluation import Feedback
@@ -108,7 +108,7 @@ ASpecificKB = TypeVar("ASpecificKB", bound=KnowledgeBase)
 
 class Trace(Generic[ASpecificScen, ASpecificKB]):
     NodeType = tuple[Experiment, ExperimentFeedback]  # Define NodeType as a new type representing the tuple
-    NEW_ROOT: Tuple = ()
+    NEW_ROOT: tuple = ()
 
     def __init__(self, scen: ASpecificScen, knowledge_base: ASpecificKB | None = None) -> None:
         self.scen: ASpecificScen = scen
@@ -163,35 +163,33 @@ class Trace(Generic[ASpecificScen, ASpecificKB]):
 
         return [self.hist[i] for i in self.get_parents(selection[0])]
 
-    def exp2idx(self, exp: Experiment | List[Experiment]) -> int | List[int] | None:
+    def exp2idx(self, exp: Experiment | list[Experiment]) -> int | list[int] | None:
         if isinstance(exp, list):
-            exps: List[Experiment] = exp
+            exps: list[Experiment] = exp
 
             # keep the order
             exp_to_index: dict[Experiment, int] = {_exp: i for i, (_exp, _) in enumerate(self.hist)}
             return [exp_to_index[_exp] for _exp in exps]
-        else:
-            for i, (_exp, _) in enumerate(self.hist):
-                if _exp == exp:
-                    return i
+        for i, (_exp, _) in enumerate(self.hist):
+            if _exp == exp:
+                return i
         return None
 
-    def idx2exp(self, idx: int | List[int]) -> Experiment | List[Experiment]:
+    def idx2exp(self, idx: int | list[int]) -> Experiment | list[Experiment]:
         if isinstance(idx, list):
-            idxs: List[int] = idx
+            idxs: list[int] = idx
             return [self.hist[_idx][0] for _idx in idxs]
-        else:
-            return self.hist[idx][0]
+        return self.hist[idx][0]
 
     def is_parent(self, parent_idx: int, child_idx: int) -> bool:
         ancestors = self.get_parents(child_idx)
         return parent_idx in ancestors
 
-    def get_parents(self, child_idx: int) -> List[int]:
+    def get_parents(self, child_idx: int) -> list[int]:
         if self.is_selection_new_tree((child_idx,)):
             return []
 
-        ancestors: List[int] = []
+        ancestors: list[int] = []
         curr = child_idx
         while True:
             ancestors.insert(0, curr)

rdagent/log/mle_summary.py

@@ -44,7 +44,7 @@ def save_grade_info(log_trace_path: Path):
 
 
 def save_all_grade_info(log_folder):
-    for log_trace_path in log_folder.iterdir():
+    for log_trace_path in Path(log_folder).iterdir():
         if is_valid_session(log_trace_path):
             try:
                 save_grade_info(log_trace_path)

rdagent/log/ui/ds_trace.py

@@ -239,10 +239,9 @@ def llm_log_win(llm_d: list):
                             st.markdown(spec)
                             rdict.pop("spec")
                         else:
-                            # show model codes
                             showed_keys = []
                             for k, v in rdict.items():
-                                if k.startswith("model_") and k.endswith(".py"):
+                                if k.endswith(".py"):
                                     st.markdown(f":red[**{k}**]")
                                     st.code(v, language="python", wrap_lines=True, line_numbers=True)
                                     showed_keys.append(k)
@@ -251,7 +250,10 @@ def llm_log_win(llm_d: list):
                         st.write(":red[**Other parts (except for the code or spec) in response dict:**]")
                         st.json(rdict)
                     except:
-                        st.json(resp)
+                        try:
+                            st.json(resp)
+                        except:
+                            show_text(resp)
                 with t2:
                     show_text(user)
                 with t3:
@@ -615,18 +617,22 @@ def stdout_win(loop_id: int):
                 st.code(v, language="log", wrap_lines=True)
 
 
-def get_folders_sorted(log_path):
-    """缓存并返回排序后的文件夹列表，并加入进度打印"""
+def get_folders_sorted(log_path, sort_by_time=False):
+    """
+    Cache and return the sorted list of folders, with progress printing.
+    :param log_path: Log path
+    :param sort_by_time: Whether to sort by time, default False (sort by name)
+    """
     if not log_path.exists():
         st.toast(f"Path {log_path} does not exist!")
         return []
-    with st.spinner("正在加载文件夹列表..."):
-        folders = sorted(
-            (folder for folder in log_path.iterdir() if is_valid_session(folder)),
-            key=lambda folder: folder.stat().st_mtime,
-            reverse=True,
-        )
-        st.write(f"找到 {len(folders)} 个文件夹")
+    with st.spinner("Loading folder list..."):
+        folders = [folder for folder in log_path.iterdir() if is_valid_session(folder)]
+        if sort_by_time:
+            folders = sorted(folders, key=lambda folder: folder.stat().st_mtime, reverse=True)
+        else:
+            folders = sorted(folders, key=lambda folder: folder.name)
+        st.write(f"Found {len(folders)} folders")
     return [folder.name for folder in folders]
 
 
@@ -657,7 +663,7 @@ with st.sidebar:
     if not state.log_folder.exists():
         st.warning(f"Path {state.log_folder} does not exist!")
     else:
-        folders = get_folders_sorted(state.log_folder)
+        folders = get_folders_sorted(state.log_folder, sort_by_time=False)
         if "selection" in st.query_params:
             default_index = (
                 folders.index(st.query_params["selection"]) if st.query_params["selection"] in folders else 0

rdagent/oai/backend/base.py

@@ -499,6 +499,13 @@ class APIBackend(ABC):
                 self.cache.embedding_set(content_to_embedding_dict)
         return [content_to_embedding_dict[content] for content in input_content_list]  # type: ignore[misc]
 
+    @abstractmethod
+    def support_function_calling(self) -> bool:
+        """
+        Check if the backend supports function calling
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
     @abstractmethod
     def _calculate_token_from_messages(self, messages: list[dict[str, Any]]) -> int:
         """

rdagent/oai/backend/deprec.py

@@ -261,6 +261,13 @@ class DeprecBackend(APIBackend):
                     raise
         return encoding
 
+    def support_function_calling(self) -> bool:
+        """
+        Check if the backend supports function calling.
+        Currently, deprec backend does not support function calling so it returns False. #FIXME: maybe a mapping to the backend class is needed.
+        """
+        return False
+
     def _create_embedding_inner_function(  # type: ignore[no-untyped-def]
         self, input_content_list: list[str], *args, **kwargs
     ) -> list[list[float]]:  # noqa: ARG002

rdagent/oai/backend/litellm.py

@@ -7,6 +7,7 @@ from litellm import (
     completion,
     completion_cost,
     embedding,
+    supports_function_calling,
     supports_response_schema,
     token_counter,
 )
@@ -93,6 +94,12 @@ class LiteLLMAPIBackend(APIBackend):
         """
         if json_mode and supports_response_schema(model=LITELLM_SETTINGS.chat_model):
             kwargs["response_format"] = {"type": "json_object"}
+        elif not supports_response_schema(model=LITELLM_SETTINGS.chat_model) and "response_format" in kwargs:
+            logger.warning(
+                f"{LogColors.RED}Model {LITELLM_SETTINGS.chat_model} does not support response schema, ignoring response_format argument.{LogColors.END}",
+                tag="llm_messages",
+            )
+            kwargs.pop("response_format")
 
         if LITELLM_SETTINGS.log_llm_chat_content:
             logger.info(self._build_log_messages(messages), tag="llm_messages")
@@ -183,3 +190,9 @@ class LiteLLMAPIBackend(APIBackend):
             tag="token_cost",
         )
         return content, finish_reason
+
+    def support_function_calling(self) -> bool:
+        """
+        Check if the backend supports function calling
+        """
+        return supports_function_calling(model=LITELLM_SETTINGS.chat_model) and LITELLM_SETTINGS.enable_function_call

rdagent/oai/llm_conf.py

@@ -16,6 +16,9 @@ class LLMSettings(ExtendedBaseSettings):
     embedding_model: str = "text-embedding-3-small"
 
     reasoning_effort: Literal["low", "medium", "high"] | None = None
+    enable_function_call: bool = (
+        True  # Whether to enable function calling in chat models. may not work for models that do not support it.
+    )
 
     # Handling format
     reasoning_think_rm: bool = False

rdagent/scenarios/data_science/dev/feedback.py

@@ -26,6 +26,7 @@ class DSExperiment2Feedback(Experiment2Feedback):
         # 3. 相对sota_exp的改动
         # 4. result 任务的结果
         # 5. sota_exp.result 之前最好的结果
+
         sota_exp = trace.sota_experiment()
         sota_desc = T("scenarios.data_science.share:describe.exp").r(
             exp=sota_exp, heading="SOTA of previous exploration of the scenario"
@@ -87,29 +88,16 @@ class DSExperiment2Feedback(Experiment2Feedback):
                     )
 
         eda_output = exp.experiment_workspace.file_dict.get("EDA.md", None)
-        if DS_RD_SETTING.proposal_version == "v3":
-            # FIXME: Some minor changes. Did not have time to test the full.
-            system_prompt = T(".prompts:exp_feedback_v3.system").r(
-                scenario=self.scen.get_scenario_all_desc(eda_output=eda_output)
-            )
-            user_prompt = T(".prompts:exp_feedback_v3.user").r(
-                sota_desc=sota_desc,
-                cur_exp=exp,
-                diff_edition=diff_edition,
-                feedback_desc=feedback_desc,
-                cur_vs_sota_score=cur_vs_sota_score,
-            )
-        else:
-            system_prompt = T(".prompts:exp_feedback.system").r(
-                scenario=self.scen.get_scenario_all_desc(eda_output=eda_output)
-            )
-            user_prompt = T(".prompts:exp_feedback.user").r(
-                sota_desc=sota_desc,
-                cur_exp=exp,
-                diff_edition=diff_edition,
-                feedback_desc=feedback_desc,
-                cur_vs_sota_score=cur_vs_sota_score,
-            )
+        system_prompt = T(".prompts:exp_feedback.system").r(
+            scenario=self.scen.get_scenario_all_desc(eda_output=eda_output)
+        )
+        user_prompt = T(".prompts:exp_feedback.user").r(
+            sota_desc=sota_desc,
+            cur_exp=exp,
+            diff_edition=diff_edition,
+            feedback_desc=feedback_desc,
+            cur_vs_sota_score=cur_vs_sota_score,
+        )
 
         resp_dict = json.loads(
             APIBackend().build_messages_and_create_chat_completion(

rdagent/scenarios/data_science/dev/prompts.yaml

@@ -15,7 +15,9 @@ exp_feedback:
       - Recommend corrective actions explicitly.
       - Set `"Replace Best Result": "no"`.
       - Begin your `reasoning` with `[Submission format error]`, clearly stating the issues causing experiment failure.
-    - If submission passes, proceed to Step 2.
+    - If submission passes the submission format check:
+      - If this is the first valid submission ever, set `"Replace Best Result": "yes"`.
+      - Otherwise, proceed to Step 2.
 
     Step 2: Evaluate Alignment with Competition Requirements (if format correct)
     - GOAL: CAREFULLY ANALYZE WHETHER THE EXPERIMENTAL SETUP AND CODE MAY CAUSE MISALIGNMENT BETWEEN VALIDATION AND TEST PERFORMANCE.
@@ -59,6 +61,8 @@ exp_feedback:
     Provide detailed and constructive feedback structured as follows:
     Example JSON Structure for Result Analysis:
     {
+      "Submission Format Check": "yes or no",
+      "First Valid Submission": "yes or no",
       "Observations": "Clearly summarize current and SOTA ensemble results with exact scores and notable patterns. Limit to no more than three concise, data-focused sentences. Your observation must be grounded by explicit evidence from scenario description or code implementation, not just validation scores.",
       "Feedback for Hypothesis": Explicitly confirm or refute the hypothesis based on specific data points or performance trends. Limit to two sentences.",
       "Evaluation Aligned With Task": "yes or no",
@@ -110,11 +114,11 @@ exp_feedback:
     {{ cur_exp.experiment_workspace.all_codes }}
 
     ## Feedback of past experiments
-    {{ feedback_desc }}
+    {{ feedback_desc or "There has not been any experiments yet." }}
     Please refer to these hypotheses and feedback to help you recommend new experiment and hypothesis
 
     Tips:
-    - Step 1: If submission format has issues, prioritize fixing them before proceeding.
+    - Step 1: If submission format has issues, prioritize fixing them before proceeding. If the format is correct and it's the first valid submission ever (there has never been valid submissions in the past), set `"Replace Best Result": "yes"`. If the format is correct and this is not the first valid submission, proceed to Step 2.
     - Step 2: If evaluation alignment issues are identified (validation approach does not follow competition requirements), address these methodological discrepancies immediately.
     - Step 3: If new results significantly worse than SOTA, or repeated hyperparameter adjustments yield no improvement, it might be time to rethink or shift focus.
 

rdagent/scenarios/data_science/dev/runner/__init__.py

@@ -114,8 +114,10 @@ class DSCoSTEERRunner(CoSTEER):
         exp.sub_tasks = [
             CoSTEERTask(
                 name="Debug running solution",
-                description=f"The whole workflow of the solution has finished with some execution error, please check the error message and debug the whole code repo.\nCurrent code repo md5: {md5_hash(exp.experiment_workspace.all_codes)}",
-            )
+                description=f"You'll be provided with the source code and the running and testing stdout. "
+                "Please check the error messages and debug the source code if any errors occur.\n"
+                f"Current code repo md5: {md5_hash(exp.experiment_workspace.all_codes)}",
+            ),
         ]
         exp = super().develop(exp)  # run strategy(code implementation & evaluation loops)
         exp.sub_tasks = bak_sub_tasks

rdagent/scenarios/data_science/dev/runner/eval.py

@@ -60,7 +60,7 @@ class DSCoSTEERCoSTEEREvaluator(CoSTEEREvaluator):
             eda_output = "No EDA output."
         implementation.inject_files(**{"EDA.md": eda_output})
         stdout = remove_eda_part(stdout)
-        stdout += f"The code executed {'successfully' if execute_ret_code == 0 else 'failed'}. {'EDA output is emmitted. ' if eda_output else ''}"
+        stdout += f"The code executed {'successfully' if execute_ret_code == 0 else 'failed'}. {'The EDA output is removed from the stdout. ' if eda_output else ''}"
 
         # Check score file
         score_fp = implementation.workspace_path / "scores.csv"
@@ -82,13 +82,13 @@ class DSCoSTEERCoSTEEREvaluator(CoSTEEREvaluator):
                 # in Pipeline task, we only check ensemble in scores.csv
                 if DS_RD_SETTING.coder_on_whole_pipeline:
                     if not score_df.index.is_unique:
-                        score_check_text += "\n[Error] The score dataframe contains duplicate model names."
+                        score_check_text += "\n[Error] The file 'scores.csv' contains duplicate model names."
                         score_ret_code = 1
                     if "ensemble" not in model_set_in_scores:
-                        score_check_text += "\n[Error] The score dataframe doesn't contain the ensemble model."
+                        score_check_text += "\n[Error] The file 'scores.csv' doesn't contain the ensemble model."
                         score_ret_code = 1
                     if score_ret_code != 0:
-                        score_check_text += f"The score_df is:\n{score_df}"
+                        score_check_text += f"The dataframe in file 'scores.csv' is:\n{score_df}"
                 else:
                     if model_set_in_scores != model_set_in_folder.union({"ensemble"}):
                         score_check_text += f"\n[Error] The scores dataframe does not contain the correct model names as index.\ncorrect model names are: {model_set_in_folder.union({'ensemble'})}\nscore_df is:\n{score_df}"

rdagent/scenarios/data_science/dev/runner/prompts.yaml

@@ -49,7 +49,7 @@ DSCoSTEER_eval:
   user: |-
     --------- code base ---------
     {{ code }}
-    --------- test stdout ---------
+    --------- the stdout of code execution and testing ---------
     {{ stdout }}
 
 DSCoSTEER_debugger:

rdagent/scenarios/data_science/experiment/experiment.py

@@ -29,3 +29,6 @@ class DSExperiment(Experiment[Task, FBWorkspace, FBWorkspace]):
         (so it is different from `trace.next_incomplete_component`.)
         """
         return self.experiment_workspace is not None and "main.py" in self.experiment_workspace.file_dict
+
+    def set_local_selection(self, local_selection: tuple[int, ...]) -> None:
+        self.local_selection = local_selection

rdagent/scenarios/data_science/loop.py

@@ -90,14 +90,11 @@ class DataScienceRDLoop(RDLoop):
         from rdagent.scenarios.data_science.proposal.exp_gen.proposal import (
             DSProposalV1ExpGen,
             DSProposalV2ExpGen,
-            DSProposalV3ExpGen,
         )
 
         if class_uri == "rdagent.scenarios.data_science.proposal.exp_gen.DSExpGen":
-            if DS_RD_SETTING.proposal_version not in ["v1", "v2", "v3"]:
+            if DS_RD_SETTING.proposal_version not in ["v1", "v2"]:
                 return import_class(DS_RD_SETTING.proposal_version)(scen=scen)
-            if DS_RD_SETTING.proposal_version == "v3":
-                return DSProposalV3ExpGen(scen=scen)
             if DS_RD_SETTING.proposal_version == "v1":
                 return DSProposalV1ExpGen(scen=scen)
             if DS_RD_SETTING.proposal_version == "v2":
@@ -145,7 +142,6 @@ class DataScienceRDLoop(RDLoop):
         super(RDLoop, self).__init__()
 
     async def direct_exp_gen(self, prev_out: dict[str, Any]):
-
         # set the SOTA experiment to submit
         sota_exp_to_submit = self.sota_exp_selector.get_sota_exp_to_submit(self.trace)
         self.trace.set_sota_exp_to_submit(sota_exp_to_submit)
@@ -154,11 +150,12 @@ class DataScienceRDLoop(RDLoop):
         selection = self.ckp_selector.get_selection(self.trace)
         # set the current selection for the trace
         self.trace.set_current_selection(selection)
-        exp = await self.exp_gen.async_gen(self.trace, self)
-        logger.log_object(exp)
 
-        # FIXME: this is for LLM debug webapp, remove this when the debugging is done.
-        logger.log_object(exp, tag="debug_exp_gen")
+        # in parallel + multi-trace mode, the above global "trace.current_selection" will not be used
+        # instead, we will use the "local_selection" attached to each exp to in async_gen().
+        exp = await self.exp_gen.async_gen(self.trace, self)
+
+        logger.log_object(exp)
         return exp
 
     def coding(self, prev_out: dict[str, Any]):
@@ -200,6 +197,11 @@ class DataScienceRDLoop(RDLoop):
         - If we come to feedback phase, the previous development steps are successful.
         """
         exp: DSExperiment = prev_out["running"]
+
+        # set the local selection to the trace after feedback
+        if exp.local_selection is not None:
+            self.trace.set_current_selection(exp.local_selection)
+
         if self.trace.next_incomplete_component() is None or DS_RD_SETTING.coder_on_whole_pipeline:
             # we have alreadly completed components in previous trace. So current loop is focusing on a new proposed idea.
             # So we need feedback for the proposal.
@@ -214,19 +216,38 @@ class DataScienceRDLoop(RDLoop):
         return feedback
 
     def record(self, prev_out: dict[str, Any]):
-        # set the DAG parent for the trace
-        self.trace.sync_dag_parent_and_hist()
+
+        exp: DSExperiment = None
 
         e = prev_out.get(self.EXCEPTION_KEY, None)
         if e is None:
-            self.trace.hist.append((prev_out["running"], prev_out["feedback"]))
+            exp = prev_out["running"]
+
+            # NOTE: we put below  operations on selections here, instead of out of the if-else block,
+            # to fit the corner case that the trace will be reset
+
+            # set the local selection to the trace as global selection, then set the DAG parent for the trace
+            if exp.local_selection is not None:
+                self.trace.set_current_selection(exp.local_selection)
+            self.trace.sync_dag_parent_and_hist()
+
+            self.trace.hist.append((exp, prev_out["feedback"]))
+
         else:
+            exp: DSExperiment = prev_out["direct_exp_gen"] if isinstance(e, CoderError) else prev_out["coding"]
+
+            # set the local selection to the trace as global selection, then set the DAG parent for the trace
+            if exp.local_selection is not None:
+                self.trace.set_current_selection(exp.local_selection)
+            self.trace.sync_dag_parent_and_hist()
+
             self.trace.hist.append(
                 (
-                    prev_out["direct_exp_gen"] if isinstance(e, CoderError) else prev_out["coding"],
+                    exp,
                     ExperimentFeedback.from_exception(e),
                 )
             )
+
             if self.trace.sota_experiment() is None:
                 if DS_RD_SETTING.coder_on_whole_pipeline:
                     #  check if feedback is not generated

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

@@ -182,27 +182,27 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):
         final_component = self.COMPLETE_ORDER[-1]
         has_final_component = True if DS_RD_SETTING.coder_on_whole_pipeline else False
         SOTA_exp_and_feedback_list = []
-        failed_exp_and_feedback_list = []
+        failed_exp_and_feedback_list_after_sota = []
         for exp, fb in search_list:
             if has_final_component:
                 if fb.decision:
                     SOTA_exp_and_feedback_list.append((exp, fb))
-                    failed_exp_and_feedback_list = []
+                    failed_exp_and_feedback_list_after_sota = []
                 else:
-                    failed_exp_and_feedback_list.append((exp, fb))
+                    failed_exp_and_feedback_list_after_sota.append((exp, fb))
             if exp.hypothesis.component == final_component and fb:
                 has_final_component = True
-        if max_retrieve_num is not None and (SOTA_exp_and_feedback_list or failed_exp_and_feedback_list):
+        if max_retrieve_num is not None and (SOTA_exp_and_feedback_list or failed_exp_and_feedback_list_after_sota):
             SOTA_exp_and_feedback_list = SOTA_exp_and_feedback_list[
                 -min(max_retrieve_num, len(SOTA_exp_and_feedback_list)) :
             ]
-            failed_exp_and_feedback_list = failed_exp_and_feedback_list[
-                -min(max_retrieve_num, len(failed_exp_and_feedback_list)) :
+            failed_exp_and_feedback_list_after_sota = failed_exp_and_feedback_list_after_sota[
+                -min(max_retrieve_num, len(failed_exp_and_feedback_list_after_sota)) :
             ]
         if return_type == "all":
-            return SOTA_exp_and_feedback_list + failed_exp_and_feedback_list
+            return SOTA_exp_and_feedback_list + failed_exp_and_feedback_list_after_sota
         elif return_type == "failed":
-            return failed_exp_and_feedback_list
+            return failed_exp_and_feedback_list_after_sota
         elif return_type == "sota":
             return SOTA_exp_and_feedback_list
         else:
@@ -266,11 +266,12 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):
     def last_exp_fb(
         self,
         search_type: Literal["all", "ancestors"] = "ancestors",
+        selection: tuple[int, ...] | None = None,
     ) -> tuple[DSExperiment, ExperimentFeedback] | None:
         """
         Access the last experiment and feedback
         """
-        search_list = self.retrieve_search_list(search_type)
+        search_list = self.retrieve_search_list(search_type, selection=selection)
         for exp, ef in search_list[::-1]:
             return exp, ef
         return None

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

@@ -119,15 +119,26 @@ class ExpGen2Hypothesis(DSProposalV2ExpGen):
         resp_dict = json.loads(response)
         return resp_dict
 
+    def get_exp_index(self, trace: DSTrace) -> int:
+        leaves: list[int] = trace.get_leaves()
+        if trace.sota_exp_to_submit is not None:
+            sota_submit_value = trace.sota_exp_to_submit.result.loc["ensemble"].iloc[0]
+            trace_scores = []
+            for i, leaf in enumerate(leaves):
+                if leaf == trace.current_selection[0]:
+                    continue
+                fb = trace.sota_experiment_fb(selection=(leaf,))
+                if fb is None:
+                    continue
+                final_score = fb[0].result.loc["ensemble"].iloc[0]
+                trace_scores.append((i, abs(final_score - sota_submit_value)))
+            if trace_scores:
+                return min(trace_scores, key=lambda item: item[1])[0]
+        return next((i for i, leaf in enumerate(leaves) if leaf != trace.current_selection[0]))
+
     def gen(self, trace: DSTrace) -> DSExperiment:
         # Ignore the selection argument and use all leaves instead.
-        leaves: list[int] = trace.get_leaves()
         sota_exp_fb = trace.sota_experiment_fb(selection=trace.current_selection)
-        exp_index = leaves[1] if trace.current_selection[0] == leaves[0] else leaves[0]
-
-        exp_to_merge_fb = trace.sota_experiment_fb(selection=(exp_index,))
-        if exp_to_merge_fb is None:
-            exp_to_merge_fb = trace.hist[exp_index]
 
         if sota_exp_fb:
             sota_exp_desc = T("scenarios.data_science.share:describe.exp").r(
@@ -139,9 +150,30 @@ class ExpGen2Hypothesis(DSProposalV2ExpGen):
             sota_exp_desc = ""
             eda_output = None
 
-        success_fb_list = trace.experiment_and_feedback_list_after_init(
-            return_type="sota", search_type="ancestors", selection=(exp_index,)
-        )
+        trace_fbs = []
+        # find the best exp to merge
+        leaves: list[int] = trace.get_leaves()
+        for leaf in leaves:
+            if leaf == trace.current_selection[0]:
+                continue
+
+            trace_fbs.append(
+                trace.experiment_and_feedback_list_after_init(
+                    return_type="sota",
+                    search_type="ancestors",
+                    selection=(leaf,),
+                )
+            )
+
+        num_to_slice = 20
+        if sum(len(fb_list) for fb_list in trace_fbs) > num_to_slice:
+            success_fb_trace_count = sum(1 for fb_list in trace_fbs if fb_list)
+            success_fb_list = [
+                fb for fb_list in trace_fbs for fb in fb_list[-(num_to_slice // success_fb_trace_count) :]
+            ]
+        else:
+            success_fb_list = [fb for fb_list in trace_fbs for fb in fb_list]
+
         if len(success_fb_list) > 0:
             exp_to_merge_fb_desc = T("scenarios.data_science.proposal.exp_gen.merge:trace").r(
                 exp_and_feedback_list=success_fb_list,
@@ -151,6 +183,11 @@ class ExpGen2Hypothesis(DSProposalV2ExpGen):
                 pipeline=DS_RD_SETTING.coder_on_whole_pipeline,
             )
         else:
+            exp_index = self.get_exp_index(trace)
+            exp_to_merge_fb = trace.sota_experiment_fb(selection=(exp_index,))
+            if exp_to_merge_fb is None:
+                exp_to_merge_fb = trace.hist[exp_index]
+
             exp_to_merge_fb_desc = T("scenarios.data_science.share:describe.feedback").r(
                 exp_and_feedback=exp_to_merge_fb,
                 heading="The feedback for the solution to be merged",
@@ -181,13 +218,14 @@ class ExpGen2Hypothesis(DSProposalV2ExpGen):
             scenario_desc=scenario_desc,
             sota_exp_desc=sota_exp_desc,
             sota_exp=sota_exp_fb[0] if sota_exp_fb else None,
-            hypothesis=new_hypothesis,
+            hypotheses=[new_hypothesis],
             pipeline=DS_RD_SETTING.coder_on_whole_pipeline,
             failed_exp_feedback_list_desc="",
         )
 
 
 class ExpGen2TraceAndMerge(ExpGen):
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.merge_exp_gen = MergeExpGen(self.scen)
@@ -396,7 +434,9 @@ class ExpGen2TraceAndMergeV3(ExpGen):
             else:
                 selection = (leaves[0],)
                 if trace.sota_exp_to_submit is not None:
-                    if trace.is_parent(trace.exp2idx(trace.sota_exp_to_submit), leaves[1]):
-                        selection = (leaves[1],)
+                    for i in range(1, len(leaves)):
+                        if trace.is_parent(trace.exp2idx(trace.sota_exp_to_submit), leaves[i]):
+                            selection = (leaves[i],)
+                            break
                 trace.set_current_selection(selection)
                 return self.merge_exp_gen.gen(trace)

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

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+import asyncio
+from datetime import timedelta
+from typing import TYPE_CHECKING
+
+from rdagent.app.data_science.conf import DS_RD_SETTING
+from rdagent.core.conf import RD_AGENT_SETTINGS
+from rdagent.core.proposal import ExpGen
+from rdagent.log import rdagent_logger as logger
+from rdagent.log.timer import RD_Agent_TIMER_wrapper, RDAgentTimer
+from rdagent.scenarios.data_science.loop import DataScienceRDLoop
+from rdagent.scenarios.data_science.proposal.exp_gen.merge import ExpGen2Hypothesis
+from rdagent.scenarios.data_science.proposal.exp_gen.trace_scheduler import (
+    RoundRobinScheduler,
+    TraceScheduler,
+)
+
+if TYPE_CHECKING:
+    from rdagent.scenarios.data_science.experiment.experiment import DSExperiment
+    from rdagent.scenarios.data_science.proposal.exp_gen.base import DSTrace, Experiment
+    from rdagent.utils.workflow.loop import LoopBase
+
+
+class ParallelMultiTraceExpGen(ExpGen):
+    """
+    An experiment generation strategy that enables parallel multi-trace exploration.
+
+    This generator is designed to work with the "Attribute Injection" model.
+    It uses a TraceScheduler to determine which parent node to expand, and
+    injects this parent context into the experiment object itself.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # The underlying generator for creating a single experiment
+        self.exp_gen = DataScienceRDLoop._get_exp_gen(
+            "rdagent.scenarios.data_science.proposal.exp_gen.DSExpGen", self.scen
+        )
+        self.merge_exp_gen = ExpGen2Hypothesis(self.scen)
+        self.trace_scheduler: TraceScheduler = RoundRobinScheduler()
+        self.max_trace_num = DS_RD_SETTING.max_trace_num
+
+    def gen(self, trace: "DSTrace") -> "Experiment":
+        raise NotImplementedError(
+            "ParallelMultiTraceExpGen is designed for async usage, please call async_gen instead."
+        )
+
+    async def async_gen(self, trace: DSTrace, loop: LoopBase) -> DSExperiment:
+        """
+        Waits for a free execution slot, selects a parent trace using the
+        scheduler, generates a new experiment, and injects the parent context
+        into it before returning.
+        """
+        timer: RDAgentTimer = RD_Agent_TIMER_wrapper.timer
+        logger.info(f"Remain time: {timer.remain_time_duration}")
+        local_selection: tuple[int, ...] = None
+
+        while True:
+
+            if timer.remain_time_duration >= timedelta(hours=DS_RD_SETTING.merge_hours):
+
+                if DS_RD_SETTING.enable_inject_knowledge_at_root:
+
+                    if len(trace.hist) == 0:
+                        # set the knowledge base option to True for the first trace
+                        DS_RD_SETTING.enable_knowledge_base = True
+
+                    else:
+                        # set the knowledge base option back to False for the other traces
+                        DS_RD_SETTING.enable_knowledge_base = False
+                # step 1: select the parant trace to expand
+                # Policy: if we have fewer traces than our target, start a new one.
+                if trace.sub_trace_count < self.max_trace_num:
+                    local_selection = trace.NEW_ROOT
+                else:
+                    # Otherwise, use the scheduler to pick an existing trace to expand.
+                    local_selection = await self.trace_scheduler.select_trace(trace)
+
+                if loop.get_unfinished_loop_cnt(loop.loop_idx) < RD_AGENT_SETTINGS.get_max_parallel():
+
+                    # set the local selection as the global current selection for the trace
+                    trace.set_current_selection(local_selection)
+                    # step 2: generate the experiment with the local selection
+                    exp = self.exp_gen.gen(trace)
+
+                    # Inject the local selection to the experiment object
+                    exp.set_local_selection(local_selection)
+
+                    return exp
+
+            else:
+                # enter the merging stage
+                # make sure the all loops are finished
+                if loop.get_unfinished_loop_cnt(loop.loop_idx) < 1:
+                    # disable reset in merging stage
+                    DS_RD_SETTING.coding_fail_reanalyze_threshold = 100000
+                    DS_RD_SETTING.consecutive_errors = 100000
+
+                    leaves: list[int] = trace.get_leaves()
+                    if len(leaves) < 2:
+                        trace.set_current_selection(selection=(-1,))
+                        return self.exp_gen.gen(trace)
+                    else:
+                        selection = (leaves[0],)
+                        if trace.sota_exp_to_submit is not None:
+                            if trace.is_parent(trace.exp2idx(trace.sota_exp_to_submit), leaves[1]):
+                                selection = (leaves[1],)
+                        trace.set_current_selection(selection)
+                        return self.merge_exp_gen.gen(trace)
+
+            await asyncio.sleep(1)

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

@@ -1,30 +1,35 @@
 scenario_problem:
   system: |-
     {% include "scenarios.data_science.share:scen.role" %}
-    The user is improving a Kaggle competition implementation iteratively through traces where each new trace is modified from the current SOTA in the trace. If new trace surpasses the current SOTA, it will be the new SOTA. If not, it will be a failed experiment.
-    You will be provide with: 
-      1. A detailed competition scenario description;
-      2. Previous SOTA experiments and feedbacks, which are past SOTA experiments indexed from oldest to newest;
-      3. Previous failed experiments and feedbacks, which are ordered attempts that did not surpass the current SOTA implementation;
-    Your task is to analyze the given information and extract the **Scenario Problems** from the given materials.
+    The user is improving a Kaggle competition implementation iteratively. Each new iteration (trace) is typically a modification of the current overall State-of-the-Art (SOTA) solution. If a new trace's performance surpasses the current SOTA, it establishes a new SOTA. Otherwise, it is considered a failed experiment.
 
-    ## Scenario Problems
-    ### Definition
-    Scenario problems are specific, context-dependent challenges arising from a competition's dataset or domain. They fall into two categories:
-    1. Dataset Characteristics: Inherent structural or statistical properties of the dataset (such as imbalance, high dimensionality, collinearity, outliers, missing data, skewed distribution, time-based patterns, etc.).
-    2. Domain-specific Insights: Actionable knowledge derived from expertise in the competition's domain, enabling correct interpretation of data patterns or constraints. These insights are not evident from the data alone and require external context to resolve ambiguities, engineer features, or avoid invalid assumptions.
+    You will be provided with:
+    1. A detailed competition scenario description;
+    2. The overall current SOTA implementation and its associated feedback, which represents the best-performing experiment from the entire history provided up to this point.
 
-    ### Specification
-    {{ problem_spec }}
+    Your task is to analyze the provided information (primarily the scenario and current SOTA, if available) and identify a concise list of **Key Challenges** or **Core Problems** relevant to achieving success in this competition and improving the target metric. Aim for **FEWER BUT BETTER** challenges (e.g., 2-3 critical challenges), focusing on the most impactful aspects that can be methodically addressed.
+
+    ### Core Analysis Dimensions for Identifying Challenges
+    1. **SOTA Alignment Analysis**: (If SOTA is provided) Systematically compare the current SOTA implementation against dataset properties and domain knowledge to identify discrepancies or areas representing core challenges to overcome for enhancement.
+    2. **Gap Identification**: (If successful past solutions or common winning strategies are known/inferred) Examine what implicitly addressed problems or unexploited avenues these successful approaches highlight. These gaps can represent current challenges.
+    3. **Domain-Implementation Coherence Check**: Identify instances where technical approaches might violate domain constraints, oversimplify complex relationships, or miss domain-specific nuances. These incoherencies are challenges.
+    4. **Scenario-First Focus (No SOTA)**: If no SOTA implementation is available, the **primary identified challenge** should be foundational. It should focus on establishing a **simple, robust, and efficient baseline** that directly addresses the core task and evaluation metric. Avoid overly complex initial challenges.
+
+    ## Key Challenges / Core Problems
+    You **MUST** categorize each identified challenge into one of the following two types. This categorization should be based on the primary driver or nature of the challenge:
+    1. **Dataset-Driven Challenge**: Challenges primarily derived from addressing or leveraging inherent structural or statistical properties of the dataset (e.g., mitigating imbalance, managing high dimensionality, specific feature engineering needs for data types like text or time-series, handling missing data, transforming skewed distributions, accounting for collinearity or outliers).
+    2. **Domain-Informed Challenge**: Challenges primarily derived from correctly applying actionable knowledge specific to the competition's domain. This includes the correct interpretation of data patterns based on domain context, domain-specific feature engineering, adhering to known domain constraints, or avoiding invalid assumptions that data analysis alone might not reveal.
+
+    ### Specification for each Identified Challenge
+    1. The challenge should be specific and fine-grained. Avoid general or vague statements.
+    2. The challenge should be technical or methodological. Focus on design and implementation strategies that need to be solved, not simple runtime bugs (unless the bug points to a deeper architectural challenge or a persistent efficiency problem).
+    3. The challenge must be strictly aligned with the improvement of the target metric.
+    4. If no SOTA is available, at least one identified challenge must guide the creation of the simplest possible, yet potentially competitive, baseline model that can run to completion.
     
-    ### Core Analysis Dimensions
-    1. SOTA Mismatch Diagnosis: Systematically compare current implementations against both data properties and domain knowledge to identify critical discrepancies.
-    2. Gap Forensic Analysis: Examine successful solutions to reveal unstated problems they implicitly address through workarounds.
-    3. Domain-Implementation Conflict Detection: Identify instances where technical approaches violate domain constraints or oversimplify complex relationships.
-    4. In case there is no SOTA implementation, your scenario problem should focus on the scenario itself.
-
+    {% if problem_output_format is not none %}
     ### Output Format
     {{ problem_output_format }}
+    {% endif %}
 
   user: |-
     # Scenario Description
@@ -36,13 +41,17 @@ scenario_problem:
 feedback_problem:
   system: |-
     {% include "scenarios.data_science.share:scen.role" %}
-    The user is improving a Kaggle competition implementation iteratively through traces where each new trace is modified from the current SOTA in the trace. If new trace surpasses the current SOTA, it will be the new SOTA. If not, it will be a failed experiment.
-    You will be provided with: 
-      1. A detailed competition scenario description;
-      2. Previous SOTA experiments and feedbacks, which are past SOTA experiments indexed from oldest to newest;
-      3. Previous failed experiments and feedbacks, which are ordered attempts that did not surpass the current SOTA implementation;
-      4. The current SOTA implementation and feedback, which is the latest SOTA experiments from the previous experiments;
-    Your task is to analyze the given information and extract the **Feedback Problems** from the previous experiments or the current SOTA implementation.
+    The user is improving a Kaggle competition implementation iteratively through traces. Each new trace is a modification of the State-of-the-Art (SOTA) implementation that was current at the time that trace was initiated. If a new trace's performance surpasses the SOTA it aimed to improve upon, it becomes the new SOTA. If not, it is considered a failed experiment.
+
+    You will be provided with:
+    1. A detailed competition scenario description;
+    2. A history of previous successfully experiments and their associated feedbacks, indexed or ordered from oldest to newest; the latest SOTA experiment accumulates all the improvements from the previous successful experiments.
+    3. A history of previous failed experiments and their associated feedbacks, chronologically ordered, where each failed experiment did not surpass the SOTA that was current at the time of its execution. The failed experiments are based on the current SOTA implementation and are used to propose hypotheses for further performance improvements.
+    4. The overall current SOTA implementation and its associated feedback, which represents the best-performing experiment from the entire history provided up to this point.
+
+    Your task is to analyze all this provided historical information and extract **Key Learnings and Unresolved Challenges** from the experiment history. These should guide concrete improvements in subsequent iterations.
+
+    ## Key Learnings and Unresolved Challenges
 
     
     {% if inject_diverse %}
@@ -53,104 +62,183 @@ feedback_problem:
     3. Do not do incremental exploration on the previous problems.
     {% endif %}
 
-    ## Feedback Problems
     ### Definition
-    Feedback problems are specific and fine-grained technical, or methodological issues within the previous experiments or the current SOTA implementation.
+    Key Learnings and Unresolved Challenges are specific, fine-grained technical or methodological observations, persistent issues, or patterns identified within previous experiments or the current SOTA implementation. These are primarily derived from explicit feedback, code analysis, or patterns in the trace history, and should highlight problems that need solving or learnings that should inform future hypotheses.
 
-    ### Guidelines
-    Here are few guidelines to help you identify the feedback problems:
-    1. Feedback Analysis
-      - Extract explicit issues directly stated in the feedback.
-      - Infer implicit issues from feedback context.
-    2. Code Review
-      - Feature Engineering. Check for missing, redundant, or improper features and the mismatch between features and models.
-      - Model Architecture. Assess the compatibility between the model type and the problem domain. Verify model architecture and hyperparameters.
-      - Ensemble. Check if ensemble methods are optimized. Identify underperforming base models to remove from the ensemble.
-      - Training. Validate hyperparameter (e.g., learning rate, batch size), loss functions, and regularization. Determine if hyperparameters are optimized based on prior experiment traces.
-    3. Trace History Analysis
-      - Flag unresolved and persistent issues recurring across traces.
-      - Highlight partial fixes (e.g., inappropriate feature engineering).
-      - Identify unexplored directions (e.g. new features, new model structures) from prior traces.
-      - Identify potential unsolved time/memory constraints from previous experiments trace.
+    ### Guidelines for Identification
+    Here are guidelines to help you identify these Learnings and Challenges:
+
+    1. **Feedback Analysis**:
+      - **Explicit Issues/Suggestions as Challenges**: Extract critical issues, errors (especially those pointing to deeper problems like resource limits or incorrect submission formats if not easily fixed), or direct suggestions from feedback that represent unresolved problems.
+      - **Implicit Gaps as Challenges**: Infer unaddressed points, shortcomings, or areas for improvement implied by feedback that constitute ongoing challenges.
+      - **Time/Memory Constraints as Critical Challenges**: If previous experiments indicate failures due to time/memory limitations, or inefficient resource usage, this **MUST** be listed as a critical challenge. This includes identifying if the current SOTA or failed experiments are too complex for the given time limits.
+
+    2. **Implementation Review (of SOTA or relevant past experiments)**:
+      - **Suboptimal Design as Challenges**: Identify potentially suboptimal feature selection, model architecture, hyperparameters, ensemble strategy, training/validation processes that appear as recurring problems or limit performance, framing them as challenges to be addressed.
+      - **Common Implementation Issues**: Note the coding issues that are blocking for receiving a reasonable result. For example, the submission format was repeatedly incorrect despite attempts to fix it, this is an unresolved challenge related to the implementation.
+
+    3. **Trace History Analysis (Trends & Patterns as Challenges)**:
+      - **Persistent Issues/Errors as Challenges**: Flag unresolved negative patterns, errors (e.g., recurrent `zipfile.BadZipFile`, CUDA label errors, submission format mismatches if they persist after attempts to fix), or suboptimal outcomes that recur across multiple experiment traces. These represent core unresolved challenges.
+      - **Ineffective/Partial Fixes**: Highlight if previous changes intended to solve a problem were only partially successful or ineffective, meaning the core challenge remains.
+      - **Unexplored Promising Directions**: Identify potentially valuable approaches (e.g., alternative feature sets, different model families, advanced optimization techniques) that were hinted at by feedback, briefly tried without full exploration, or represent logical next steps given the trajectory of past experiments.
+      - **Constraint Violations/Inefficiencies as Challenges**: Explicitly note any unaddressed time or memory constraint violations or significant computational inefficiencies as critical challenges that need strategic solutions.
+
+    ### Specification for each Learning/Challenge
+    1. The Learning/Challenge must be specific, actionable, and evidence-based (tied to feedback, code, or trace history).
+    2. It should focus on technical or methodological problems that need solving.
+    3. Clearly state the learning or articulate the challenge.
+    4. Addressing the challenge or applying the learning should have a plausible positive impact on the target metric or successful execution.
+    5. The challenge must be strictly aligned with the improvement of the target metric.
     
-    ### Specification
-    {{ problem_spec }}
-
+    {% if problem_output_format is not none %}
     ### Output Format
     {{ problem_output_format }}
+    {% endif %}
 
   user: |-
     # Scenario Description
     {{ scenario_desc }}
-    
+
     # Previous Experiments and Feedbacks
     {{ exp_and_feedback_list_desc }}    
 
     # Current SOTA Implementation
     {{ sota_exp_desc }}
 
+scenario_description: |-
+  {% if use_raw_description -%}
+  ====== Background ======
+  {{ raw_description }}
+
+  {% else %}
+  ====== Background ======
+  {{ background }}
+
+  {% if eda_output is not none %}
+  ====== Data Overview (EDA) ======
+  {{ eda_output }}
+  {% endif %}
+
+  ====== Submission Format ======
+  Please ensure your submission adheres to the following specifications:
+  {{ submission_specifications }}
+
+  ====== Important Guidelines ======
+  Before submitting your results, please note the following:
+  - We have numerous tests in place to check your code.
+  - Ensure your submission is genuine.
+  - Do not manipulate data or return values solely to pass preliminary tests, as this will not lead to successful final evaluation.
+
+  {% endif %}
+
+  ====== Evaluation ======
+  {% if not use_raw_description and metric_name %}
+  The primary evaluation metric for this task is: **{{ metric_name }}**.
+  {% endif %}
+  This metric is considered better when it is **{% if metric_direction %}larger{% else %}smaller{% endif %}**.
+
+  {% if evaluation is not none %}
+  Additional Evaluation Details:
+  {{ evaluation }}
+  {% endif %}
+
+  {% if time_limit %}
+  ====== Time Limit ======
+  Your code's execution is limited to **{{ time_limit }}**.
+  Please optimize your model and parameters to ensure your code runs within this specified time constraint.
+  {% endif %}
+
 hypothesis_gen:
   system: |-
     {% include "scenarios.data_science.share:scen.role" %}
-    The user is improving a Kaggle competition implementation iteratively through traces where each new trace is modified from the current SOTA in the trace. If new trace surpasses the current SOTA, it will be the new SOTA. If not, it will be a failed experiment.
-    You will be provided with: 
-      1. A detailed competition scenario description;
-      2. Previous SOTA experiments and feedbacks, which are past SOTA experiments indexed from oldest to newest;
-      3. Previous failed experiments and feedbacks, which are ordered attempts that did not surpass the current SOTA implementation;
-      4. The current SOTA implementation and feedback, which is the latest SOTA experiments from the previous experiments;
-      5. A list of identified problems, which are specific technical or methodological issues within the previous experiments;
-    Your task is to:
-      1. **Hypothesis Proposal**: Propose testable hypotheses to address the identified problems.
-      2. **Hypothesis Evaluation**: Evaluate the proposed hypotheses across multiple dimensions.
+    The user is iteratively improving a Kaggle competition implementation. Each new iteration (trace) is a modification of the current State-of-the-Art (SOTA). If a new trace surpasses the current SOTA, it becomes the new SOTA. Otherwise, it's a failed experiment.
+    You will be provided with:
+    1. A detailed competition scenario description.
+    2. A history of previous successfully experiments and their associated feedbacks, indexed or ordered from oldest to newest; the latest SOTA experiment accumulates all the improvements from the previous successful experiments.
+    3. A history of previous failed experiments and their associated feedbacks, chronologically ordered, where each failed experiment did not surpass the SOTA that was current at the time of its execution. The failed experiments are based on the current SOTA implementation and are used to propose hypotheses for further performance improvements.
+    4. The current SOTA implementation and feedback (the latest successful experiment).
+    5. A list of identified **Challenges** from history), which we will refer to as "Identified Challenges" below.
+
+    Your task is to perform two main steps:
+    1. **Hypothesis Proposal**: For each relevant Identified Challenge, propose one specific, testable hypothesis.
+    2. **Hypothesis Evaluation**: Evaluate each proposed hypothesis 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.
+    To help you propose hypotheses, the user may provide a list of ideas for each Identified Challenge. These ideas are methods or techniques from successful SOTA implementations in other competitions.
+    Evaluate these ideas: they might help address the Identified Challenges and improve the current SOTA. You must decide whether to use them. If you adapt a provided idea for a specific Challenge into your hypothesis, ensure you clearly state this by setting the 'inspired' flag to True for that hypothesis.
     {% endif %}
 
     # Task 1: Hypothesis Proposal
-    For each identified problem, propose a hypothesis to improve the current SOTA implementation.
+    First note that the user might provide a list of challenges containing duplicates. You should only propose one hypothesis for each unique challenge. If a challenge is a duplicate of a previous one, you can skip it.
+    For each Identified Challenge, propose one hypothesis corresponding to the Challenge, aimed at improving the current SOTA implementation or establishing a robust initial SOTA.
 
-    ## Hypothesis Guidelines
-    Here are few guidelines to help you formulate hypotheses:
-    1. Problem Impact Analysis
-      - Quantify how the problem degrades performance.
-    2. Previous Experiments Analysis
-      - For previous SOTA experiments, analyze insights and implicit patterns that can be leveraged to improve the current SOTA implementation.
-      - For failed experiments, think about the persistent problems they facing. If these experiments consistently failed due to time/memory constraints, prioritize changes on efficiency.
-    3. Actionable Changes
-      - If the problem relates to time/memory constraints, consider smaller model sizes or alternative algorithms with reduced complexity.
-      - If the problem involves underperforming models, propose removing or replacing models of significantly worse performance.
-      - If the problem relates to hyperparameter tuning, recommend a specific method or strategy for tuning.
-    4. Note on Time/Memory Constraints
-      - If prior experiments failed due to time/memory limitations, assume your new hypothesis will face the same constraints. In this case, prioritize efficiency and **ONLY** response to the problems related to time/memory constraints in your response dictionary.
-      - Besides, do not compromise performance merely for efficiency since the current SOTA implementation do not encounter the constraints. You should think about how to balance the efficiency and performance so that your new hypothesis can be executed successfully and achieve satisfactory performance. 
-    5. Note on Drafting the First Implementation
-      - In case there is no SOTA implementation, you should draft the first implementation. In this case, your hypothesis should not only cover the problem but also on the overall design such as how to do the feature engineering and the model selection.
+    ## 1.1. Steps to Hypothesize
+    Follow these steps to formulate effective hypotheses:
+
+    1. **Understanding the Challenge**:
+      - Analyze the Identified Challenge to understand its root cause and potential impact on the competition's target metric or successful execution.
+      - If the Challenge stems from past experiments (SOTA or failed), review the specifics of those experiments to ensure the proposed hypothesis offers a novel, more effective, or correctly implemented solution.
+      - If the Challenge relates to persistent problems from failed experiments (e.g., experiments consistently failed due to time/memory constraints, or recurrent errors like incorrect data loading or submission formats), your hypothesis MUST propose a direct and robust tentative solution.
+    2. **Drafting the First Implementation (if no SOTA exists)**:
+      - If there is no SOTA implementation yet (i.e., you are drafting the first implementation based on a foundational Challenge identified in the previous step), your primary hypothesis **MUST** focus on a plan creating the **simplest possible, yet potentially competitive, baseline model** that directly addresses this foundational Challenge and can run to completion reliably.
+      - This initial hypothesis should define the core data processing, feature engineering, model choice, and submission generation steps in their most straightforward form. Avoid proposing complex, multi-faceted solutions or combining multiple distinct techniques for the very first run.
+    3. **Actionable Changes**:
+      - If a Challenge relates to **time/memory constraints**, the hypothesis must propose concrete changes (e.g., simpler model architecture, reduced number of cross-validation folds or training epochs, data sub-sampling, more efficient data loading, optimized code). The aim is to create a solution that can run to completion reliably. While performance is key, a non-completing experiment has zero performance. Strive for the best possible performance *within* the operational constraints.
+      - If a Challenge involves underperforming models (e.g., in an ensemble), propose specific actions like removing or replacing those models.
+      - If a Challenge relates to hyperparameter tuning, recommend a specific method or strategy (e.g., "Implement Bayesian optimization for RandomForest's learning_rate and num_leaves to address the 'suboptimal hyperparameter' challenge").
+      - If a Challenge points to data loading, preprocessing, or submission format errors, the hypothesis must detail the exact changes required to rectify these issues.
     {% if enable_idea_pool %}
-    6. 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.
+    4. **Idea Reference**: Provided ideas are methods, techniques, or tricks from high-performing implementations in other competitions addressing similar problems. Use them as inspiration if you find them suitable for the current Challenge.
     {% endif %}
 
-    ## Hypothesis Specification
-    {{ hypothesis_spec }}
+    ## 1.2. Guidelines for Writing Hypotheses
 
+    1. **Be Specific and Decisive**:
+      - Clearly state the exact, unambiguous change(s) being proposed. Avoid vague goals like "improve the model" or "optimize the pipeline."
+      - The hypothesis must propose a single, clear course of action. Do not suggest alternatives (e.g., "try method A or method B").
+      - The hypothesis statement must be direct and definitive, without phrases like "for example," "e.g.," or "might involve."
+      - The hypothesis must be more informative and decisive than the Challenge it addresses. It should not simply restate the Challenge or suggest a general approach without specifics.
+    2. **Ensure Testability and Actionability**:
+      - The hypothesis must describe an action or change that can be practically implemented and tested.
+      - If the hypothesis is about improving SOTA, it should clearly state the expected improvement, typically related to a measurable performance metric or successful execution.
+        - *Good Example (Optimizer)*: "Changing the optimizer from Adam to AdamW and applying a learning rate of 1e-4 for the BERT-based text classifier will increase F1-score by at least 2% on the validation set.""
+        - *Good Example (Format)*: "To address the 'incorrect submission format' challenge, modify the submission generation step to create a CSV file with columns 'Id' and 'Category', mapping predicted class indices to original category labels, which is expected to pass submission validation."
+        - *Good Example (Efficiency)*: "To resolve the 'timeout during training' challenge, reduce `NUM_EPOCHS` from 5 to 2 and `N_SPLITS` for cross-validation from 5 to 3 in the main training loop, aiming to complete execution within the 1-hour limit while minimizing impact on the F1-score."
+        - *Poor Example*: "Tune the model for better results."
+      - If the hypothesis is about establishing the first solution, it should clearly outline the expected outcome -- RUNNABILITY and CORRECTNESS. Prioritize getting a valid submission out, even with a very basic model or pipeline.
+        - *Good Example*: "Implement a simple RandomForest classifier with default parameters, using 5-fold cross-validation for model evaluation. This will lead to a decent baseline model that can run to completion and generate a valid submission file."
+    3. **Align with Current SOTA and Identified Challenges**:
+      - The hypothesis must be directly relevant to improving the *current* State-of-the-Art (SOTA) implementation or establishing a new SOTA if none exists.
+      - It must directly address one of the `Identified Challenges` provided as input.
+    4. **Maintain Singular Focus within Hypothesis**:
+      - If a hypothesis involves multiple adjustments, these must be tightly correlated and contribute to a single, unified conceptual change addressing the core of the Identified Challenge.
+      - Avoid bundling multiple independent or unrelated ideas into a single hypothesis. Each hypothesis should test one core concept.
+    5. **Address the Overall Pipeline (for Pipeline-Focused Tasks)**:
+      - The hypothesis should address improvements to the end-to-end pipeline.
+      - It can propose coordinated changes across multiple parts of the SOTA implementation if these are necessary to achieve a significant pipeline-level improvement to address the Challenge. (Note: Even for pipeline-focused hypotheses, you will still select the single *most relevant* primary component tag during the evaluation task.)
 
     # Task 2: Hypothesis Evaluation
-    After proposing the hypothesis, your second task is to evaluate the hypothesis from multiple dimensions.
+    After proposing one hypothesis for each relevant Identified Challenge, evaluate each one.
 
-    ## Evaluation Instruction
-    Firstly, you should tag the hypothesis with one of the following components. If the hypothesis is related to multiple components, you should choose the most relevant one.
-    {{ component_desc }}
+    ## 2.1. Evaluation Instruction
+    For each individual hypothesis you proposed in Task 1, perform the following two evaluation steps:
 
-    Secondly, please score the proposed hypothesis from 1 to 10 for each of the following dimensions (where 1 means lowest and 10 means highest):
-    1. Problem-Hypothesis Alignment: How well the hypothesis addresses the identified problem.
-    2. Expected Impact: The estimated improvement after applying the hypothesis to current SOTA implementation.
-    3. Novelty: Degree of innovation compared to previous attempts. If the proposed hypothesis is similar to previous experiments' hypothesis, assign novelty score to one.
-    4. Feasibility: The ease of implementing the proposed hypothesis in the current SOTA implementation.
-    5. Risk-Reward Balance: The exploration-exploitation balance of the proposed hypothesis.
+    1. **Assign a Component Tag:** Assign a single component tag to the hypothesis. Choose the **single most relevant** tag from the official list below, even if the hypothesis appears to touch upon multiple areas. Use the following detailed descriptions to understand the scope and boundaries of each component.
+
+      - **`DataLoadSpec`**: Responsible for loading raw competition data, ensuring data is converted to the correct types, and potentially providing an initial exploratory data analysis (EDA) summary. (e.g., fixing `zipfile.BadZipFile` by improving loading logic).
+      - **`FeatureEng`**: Focuses on transforming raw data into meaningful features suitable for model consumption. Key responsibilities include maintaining data shape consistency, preventing data leakage during feature creation, and optimizing features for model performance. Feature engineering should be model-agnostic.
+      - **`Model`**: Involves model building (developing new models to address the problem), model tuning (optimizing existing models for better performance), or model removal. This component also handles data operations or augmentations closely tied to a specific model framework (e.g., PyTorch `Datasets` & `DataLoaders`, TensorFlow `tf.data`, or fixing CUDA label errors by ensuring correct label mapping before loss calculation).
+      - **`Ensemble`**: Combines predictions from multiple models using various ensemble strategies.
+      - **`Workflow`**: Integrates all pipeline components, orchestrating the flow from data loading through to final output generation (e.g., correcting `submission.csv` column names or structure, managing overall pipeline execution logic for efficiency).
+
+    2. **Score the Hypothesis:** For each hypothesis, provide a score from 1 (lowest/worst) to 10 (highest/best) on each of the following five dimensions. Base your scores on all provided information.
+
+      - **Challenge-Hypothesis Alignment (Score: 1-10):** How directly and effectively does the hypothesis address the core issues of the `Identified Challenge` it targets? A higher score means a stronger, more direct alignment.
+      - **Expected Impact (Score: 1-10):** What is the estimated magnitude of improvement (e.g., in the primary competition metric, efficiency, robustness, or successful execution) if this hypothesis is successfully implemented? Higher scores for greater positive impact.
+      - **Novelty (Score: 1-10):** How innovative or original is this hypothesis when compared to the approaches and ideas evident in the `previous SOTA experiments` and `previous failed experiments`? Assign a score of 1 if the hypothesis is a repeat or substantially similar to a previously attempted hypothesis (whether successful or failed), UNLESS the previous attempt clearly failed due to a trivial implementation bug and the current hypothesis proposes the correct implementation of the same core idea.
+      - **Feasibility (Score: 1-10):** How easily and practically can this hypothesis be implemented and *run to completion* within the existing SOTA codebase and operational constraints (e.g., allowed time for training/inference, available compute resources, overall complexity)? Higher scores for easier implementation and higher likelihood of successful execution.
+      - **Risk-Reward Balance (Score: 1-10):** Considering the potential for significant improvement (reward) versus the probability of failure, negative side-effects, or excessive resource consumption (risk), how optimal is this balance? A high score indicates a favorable balance.
+      - **Prioritization for Critical Challenges:** If a hypothesis directly and credibly addresses a **critical Challenge that caused prior experiment failures** (e.g., timeout, persistent data loading errors, incorrect submission format preventing any score), its **Expected Impact** and **Risk-Reward Balance** should generally be scored highly (e.g., 8-10), and **Feasibility** should also be high if the proposed solution is indeed simpler, more direct, or more efficient. This ensures such critical hypotheses are prioritized.
 
     {% if inject_diverse %}
     # Focus on Diversity!!
@@ -160,10 +248,11 @@ hypothesis_gen:
     3. Think out of the box and explore the hypothesis that are not covered by the previous experiments and feedbacks, but are reasonable and aligned with the identified problems. 
     4. Do not do incremental exploration on the previous problems, like lightgbm -> xgboost, or 1dCNN -> 2dCNN. Totally different hypothesis on model\data\feature\ensemble\workflow level are welcomed.
     {% endif %}
-
-
+    
+    {% if hypothesis_output_format is not none %}
     ## Final Output Format in JSON Schema:
     {{ hypothesis_output_format }}
+    {% endif %}
     
   user: |-
     # Scenario Description
@@ -175,40 +264,97 @@ hypothesis_gen:
     # Current SOTA Implementation
     {{ sota_exp_desc }}
 
-    # Identified Problems{% if enable_idea_pool %} with Sampled Ideas{% endif %}
+    # Identified Challenges{% if enable_idea_pool %} with Sampled Ideas{% endif %}
     {{ problems }}
 
 task_gen:
   system: |-
     {% include "scenarios.data_science.share:scen.role" %}
-    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 provided with: 
-      1. A detailed competition scenario description;
-      2. Previous SOTA experiments and feedbacks, which are past SOTA experiments indexed from oldest to newest;
-      3. Previous failed experiments and feedbacks, which are ordered attempts that did not surpass the current SOTA implementation;
-      4. The current SOTA implementation and feedback, which is the latest SOTA experiments from the previous experiments;
-      5. A proposed hypothesis to improve the current SOTA implementation;
+    The user is iteratively developing a Kaggle competition solution. Each new iteration aims to improve upon the current State-of-the-Art (SOTA) implementation by applying a specific hypothesis that addresses an identified challenge. The new trace is based on the current SOTA; the SOTA itself evolves.
 
-    # Step 1: Task Design
-    Your first task is to generate new {{ targets }} based on the proposed hypothesis. Your task should very detailed with specific steps and instructions. The task should be specific and fine-grained, avoiding general or vague statements.
+    You will be provided with the following inputs:
+    1. **Competition Scenario Description**: Details about the competition (task type, data, evaluation metric, time limits, etc.).
+    2. **Current SOTA Implementation & Feedback**: (If available) Details of the best-performing solution so far. **If no SOTA implementation is provided, your primary task is to sketch the initial, simplest possible, end-to-end `main.py` workflow.**
+    3. **Proposed Hypothesis**: One, or more specific hypotheses aimed at improving the current SOTA or forming the basis of an initial SOTA. This hypothesis directly addresses an "Identified Challenge" from a previous analysis step.
+    4. **Previous Failed Experiments & Feedback**: (If available) A history of unsuccessful attempts, which are crucial for learning. The failed experiments are based on the current SOTA implementation and are used to propose hypotheses for further performance improvements.
 
-    ## Specification
-    {{ task_specification }}
+    Your primary goal is to generate a detailed, step-by-step **sketch or refinement plan** for a new data processing and modeling pipeline, specifically for the main workflow script (`main.py`), that effectively implements the `Proposed Hypothesis`. This sketch will guide a developer to write the code correctly.
 
-    ## Task Design Guidelines
-    1. The task should be concise with several steps each only in a few sentences. 
-    2. DO NOT repeat the details which has already included in the SOTA code. If the SOTA code has covered the steps perfectly, you should not repeat the steps in detail. 
-    3. DO NOT write any code in the task description!
-    4. Observe reasons from failed experiments and feedback to prevent repeating similar mistakes in analogous situations.
-    5. Specific and Non-Vague
-      - Avoid vague statements like "choose a proper model" Instead, specify the exact task to be made.
-      - No phrases like "for example" or "eg.," should be used in the task. Give a clear decision in the task.
-    6. Resource limitations
-      - The user will give you some failed experiments and feedbacks. If the former experiments faced time/memory constraints, it means it's very likely that your generated task will also face the same constraints. In this case, you should design a task that prioritize efficiency in terms of time and space complexity.
+    ### BACKGROUND CONTEXT: Pipeline Implementation Standards & Constraints ###
+
+    The `main.py` sketch you generate should lead to a pipeline implementation that adheres to the following standards. These are guiding principles for the final *outcome* of your sketch:
+
+    1. **Program Execution**: The resulting `main.py` script must be executable via `python main.py` without command-line parameters. Configurations should be hardcoded for simplicity.
+    2. **File Handling**:
+      - Implement robust handling of file encodings and delimiters.
+      - Input files are under `{% include "scenarios.data_science.share:scen.input_path" %}`. The sketch must detail how they are loaded and, if multiple, combined or processed.
+      - Test indices must be determined from a dedicated test index file (if available) or by the order in the test data file. **Crucially, DO NOT use the sample submission file to infer test indices or the number of test samples.**
+      - Ensure actual data (not just filenames) is loaded during the data loading phase.
+      - If data is in zip files, the sketch should advise on robust loading, e.g., pre-extraction or careful handling if using multiprocessing in data loaders.
+    3. **Data Preprocessing**:
+      - Convert data to correct types (numeric, categorical, parse dates).
+      - Optimize memory usage (e.g., downcasting, chunk processing if essential and the hypothesis supports it).
+      - Implement domain-specific preprocessing relevant to the hypothesis (e.g., text tokenization, image resizing/augmentation).
+    4. **Code Standards**:
+      - The pipeline must **NOT** use progress bars (e.g., `tqdm`) in the submission code.
+      - Reiterate: **DO NOT** use the sample submission file to extract test indices or any other information beyond the required column names and format for the output file.
+      - Ensure no features are inadvertently excluded during processing.
+    5. **Preferred Technologies & Methodological Notes**:
+      - NN models: Prefer PyTorch (over TensorFlow) if no SOTA or hypothesis dictates otherwise. Prioritize fine-tuning pre-trained models.
+      - Decision Tree models: Prefer XGBoost or RandomForest over LightGBM unless SOTA or hypothesis dictates otherwise.
+    6. **General Data Science Considerations**:
+      - Design for scalability.
+      - Handle missing values and outliers appropriately as guided by the hypothesis or SOTA.
+      - Ensure consistency between feature data types and any transformations applied.
+      - Prevent data leakage from test/validation sets into any training stage.
+    7. **Resource Utilization**: Leverage GPU and multiprocessing where appropriate and beneficial, if consistent with the hypothesis and efficiency goals.
+    8. **Metric Calculation and Storage (`scores.csv`)**:
+      - Calculate the official competition metric on a proper validation set (e.g., K-fold CV, typically 3-5 folds unless efficiency dictates fewer). Save results to `scores.csv`.
+      - The sketch must ensure this step is included. A successful run should always produce scores.
+      - `scores.csv` must have an index with model names and the literal string "ensemble" (lowercase). Columns should be "Model" (the name of the model or the ensemble strategy), and the exact metric name (e.g., "AUC").
+      - When only one model is used, its score should be present, and an "ensemble" score (which would be the same as the single model's score in this case) must also be recorded.
+      - Ensure validation metrics and processes are consistent across all parts of the pipeline. Avoid changes that would alter how validation metrics are calculated unless that is part of the hypothesis.
+    9. **Submission File (`submission.csv`)**: Generate `submission.csv` in the **exact format** required (column names, order, data types), as detailed by `sample_submission.csv` in the `Competition Scenario Description`. This is a critical step.
+
+    ### END OF BACKGROUND CONTEXT ###
+
+    # Guidelines for Sketching the `main.py` Workflow
+
+    YOUR TASK IS TO create a conceptual sketch for drafting or updating the `main.py` workflow. This is a plan, not code.
+
+    1. **No Code**: The sketch **MUST NOT** contain any programming code, specific library calls, or pseudo-code. Describe steps conceptually (e.g., "Load training data from {% include "scenarios.data_science.share:scen.input_path" %}/train.csv"). List specific algorithm names where appropriate (e.g., "Apply XGBoost classifier," "Use Isotonic Regression for calibration").
+    2. **Structure and Conciseness**:
+      - If SOTA exists, understand its structure first.
+      - If no SOTA, outline a clear, logical sequence of steps for the new `main.py`.
+    3. **Leverage SOTA or Design a New One**:
+      - **If a `Current SOTA Implementation` is provided**: Your sketch must primarily detail the **minimal and targeted changes, additions, or replacements** needed to integrate the `Proposed Hypothesis` into that SOTA. Focus only on what needs to change.
+      - **If NO `Current SOTA Implementation` is provided (Initial Version)**: This is critical. Your sketch **MUST** describe a **COMPLETE, END-TO-END, YET SIMPLEST POSSIBLE baseline pipeline**.
+        - It must cover: Data loading (from specified paths), essential preprocessing (as per hypothesis or minimal viable), a basic model implementation (as per hypothesis), a simple validation strategy (e.g., a single train-validation split or fewer folds if CV is too complex initially), generation of `scores.csv`, and `submission.csv` in the correct format.
+        - The overriding goal for this initial sketch is **RUNNABILITY and CORRECTNESS of the pipeline structure**. Prioritize getting a valid submission out, even with a very basic model. Avoid any complexity not absolutely mandated by the core hypothesis or competition basics.
+    4. **Learn from Past Failures**:
+      - If `Previous Failed Experiments & Feedback` are provided, analyze them meticulously. Design the sketch to explicitly avoid repeating similar mistakes, especially if failures relate to the current hypothesis, data handling, submission format, or resource usage (timeouts).
+      - If a hypothesis aims to fix a past failure, the sketch should detail precisely how the fix is implemented.
+    5. **Specificity and Clarity**:
+      - Be unambiguous. Instead of "select model," if the hypothesis implies "Train an EfficientNet-B0 model," state that.
+      - The sketch must be definitive. No open-ended options or phrases like "for example," or "e.g.," within a step's action.
+    6. **Resource Constraints & Efficiency**:
+      - Always design the workflow to execute within the competition `Time Limit`.
+      - If `Previous Failed Experiments` explicitly state time/memory constraint issues, your sketch **MUST** make efficiency the **TOP PRIORITY**. Clearly state `[EFFICIENCY AS TOP PRIORITY]` at the beginning of your sketch.
+      - The sketch must then detail *specific measures* to achieve this (e.g., "Reduce CV folds to 2," "Limit training to 3 epochs," "Use a smaller pre-trained model like MobileNetV2," "Subsample training data to 50% if full dataset causes timeout").
+      - Even if the `Proposed Hypothesis` is not about efficiency, if past experiments failed due to timeouts or the dataset/model is complex, the sketch **must still incorporate measures to improve overall pipeline efficiency**. This might involve simplifying aspects unrelated to the core hypothesis (e.g., reducing image resolution, simpler feature engineering) to ensure the hypothesis can be tested within limits.
+      - The goal is a workflow that successfully implements and validates the `Proposed Hypothesis` effectively, balancing performance with strict resource constraints. An experiment that times out provides no information.
       - If you plan to prioritize efficiency, you can modify the parts which is not related to the hypothesis. Which means your task should still able to validate the hypothesis.
       - Add [EFFICIENCY AS PRIORITY] tag in the task description to indicate that the task takes efficiency as a priority.
       - Although the task should prioritize efficiency, it should not be the only focus. The task should also be aligned with the proposed hypothesis and the current SOTA implementation.
-
+    7. **Reminders of Common Mistakes (Especially for New `main.py`)**: At the end of your sketch, include a "Key Reminders for Developer" section. Add the following reminders if appropriate.
+      - Ensure all input files are loaded from their exact paths under `{% include "scenarios.data_science.share:scen.input_path" %}` (e.g., `{% include "scenarios.data_science.share:scen.input_path" %}<competition_name>/train.csv`)."
+      - Verify `submission.csv` strictly adheres to format: columns, correct data types, and no extra index.
+      - "Implement correct label mapping for classification tasks (e.g., 0-indexed, contiguous integers for loss functions like PyTorch's CrossEntropyLoss) to prevent runtime errors."
+      - Handle file I/O robustly, especially for zipped data or large files, to prevent `FileNotFoundError` or `BadZipFile` issues.
+      - Confirm no `tqdm` or other progress bars are in the final script.
+      - Double-check that validation scores are saved correctly to `scores.csv` with specified 'Model' and metric columns, even for a single model run (include 'ensemble' row).
+    
+    {% if task_output_format is not none %}
     ## [Partial Response Format 1] Task Output Format:
     {{ task_output_format }}
 
@@ -224,18 +370,29 @@ task_gen:
       "task_design": ---The dict corresponding to task output format---,
       {% if workflow_check %}"workflow_update": ---A string corresponding to workflow description--- {% endif %}
     }
+    {% endif %}
     
   user: |-
-    # Scenario Description
+    # Competition Scenario Description
     {{ scenario_desc }}
 
-    # Current SOTA Implementation
+    # Data Folder Structure (All files are under {% include "scenarios.data_science.share:scen.input_path" %})
+    {{ data_folder_info }}
+
+    # Current SOTA Implementation & Feedback
     {{ sota_exp_desc }}
 
-    # Proposed Hypothesis you should strictly follow:
-    {{ hypothesis }}
+    # Proposed Hypothesis
+    This sketch should implement the following hypotheses:
+
+    {% for hypothesis in hypotheses %}
+    ## {{ hypothesis.problem_name }}
+    **Why:** {{ hypothesis.problem_desc }}
+    **Hypothesis:** {{ hypothesis.hypothesis }}
+
+    {% endfor %}
+    # Previous Failed Experiments & Feedback (e.g., experiments that did not pass evaluation, encountered bugs, or failed to surpass SOTA performance)
 
-    # 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:
@@ -268,11 +425,6 @@ idea_sample:
     {{ problem_ideas }}
 
 specification:
-  problem: |-
-    1. The problem should be specific and fine-grained. Avoid general or vague statements. 
-    2. The problem should technical or methodological. Focus on design and implementation flaws, not runtime errors.
-    3. The problem should be strictly aligned with the improvement of target metric. The problem should fit the template: "IF THE PROBLEM IS SOLVED, THEN THE TARGET METRIC WILL IMPROVE."
-  
   hypothesis: |-
     1. Each hypothesis should be specific and non-vague.
       - Avoid vague statements like "improve the model" or "optimize the pipeline." Instead, specify the exact changes to be made. Do not use ambiguous changes like "try method A or method B". 

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

@@ -87,47 +87,7 @@ feedback_problem:
     {{ sota_exp_desc }}
 
 scenario_description: |-
-  {% if use_raw_description -%}
-  ====== Background ======
-  {{ raw_description }}
-
-  {% else %}
-  ====== Background ======
-  {{ background }}
-
-  {% if eda_output is not none %}
-  ====== Data Overview (EDA) ======
-  {{ eda_output }}
-  {% endif %}
-
-  ====== Submission Format ======
-  Please ensure your submission adheres to the following specifications:
-  {{ submission_specifications }}
-
-  ====== Important Guidelines ======
-  Before submitting your results, please note the following:
-  - We have numerous tests in place to check your code.
-  - Ensure your submission is genuine.
-  - Do not manipulate data or return values solely to pass preliminary tests, as this will not lead to successful final evaluation.
-
-  {% endif %}
-
-  ====== Evaluation ======
-  {% if not use_raw_description and metric_name %}
-  The primary evaluation metric for this task is: **{{ metric_name }}**.
-  {% endif %}
-  This metric is considered better when it is **{% if metric_direction %}larger{% else %}smaller{% endif %}**.
-
-  {% if evaluation is not none %}
-  Additional Evaluation Details:
-  {{ evaluation }}
-  {% endif %}
-
-  {% if time_limit %}
-  ====== Time Limit ======
-  Your code's execution is limited to **{{ time_limit }}**.
-  Please optimize your model and parameters to ensure your code runs within this specified time constraint.
-  {% endif %}
+  {% include "scenarios.data_science.proposal.exp_gen.prompts_v2:scenario_description" %}
 
 hypothesis_gen:
   system: |-
@@ -320,24 +280,4 @@ task_gen:
       - Double-check that validation scores are saved correctly to `scores.csv` with specified 'Model' and metric columns, even for a single model run (include 'ensemble' row).
 
   user: |-
-    # Competition Scenario Description
-    {{ scenario_desc }}
-
-    # Data Folder Structure (All files are under {% include "scenarios.data_science.share:scen.input_path" %})
-    {{ data_folder_info }}
-
-    # Current SOTA Implementation
-    {{ sota_exp_desc }}
-
-    # Proposed Hypothesis
-    This sketch should implement the following hypotheses:
-
-    {% for hypothesis in hypotheses %}
-    ## {{ hypothesis.problem_name }}
-    **Why:** {{ hypothesis.problem_desc }}
-    **Hypothesis:** {{ hypothesis.hypothesis }}
-
-    {% endfor %}
-    # 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 }}
+    {% include "scenarios.data_science.proposal.exp_gen.prompts_v2:task_gen.user" %}

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

@@ -454,10 +454,15 @@ class DSProposalV1ExpGen(ExpGen):
 
 
 class DSProposalV2ExpGen(ExpGen):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.support_function_calling = APIBackend().support_function_calling()
+
     def identify_scenario_problem(self, scenario_desc: str, sota_exp_desc: str) -> Dict:
         sys_prompt = T(".prompts_v2:scenario_problem.system").r(
-            problem_spec=T(".prompts_v2:specification.problem").r(),
-            problem_output_format=T(".prompts_v2:output_format.problem").r(),
+            problem_output_format=(
+                T(".prompts_v2:output_format.problem").r() if not self.support_function_calling else None
+            ),
         )
         user_prompt = T(".prompts_v2:scenario_problem.user").r(
             scenario_desc=scenario_desc,
@@ -466,17 +471,26 @@ class DSProposalV2ExpGen(ExpGen):
         response = APIBackend().build_messages_and_create_chat_completion(
             user_prompt=user_prompt,
             system_prompt=sys_prompt,
-            json_mode=True,
-            json_target_type=Dict[str, Dict[str, str]],
+            response_format=ScenarioChallenges if self.support_function_calling else {"type": "json_object"},
+            json_target_type=Dict[str, Dict[str, str]] if not self.support_function_calling else None,
         )
-        return json.loads(response)
+        if self.support_function_calling:
+            challenges = ScenarioChallenges(**json.loads(response))
+            # Translate to problems
+            problems = {o.caption: {"problem": o.statement, "reason": o.reasoning} for o in challenges.challenges}
+            logger.info(f"Identified scenario problems:\n" + json.dumps(problems))
+        else:
+            problems = json.loads(response)
+            logger.info(f"Identified scenario problems:\n" + json.dumps(problems))
+        return problems
 
     def identify_feedback_problem(
         self, scenario_desc: str, exp_feedback_list_desc: str, sota_exp_desc: str, inject_diverse: bool = False
     ) -> Dict:
         sys_prompt = T(".prompts_v2:feedback_problem.system").r(
-            problem_spec=T(".prompts_v2:specification.problem").r(),
-            problem_output_format=T(".prompts_v2:output_format.problem").r(),
+            problem_output_format=(
+                T(".prompts_v2:output_format.problem").r() if not self.support_function_calling else None
+            ),
             inject_diverse=inject_diverse,
         )
         user_prompt = T(".prompts_v2:feedback_problem.user").r(
@@ -487,10 +501,18 @@ class DSProposalV2ExpGen(ExpGen):
         response = APIBackend().build_messages_and_create_chat_completion(
             user_prompt=user_prompt,
             system_prompt=sys_prompt,
-            json_mode=True,
-            json_target_type=Dict[str, Dict[str, str]],
+            response_format=TraceChallenges if self.support_function_calling else {"type": "json_object"},
+            json_target_type=Dict[str, Dict[str, str]] if not self.support_function_calling else None,
         )
-        return json.loads(response)
+        if self.support_function_calling:
+            challenges = TraceChallenges(**json.loads(response))
+            # Translate to problems
+            problems = {o.caption: {"problem": o.statement, "reason": o.reasoning} for o in challenges.challenges}
+            logger.info(f"Identified feedback problems:\n" + json.dumps(problems))
+        else:
+            problems = json.loads(response)
+            logger.info(f"Identified feedback problems:\n" + json.dumps(problems))
+        return problems
 
     def identify_problem(
         self, current_sub_trace, scenario_desc, sota_exp_desc, exp_feedback_list_desc, inject_diverse
@@ -535,21 +557,19 @@ class DSProposalV2ExpGen(ExpGen):
         inject_diverse: bool = False,
     ) -> Dict:
         problem_formatted_str = ""
-        for problem_name, problem_dict in problems.items():
-            if "problem" not in problem_dict:
-                continue
-            problem_formatted_str += f"Problem Name: {problem_name}\n"
-            problem_formatted_str += f"- Problem Description: {problem_dict['problem']}\n"
+        for i, (problem_name, problem_dict) in enumerate(problems.items()):
+            problem_formatted_str += f"## {i+1}. {problem_name}\n"
+            problem_formatted_str += f"{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 += 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(pipeline=pipeline),
-            hypothesis_output_format=T(".prompts_v2:output_format.hypothesis").r(
-                pipeline=pipeline, enable_idea_pool=enable_idea_pool
+            hypothesis_output_format=(
+                T(".prompts_v2:output_format.hypothesis").r(pipeline=pipeline, enable_idea_pool=enable_idea_pool)
+                if not self.support_function_calling
+                else None
             ),
             pipeline=pipeline,
             enable_idea_pool=enable_idea_pool,
@@ -565,10 +585,31 @@ class DSProposalV2ExpGen(ExpGen):
         response = APIBackend().build_messages_and_create_chat_completion(
             user_prompt=user_prompt,
             system_prompt=sys_prompt,
-            json_mode=True,
-            json_target_type=Dict[str, Dict[str, str | Dict[str, str | int]]],
+            response_format=HypothesisList if self.support_function_calling else {"type": "json_object"},
+            json_target_type=(
+                Dict[str, Dict[str, str | Dict[str, str | int]]] if not self.support_function_calling else None
+            ),
         )
-        resp_dict = json.loads(response)
+        if self.support_function_calling:
+            hypotheses = HypothesisList(**json.loads(response))
+            resp_dict = {
+                h.caption: {
+                    "reason": h.challenge,
+                    "component": h.component,
+                    "hypothesis": h.hypothesis,
+                    "evaluation": {
+                        "alignment_score": h.evaluation.alignment.score,
+                        "impact_score": h.evaluation.impact.score,
+                        "novelty_score": h.evaluation.novelty.score,
+                        "feasibility_score": h.evaluation.feasibility.score,
+                        "risk_reward_balance_score": h.evaluation.risk_reward_balance.score,
+                    },
+                }
+                for h in hypotheses.hypotheses
+            }
+        else:
+            resp_dict = json.loads(response)
+        logger.info(f"Generated hypotheses:\n" + json.dumps(resp_dict, indent=2))
 
         # make sure the problem name is aligned
         problem_keys = set(problems.keys())
@@ -675,44 +716,40 @@ class DSProposalV2ExpGen(ExpGen):
         scenario_desc: str,
         sota_exp_desc: str,
         sota_exp: DSExperiment,
-        hypothesis: DSHypothesis,
+        hypotheses: list[DSHypothesis],
         pipeline: bool,
         failed_exp_feedback_list_desc: str,
     ) -> DSExperiment:
         if pipeline:
             component_info = get_component("Pipeline")
         else:
-            component_info = get_component(hypothesis.component)
-        if pipeline:
-            task_spec = T(f"scenarios.data_science.share:component_spec.Pipeline").r()
-        elif DS_RD_SETTING.spec_enabled and sota_exp is not None:
-            task_spec = sota_exp.experiment_workspace.file_dict[component_info["spec_file"]]
-        else:
-            task_spec = T(f"scenarios.data_science.share:component_spec.{hypothesis.component}").r()
+            component_info = get_component(hypotheses[0].component)
+        data_folder_info = self.scen.processed_data_folder_description
         sys_prompt = T(".prompts_v2:task_gen.system").r(
-            targets=component_info["target_name"],
-            task_specification=task_spec,
-            task_output_format=component_info["task_output_format"],
+            task_output_format=component_info["task_output_format"] if not self.support_function_calling else None,
+            # task_output_format=component_info["task_output_format"],
             component_desc=component_desc,
-            workflow_check=not pipeline and hypothesis.component != "Workflow",
+            workflow_check=not pipeline and hypotheses[0].component != "Workflow",
         )
         user_prompt = T(".prompts_v2:task_gen.user").r(
             scenario_desc=scenario_desc,
+            data_folder_info=data_folder_info,
             sota_exp_desc=sota_exp_desc,
-            hypothesis=str(hypothesis),
+            hypotheses=hypotheses,
             failed_exp_and_feedback_list_desc=failed_exp_feedback_list_desc,
         )
         response = APIBackend().build_messages_and_create_chat_completion(
             user_prompt=user_prompt,
             system_prompt=sys_prompt,
-            json_mode=True,
-            json_target_type=Dict[str, str | Dict[str, str]],
+            response_format=CodingSketch if self.support_function_calling else {"type": "json_object"},
+            json_target_type=Dict[str, str | Dict[str, str]] if not self.support_function_calling else None,
         )
         task_dict = json.loads(response)
-        task_design = task_dict.get("task_design", {})
-        task_name = (
-            task_design["model_name"] if (hypothesis.component == "Model" and not pipeline) else hypothesis.component
+        task_design = (
+            task_dict.get("task_design", {}) if not self.support_function_calling else task_dict.get("sketch", {})
         )
+        logger.info(f"Task design:\n{task_design}")
+        task_name = hypotheses[0].component
         description = (
             task_design
             if isinstance(task_design, str)
@@ -724,7 +761,7 @@ class DSProposalV2ExpGen(ExpGen):
             description=description,
         )
         new_workflow_desc = task_dict.get("workflow_update", "No update needed")
-        exp = DSExperiment(pending_tasks_list=[[task]], hypothesis=hypothesis)
+        exp = DSExperiment(pending_tasks_list=[[task]], hypothesis=hypotheses[0])
         # exp.experiment_workspace.inject_code_from_folder(sota_exp.experiment_workspace.workspace_path)
         if sota_exp is not None:
             exp.experiment_workspace.inject_code_from_file_dict(sota_exp.experiment_workspace)
@@ -736,7 +773,40 @@ class DSProposalV2ExpGen(ExpGen):
             exp.pending_tasks_list.append([workflow_task])
         return exp
 
-    def gen(self, trace: DSTrace) -> DSExperiment:
+    def get_scenario_all_desc(self, trace: DSTrace, eda_output=None) -> str:
+        return T(".prompts_v2:scenario_description").r(
+            background=trace.scen.background,
+            submission_specifications=trace.scen.submission_specifications,
+            evaluation=trace.scen.metric_description,
+            metric_name=trace.scen.metric_name,
+            metric_direction=trace.scen.metric_direction,
+            raw_description=trace.scen.raw_description,
+            use_raw_description=DS_RD_SETTING.use_raw_description,
+            time_limit=f"{DS_RD_SETTING.full_timeout / 60 / 60 : .2f} hours",
+            eda_output=eda_output,
+        )
+
+    def get_all_hypotheses(self, problem_dict: dict, hypothesis_dict: dict) -> list[DSHypothesis]:
+        result = []
+        for name, data in hypothesis_dict.items():
+            problem_data = problem_dict.get(name, {})
+            result.append(
+                DSHypothesis(
+                    component=data.get("component", "Model"),
+                    hypothesis=data.get("hypothesis", "Hypothesis not provided"),
+                    reason=data.get("reason", "Reason not provided"),
+                    problem_name=name,
+                    problem_desc=problem_data.get("problem", "Problem description not provided"),
+                    problem_label=problem_data.get("label", "FEEDBACK_PROBLEM"),
+                )
+            )
+        return result
+
+    def gen(
+        self,
+        trace: DSTrace,
+    ) -> DSExperiment:
+
         pipeline = DS_RD_SETTING.coder_on_whole_pipeline
         if not pipeline and (draft_exp := draft_exp_in_decomposition(self.scen, trace)):
             return draft_exp
@@ -756,7 +826,7 @@ class DSProposalV2ExpGen(ExpGen):
             eda_output = None
         else:
             eda_output = sota_exp.experiment_workspace.file_dict.get("EDA.md", None)
-        scenario_desc = trace.scen.get_scenario_all_desc(eda_output=eda_output)
+        scenario_desc = self.get_scenario_all_desc(trace, eda_output=eda_output)
 
         sota_exp_desc = T("scenarios.data_science.share:describe.exp").r(
             exp=sota_exp, heading="Best of previous exploration of the scenario"
@@ -773,6 +843,7 @@ class DSProposalV2ExpGen(ExpGen):
             pipeline=pipeline,
         )
 
+        # NOTE: we currently don't support inject diverse problems for the parallel + multi-trace mode,
         if DS_RD_SETTING.enable_inject_diverse and len(trace.hist) > 0:
             if len(trace.current_selection) == 0:
                 # start a new sub-trace, and inject diverse problems.
@@ -838,374 +909,6 @@ class DSProposalV2ExpGen(ExpGen):
         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,
-            scenario_desc=scenario_desc,
-            sota_exp_desc=sota_exp_desc,
-            sota_exp=sota_exp,
-            hypothesis=new_hypothesis,
-            pipeline=pipeline,
-            failed_exp_feedback_list_desc=failed_exp_feedback_list_desc,
-        )
-
-
-class DSProposalV3ExpGen(DSProposalV2ExpGen):
-    def identify_scenario_problem(self, scenario_desc: str, sota_exp_desc: str) -> Dict:
-        sys_prompt = T(".prompts_v3:scenario_problem.system").r()
-        user_prompt = T(".prompts_v3:scenario_problem.user").r(
-            scenario_desc=scenario_desc,
-            sota_exp_desc=sota_exp_desc,
-        )
-        response = APIBackend().build_messages_and_create_chat_completion(
-            user_prompt=user_prompt,
-            system_prompt=sys_prompt,
-            response_format=ScenarioChallenges,
-            # json_mode=True,
-            # json_target_type=Dict[str, Dict[str, str]],
-        )
-        challenges = ScenarioChallenges(**json.loads(response))
-        # Translate to problems
-        problems = {o.caption: {"problem": o.statement, "reason": o.reasoning} for o in challenges.challenges}
-        logger.info(f"Identified scenario problems:\n" + json.dumps(problems))
-        return problems
-
-    def identify_feedback_problem(self, scenario_desc: str, exp_feedback_list_desc: str, sota_exp_desc: str) -> Dict:
-        sys_prompt = T(".prompts_v3:feedback_problem.system").r()
-        user_prompt = T(".prompts_v3:feedback_problem.user").r(
-            scenario_desc=scenario_desc,
-            exp_and_feedback_list_desc=exp_feedback_list_desc,
-            sota_exp_desc=sota_exp_desc,
-        )
-        response = APIBackend().build_messages_and_create_chat_completion(
-            user_prompt=user_prompt,
-            system_prompt=sys_prompt,
-            response_format=TraceChallenges,
-            # json_mode=True,
-            # json_target_type=Dict[str, Dict[str, str]],
-        )
-        challenges = TraceChallenges(**json.loads(response))
-        # Translate to problems
-        problems = {o.caption: {"problem": o.statement, "reason": o.reasoning} for o in challenges.challenges}
-        logger.info(f"Identified feedback problems:\n" + json.dumps(problems))
-        return problems
-
-    def get_scenario_all_desc_v3(self, trace: DSTrace, eda_output=None) -> str:
-        return T(".prompts_v3:scenario_description").r(
-            background=trace.scen.background,
-            submission_specifications=trace.scen.submission_specifications,
-            evaluation=trace.scen.metric_description,
-            metric_name=trace.scen.metric_name,
-            metric_direction=trace.scen.metric_direction,
-            raw_description=trace.scen.raw_description,
-            use_raw_description=DS_RD_SETTING.use_raw_description,
-            time_limit=f"{DS_RD_SETTING.full_timeout / 60 / 60 : .2f} hours",
-            eda_output=eda_output,
-        )
-
-    @wait_retry(retry_n=5)
-    def hypothesis_gen(
-        self,
-        component_desc: str,
-        scenario_desc: str,
-        exp_feedback_list_desc: str,
-        sota_exp_desc: str,
-        problems: dict,
-        pipeline: bool,
-        enable_idea_pool: bool,
-    ) -> Dict:
-        problem_formatted_str = ""
-        for i, (problem_name, problem_dict) in enumerate(problems.items()):
-            problem_formatted_str += f"## {i+1}. {problem_name}\n"
-            problem_formatted_str += f"{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_v3:hypothesis_gen.system").r(
-            pipeline=pipeline,
-            enable_idea_pool=enable_idea_pool,
-        )
-        user_prompt = T(".prompts_v3:hypothesis_gen.user").r(
-            scenario_desc=scenario_desc,
-            exp_and_feedback_list_desc=exp_feedback_list_desc,
-            sota_exp_desc=sota_exp_desc,
-            problems=problem_formatted_str,
-            enable_idea_pool=enable_idea_pool,
-        )
-        response = APIBackend().build_messages_and_create_chat_completion(
-            user_prompt=user_prompt, system_prompt=sys_prompt, response_format=HypothesisList
-        )
-        hypotheses = HypothesisList(**json.loads(response))
-        resp_dict = {
-            h.caption: {
-                "reason": h.challenge,
-                "component": h.component,
-                "hypothesis": h.hypothesis,
-                "evaluation": {
-                    "alignment_score": h.evaluation.alignment.score,
-                    "impact_score": h.evaluation.impact.score,
-                    "novelty_score": h.evaluation.novelty.score,
-                    "feasibility_score": h.evaluation.feasibility.score,
-                    "risk_reward_balance_score": h.evaluation.risk_reward_balance.score,
-                },
-            }
-            for h in hypotheses.hypotheses
-        }
-
-        logger.info(f"Generated hypotheses:\n" + json.dumps(resp_dict, indent=2))
-
-        if len(resp_dict) == 0:
-            logger.error("No hypothesis generated. Retrying...")
-            raise ValueError("No hypothesis generated.")
-
-        # make sure the problem name is aligned
-        problem_keys = set(problems.keys())
-        resp_keys = set(resp_dict.keys())
-        if not resp_keys.issubset(problem_keys):
-            logger.error("Problem names are not fully aligned. Retrying...")
-            raise ValueError("Problem names are not fully aligned.")
-
-        return resp_dict
-
-    def task_gen(
-        self,
-        component_desc: str,
-        scenario_desc: str,
-        sota_exp_desc: str,
-        sota_exp: DSExperiment,
-        hypotheses: list[DSHypothesis],
-        pipeline: bool,
-        failed_exp_feedback_list_desc: str,
-    ) -> DSExperiment:
-        if pipeline:
-            component_info = get_component("Pipeline")
-        else:
-            component_info = get_component(hypotheses[0].component)
-        data_folder_info = self.scen.processed_data_folder_description
-        sys_prompt = T(".prompts_v3:task_gen.system").r(
-            # targets=component_info["target_name"],
-            # task_output_format=component_info["task_output_format"],
-            # component_desc=component_desc,
-            # workflow_check=not pipeline and hypothesis.component != "Workflow",
-        )
-        user_prompt = T(".prompts_v3:task_gen.user").r(
-            scenario_desc=scenario_desc,
-            data_folder_info=data_folder_info,
-            sota_exp_desc=sota_exp_desc,
-            hypotheses=hypotheses,
-            failed_exp_and_feedback_list_desc=failed_exp_feedback_list_desc,
-        )
-        response = APIBackend().build_messages_and_create_chat_completion(
-            user_prompt=user_prompt,
-            system_prompt=sys_prompt,
-            response_format=CodingSketch,
-            # json_mode=True,
-            # json_target_type=Dict[str, str | Dict[str, str]],
-        )
-        task_dict = json.loads(response)
-        task_design = task_dict.get("sketch", {})
-        logger.info("Task design:\n" + task_design)
-        task_name = hypotheses[0].component
-        description = (
-            task_design
-            if isinstance(task_design, str)
-            else task_design.get("description", f"{component_info['target_name']} description not provided")
-        )
-        task_class = component_info["task_class"]
-        task = task_class(
-            name=task_name,
-            description=description,
-        )
-        new_workflow_desc = task_dict.get("workflow_update", "No update needed")
-        exp = DSExperiment(pending_tasks_list=[[task]], hypothesis=hypotheses[0])
-        # exp.experiment_workspace.inject_code_from_folder(sota_exp.experiment_workspace.workspace_path)
-        if sota_exp is not None:
-            exp.experiment_workspace.inject_code_from_file_dict(sota_exp.experiment_workspace)
-        if not pipeline and new_workflow_desc != "No update needed":
-            workflow_task = WorkflowTask(
-                name="Workflow",
-                description=new_workflow_desc,
-            )
-            exp.pending_tasks_list.append([workflow_task])
-        return exp
-
-    def get_all_hypotheses(self, problem_dict: dict, hypothesis_dict: dict) -> list[DSHypothesis]:
-        result = []
-        for name, data in hypothesis_dict.items():
-            problem_data = problem_dict.get(name, {})
-            result.append(
-                DSHypothesis(
-                    component=data.get("component", "Model"),
-                    hypothesis=data.get("hypothesis", "Hypothesis not provided"),
-                    reason=data.get("reason", "Reason not provided"),
-                    problem_name=name,
-                    problem_desc=problem_data.get("problem", "Problem description not provided"),
-                    problem_label=problem_data.get("label", "FEEDBACK_PROBLEM"),
-                )
-            )
-        return result
-
-    # FIXME: remove this, dump solution, should be merged into identify_problem in V2
-    def identify_problems_v3(
-        self, trace: DSTrace, scenario_desc: str, sota_exp_desc: str, exp_feedback_list_desc: str
-    ) -> Dict:
-        sub_trace = trace.get_parent_exps()
-        trace_length = len(trace.hist)
-        all_problems = {}
-
-        # 阶段一：探索期（主要场景问题）
-        if trace_length <= 3:
-            scen_problems = self.identify_scenario_problem(scenario_desc, sota_exp_desc)
-            for problem_name in scen_problems:
-                scen_problems[problem_name]["label"] = "SCENARIO_PROBLEM"
-                all_problems[problem_name] = scen_problems[problem_name]
-            self.scen_prob_multiplier = 3
-
-        # 阶段二：混合期（两种问题都考虑）
-        elif trace_length <= 6:
-            # 优先场景问题，但也考虑反馈
-            scen_problems = self.identify_scenario_problem(scenario_desc, sota_exp_desc)
-            for problem_name in scen_problems:
-                scen_problems[problem_name]["label"] = "SCENARIO_PROBLEM"
-                all_problems[problem_name] = scen_problems[problem_name]
-
-            fb_problems = self.identify_feedback_problem(scenario_desc, exp_feedback_list_desc, sota_exp_desc)
-            for problem_name in fb_problems:
-                fb_problems[problem_name]["label"] = "FEEDBACK_PROBLEM"
-                all_problems[problem_name] = fb_problems[problem_name]
-            self.scen_prob_multiplier = 2
-
-        # 阶段三：优化期（主要反馈问题）
-        else:
-            fb_problems = self.identify_feedback_problem(scenario_desc, exp_feedback_list_desc, sota_exp_desc)
-            for problem_name in fb_problems:
-                fb_problems[problem_name]["label"] = "FEEDBACK_PROBLEM"
-                all_problems[problem_name] = fb_problems[problem_name]
-            self.scen_prob_multiplier = 1
-
-        return all_problems
-
-    def gen(self, trace: DSTrace) -> DSExperiment:
-        pipeline = DS_RD_SETTING.coder_on_whole_pipeline
-        if not pipeline and (draft_exp := draft_exp_in_decomposition(self.scen, trace)):
-            return draft_exp
-
-        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):
-            eda_output = None
-        else:
-            eda_output = sota_exp.experiment_workspace.file_dict.get("EDA.md", None)
-        scenario_desc = self.get_scenario_all_desc_v3(trace, eda_output=eda_output)
-
-        sota_exp_desc = T("scenarios.data_science.share:describe.exp").r(
-            exp=sota_exp, heading="Best of previous exploration of the scenario"
-        )
-
-        exp_feedback_list_desc = T("scenarios.data_science.share:describe.trace").r(
-            exp_and_feedback_list=trace.experiment_and_feedback_list_after_init(return_type="all"),
-            type="all",
-            pipeline=pipeline,
-        )
-        failed_exp_feedback_list_desc = T("scenarios.data_science.share:describe.trace").r(
-            exp_and_feedback_list=trace.experiment_and_feedback_list_after_init(return_type="failed"),
-            type="failed",
-            pipeline=pipeline,
-        )
-
-        if DS_RD_SETTING.enable_inject_diverse and len(trace.hist) > 0:
-            if len(trace.current_selection) == 0:
-                # start a new sub-trace, and inject diverse problems.
-                inject_diverse = True
-                logger.info("Start a new sub-trace, and inject diverse problems.")
-            else:
-                inject_diverse = False
-        else:
-            inject_diverse = False
-        # Step 1: Identify problems
-        all_problems = {}
-
-        all_problems = self.identify_problems_v3(
-            trace=trace,
-            scenario_desc=scenario_desc,
-            sota_exp_desc=sota_exp_desc,
-            exp_feedback_list_desc=exp_feedback_list_desc,
-        )
-
-        # if len(trace.hist) > 3:
-        #     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[problem_name] = fb_problems[problem_name]
-
-        # if len(trace.hist) < 9:
-        #     scen_problems = self.identify_scenario_problem(
-        #         scenario_desc=scenario_desc,
-        #         sota_exp_desc=sota_exp_desc,
-        #     )
-        #     for problem_name in scen_problems:
-        #         scen_problems[problem_name]["label"] = "SCENARIO_PROBLEM"
-        #         all_problems[problem_name] = scen_problems[problem_name]
-
-        # 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,
-            exp_feedback_list_desc=exp_feedback_list_desc,
-            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(
-                [
-                    k
-                    for k in sota_exp.experiment_workspace.file_dict.keys()
-                    if k.endswith(".py") and "test" not in k and k.startswith("model")
-                ]
-            )
-            if sota_exp_model_file_count <= 1:
-                pop_names = []
-                for problem_name in hypothesis_dict:
-                    if hypothesis_dict[problem_name].get("component", "") == "Ensemble":
-                        pop_names.append(problem_name)
-                for name in pop_names:
-                    hypothesis_dict.pop(name)
-
-        # Step 3: Select the best hypothesis
-        pickled_problem_name, new_hypothesis = self.hypothesis_rank(
-            hypothesis_dict=hypothesis_dict,
-            problem_dict=all_problems,
-        )
-        # 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,
             scenario_desc=scenario_desc,

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

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import asyncio
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from rdagent.scenarios.data_science.proposal.exp_gen.base import DSTrace
+
+
+class TraceScheduler(ABC):
+    """
+    An abstract base class for trace scheduling strategies.
+    Determines which active trace to expand next during parallel exploration.
+    """
+
+    @abstractmethod
+    async def select_trace(self, trace: DSTrace) -> tuple[int, ...]:
+        """
+        Selects the next trace to expand.
+
+        This method must be async to allow for safe concurrent access.
+
+        Args:
+            trace: The DSTrace object containing the full experiment history.
+
+        Returns:
+            A tuple representing the selection of the parent node for the new experiment.
+            e.g., (leaf_idx,) for an existing trace, or trace.NEW_ROOT for a new one.
+        """
+        raise NotImplementedError
+
+
+class RoundRobinScheduler(TraceScheduler):
+    """
+    A concurrency-safe scheduling strategy that cycles through active traces
+    in a round-robin fashion.
+
+    NOTE: we don't need to use asyncio.Lock here as the kickoff_loop ensures the ExpGen is always sequential, instead of parallel.
+    """
+
+    def __init__(self):
+        self._last_selected_leaf_id = -1
+
+    async def select_trace(self, trace: DSTrace) -> tuple[int, ...]:
+        """
+        Atomically selects the next leaf node from the trace in order.
+        """
+
+        leaves = trace.get_leaves()
+        if not leaves:
+            # This is the very first experiment in a new tree.
+            return trace.NEW_ROOT
+
+        # Find the index of the last selected leaf in the current list of leaves
+        try:
+            current_position = leaves.index(self._last_selected_leaf_id)
+            # Move to the next position, wrapping around if necessary
+            next_position = (current_position + 1) % len(leaves)
+        except ValueError:
+            # This can happen if the last selected leaf is no longer a leaf
+            # (it has been expanded) or if this is the first selection.
+            # In either case, start from the beginning.
+            next_position = 0
+
+        selected_leaf = leaves[next_position]
+        self._last_selected_leaf_id = selected_leaf
+
+        return (selected_leaf,)

rdagent/scenarios/data_science/scen/utils.py

@@ -306,19 +306,17 @@ def _walk(path: Path):
             yield p
 
 
-def preview_csv(p: Path, file_name: str, simple=True, show_nan_columns=False) -> str:
-    """Generate a textual preview of a csv file
+def preview_df(df: pd.DataFrame, file_name: str, simple=True, show_nan_columns=False) -> str:
+    """Generate a textual preview of a dataframe
 
     Args:
-        p (Path): the path to the csv file
+        df (pd.DataFrame): the dataframe to preview
         file_name (str): the file name to use in the preview
         simple (bool, optional): whether to use a simplified version of the preview. Defaults to True.
 
     Returns:
         str: the textual preview
     """
-    df = pd.read_csv(p)
-
     out = []
 
     out.append(f"-> {file_name} has {df.shape[0]} rows and {df.shape[1]} columns.")
@@ -358,6 +356,18 @@ def preview_csv(p: Path, file_name: str, simple=True, show_nan_columns=False) ->
     return "\n".join(out)
 
 
+def preview_csv(p: Path, file_name: str, simple=True, show_nan_columns=False) -> str:
+    """Generate a textual preview of a csv file"""
+    df = pd.read_csv(p)
+    return preview_df(df, file_name, simple=simple, show_nan_columns=show_nan_columns)
+
+
+def preview_parquet(p: Path, file_name: str, simple=True, show_nan_columns=False) -> str:
+    """Generate a textual preview of a parquet file"""
+    df = pd.read_parquet(p)
+    return preview_df(df, file_name, simple=simple, show_nan_columns=show_nan_columns)
+
+
 def preview_json(p: Path, file_name: str):
     """Generate a textual preview of a json file using a generated json schema"""
     builder = SchemaBuilder()
@@ -388,7 +398,9 @@ def preview_json(p: Path, file_name: str):
     return f"-> {file_name} has auto-generated json schema:\n" + builder.to_json(indent=2)
 
 
-def describe_data_folder_v2(base_path, include_file_details=True, simple=False, show_nan_columns=False):
+def describe_data_folder_v2(
+    base_path, include_file_details=True, simple=False, show_nan_columns=False, max_length: int = 6_000
+):
     """
     Generate a textual preview of a directory, including an overview of the directory
     structure and previews of individual files
@@ -404,6 +416,8 @@ def describe_data_folder_v2(base_path, include_file_details=True, simple=False,
                 out.append(preview_csv(fn, file_name, simple=simple, show_nan_columns=show_nan_columns))
             elif fn.suffix == ".json":
                 out.append(preview_json(fn, file_name))
+            elif fn.suffix == ".parquet":
+                out.append(preview_parquet(fn, file_name, simple=simple, show_nan_columns=show_nan_columns))
             elif fn.suffix in plaintext_files:
                 if get_file_len_size(fn)[0] < 30:
                     with open(fn) as f:
@@ -411,16 +425,22 @@ def describe_data_folder_v2(base_path, include_file_details=True, simple=False,
                         if fn.suffix in code_files:
                             content = f"```\n{content}\n```"
                         out.append(f"-> {file_name} has content:\n\n{content}")
+            if len("\n\n".join(out)) > max_length:
+                break
 
     result = "\n\n".join(out)
 
     # if the result is very long we generate a simpler version
-    if len(result) > 6_000 and not simple:
+    if len(result) > max_length and not simple:
         return describe_data_folder_v2(
-            base_path, include_file_details=include_file_details, simple=True, show_nan_columns=show_nan_columns
+            base_path,
+            include_file_details=include_file_details,
+            simple=True,
+            show_nan_columns=show_nan_columns,
+            max_length=max_length,
         )
     # if still too long, we truncate
-    if len(result) > 6_000 and simple:
-        return result[:6_000] + "\n... (truncated)"
+    if len(result) > max_length and simple:
+        return result[:max_length] + "\n... (truncated)"
 
     return result

rdagent/scenarios/kaggle/kaggle_crawler.py

@@ -3,6 +3,7 @@ import bisect
 import json
 import shutil
 import subprocess
+import tarfile
 import time
 import zipfile
 from itertools import chain
@@ -111,6 +112,7 @@ def download_data(competition: str, settings: ExtendedBaseSettings, enable_creat
     if settings.if_using_mle_data:
         zipfile_path = f"{local_path}/zip_files"
         zip_competition_path = Path(zipfile_path) / competition
+        competition_local_path = Path(local_path) / competition
 
         if not zip_competition_path.exists():
             mleb_env = MLEBDockerEnv()
@@ -122,24 +124,41 @@ def download_data(competition: str, settings: ExtendedBaseSettings, enable_creat
                 running_extra_volume={str(Path("~/.kaggle").expanduser().absolute()): "/root/.kaggle"},
             )
 
-        if not (Path(local_path) / competition).exists() or list((Path(local_path) / competition).iterdir()) == []:
-            (Path(local_path) / competition).mkdir(parents=True, exist_ok=True)
+        if not competition_local_path.exists() or list(competition_local_path.iterdir()) == []:
+            competition_local_path.mkdir(parents=True, exist_ok=True)
 
             mleb_env = MLEBDockerEnv()
             mleb_env.prepare()
             mleb_env.run(f"cp -r ./zip_files/{competition}/prepared/public/* ./{competition}", local_path=local_path)
 
-            for zip_path in (Path(local_path) / competition).rglob("*.zip"):
+            for zip_path in competition_local_path.rglob("*.zip"):
                 with zipfile.ZipFile(zip_path, "r") as zip_ref:
                     if len(zip_ref.namelist()) == 1:
                         mleb_env.run(
-                            f"unzip -o ./{zip_path.relative_to(local_path)} -d {zip_path.parent.relative_to(local_path)}",
-                            local_path=local_path,
+                            f"unzip -o ./{zip_path.relative_to(competition_local_path)} -d {zip_path.parent.relative_to(competition_local_path)}",
+                            local_path=competition_local_path,
                         )
                     else:
                         mleb_env.run(
-                            f"mkdir -p ./{zip_path.parent.relative_to(local_path)}/{zip_path.stem}; unzip -o ./{zip_path.relative_to(local_path)} -d ./{zip_path.parent.relative_to(local_path)}/{zip_path.stem}",
-                            local_path=local_path,
+                            f"mkdir -p ./{zip_path.parent.relative_to(competition_local_path)}/{zip_path.stem}; unzip -o ./{zip_path.relative_to(competition_local_path)} -d ./{zip_path.parent.relative_to(competition_local_path)}/{zip_path.stem}",
+                            local_path=competition_local_path,
+                        )
+            for tar_path in competition_local_path.rglob("*.tar*"):
+                if not tarfile.is_tarfile(tar_path):
+                    logger.error(f"{tar_path} is not a valid tar file.")
+                    continue
+                is_gzip_file = open(tar_path, "rb").read(2) == b"\x1f\x8b"
+                with tarfile.open(tar_path, "r:gz") if is_gzip_file else tarfile.open(tar_path, "r") as tar_ref:
+                    if len(tar_ref.getmembers()) == 1:
+                        mleb_env.run(
+                            f"tar -{'xzf' if is_gzip_file else 'xf'} ./{tar_path.relative_to(competition_local_path)} -C {tar_path.parent.relative_to(competition_local_path)}",
+                            local_path=competition_local_path,
+                        )
+                    else:
+                        folder_name = tar_path.name.replace(".tar", "").replace(".gz", "")
+                        mleb_env.run(
+                            f"mkdir -p ./{tar_path.parent.relative_to(competition_local_path)}/{folder_name}; tar -{'xzf' if is_gzip_file else 'xf'} ./{tar_path.relative_to(competition_local_path)} -C ./{tar_path.parent.relative_to(competition_local_path)}/{folder_name}",
+                            local_path=competition_local_path,
                         )
             # NOTE:
             # Patching:  due to mle has special renaming mechanism for different competition;

rdagent/utils/agent/ret.py

@@ -31,7 +31,8 @@ class PythonAgentOut(AgentOut):
 
     @classmethod
     def extract_output(cls, resp: str):
-        match = re.search(r".*```[Pp]ython\n(.*)\n```.*", resp, re.DOTALL)
+        # We use lazy mode (.*?) to only extract the first code block in the response.
+        match = re.search(r".*```[Pp]ython\n(.*?)\n```.*", resp, re.DOTALL)
         if match:
             code = match.group(1)
             code = re.sub(r"</?code>", "", code, flags=re.IGNORECASE)

rdagent/utils/env.py

@@ -44,6 +44,29 @@ from rdagent.utils.agent.tpl import T
 from rdagent.utils.workflow import wait_retry
 
 
+def cleanup_container(container: docker.models.containers.Container | None, context: str = "") -> None:  # type: ignore[no-any-unimported]
+    """
+    Shared helper function to clean up a Docker container.
+    Always stops the container before removing it.
+
+    Parameters
+    ----------
+    container : docker container object or None
+        The container to clean up, or None if no container to clean up
+    context : str
+        Additional context for logging (e.g., "health check", "GPU test")
+    """
+    if container is not None:
+        try:
+            # Always stop first - stop() doesn't raise error if already stopped
+            container.stop()
+            container.remove()
+        except Exception as cleanup_error:
+            # Log cleanup error but don't mask the original exception
+            context_str = f" {context}" if context else ""
+            logger.warning(f"Failed to cleanup{context_str} container {container.id}: {cleanup_error}")
+
+
 # Normalize all bind paths in volumes to absolute paths using the workspace (working_dir).
 def normalize_volumes(vols: dict[str, str | dict[str, str]], working_dir: str) -> dict:
     abs_vols: dict[str, str | dict[str, str]] = {}
@@ -93,6 +116,7 @@ def pull_image_with_progress(image: str) -> None:
 
 
 class EnvConf(ExtendedBaseSettings):
+    # TODO: add prefix ....
     default_entry: str
     extra_volumes: dict = {}
     running_timeout_period: int = 3600  # 10 minutes
@@ -578,6 +602,10 @@ class DockerConf(EnvConf):
     default_entry: str  # the entry point of the image
 
     extra_volumes: dict = {}
+    """It accept a dict of volumes, which can be either
+    {<host_path>: <container_path>} or
+    {<host_path>: {"bind": <container_path>, "mode": <mode, ro/rw/default is extra_volume_mode>}}
+    """
     extra_volume_mode: str = "ro"  # by default. only the mount_path should be writable, others are changed to read-only
     # Sometime, we need maintain some extra data for the workspace.
     # And the extra data may be shared and the downloading can be time consuming.
@@ -638,7 +666,9 @@ class QlibDockerConf(DockerConf):
     image: str = "local_qlib:latest"
     mount_path: str = "/workspace/qlib_workspace/"
     default_entry: str = "qrun conf.yaml"
-    extra_volumes: dict = {str(Path("~/.qlib/").expanduser().resolve().absolute()): "/root/.qlib/"}
+    extra_volumes: dict = {
+        str(Path("~/.qlib/").expanduser().resolve().absolute()): {"bind": "/root/.qlib/", "mode": "rw"}
+    }
     shm_size: str | None = "16g"
     enable_gpu: bool = True
     enable_cache: bool = False
@@ -785,12 +815,17 @@ class DockerEnv(Env[DockerConf]):
 
         @wait_retry(5, 10)
         def _f() -> dict:
+            container = None
             try:
                 get_image(self.conf.image)
-                client.containers.run(self.conf.image, "nvidia-smi", **gpu_kwargs)
+                container = client.containers.run(self.conf.image, "nvidia-smi", detach=True, **gpu_kwargs)
+                # Wait for container to complete
+                container.wait()
                 logger.info("GPU Devices are available.")
             except docker.errors.APIError:
                 return {}
+            finally:
+                cleanup_container(container, context="GPU test")
             return gpu_kwargs
 
         return _f()
@@ -825,19 +860,20 @@ class DockerEnv(Env[DockerConf]):
 
         if self.conf.extra_volumes is not None:
             for lp, rp in self.conf.extra_volumes.items():
-                volumes[lp] = {"bind": rp, "mode": self.conf.extra_volume_mode}
+                volumes[lp] = rp if isinstance(rp, dict) else {"bind": rp, "mode": self.conf.extra_volume_mode}
             cache_path = "/tmp/sample" if "/sample/" in "".join(self.conf.extra_volumes.keys()) else "/tmp/full"
             Path(cache_path).mkdir(parents=True, exist_ok=True)
             volumes[cache_path] = {"bind": T("scenarios.data_science.share:scen.cache_path").r(), "mode": "rw"}
         for lp, rp in running_extra_volume.items():
-            volumes[lp] = {"bind": rp, "mode": self.conf.extra_volume_mode}
+            volumes[lp] = rp if isinstance(rp, dict) else {"bind": rp, "mode": self.conf.extra_volume_mode}
 
         volumes = normalize_volumes(cast(dict[str, str | dict[str, str]], volumes), self.conf.mount_path)
 
         log_output = ""
+        container: docker.models.containers.Container | None = None  # type: ignore[no-any-unimported]
 
         try:
-            container: docker.models.containers.Container = client.containers.run(  # type: ignore[no-any-unimported]
+            container = client.containers.run(
                 image=self.conf.image,
                 command=entry,
                 volumes=volumes,
@@ -851,6 +887,7 @@ class DockerEnv(Env[DockerConf]):
                 cpu_count=self.conf.cpu_count,  # Set CPU limit
                 **self._gpu_kwargs(client),
             )
+            assert container is not None  # Ensure container was created successfully
             logs = container.logs(stream=True)
             print(Rule("[bold green]Docker Logs Begin[/bold green]", style="dark_orange"))
             table = Table(title="Run Info", show_header=False)
@@ -869,8 +906,6 @@ class DockerEnv(Env[DockerConf]):
                 Console().print(decoded_log, markup=False)
                 log_output += decoded_log + "\n"
             exit_status = container.wait()["StatusCode"]
-            container.stop()
-            container.remove()
             print(Rule("[bold green]Docker Logs End[/bold green]", style="dark_orange"))
             return log_output, exit_status
         except docker.errors.ContainerError as e:
@@ -879,6 +914,8 @@ class DockerEnv(Env[DockerConf]):
             raise RuntimeError("Docker image not found.")
         except docker.errors.APIError as e:
             raise RuntimeError(f"Error while running the container: {e}")
+        finally:
+            cleanup_container(container)
 
 
 class QTDockerEnv(DockerEnv):

rdagent/utils/workflow/loop.py

@@ -136,6 +136,11 @@ class LoopBase:
         if isinstance(limit := RD_AGENT_SETTINGS.step_semaphore, dict):
             limit = limit.get(step_name, 1)  # default to 1 if not specified
 
+        # NOTE: we assume the record step is always the last step to modify the global environment,
+        # so we set the limit to 1 to avoid race condition
+        if step_name == "record":
+            limit = 1
+
         if step_name not in self.semaphores:
             self.semaphores[step_name] = asyncio.Semaphore(limit)
         return self.semaphores[step_name]
@@ -219,6 +224,10 @@ class LoopBase:
                             result = func(self.loop_prev_out[li])
                     # Store result in the nested dictionary
                     self.loop_prev_out[li][name] = result
+
+                    # Record the trace
+                    end = datetime.datetime.now(datetime.timezone.utc)
+                    self.loop_trace[li].append(LoopTrace(start, end, step_idx=si))
                     # Save snapshot after completing the step
                     self.dump(self.session_folder / f"{li}" / f"{si}_{name}")
                 except Exception as e:
@@ -239,10 +248,6 @@ class LoopBase:
                         raise  # re-raise unhandled exceptions
                 finally:
                     if step_forward:
-                        # Record execution trace and update progress bar
-                        end = datetime.datetime.now(datetime.timezone.utc)
-                        self.loop_trace[li].append(LoopTrace(start, end, step_idx=si))
-
                         # Increment step index
                         self.step_idx[li] = next_step_idx
 
@@ -250,7 +255,11 @@ class LoopBase:
                         current_step = self.step_idx[li]
                         self.pbar.n = current_step
                         next_step = self.step_idx[li] % len(self.steps)
-                        self.pbar.set_postfix(loop_index=li, step_index=next_step, step_name=self.steps[next_step])
+                        self.pbar.set_postfix(
+                            loop_index=li + next_step_idx // len(self.steps),
+                            step_index=next_step,
+                            step_name=self.steps[next_step],
+                        )
                         self._check_exit_conditions_on_step()
                     else:
                         logger.warning(f"Step forward {si} of loop {li} is skipped.")
@@ -411,6 +420,18 @@ class LoopBase:
             An instance of LoopBase with the loaded session.
         """
         path = Path(path)
+        # if the path is a directory, load the latest session
+        if path.is_dir():
+            if path.name != "__session__":
+                path = path / "__session__"
+
+            if not path.exists():
+                raise FileNotFoundError(f"No session file found in {path}")
+
+            # iterate the dump steps in increasing order
+            files = sorted(path.glob("*/*_*"), key=lambda f: (int(f.parent.name), int(f.name.split("_")[0])))
+            path = files[-1]
+            logger.info(f"Loading latest session from {path}")
         with path.open("rb") as f:
             session = cast(LoopBase, pickle.load(f))
 

test/utils/test_env.py

@@ -13,6 +13,7 @@ from rdagent.utils.env import (
     LocalEnv,
     QlibDockerConf,
     QTDockerEnv,
+    cleanup_container,
 )
 
 DIRNAME = Path(__file__).absolute().resolve().parent
@@ -142,6 +143,18 @@ class EnvUtils(unittest.TestCase):
         # docker run  --memory=10m  -it --rm local_qlib:latest python -c 'import numpy as np; print(123);  size_mb = 1; size = size_mb * 1024 * 1024 // 8; array = np.random.randn(size).astype(np.float64); array[0], array[-1] = 1.0, 1.0; print(321)'
         # docker run  --memory=10g  -it --rm local_qlib:latest python -c 'import numpy as np; print(123);  size_mb = 1; size = size_mb * 1024 * 1024 // 8; array = np.random.randn(size).astype(np.float64); array[0], array[-1] = 1.0, 1.0; print(321)'
 
+    def test_cleanup_container_import(self):
+        """Test that cleanup_container function can be imported and has correct interface."""
+        # Test that the function exists and can be called
+        self.assertTrue(callable(cleanup_container))
+
+        # Test with None (should not raise an exception)
+        cleanup_container(None, "test context")
+
+        # The function should accept positional and keyword arguments
+        cleanup_container(None)
+        cleanup_container(None, context="test")
+
 
 if __name__ == "__main__":
     unittest.main()