feat: parallel loop running based on asyncio (#932)

* refactor: split workflow into pkg, add WorkflowTracker & wait_retry

* feat: add async LoopBase with parallel workers and step semaphores

* fix: replace pickle with dill and run blocking tasks via joblib wrapper

* feat: add log format settings, dynamic parallelism & pickle-based snapshot

* fix: default step semaphore to 1 and avoid subprocess when single worker

* merge bowen's changes

* merge tim's changes

* refactor: extract component task mapping, add conditional logger setup

* lint

* refactor: add type hints and safer remain_time metric logging in workflow

* lint

* fix: allow BadRequestError to be pickled via custom copyreg reducer

* fix: stop loop when LoopTerminationError is raised in LoopBase

* lint

* refactor: make log tag context-local using ContextVar for thread safety

* feat: add subproc_step flag and helper to decide subprocess execution

* fix: use ./cache path and normalize relative volume bind paths

* fix: reset loop_idx to 0 on loop restart/resume to ensure correct flow

* fix: avoid chmod on cache and input dirs in Env timeout wrapper

* fix: skip chmod on 'cache' and 'input' dirs using find -prune

* fix: restrict chmod to immediate mount dirs excluding cache/input

* fix: chmod cache and input dirs alongside their contents after entry run

* fix: guard chmod with directory checks for cache and input

* fix: prefix mount_path in chmod command for cache/input dirs

* fix: drop quotes from find exclude patterns to ensure chmod executes

* fix: skip chmod on cache/input directories to avoid warning spam

* feat: support string volume mappings and poll subprocess stdout/stderr

* support remove symbolic link

* test: use dynamic home path and code volume in LocalEnv local_simple

* fix: skip trace and progress update when loop step is withdrawn

* refactor: add clean_workspace util and non-destructive workspace backup

* fix: preserve symlinks when backing up workspace with copytree

* fix: prevent AttributeError when _pbar not yet initialized in LoopBase

* perf: replace shutil.copytree with rsync for faster workspace backup

* fix: cast log directory Path to str in tar command of data science loop

* fix: use portable 'cp -r -P' instead of rsync for workspace backup

* fix: add retry and logging to workspace backup for robustness

* refactor: extract backup_folder helper and reuse in DataScienceRDLoop

* fix: propagate backup errors & default _pbar getattr to avoid error

* fix the division by zero bug

* refactor: execute RD loops via asyncio.run and add necessary imports

* lint

* lint

* lint

---------

Co-authored-by: Xu <v-xuminrui@microsoft.com>

rdagent/app/data_mining/model.py

@@ -1,3 +1,5 @@
+import asyncio
+
 import fire
 
 from rdagent.app.data_mining.conf import MED_PROP_SETTING
@@ -24,7 +26,7 @@ def main(path=None, step_n=None, loop_n=None, all_duration=None, checkout=True):
         model_loop = ModelRDLoop(MED_PROP_SETTING)
     else:
         model_loop = ModelRDLoop.load(path, checkout=checkout)
-    model_loop.run(step_n=step_n, loop_n=loop_n, all_duration=all_duration)
+    asyncio.run(model_loop.run(step_n=step_n, loop_n=loop_n, all_duration=all_duration))
 
 
 if __name__ == "__main__":

rdagent/app/data_science/loop.py

@@ -1,3 +1,4 @@
+import asyncio
 from pathlib import Path
 
 import fire
@@ -66,7 +67,7 @@ def main(
     if exp_gen_cls is not None:
         kaggle_loop.exp_gen = import_class(exp_gen_cls)(kaggle_loop.exp_gen.scen)
 
-    kaggle_loop.run(step_n=step_n, loop_n=loop_n, all_duration=timeout)
+    asyncio.run(kaggle_loop.run(step_n=step_n, loop_n=loop_n, all_duration=timeout))
 
 
 if __name__ == "__main__":

rdagent/app/qlib_rd_loop/factor.py

@@ -2,6 +2,7 @@
 Factor workflow with session control
 """
 
+import asyncio
 from typing import Any
 
 import fire
@@ -40,7 +41,7 @@ def main(path=None, step_n=None, loop_n=None, all_duration=None, checkout=True):
         model_loop = FactorRDLoop(FACTOR_PROP_SETTING)
     else:
         model_loop = FactorRDLoop.load(path, checkout=checkout)
-    model_loop.run(step_n=step_n, loop_n=loop_n, all_duration=all_duration)
+    asyncio.run(model_loop.run(step_n=step_n, loop_n=loop_n, all_duration=all_duration))
 
 
 if __name__ == "__main__":

rdagent/app/qlib_rd_loop/factor_from_report.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 from pathlib import Path
 from typing import Any, Dict, Tuple
@@ -162,7 +163,7 @@ def main(report_folder=None, path=None, step_n=None, loop_n=None, all_duration=N
     else:
         model_loop = FactorReportLoop(report_folder=report_folder)
 
-    model_loop.run(step_n=step_n, loop_n=loop_n, all_duration=all_duration)
+    asyncio.run(model_loop.run(step_n=step_n, loop_n=loop_n, all_duration=all_duration))
 
 
 if __name__ == "__main__":

rdagent/app/qlib_rd_loop/model.py

@@ -2,6 +2,8 @@
 Model workflow with session control
 """
 
+import asyncio
+
 import fire
 
 from rdagent.app.qlib_rd_loop.conf import MODEL_PROP_SETTING
@@ -28,7 +30,7 @@ def main(path=None, step_n=None, loop_n=None, all_duration=None, checkout=True):
         model_loop = ModelRDLoop(MODEL_PROP_SETTING)
     else:
         model_loop = ModelRDLoop.load(path, checkout=checkout)
-    model_loop.run(step_n=step_n, loop_n=loop_n, all_duration=all_duration)
+    asyncio.run(model_loop.run(step_n=step_n, loop_n=loop_n, all_duration=all_duration))
 
 
 if __name__ == "__main__":

rdagent/app/qlib_rd_loop/quant.py

@@ -2,6 +2,7 @@
 Quant (Factor & Model) workflow with session control
 """
 
+import asyncio
 from typing import Any
 
 import fire
@@ -130,7 +131,8 @@ def main(path=None, step_n=None, loop_n=None, all_duration=None, checkout=True):
         quant_loop = QuantRDLoop(QUANT_PROP_SETTING)
     else:
         quant_loop = QuantRDLoop.load(path, checkout=checkout)
-    quant_loop.run(step_n=step_n, loop_n=loop_n, all_duration=all_duration)
+
+    asyncio.run(quant_loop.run(step_n=step_n, loop_n=loop_n, all_duration=all_duration))
 
 
 if __name__ == "__main__":

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

@@ -46,10 +46,10 @@ feature_coder:
     5. You should use the following cache decorator to cache the results of the function:
     ```python
     from joblib import Memory
-    memory = Memory(location='/tmp/cache', verbose=0)
+    memory = Memory(location='./cache', verbose=0)
     @memory.cache```
     6. Coding tricks:
-      - If the input consists of a batch of file paths and you need to modify the file contents to complete your feature engineering task, you can accomplish your feature engineering task by modifying these files and creating new files in a subfolder within "/tmp/cache" (this path is persistent, otherwise you may lose your created file). Then the new file paths are returned.
+      - If the input consists of a batch of file paths and you need to modify the file contents to complete your feature engineering task, you can accomplish your feature engineering task by modifying these files and creating new files in a subfolder within "./cache" (this path is persistent, otherwise you may lose your created file). Then the new file paths are returned.
 
     {% include "scenarios.data_science.share:guidelines.coding" %}
 

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

@@ -43,7 +43,7 @@ model_coder:
     4. You should use the following cache decorator to cache the results of the function:
     ```python
     from joblib import Memory
-    memory = Memory(location='/tmp/cache', verbose=0)
+    memory = Memory(location='./cache', verbose=0)
     @memory.cache``
     {% include "scenarios.data_science.share:guidelines.coding" %}
 

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

@@ -273,7 +273,7 @@ data_loader_coder:
     3. You should use the following cache decorator to cache the results of the function:
     ```python
     from joblib import Memory
-    memory = Memory(location='/tmp/cache', verbose=0)
+    memory = Memory(location='./cache', verbose=0)
     @memory.cache```
     {% include "scenarios.data_science.share:guidelines.coding" %}
     

rdagent/core/conf.py

@@ -78,5 +78,24 @@ class RDAgentSettings(ExtendedBaseSettings):
 
     initial_fator_library_size: int = 20
 
+    # parallel loop
+    step_semaphore: int | dict[str, int] = 1
+    """the semaphore for each step;  you can specify a overall semaphore
+    or a step-wise semaphore like {"coding": 3, "running": 2}"""
+
+    def get_max_parallel(self) -> int:
+        """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())
+
+    # NOTE: for debug
+    # the following function only serves as debugging and is necessary in main logic.
+    subproc_step: bool = False
+
+    def is_force_subproc(self) -> bool:
+        return self.subproc_step or self.get_max_parallel() > 1
+
 
 RD_AGENT_SETTINGS = RDAgentSettings()

rdagent/core/proposal.py

@@ -2,14 +2,19 @@
 
 from __future__ import annotations
 
+import asyncio
 from abc import ABC, abstractmethod
-from typing import Generic, List, Tuple, TypeVar
+from typing import TYPE_CHECKING, Generic, List, Tuple, TypeVar
 
+from rdagent.core.conf import RD_AGENT_SETTINGS
 from rdagent.core.evaluation import Feedback
 from rdagent.core.experiment import ASpecificExp, Experiment
 from rdagent.core.knowledge_base import KnowledgeBase
 from rdagent.core.scenario import Scenario
 
+if TYPE_CHECKING:
+    from rdagent.utils.workflow.loop import LoopBase
+
 
 class Hypothesis:
     """
@@ -248,6 +253,17 @@ class ExpGen(ABC):
             )
         """
 
+    async def async_gen(self, trace: Trace, loop: LoopBase) -> Experiment:
+        """
+        generate the experiment and decide whether to stop yield generation and give up control to other routines.
+        """
+        # we give a default implementation here.
+        # The proposal is set to try best to generate the experiment in max-parallel level.
+        while True:
+            if loop.get_unfinished_loop_cnt(loop.loop_idx) < RD_AGENT_SETTINGS.get_max_parallel():
+                return self.gen(trace)
+            await asyncio.sleep(1)
+
 
 class HypothesisGen(ABC):
 

rdagent/log/conf.py

@@ -12,6 +12,9 @@ class LogSettings(ExtendedBaseSettings):
 
     trace_path: str = str(Path.cwd() / "log" / datetime.now(timezone.utc).strftime("%Y-%m-%d_%H-%M-%S-%f"))
 
+    format_console: str | None = None
+    """"If it is None, leave it as the default"""
+
     ui_server_port: int | None = None
 
     storages: dict[str, list[int | str]] = {}

rdagent/log/logger.py

@@ -1,17 +1,24 @@
 import os
 import sys
 from contextlib import contextmanager
+from contextvars import ContextVar
 from datetime import datetime
 from pathlib import Path
 from typing import Generator
 
 from loguru import logger
+
+from .conf import LOG_SETTINGS
+
+if LOG_SETTINGS.format_console is not None:
+    logger.remove()
+    logger.add(sys.stdout, format=LOG_SETTINGS.format_console)
+
 from psutil import Process
 
 from rdagent.core.utils import SingletonBaseClass, import_class
 
 from .base import Storage
-from .conf import LOG_SETTINGS
 from .storage import FileStorage
 from .utils import get_caller_info
 
@@ -39,14 +46,16 @@ class RDAgentLog(SingletonBaseClass):
 
     """
 
-    # TODO: Simplify it to introduce less concepts ( We may merge RDAgentLog, Storage &)
-    # Solution:  Storage => PipeLog, View => PipeLogView, RDAgentLog is an instance of PipeLogger
-    # PipeLogger.info(...) ,  PipeLogger.get_resp() to get feedback from frontend.
-    # def f():
-    #   logger = PipeLog()
-    #   logger.info("<code>")
-    #   feedback = logger.get_reps()
-    _tag: str = ""
+    # Thread-/coroutine-local tag;  In Linux forked subprocess, it will be copied to the subprocess.
+    _tag_ctx: ContextVar[str] = ContextVar("_tag_ctx", default="")
+
+    @property
+    def _tag(self) -> str:  # Get current tag
+        return self._tag_ctx.get()
+
+    @_tag.setter  # Set current tag
+    def _tag(self, value: str) -> None:
+        self._tag_ctx.set(value)
 
     def __init__(self) -> None:
         self.storage = FileStorage(LOG_SETTINGS.trace_path)
@@ -61,15 +70,16 @@ class RDAgentLog(SingletonBaseClass):
     def tag(self, tag: str) -> Generator[None, None, None]:
         if tag.strip() == "":
             raise ValueError("Tag cannot be empty.")
-        if self._tag != "":
-            tag = "." + tag
-
-        # TODO: It may result in error in mutithreading or co-routine
-        self._tag = self._tag + tag
+        # Generate a new complete tag
+        current_tag = self._tag_ctx.get()
+        new_tag = tag if current_tag == "" else f"{current_tag}.{tag}"
+        # Set and save token for later restore
+        token = self._tag_ctx.set(new_tag)
         try:
             yield
         finally:
-            self._tag = self._tag[: -len(tag)]
+            # Restore previous tag (thread/coroutine safe)
+            self._tag_ctx.reset(token)
 
     def set_storages_path(self, path: str | Path) -> None:
         for storage in [self.storage] + self.other_storages:
@@ -96,7 +106,6 @@ class RDAgentLog(SingletonBaseClass):
         return pid_chain
 
     def log_object(self, obj: object, *, tag: str = "") -> None:
-        # TODO: I think we can merge the log_object function with other normal log methods to make the interface simpler.
         caller_info = get_caller_info()
         tag = f"{self._tag}.{tag}.{self.get_pids()}".strip(".")
 

rdagent/oai/backend/litellm.py

@@ -1,7 +1,9 @@
+import copyreg
 from typing import Any, Literal, cast
 
 import numpy as np
 from litellm import (
+    BadRequestError,
     completion,
     completion_cost,
     embedding,
@@ -15,6 +17,17 @@ from rdagent.oai.backend.base import APIBackend
 from rdagent.oai.llm_conf import LLMSettings
 
 
+# NOTE: Patching! Otherwise, the exception will call the constructor and with following error:
+# `BadRequestError.__init__() missing 2 required positional arguments: 'model' and 'llm_provider'`
+def _reduce_no_init(exc: Exception) -> tuple:
+    cls = exc.__class__
+    return (cls.__new__, (cls,), exc.__dict__)
+
+
+# suppose you want to apply this to MyError
+copyreg.pickle(BadRequestError, _reduce_no_init)
+
+
 class LiteLLMSettings(LLMSettings):
 
     class Config:

rdagent/scenarios/data_science/loop.py

@@ -1,3 +1,4 @@
+import asyncio
 import shutil
 import subprocess
 from datetime import datetime
@@ -31,6 +32,47 @@ from rdagent.scenarios.data_science.dev.runner import DSCoSTEERRunner
 from rdagent.scenarios.data_science.experiment.experiment import DSExperiment
 from rdagent.scenarios.data_science.proposal.exp_gen import DSTrace
 from rdagent.scenarios.data_science.proposal.exp_gen.idea_pool import DSKnowledgeBase
+from rdagent.utils.workflow.misc import wait_retry
+
+
+def clean_workspace(workspace_root: Path) -> None:
+    """
+    Clean the workspace folder and only keep the essential files to save more space.
+
+    # remove all files and folders in the workspace except for .py, .md, and .csv files to avoid large workspace dump
+    """
+    for file_and_folder in workspace_root.iterdir():
+        if file_and_folder.is_dir():
+            if file_and_folder.is_symlink():
+                file_and_folder.unlink()
+            else:
+                shutil.rmtree(file_and_folder)
+        elif file_and_folder.is_file() and file_and_folder.suffix not in [".py", ".md", ".csv"]:
+            file_and_folder.unlink()
+
+
+@wait_retry()
+def backup_folder(path: str | Path) -> Path:
+    path = Path(path)
+    workspace_bak_path = path.with_name(path.name + ".bak")
+    if workspace_bak_path.exists():
+        shutil.rmtree(workspace_bak_path)
+
+    try:
+        # `cp` may raise error if the workspace is beiing modified.
+        # rsync is more robust choice, but it is not installed in some docker images.
+        # use shutil.copytree(..., symlinks=True) should be more elegant, but it has more changes to raise error.
+        subprocess.run(
+            ["cp", "-r", "-P", str(path), str(workspace_bak_path)],
+            check=True,
+            capture_output=True,
+        )
+    except subprocess.CalledProcessError as e:
+        logger.error(f"Error copying {path} to {workspace_bak_path}: {e}")
+        logger.error(f"Stdout: {e.stdout.decode() if e.stdout else ''}")
+        logger.error(f"Stderr: {e.stderr.decode() if e.stderr else ''}")
+        raise
+    return workspace_bak_path
 
 
 class DataScienceRDLoop(RDLoop):
@@ -102,7 +144,7 @@ class DataScienceRDLoop(RDLoop):
         self.summarizer = DSExperiment2Feedback(scen)
         super(RDLoop, self).__init__()
 
-    def direct_exp_gen(self, prev_out: dict[str, Any]):
+    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)
@@ -112,8 +154,7 @@ 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 = self.exp_gen.gen(self.trace)
+        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.
@@ -241,19 +282,29 @@ class DataScienceRDLoop(RDLoop):
                 )
                 / "mid_workspace.tar"
             )
-            subprocess.run(["tar", "-cf", str(mid_log_tar_path), "-C", (Path().cwd() / "log"), "."], check=True)
+            log_back_path = backup_folder(Path().cwd() / "log")
+            subprocess.run(["tar", "-cf", str(mid_log_tar_path), "-C", str(log_back_path), "."], check=True)
 
-            # remove all files and folders in the workspace except for .py, .md, and .csv files to avoid large workspace dump
-            for workspace_id in Path(RD_AGENT_SETTINGS.workspace_path).iterdir():
-                for file_and_folder in workspace_id.iterdir():
-                    if file_and_folder.is_dir():
-                        shutil.rmtree(file_and_folder)
-                    elif file_and_folder.is_file() and file_and_folder.suffix not in [".py", ".md", ".csv"]:
-                        file_and_folder.unlink()
+            # only clean current workspace without affecting other loops.
+            for k in "direct_exp_gen", "coding", "running":
+                if k in prev_out:
+                    assert isinstance(prev_out[k], DSExperiment)
+                    clean_workspace(prev_out[k].experiment_workspace.workspace_path)
+
+            # Backup the workspace (only necessary files are included)
+            # - Step 1: Copy the workspace to a .bak package
+            workspace_bak_path = backup_folder(RD_AGENT_SETTINGS.workspace_path)
+
+            # - Step 2: Clean .bak package
+            for bak_workspace in workspace_bak_path.iterdir():
+                clean_workspace(bak_workspace)
+
+            # - Step 3: Create tarball from the cleaned .bak workspace
+            subprocess.run(["tar", "-cf", str(mid_workspace_tar_path), "-C", str(workspace_bak_path), "."], check=True)
+
+            # - Step 4: Remove .bak package
+            shutil.rmtree(workspace_bak_path)
 
-            subprocess.run(
-                ["tar", "-cf", str(mid_workspace_tar_path), "-C", (RD_AGENT_SETTINGS.workspace_path), "."], check=True
-            )
             if DS_RD_SETTING.log_archive_temp_path is not None:
                 shutil.move(mid_log_tar_path, Path(DS_RD_SETTING.log_archive_path) / "mid_log.tar")
                 mid_log_tar_path = Path(DS_RD_SETTING.log_archive_path) / "mid_log.tar"

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

@@ -297,11 +297,11 @@ output_format:
     Your final output should be a dict containing all the identified problem without anything else.
     Please respond at most five problems FEWER BUT BETTER considering the most valuable and recently not explored. Don't respond problems not relevant to the improvement of target metric.
     {
-      "problem name 1": {
+      "problem name 1 (name of the identified problem without anything else)": {
         "problem": "Description of the first issue in no more than three sentences.",
         "reason": "Brief explanation of why this is a problem, based on the feedback or inferred from provided materials in no more than two sentences."
       },
-      "problem name 2": {
+      "problem name 2 (name of the identified problem without anything else)": {
         "problem": "Description of the second issue in no more than three sentences.",
         "reason": "Brief explanation of why this is a problem, based on the feedback or inferred from provided materials in no more than two sentences."
       }

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

@@ -255,7 +255,7 @@ task_gen:
     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 `/kaggle/input/`. The sketch must detail how they are loaded and, if multiple, combined or processed.
+      - 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.
@@ -290,7 +290,7 @@ task_gen:
 
     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 /kaggle/input/train.csv"). List specific algorithm names where appropriate (e.g., "Apply XGBoost classifier," "Use Isotonic Regression for calibration").
+    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`.
@@ -312,7 +312,7 @@ task_gen:
       - 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.
     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 `/kaggle/input/` (e.g., `/kaggle/input/<competition_name>/train.csv`)."
+      - 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.
@@ -323,7 +323,7 @@ task_gen:
     # Competition Scenario Description
     {{ scenario_desc }}
 
-    # Data Folder Structure (All files are under /kaggle/input/)
+    # Data Folder Structure (All files are under {% include "scenarios.data_science.share:scen.input_path" %})
     {{ data_folder_info }}
 
     # Current SOTA Implementation

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

@@ -285,7 +285,6 @@ def draft_exp_in_decomposition(scen: Scenario, trace: DSTrace) -> None | DSDraft
 
 
 class DSProposalV1ExpGen(ExpGen):
-
     def gen(self, trace: DSTrace) -> DSExperiment:
         # Drafting Stage
         if draft_exp := draft_exp_in_decomposition(self.scen, trace):
@@ -570,6 +569,14 @@ class DSProposalV2ExpGen(ExpGen):
             json_target_type=Dict[str, Dict[str, str | Dict[str, str | int]]],
         )
         resp_dict = json.loads(response)
+
+        # 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 compute_top_scores(
@@ -623,10 +630,12 @@ class DSProposalV2ExpGen(ExpGen):
         for j, problem_name in enumerate(scores_sorted.index):
             if hypothesis_dict[problem_name].get("inspired", False):
                 index_to_pick_pool_list.extend([j] * 2)
-            if problem_dict.get(problem_name, {}).get("label", "") == "SCENARIO_PROBLEM":
+            if problem_dict[problem_name]["label"] == "SCENARIO_PROBLEM":
                 index_to_pick_pool_list.extend([j] * self.scen_prob_multiplier)
-            else:
+            elif problem_dict[problem_name]["label"] == "FEEDBACK_PROBLEM":
                 index_to_pick_pool_list.extend([j] * (3 - self.scen_prob_multiplier))
+            else:
+                index_to_pick_pool_list.extend([j] * 1)
         logger.info(f"index_to_pick_pool_list: {index_to_pick_pool_list}")
 
         # Create a random but reproducible integer
@@ -950,6 +959,13 @@ class DSProposalV3ExpGen(DSProposalV2ExpGen):
             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(

rdagent/scenarios/data_science/share.yaml

@@ -61,7 +61,7 @@ describe: # some template to describe some object
 scen:  # customizable
   role: |-
     You are a Kaggle Grandmaster and expert ML engineer with deep expertise in statistics, machine learning, and competition optimization.
-  input_path: "/kaggle/input"
+  input_path: "./input/"
 
 component_description:
   DataLoadSpec: |-

rdagent/utils/env.py

@@ -7,6 +7,7 @@ Tries to create uniform environment for the agent to run;
 
 # TODO: move the scenario specific docker env into other folders.
 
+import contextlib
 import json
 import os
 import pickle
@@ -20,7 +21,7 @@ import zipfile
 from abc import abstractmethod
 from pathlib import Path
 from types import MappingProxyType
-from typing import Any, Generic, Mapping, Optional, TypeVar
+from typing import Any, Generator, Generic, Mapping, Optional, TypeVar, cast
 
 import docker  # type: ignore[import-untyped]
 import docker.models  # type: ignore[import-untyped]
@@ -42,6 +43,29 @@ from rdagent.oai.llm_utils import md5_hash
 from rdagent.utils.workflow import wait_retry
 
 
+# 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]] = {}
+
+    def to_abs(path: str) -> str:
+        # Converts a relative path to an absolute path using the workspace (working_dir).
+        return os.path.abspath(os.path.join(working_dir, path)) if not os.path.isabs(path) else path
+
+    for lp, vinfo in vols.items():
+        # Support both:
+        # 1. {'host_path': {'bind': 'container_path', ...}}
+        # 2. {'host_path': 'container_path'}
+        if isinstance(vinfo, dict):
+            # abs_vols = cast(dict[str, dict[str, str]], abs_vols)
+            vinfo = vinfo.copy()
+            vinfo["bind"] = to_abs(vinfo["bind"])
+            abs_vols[lp] = vinfo
+        else:
+            # abs_vols = cast(dict[str, str], abs_vols)
+            abs_vols[lp] = to_abs(vinfo)
+    return abs_vols
+
+
 def pull_image_with_progress(image: str) -> None:
     client = docker.APIClient(base_url="unix://var/run/docker.sock")
     pull_logs = client.pull(image, stream=True, decode=True)
@@ -213,10 +237,20 @@ class Env(Generic[ASpecificEnvConf]):
                 "the last command in the pipeline.",
             )
 
+        # FIXME: the input path and cache path is hard coded here.
+        # We don't want to change the content in input and cache path.
+        # Otherwise, it may produce large amount of warnings.
         entry_add_timeout = (
             f"/bin/sh -c 'timeout --kill-after=10 {self.conf.running_timeout_period} {entry}; "
             + "entry_exit_code=$?; "
-            + (f"chmod -R 777 {self.conf.mount_path}; " if hasattr(self.conf, "mount_path") else "")
+            + (
+                f"chmod -R 777 $(find {self.conf.mount_path} -mindepth 1 -maxdepth 1 ! -name cache ! -name input); "
+                # We don't have to change the permission of the cache and input folder to remove it
+                # + f"if [ -d {self.conf.mount_path}/cache ]; then chmod 777 {self.conf.mount_path}/cache; fi; " +
+                #     f"if [ -d {self.conf.mount_path}/input ]; then chmod 777 {self.conf.mount_path}/input; fi; "
+                if hasattr(self.conf, "mount_path")
+                else ""
+            )
             + "exit $entry_exit_code'"
         )
 
@@ -375,98 +409,116 @@ class LocalEnv(Env[ASpecificLocalConf]):
                 volumes[lp] = rp
             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] = "/tmp/cache"
+            volumes[cache_path] = "./cache"
         for lp, rp in running_extra_volume.items():
             volumes[lp] = rp
-        for rp, lp in volumes.items():
-            link_path = Path(lp)
-            real_path = Path(rp)
-            if not link_path.parent.exists():
-                link_path.parent.mkdir(parents=True, exist_ok=True)
-            if link_path.exists() or link_path.is_symlink():
-                link_path.unlink()
-            link_path.symlink_to(real_path)
 
-        # Setup environment
-        if env is None:
-            env = {}
-        path = [*self.conf.bin_path.split(":"), "/bin/", "/usr/bin/", *env.get("PATH", "").split(":")]
-        env["PATH"] = ":".join(path)
+        assert local_path is not None, "local_path should not be None"
+        volumes = normalize_volumes(volumes, local_path)
 
-        if entry is None:
-            entry = self.conf.default_entry
+        @contextlib.contextmanager
+        def _symlink_ctx(vol_map: Mapping[str, str]) -> Generator[None, None, None]:
+            created_links: list[Path] = []
+            try:
+                for real, link in vol_map.items():
+                    link_path = Path(link)
+                    real_path = Path(real)
+                    if not link_path.parent.exists():
+                        link_path.parent.mkdir(parents=True, exist_ok=True)
+                    if link_path.exists() or link_path.is_symlink():
+                        link_path.unlink()
+                    link_path.symlink_to(real_path)
+                    created_links.append(link_path)
+                yield
+            finally:
+                for p in created_links:
+                    try:
+                        if p.is_symlink() or p.exists():
+                            p.unlink()
+                    except FileNotFoundError:
+                        pass
 
-        print(Rule("[bold green]LocalEnv Logs Begin[/bold green]", style="dark_orange"))
-        table = Table(title="Run Info", show_header=False)
-        table.add_column("Key", style="bold cyan")
-        table.add_column("Value", style="bold magenta")
-        table.add_row("Entry", entry)
-        table.add_row("Local Path", local_path or "")
-        table.add_row("Env", "\n".join(f"{k}:{v}" for k, v in env.items()))
-        table.add_row("Volumes", "\n".join(f"{k}:{v}" for k, v in volumes.items()))
-        print(table)
+        with _symlink_ctx(volumes):
+            # Setup environment
+            if env is None:
+                env = {}
+            path = [*self.conf.bin_path.split(":"), "/bin/", "/usr/bin/", *env.get("PATH", "").split(":")]
+            env["PATH"] = ":".join(path)
 
-        cwd = Path(local_path).resolve() if local_path else None
-        env = {k: str(v) if isinstance(v, int) else v for k, v in env.items()}
+            if entry is None:
+                entry = self.conf.default_entry
 
-        process = subprocess.Popen(
-            entry,
-            cwd=cwd,
-            env={**os.environ, **env},
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            text=True,
-            shell=True,
-            bufsize=1,
-            universal_newlines=True,
-        )
+            print(Rule("[bold green]LocalEnv Logs Begin[/bold green]", style="dark_orange"))
+            table = Table(title="Run Info", show_header=False)
+            table.add_column("Key", style="bold cyan")
+            table.add_column("Value", style="bold magenta")
+            table.add_row("Entry", entry)
+            table.add_row("Local Path", local_path or "")
+            table.add_row("Env", "\n".join(f"{k}:{v}" for k, v in env.items()))
+            table.add_row("Volumes", "\n".join(f"{k}:{v}" for k, v in volumes.items()))
+            print(table)
 
-        # Setup polling
-        if process.stdout is None or process.stderr is None:
-            raise RuntimeError("The subprocess did not correctly create stdout/stderr pipes")
+            cwd = Path(local_path).resolve() if local_path else None
+            env = {k: str(v) if isinstance(v, int) else v for k, v in env.items()}
 
-        stdout_fd = process.stdout.fileno()
-        stderr_fd = process.stderr.fileno()
+            process = subprocess.Popen(
+                entry,
+                cwd=cwd,
+                env={**os.environ, **env},
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                shell=True,
+                bufsize=1,
+                universal_newlines=True,
+            )
 
-        poller = select.poll()
-        poller.register(stdout_fd, select.POLLIN)
-        poller.register(stderr_fd, select.POLLIN)
+            # Setup polling
+            if process.stdout is None or process.stderr is None:
+                raise RuntimeError("The subprocess did not correctly create stdout/stderr pipes")
 
-        combined_output = ""
-        while True:
-            if process.poll() is not None:
-                break
-            events = poller.poll(100)
-            for fd, event in events:
-                if event & select.POLLIN:
-                    if fd == stdout_fd:
-                        while True:
-                            output = process.stdout.readline()
-                            if output == "":
-                                break
-                            Console().print(output.strip(), markup=False)
-                            combined_output += output
-                    elif fd == stderr_fd:
-                        while True:
-                            error = process.stderr.readline()
-                            if error == "":
-                                break
-                            Console().print(error.strip(), markup=False)
-                            combined_output += error
+            stdout_fd = process.stdout.fileno()
+            stderr_fd = process.stderr.fileno()
 
-        # Capture any final output
-        remaining_output, remaining_error = process.communicate()
-        if remaining_output:
-            Console().print(remaining_output.strip(), markup=False)
-            combined_output += remaining_output
-        if remaining_error:
-            Console().print(remaining_error.strip(), markup=False)
-            combined_output += remaining_error
+            poller = select.poll()
+            poller.register(stdout_fd, select.POLLIN)
+            poller.register(stderr_fd, select.POLLIN)
 
-        return_code = process.returncode
-        print(Rule("[bold green]LocalEnv Logs End[/bold green]", style="dark_orange"))
+            combined_output = ""
+            while True:
+                if process.poll() is not None:
+                    break
+                events = poller.poll(100)
+                for fd, event in events:
+                    if event & select.POLLIN:
+                        if fd == stdout_fd:
+                            while True:
+                                output = process.stdout.readline()
+                                if output == "":
+                                    break
+                                Console().print(output.strip(), markup=False)
+                                combined_output += output
+                        elif fd == stderr_fd:
+                            while True:
+                                error = process.stderr.readline()
+                                if error == "":
+                                    break
+                                Console().print(error.strip(), markup=False)
+                                combined_output += error
 
-        return combined_output, return_code
+            # Capture any final output
+            remaining_output, remaining_error = process.communicate()
+            if remaining_output:
+                Console().print(remaining_output.strip(), markup=False)
+                combined_output += remaining_output
+            if remaining_error:
+                Console().print(remaining_error.strip(), markup=False)
+                combined_output += remaining_error
+
+            return_code = process.returncode
+            print(Rule("[bold green]LocalEnv Logs End[/bold green]", style="dark_orange"))
+
+            return combined_output, return_code
 
 
 class CondaConf(LocalConf):
@@ -769,10 +821,12 @@ class DockerEnv(Env[DockerConf]):
                 volumes[lp] = {"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": "/tmp/cache", "mode": "rw"}
+            volumes[cache_path] = {"bind": "./cache", "mode": "rw"}
         for lp, rp in running_extra_volume.items():
             volumes[lp] = {"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 = ""
 
         try:

rdagent/utils/workflow.py

@@ -1,370 +0,0 @@
-"""
-This is a class that try to store/resume/traceback the workflow session
-
-
-Postscripts:
-- Originally, I want to implement it in a more general way with python generator.
-  However, Python generator is not picklable (dill does not support pickle as well)
-
-"""
-
-import datetime
-import pickle
-import time
-from collections import defaultdict
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any, Callable, TypeVar, cast
-
-import pytz
-from tqdm.auto import tqdm
-
-from rdagent.core.conf import RD_AGENT_SETTINGS
-from rdagent.log import rdagent_logger as logger
-from rdagent.log.conf import LOG_SETTINGS
-from rdagent.log.timer import RD_Agent_TIMER_wrapper, RDAgentTimer
-
-if RD_AGENT_SETTINGS.enable_mlflow:
-    import mlflow
-
-
-class LoopMeta(type):
-    @staticmethod
-    def _get_steps(bases: tuple[type, ...]) -> list[str]:
-        """
-        Recursively get all the `steps` from the base classes and combine them into a single list.
-
-        Args:
-            bases (tuple): A tuple of base classes.
-
-        Returns:
-            List[Callable]: A list of steps combined from all base classes.
-        """
-        steps = []
-        for base in bases:
-            for step in LoopMeta._get_steps(base.__bases__) + getattr(base, "steps", []):
-                if step not in steps and step not in ["load", "dump"]:  # incase user override the load/dump method
-                    steps.append(step)
-        return steps
-
-    def __new__(mcs, clsname: str, bases: tuple[type, ...], attrs: dict[str, Any]) -> Any:
-        """
-        Create a new class with combined steps from base classes and current class.
-
-        Args:
-            clsname (str): Name of the new class.
-            bases (tuple): Base classes.
-            attrs (dict): Attributes of the new class.
-
-        Returns:
-            LoopMeta: A new instance of LoopMeta.
-        """
-        steps = LoopMeta._get_steps(bases)  # all the base classes of parents
-        for name, attr in attrs.items():
-            if not name.startswith("_") and callable(attr):
-                if name not in steps and name not in ["load", "dump"]:  # incase user override the load/dump method
-                    # NOTE: if we override the step in the subclass
-                    # Then it is not the new step. So we skip it.
-                    steps.append(name)
-        attrs["steps"] = steps
-        return super().__new__(mcs, clsname, bases, attrs)
-
-
-@dataclass
-class LoopTrace:
-    start: datetime.datetime  # the start time of the trace
-    end: datetime.datetime  # the end time of the trace
-    step_idx: int
-    # TODO: more information about the trace
-
-
-class LoopBase:
-    """
-    Assumption:
-    - The last step is responsible for recording information!!!!
-    """
-
-    steps: list[str]  # a list of steps to work on
-    loop_trace: dict[int, list[LoopTrace]]
-
-    skip_loop_error: tuple[type[BaseException], ...] = ()  # you can define a list of error that will skip current loop
-    withdraw_loop_error: tuple[
-        type[BaseException], ...
-    ] = ()  # you can define a list of error that will withdraw current loop
-
-    EXCEPTION_KEY = "_EXCEPTION"
-
-    def __init__(self) -> None:
-        self.loop_idx = 0  # current loop index
-        self.step_idx = 0  # the index of next step to be run
-        self.loop_prev_out: dict[str, Any] = {}  # the step results of current loop
-        self.loop_trace = defaultdict(list[LoopTrace])  # the key is the number of loop
-        self.session_folder = Path(LOG_SETTINGS.trace_path) / "__session__"
-        self.timer: RDAgentTimer = RD_Agent_TIMER_wrapper.timer
-
-    def run(self, step_n: int | None = None, loop_n: int | None = None, all_duration: str | None = None) -> None:
-        """
-
-        Parameters
-        ----------
-        step_n : int | None
-            How many steps to run;
-            `None` indicates to run forever until error or KeyboardInterrupt
-        loop_n: int | None
-            How many steps to run; if current loop is incomplete, it will be counted as the first loop for completion
-            `None` indicates to run forever until error or KeyboardInterrupt
-        """
-
-        if all_duration is not None and not self.timer.started:
-            self.timer.reset(all_duration=all_duration)
-
-        with tqdm(total=len(self.steps), desc="Workflow Progress", unit="step") as pbar:
-            while True:
-                if step_n is not None:
-                    if step_n <= 0:
-                        break
-                    step_n -= 1
-                if loop_n is not None:
-                    if loop_n <= 0:
-                        break
-
-                if RD_AGENT_SETTINGS.enable_mlflow:
-                    mlflow.log_metric("loop_index", self.loop_idx)
-                    mlflow.log_metric("step_index", self.step_idx)
-                    current_local_datetime = datetime.datetime.now(pytz.timezone("Asia/Shanghai"))
-                    float_like_datetime = (
-                        current_local_datetime.second
-                        + current_local_datetime.minute * 1e2
-                        + current_local_datetime.hour * 1e4
-                        + current_local_datetime.day * 1e6
-                        + current_local_datetime.month * 1e8
-                        + current_local_datetime.year * 1e10
-                    )
-                    mlflow.log_metric("current_datetime", float_like_datetime)
-                    mlflow.log_metric("api_fail_count", RD_Agent_TIMER_wrapper.api_fail_count)
-                    lastest_api_fail_time = RD_Agent_TIMER_wrapper.latest_api_fail_time
-                    if lastest_api_fail_time is not None:
-                        mlflow.log_metric(
-                            "lastest_api_fail_time",
-                            (
-                                lastest_api_fail_time.second
-                                + lastest_api_fail_time.minute * 1e2
-                                + lastest_api_fail_time.hour * 1e4
-                                + lastest_api_fail_time.day * 1e6
-                                + lastest_api_fail_time.month * 1e8
-                                + lastest_api_fail_time.year * 1e10
-                            ),
-                        )
-
-                if self.timer.started:
-                    if RD_AGENT_SETTINGS.enable_mlflow:
-                        mlflow.log_metric("remain_time", self.timer.remain_time().seconds)  # type: ignore[union-attr]
-                        mlflow.log_metric(
-                            "remain_percent", self.timer.remain_time() / self.timer.all_duration * 100  # type: ignore[operator]
-                        )
-
-                    if self.timer.is_timeout():
-                        logger.warning("Timeout, exiting the loop.")
-                        break
-                    else:
-                        logger.info(f"Timer remaining time: {self.timer.remain_time()}")
-
-                li, si = self.loop_idx, self.step_idx
-                name = self.steps[si]
-                logger.info(f"Start Loop {li}, Step {si}: {name}")
-                with logger.tag(f"Loop_{li}.{name}"):
-                    start = datetime.datetime.now(datetime.timezone.utc)
-                    func: Callable[..., Any] = cast(Callable[..., Any], getattr(self, name))
-                    try:
-                        self.loop_prev_out[name] = func(self.loop_prev_out)
-                        # TODO: Fix the error logger.exception(f"Skip loop {li} due to {e}")
-                    except Exception as e:
-                        if isinstance(e, self.skip_loop_error):
-                            # FIXME: This does not support previous demo (due to their last step is not for recording)
-                            logger.warning(f"Skip loop {li} due to {e}")
-                            # NOTE: strong assumption!  The last step is responsible for recording information
-                            self.step_idx = len(self.steps) - 1  # directly jump to the last step.
-                            self.loop_prev_out[self.EXCEPTION_KEY] = e
-                            continue
-                        elif isinstance(e, self.withdraw_loop_error):
-                            logger.warning(f"Withdraw loop {li} due to {e}")
-                            # Back to previous loop
-                            self.withdraw_loop(li)
-                            continue
-                        else:
-                            raise
-                    finally:
-                        # make sure failure steps are displayed correclty
-                        end = datetime.datetime.now(datetime.timezone.utc)
-                        self.loop_trace[li].append(LoopTrace(start, end, step_idx=si))
-
-                        # Update tqdm progress bar directly to step_idx
-                        pbar.n = si + 1
-                        pbar.set_postfix(
-                            loop_index=li, step_index=si + 1, step_name=name
-                        )  # step_name indicate  last finished step_name
-
-                # index increase and save session
-                self.step_idx = (self.step_idx + 1) % len(self.steps)
-                if self.step_idx == 0:  # reset to step 0 in next round
-                    self.loop_idx += 1
-                    if loop_n is not None:
-                        loop_n -= 1
-                    self.loop_prev_out = {}
-                    pbar.reset()  # reset the progress bar for the next loop
-
-                self.dump(self.session_folder / f"{li}" / f"{si}_{name}")  # save a snapshot after the session
-
-    def withdraw_loop(self, loop_idx: int) -> None:
-        prev_session_dir = self.session_folder / str(loop_idx - 1)
-        prev_path = min(
-            (p for p in prev_session_dir.glob("*_*") if p.is_file()),
-            key=lambda item: int(item.name.split("_", 1)[0]),
-            default=None,
-        )
-        if prev_path:
-            loaded = type(self).load(
-                prev_path,
-                checkout=True,
-                replace_timer=True,
-            )
-            logger.info(f"Load previous session from {prev_path}")
-            # Overwrite current instance state
-            self.__dict__ = loaded.__dict__
-        else:
-            logger.error(f"No previous dump found at {prev_session_dir}, cannot withdraw loop {loop_idx}")
-            raise
-
-    def dump(self, path: str | Path) -> None:
-        if RD_Agent_TIMER_wrapper.timer.started:
-            RD_Agent_TIMER_wrapper.timer.update_remain_time()
-        path = Path(path)
-        path.parent.mkdir(parents=True, exist_ok=True)
-        with path.open("wb") as f:
-            pickle.dump(self, f)
-
-    def truncate_session_folder(self, li: int, si: int) -> None:
-        """
-        Clear the session folder by removing all session objects after the given loop index (li) and step index (si).
-        """
-        # clear session folders after the li
-        for sf in self.session_folder.iterdir():
-            if sf.is_dir() and int(sf.name) > li:
-                for file in sf.iterdir():
-                    file.unlink()
-                sf.rmdir()
-
-        # clear step session objects in the li
-        final_loop_session_folder = self.session_folder / str(li)
-        for step_session in final_loop_session_folder.glob("*_*"):
-            if step_session.is_file():
-                step_id = int(step_session.name.split("_", 1)[0])
-                if step_id > si:
-                    step_session.unlink()
-
-    @classmethod
-    def load(
-        cls,
-        path: str | Path,
-        checkout: bool | Path | str = False,
-        replace_timer: bool = True,
-    ) -> "LoopBase":
-        """
-        Load a session from a given path.
-        Parameters
-        ----------
-        path : str | Path
-            The path to the session file.
-        checkout : bool | Path | str
-            If True, the new loop will use the existing folder and clear logs for sessions after the one corresponding to the given path.
-            If False, the new loop will use the existing folder but keep the logs for sessions after the one corresponding to the given path.
-            If a path (or a str like Path) is provided, the new loop will be saved to that path, leaving the original path unchanged.
-        replace_timer : bool
-            If a session is loaded, determines whether to replace the timer with session.timer.
-            Default is True, which means the session timer will be replaced with the current timer.
-            If False, the session timer will not be replaced.
-        Returns
-        -------
-        LoopBase
-            An instance of LoopBase with the loaded session.
-        """
-        path = Path(path)
-        with path.open("rb") as f:
-            session = cast(LoopBase, pickle.load(f))
-
-        # set session folder
-        if checkout:
-            if checkout is True:
-                logger.set_storages_path(session.session_folder.parent)
-                max_loop = max(session.loop_trace.keys())
-
-                # truncate log storages after the max loop
-                session.truncate_session_folder(max_loop, len(session.loop_trace[max_loop]) - 1)
-                logger.truncate_storages(session.loop_trace[max_loop][-1].end)
-            else:
-                checkout = Path(checkout)
-                checkout.mkdir(parents=True, exist_ok=True)
-                session.session_folder = checkout / "__session__"
-                logger.set_storages_path(checkout)
-
-        if session.timer.started:
-            if replace_timer:
-                RD_Agent_TIMER_wrapper.replace_timer(session.timer)
-                RD_Agent_TIMER_wrapper.timer.restart_by_remain_time()
-            else:
-                # Use the default timer to replace the session timer
-                session.timer = RD_Agent_TIMER_wrapper.timer
-
-        return session
-
-
-ASpecificRet = TypeVar("ASpecificRet")
-
-
-def wait_retry(
-    retry_n: int = 3, sleep_time: int = 1, transform_args_fn: Callable[[tuple, dict], tuple[tuple, dict]] | None = None
-) -> Callable[[Callable[..., ASpecificRet]], Callable[..., ASpecificRet]]:
-    """Decorator to wait and retry the function for retry_n times.
-
-    Example:
-    >>> import time
-    >>> @wait_retry(retry_n=2, sleep_time=1)
-    ... def test_func():
-    ...     global counter
-    ...     counter += 1
-    ...     if counter < 3:
-    ...         raise ValueError("Counter is less than 3")
-    ...     return counter
-    >>> counter = 0
-    >>> try:
-    ...     test_func()
-    ... except ValueError as e:
-    ...     print(f"Caught an exception: {e}")
-    Error: Counter is less than 3
-    Error: Counter is less than 3
-    Caught an exception: Counter is less than 3
-    >>> counter
-    2
-    """
-    assert retry_n > 0, "retry_n should be greater than 0"
-
-    def decorator(f: Callable[..., ASpecificRet]) -> Callable[..., ASpecificRet]:
-        def wrapper(*args: Any, **kwargs: Any) -> ASpecificRet:
-            for i in range(retry_n + 1):
-                try:
-                    return f(*args, **kwargs)
-                except Exception as e:
-                    print(f"Error: {e}")
-                    time.sleep(sleep_time)
-                    if i == retry_n:
-                        raise
-                    # Update args and kwargs using the transform function if provided.
-                    if transform_args_fn is not None:
-                        args, kwargs = transform_args_fn(args, kwargs)
-            else:
-                # just for passing mypy CI.
-                return f(*args, **kwargs)
-
-        return wrapper
-
-    return decorator

rdagent/utils/workflow/__init__.py

@@ -0,0 +1,5 @@
+from .loop import LoopBase, LoopMeta
+from .misc import wait_retry
+from .tracking import WorkflowTracker
+
+__all__ = ["LoopBase", "LoopMeta", "WorkflowTracker", "wait_retry"]

rdagent/utils/workflow/loop.py

@@ -0,0 +1,444 @@
+"""
+This is a class that try to store/resume/traceback the workflow session
+
+
+Postscripts:
+- Originally, I want to implement it in a more general way with python generator.
+  However, Python generator is not picklable (dill does not support pickle as well)
+
+"""
+
+import asyncio
+import concurrent.futures
+import datetime
+import pickle
+from collections import defaultdict
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable, Optional, Union, cast
+
+from tqdm.auto import tqdm
+
+from rdagent.core.conf import RD_AGENT_SETTINGS
+from rdagent.log import rdagent_logger as logger
+from rdagent.log.conf import LOG_SETTINGS
+from rdagent.log.timer import RD_Agent_TIMER_wrapper, RDAgentTimer
+from rdagent.utils.workflow.tracking import WorkflowTracker
+
+
+class LoopMeta(type):
+
+    @staticmethod
+    def _get_steps(bases: tuple[type, ...]) -> list[str]:
+        """
+        Recursively get all the `steps` from the base classes and combine them into a single list.
+
+        Args:
+            bases (tuple): A tuple of base classes.
+
+        Returns:
+            List[Callable]: A list of steps combined from all base classes.
+        """
+        steps = []
+        for base in bases:
+            for step in LoopMeta._get_steps(base.__bases__) + getattr(base, "steps", []):
+                if step not in steps and step not in ["load", "dump"]:  # incase user override the load/dump method
+                    steps.append(step)
+        return steps
+
+    def __new__(mcs, clsname: str, bases: tuple[type, ...], attrs: dict[str, Any]) -> Any:
+        """
+        Create a new class with combined steps from base classes and current class.
+
+        Args:
+            clsname (str): Name of the new class.
+            bases (tuple): Base classes.
+            attrs (dict): Attributes of the new class.
+
+        Returns:
+            LoopMeta: A new instance of LoopMeta.
+        """
+        steps = LoopMeta._get_steps(bases)  # all the base classes of parents
+        for name, attr in attrs.items():
+            if not name.startswith("_") and callable(attr):
+                if name not in steps and name not in ["load", "dump"]:  # incase user override the load/dump method
+                    # NOTE: if we override the step in the subclass
+                    # Then it is not the new step. So we skip it.
+                    steps.append(name)
+        attrs["steps"] = steps
+        return super().__new__(mcs, clsname, bases, attrs)
+
+
+@dataclass
+class LoopTrace:
+    start: datetime.datetime  # the start time of the trace
+    end: datetime.datetime  # the end time of the trace
+    step_idx: int
+    # TODO: more information about the trace
+
+
+class LoopBase:
+    """
+    Assumption:
+    - The last step is responsible for recording information!!!!
+
+    Unsolved problem:
+    - Global variable synchronization when `force_subproc` is True
+        - Timer
+    """
+
+    steps: list[str]  # a list of steps to work on
+    loop_trace: dict[int, list[LoopTrace]]
+
+    skip_loop_error: tuple[type[BaseException], ...] = ()  # you can define a list of error that will skip current loop
+    withdraw_loop_error: tuple[
+        type[BaseException], ...
+    ] = ()  # you can define a list of error that will withdraw current loop
+
+    EXCEPTION_KEY = "_EXCEPTION"
+
+    _pbar: tqdm  # progress bar instance
+
+    class LoopTerminationError(Exception):
+        """Exception raised when loop conditions indicate the loop should terminate"""
+
+    class LoopResumeError(Exception):
+        """Exception raised when loop conditions indicate the loop should stop all coroutines and resume"""
+
+    def __init__(self) -> None:
+        # progress control
+        self.loop_idx: int = 0  # current loop index / next loop index to kickoff
+        self.step_idx: defaultdict[int, int] = defaultdict(int)  # dict from loop index to next step index
+        self.queue: asyncio.Queue[Any] = asyncio.Queue()
+
+        # Store step results for all loops in a nested dictionary: loop_prev_out[loop_index][step_name]
+        self.loop_prev_out: dict[int, dict[str, Any]] = defaultdict(dict)
+        self.loop_trace = defaultdict(list[LoopTrace])  # the key is the number of loop
+        self.session_folder = Path(LOG_SETTINGS.trace_path) / "__session__"
+        self.timer: RDAgentTimer = RD_Agent_TIMER_wrapper.timer
+        self.tracker = WorkflowTracker(self)  # Initialize tracker with this LoopBase instance
+
+        # progress control
+        self.loop_n: Optional[int] = None  # remain loop count
+        self.step_n: Optional[int] = None  # remain step count
+
+        self.semaphores: dict[str, asyncio.Semaphore] = {}
+
+    def get_unfinished_loop_cnt(self, next_loop: int) -> int:
+        n = 0
+        for li in range(next_loop):
+            if self.step_idx[li] < len(self.steps):  # unfinished loop
+                n += 1
+        return n
+
+    def get_semaphore(self, step_name: str) -> asyncio.Semaphore:
+        if isinstance(limit := RD_AGENT_SETTINGS.step_semaphore, dict):
+            limit = limit.get(step_name, 1)  # default to 1 if not specified
+
+        if step_name not in self.semaphores:
+            self.semaphores[step_name] = asyncio.Semaphore(limit)
+        return self.semaphores[step_name]
+
+    @property
+    def pbar(self) -> tqdm:
+        """Progress bar property that initializes itself if it doesn't exist."""
+        if getattr(self, "_pbar", None) is None:
+            self._pbar = tqdm(total=len(self.steps), desc="Workflow Progress", unit="step")
+        return self._pbar
+
+    def close_pbar(self) -> None:
+        if getattr(self, "_pbar", None) is not None:
+            self._pbar.close()
+            del self._pbar
+
+    def _check_exit_conditions_on_step(self) -> None:
+        """Check if the loop should continue or terminate.
+
+        Raises
+        ------
+        LoopTerminationException
+            When conditions indicate that the loop should terminate
+        """
+        # Check step count limitation
+        if self.step_n is not None:
+            if self.step_n <= 0:
+                raise self.LoopTerminationError("Step count reached")
+            self.step_n -= 1
+
+        # Check timer timeout
+        if self.timer.started:
+            if self.timer.is_timeout():
+                logger.warning("Timeout, exiting the loop.")
+                raise self.LoopTerminationError("Timer timeout")
+            else:
+                logger.info(f"Timer remaining time: {self.timer.remain_time()}")
+
+    async def _run_step(self, li: int, force_subproc: bool = False) -> None:
+        """Execute a single step (next unrun step) in the workflow (async version with force_subproc option).
+
+        Parameters
+        ----------
+        li : int
+            Loop index
+
+        force_subproc : bool
+            Whether to force the step to run in a subprocess in asyncio
+
+        Returns
+        -------
+        Any
+            The result of the step function
+        """
+        si = self.step_idx[li]
+        name = self.steps[si]
+
+        async with self.get_semaphore(name):
+
+            logger.info(f"Start Loop {li}, Step {si}: {name}")
+            self.tracker.log_workflow_state()
+
+            with logger.tag(f"Loop_{li}.{name}"):
+                start = datetime.datetime.now(datetime.timezone.utc)
+                func: Callable[..., Any] = cast(Callable[..., Any], getattr(self, name))
+
+                next_step_idx = si + 1
+                step_forward = True
+                try:
+                    # Call function with current loop's output, await if coroutine or use ProcessPoolExecutor for sync if required
+                    if force_subproc:
+                        curr_loop = asyncio.get_running_loop()
+                        with concurrent.futures.ProcessPoolExecutor() as pool:
+                            result = await curr_loop.run_in_executor(pool, func, self.loop_prev_out[li])
+                    else:
+                        # auto determine whether to run async or sync
+                        if asyncio.iscoroutinefunction(func):
+                            result = await func(self.loop_prev_out[li])
+                        else:
+                            # Default: run sync function directly
+                            result = func(self.loop_prev_out[li])
+                    # Store result in the nested dictionary
+                    self.loop_prev_out[li][name] = result
+                    # Save snapshot after completing the step
+                    self.dump(self.session_folder / f"{li}" / f"{si}_{name}")
+                except Exception as e:
+                    if isinstance(e, self.skip_loop_error):
+                        logger.warning(f"Skip loop {li} due to {e}")
+                        # Jump to the last step (assuming last step is for recording)
+                        next_step_idx = len(self.steps) - 1
+                        self.loop_prev_out[li][self.EXCEPTION_KEY] = e
+                    elif isinstance(e, self.withdraw_loop_error):
+                        logger.warning(f"Withdraw loop {li} due to {e}")
+                        # Back to previous loop
+                        self.withdraw_loop(li)
+                        step_forward = False
+
+                        msg = "We have reset the loop instance, stop all the routines and resume."
+                        raise self.LoopResumeError(msg) from e
+                    else:
+                        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
+
+                        # Update progress bar
+                        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._check_exit_conditions_on_step()
+                    else:
+                        logger.warning(f"Step forward {si} of loop {li} is skipped.")
+
+    async def kickoff_loop(self) -> None:
+        while True:
+            li = self.loop_idx
+
+            # exit on loop limitation
+            if self.loop_n is not None:
+                if self.loop_n <= 0:
+                    break
+                self.loop_n -= 1
+
+            # NOTE:
+            # Try best to kick off the first step; the first step is always the ExpGen;
+            # it have the right to decide when to stop yield new Experiment
+            if self.step_idx[li] == 0:
+                # Assume the first step is ExpGen
+                # Only kick off ExpGen when it is never kicked off before
+                await self._run_step(li)
+            self.queue.put_nowait(li)  # the loop `li` has been kicked off, waiting for workers to pick it up
+            self.loop_idx += 1
+
+    async def execute_loop(self) -> None:
+        while True:
+            # 1) get the tasks to goon loop `li`
+            li = await self.queue.get()
+            # 2) run the unfinished steps
+            while self.step_idx[li] < len(self.steps):
+                if self.step_idx[li] == len(self.steps) - 1:
+                    # NOTE: assume the last step is record, it will be fast and affect the global environment
+                    # if it is the last step, run it directly ()
+                    await self._run_step(li)
+                else:
+                    # await the step; parallel running happens here!
+                    # Only trigger subprocess if we have more than one process.
+                    await self._run_step(li, force_subproc=RD_AGENT_SETTINGS.is_force_subproc())
+
+    async def run(self, step_n: int | None = None, loop_n: int | None = None, all_duration: str | None = None) -> None:
+        """Run the workflow loop.
+
+        Parameters
+        ----------
+        loop_n: int | None
+            How many loops to run; if current loop is incomplete, it will be counted as the first loop for completion
+            `None` indicates to run forever until error or KeyboardInterrupt
+        all_duration : str | None
+            Maximum duration to run, in format accepted by the timer
+        """
+        # Initialize timer if duration is provided
+        if all_duration is not None and not self.timer.started:
+            self.timer.reset(all_duration=all_duration)
+
+        self.step_n, self.loop_n = step_n, loop_n
+
+        # empty the queue when restarting
+        while not self.queue.empty():
+            self.queue.get_nowait()
+        self.loop_idx = (
+            0  # if we rerun the loop, we should revert the loop index to 0 to make sure every loop is correctly kicked
+        )
+
+        while True:
+            try:
+                # run one kickoff_loop and execute_loop
+                await asyncio.gather(
+                    self.kickoff_loop(), *[self.execute_loop() for _ in range(RD_AGENT_SETTINGS.get_max_parallel())]
+                )
+                break
+            except self.LoopResumeError as e:
+                logger.warning(f"Stop all the routines and resume loop: {e}")
+                self.loop_idx = 0
+            except self.LoopTerminationError as e:
+                logger.warning(f"Reach stop criterion and stop loop: {e}")
+                break
+            finally:
+                self.close_pbar()
+
+    def withdraw_loop(self, loop_idx: int) -> None:
+        prev_session_dir = self.session_folder / str(loop_idx - 1)
+        prev_path = min(
+            (p for p in prev_session_dir.glob("*_*") if p.is_file()),
+            key=lambda item: int(item.name.split("_", 1)[0]),
+            default=None,
+        )
+        if prev_path:
+            loaded = type(self).load(
+                prev_path,
+                checkout=True,
+                replace_timer=True,
+            )
+            logger.info(f"Load previous session from {prev_path}")
+            # Overwrite current instance state
+            self.__dict__ = loaded.__dict__
+        else:
+            logger.error(f"No previous dump found at {prev_session_dir}, cannot withdraw loop {loop_idx}")
+            raise
+
+    def dump(self, path: str | Path) -> None:
+        if RD_Agent_TIMER_wrapper.timer.started:
+            RD_Agent_TIMER_wrapper.timer.update_remain_time()
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with path.open("wb") as f:
+            pickle.dump(self, f)
+
+    def truncate_session_folder(self, li: int, si: int) -> None:
+        """
+        Clear the session folder by removing all session objects after the given loop index (li) and step index (si).
+        """
+        # clear session folders after the li
+        for sf in self.session_folder.iterdir():
+            if sf.is_dir() and int(sf.name) > li:
+                for file in sf.iterdir():
+                    file.unlink()
+                sf.rmdir()
+
+        # clear step session objects in the li
+        final_loop_session_folder = self.session_folder / str(li)
+        for step_session in final_loop_session_folder.glob("*_*"):
+            if step_session.is_file():
+                step_id = int(step_session.name.split("_", 1)[0])
+                if step_id > si:
+                    step_session.unlink()
+
+    @classmethod
+    def load(
+        cls,
+        path: str | Path,
+        checkout: bool | Path | str = False,
+        replace_timer: bool = True,
+    ) -> "LoopBase":
+        """
+        Load a session from a given path.
+        Parameters
+        ----------
+        path : str | Path
+            The path to the session file.
+        checkout : bool | Path | str
+            If True, the new loop will use the existing folder and clear logs for sessions after the one corresponding to the given path.
+            If False, the new loop will use the existing folder but keep the logs for sessions after the one corresponding to the given path.
+            If a path (or a str like Path) is provided, the new loop will be saved to that path, leaving the original path unchanged.
+        replace_timer : bool
+            If a session is loaded, determines whether to replace the timer with session.timer.
+            Default is True, which means the session timer will be replaced with the current timer.
+            If False, the session timer will not be replaced.
+        Returns
+        -------
+        LoopBase
+            An instance of LoopBase with the loaded session.
+        """
+        path = Path(path)
+        with path.open("rb") as f:
+            session = cast(LoopBase, pickle.load(f))
+
+        # set session folder
+        if checkout:
+            if checkout is True:
+                logger.set_storages_path(session.session_folder.parent)
+                max_loop = max(session.loop_trace.keys())
+
+                # truncate log storages after the max loop
+                session.truncate_session_folder(max_loop, len(session.loop_trace[max_loop]) - 1)
+                logger.truncate_storages(session.loop_trace[max_loop][-1].end)
+            else:
+                checkout = Path(checkout)
+                checkout.mkdir(parents=True, exist_ok=True)
+                session.session_folder = checkout / "__session__"
+                logger.set_storages_path(checkout)
+
+        if session.timer.started:
+            if replace_timer:
+                RD_Agent_TIMER_wrapper.replace_timer(session.timer)
+                RD_Agent_TIMER_wrapper.timer.restart_by_remain_time()
+            else:
+                # Use the default timer to replace the session timer
+                session.timer = RD_Agent_TIMER_wrapper.timer
+
+        return session
+
+    def __getstate__(self) -> dict[str, Any]:
+        res = {}
+        for k, v in self.__dict__.items():
+            if k not in ["queue", "semaphores", "_pbar"]:
+                res[k] = v
+        return res
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        self.queue = asyncio.Queue()
+        self.semaphores = {}

rdagent/utils/workflow/misc.py

@@ -0,0 +1,54 @@
+import time
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+ASpecificRet = TypeVar("ASpecificRet")
+
+
+def wait_retry(
+    retry_n: int = 3, sleep_time: int = 1, transform_args_fn: Callable[[tuple, dict], tuple[tuple, dict]] | None = None
+) -> Callable[[Callable[..., ASpecificRet]], Callable[..., ASpecificRet]]:
+    """Decorator to wait and retry the function for retry_n times.
+
+    Example:
+    >>> import time
+    >>> @wait_retry(retry_n=2, sleep_time=1)
+    ... def test_func():
+    ...     global counter
+    ...     counter += 1
+    ...     if counter < 3:
+    ...         raise ValueError("Counter is less than 3")
+    ...     return counter
+    >>> counter = 0
+    >>> try:
+    ...     test_func()
+    ... except ValueError as e:
+    ...     print(f"Caught an exception: {e}")
+    Error: Counter is less than 3
+    Error: Counter is less than 3
+    Caught an exception: Counter is less than 3
+    >>> counter
+    2
+    """
+    assert retry_n > 0, "retry_n should be greater than 0"
+
+    def decorator(f: Callable[..., ASpecificRet]) -> Callable[..., ASpecificRet]:
+        def wrapper(*args: Any, **kwargs: Any) -> ASpecificRet:
+            for i in range(retry_n + 1):
+                try:
+                    return f(*args, **kwargs)
+                except Exception as e:
+                    print(f"Error: {e}")
+                    time.sleep(sleep_time)
+                    if i == retry_n:
+                        raise
+                    # Update args and kwargs using the transform function if provided.
+                    if transform_args_fn is not None:
+                        args, kwargs = transform_args_fn(args, kwargs)
+            else:
+                # just for passing mypy CI.
+                return f(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

rdagent/utils/workflow/tracking.py

@@ -0,0 +1,93 @@
+"""
+Tracking module for experiment tracking using MLflow.
+
+This module provides a clean interface for tracking metrics and parameters
+while keeping the MLflow dependency optional based on configuration.
+"""
+
+import datetime
+from typing import TYPE_CHECKING
+
+import pytz
+
+from rdagent.core.conf import RD_AGENT_SETTINGS
+from rdagent.log.timer import RD_Agent_TIMER_wrapper
+
+if TYPE_CHECKING:
+    # Import here to avoid circular dependency
+    from rdagent.utils.workflow.loop import LoopBase
+
+from rdagent.log import rdagent_logger as logger
+
+# Define a placeholder for mlflow if it's not available
+mlflow = None
+
+# Conditional import to make MLflow optional
+if RD_AGENT_SETTINGS.enable_mlflow:
+    try:
+        import mlflow  # type: ignore[assignment]
+    except ImportError:
+        logger.warning("MLflow is enabled in settings but could not be imported.")
+        RD_AGENT_SETTINGS.enable_mlflow = False
+
+
+class WorkflowTracker:
+    """
+    A workflow-specific tracking system that logs metrics related to workflow execution.
+
+    This class handles metric logging while keeping the MLflow dependency optional.
+    If MLflow is not enabled in settings, tracking calls become no-ops.
+    """
+
+    def __init__(self, loop_base: "LoopBase"):
+        """
+        Initialize a WorkflowTracker with a LoopBase instance.
+
+        Args:
+            loop_base: The LoopBase instance to track metrics for
+        """
+        self.loop_base = loop_base
+
+    @staticmethod
+    def is_enabled() -> bool:
+        """Check if tracking is enabled."""
+        return RD_AGENT_SETTINGS.enable_mlflow
+
+    @staticmethod
+    def _datetime_to_float(dt: datetime.datetime) -> float:
+        """Convert datetime to a structured float representation."""
+        return dt.second + dt.minute * 1e2 + dt.hour * 1e4 + dt.day * 1e6 + dt.month * 1e8 + dt.year * 1e10
+
+    def log_workflow_state(self) -> None:
+        """
+        Log all workflow state metrics from the associated LoopBase instance.
+        """
+        if not RD_AGENT_SETTINGS.enable_mlflow or mlflow is None:
+            return
+
+        # Log workflow progress
+        mlflow.log_metric("loop_index", self.loop_base.loop_idx)
+        mlflow.log_metric("step_index", self.loop_base.step_idx[self.loop_base.loop_idx])
+
+        current_local_datetime = datetime.datetime.now(pytz.timezone("Asia/Shanghai"))
+        float_like_datetime = self._datetime_to_float(current_local_datetime)
+        mlflow.log_metric("current_datetime", float_like_datetime)
+
+        # Log API status
+        mlflow.log_metric("api_fail_count", RD_Agent_TIMER_wrapper.api_fail_count)
+        latest_api_fail_time = RD_Agent_TIMER_wrapper.latest_api_fail_time
+        if latest_api_fail_time is not None:
+            float_like_datetime = self._datetime_to_float(latest_api_fail_time)
+            mlflow.log_metric("lastest_api_fail_time", float_like_datetime)
+
+        # Log timer status if timer is started
+        if self.loop_base.timer.started:
+            remain_time = self.loop_base.timer.remain_time()
+            assert remain_time is not None
+            mlflow.log_metric("remain_time", remain_time.seconds)
+            mlflow.log_metric(
+                "remain_percent",
+                remain_time / self.loop_base.timer.all_duration * 100,
+            )
+
+    # Keep only the log_workflow_state method as it's the primary entry point now

test/utils/test_env.py

@@ -52,12 +52,16 @@ class EnvUtils(unittest.TestCase):
         self.assertTrue(mlrun_p.exists(), f"Expected output file {mlrun_p} not found")
 
     def test_local_simple(self):
-        local_conf = LocalConf(bin_path="/home/xiaoyang/miniconda3/bin/", default_entry="which python")
-        le = LocalEnv(conf=local_conf)
-        print(local_conf)
-        le.prepare()
         code_path = DIRNAME / "tmp_code"
         code_path.mkdir(exist_ok=True)
+        # Get user home dynamically
+        home_bin = str(Path.home() / "miniconda3/bin/")
+        local_conf = LocalConf(bin_path=home_bin, default_entry="which python")
+
+        local_conf.extra_volumes = {str(code_path): "./code"}
+        print(local_conf)
+        le = LocalEnv(conf=local_conf)
+        le.prepare()
         res, code = le.run_ret_code(local_path=str(code_path))
         print(res, code)