gavindiaz/NexQuant
1
0
Fork 0
mirror of https://github.com/NicolasBohn/NexQuant.git synced 2026-08-01 17:37:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
* 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>
This commit is contained in:
you-n-g
2025-06-12 11:44:14 +08:00
committed by GitHub
parent 235fcd308a
commit 09be71d586
26 changed files with 926 additions and 506 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rdagent/utils/env.py
+135 -81
View File
@@ -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
-370
View File
@@ -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
+5
View File
@@ -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
+444
View File
@@ -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
+54
View File
@@ -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
+93
View File
@@ -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