mirror of
https://github.com/NicolasBohn/NexQuant.git
synced 2026-07-28 16:07:46 +00:00
you-n-g
ec51bb94b6
* feat: runnalbe -- add exp_gen_cls param, get_leaves and merge exp gen functionalities * fix: remove unused scenario_desc and update YAML task labels * feat: override selection and update merge task description * lint * lint * lint * lint * lint * fix: log competition setting to enable mle_summary * fix name error
294 lines
12 KiB
Python
294 lines
12 KiB
Python
"""
|
|
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 os
|
|
import pickle
|
|
import time
|
|
from collections import defaultdict
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Optional, TypeVar, Union, 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.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
|
|
|
|
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 = logger.log_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)
|
|
|
|
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
|
|
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 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)
|
|
|
|
@classmethod
|
|
def load(
|
|
cls,
|
|
path: Union[str, Path],
|
|
output_path: Optional[Union[str, Path]] = None,
|
|
do_truncate: bool = False,
|
|
replace_timer: bool = True,
|
|
) -> "LoopBase":
|
|
path = Path(path)
|
|
with path.open("rb") as f:
|
|
session = cast(LoopBase, pickle.load(f))
|
|
|
|
# set session folder
|
|
# - P1: if output_path explicitly specified.
|
|
# - P2: RD_AGENT_SETTINGS.log_trace_path
|
|
output_path_value = output_path if output_path is not None else RD_AGENT_SETTINGS.log_trace_path
|
|
if output_path_value is not None:
|
|
output_path_path = Path(output_path_value)
|
|
output_path_path.mkdir(parents=True, exist_ok=True)
|
|
session.session_folder = output_path_path / "__session__"
|
|
|
|
# set trace path
|
|
logger.set_trace_path(session.session_folder.parent)
|
|
|
|
# truncate future message
|
|
if do_truncate:
|
|
max_loop = max(session.loop_trace.keys())
|
|
logger.storage.truncate(time=session.loop_trace[max_loop][-1].end)
|
|
|
|
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
|