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

Author: 2 people

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()