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/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(".")