diff --git a/rdagent/app/data_mining/model.py b/rdagent/app/data_mining/model.py index d0ea970c..ba83b553 100644 --- a/rdagent/app/data_mining/model.py +++ b/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__": diff --git a/rdagent/app/data_science/loop.py b/rdagent/app/data_science/loop.py index 18addef2..8df7853c 100644 --- a/rdagent/app/data_science/loop.py +++ b/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__": diff --git a/rdagent/app/qlib_rd_loop/factor.py b/rdagent/app/qlib_rd_loop/factor.py index 5c08a091..f13b3d8a 100755 --- a/rdagent/app/qlib_rd_loop/factor.py +++ b/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__": diff --git a/rdagent/app/qlib_rd_loop/factor_from_report.py b/rdagent/app/qlib_rd_loop/factor_from_report.py index 2cff1ca7..24213b49 100644 --- a/rdagent/app/qlib_rd_loop/factor_from_report.py +++ b/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__": diff --git a/rdagent/app/qlib_rd_loop/model.py b/rdagent/app/qlib_rd_loop/model.py index 40183360..d66a5a1e 100644 --- a/rdagent/app/qlib_rd_loop/model.py +++ b/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__": diff --git a/rdagent/app/qlib_rd_loop/quant.py b/rdagent/app/qlib_rd_loop/quant.py index c7778ed4..6573a537 100644 --- a/rdagent/app/qlib_rd_loop/quant.py +++ b/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__": diff --git a/rdagent/components/coder/data_science/feature/prompts.yaml b/rdagent/components/coder/data_science/feature/prompts.yaml index 047d9c8a..56a384a0 100644 --- a/rdagent/components/coder/data_science/feature/prompts.yaml +++ b/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" %} diff --git a/rdagent/components/coder/data_science/model/prompts.yaml b/rdagent/components/coder/data_science/model/prompts.yaml index c0af430e..9877cfca 100644 --- a/rdagent/components/coder/data_science/model/prompts.yaml +++ b/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" %} diff --git a/rdagent/components/coder/data_science/raw_data_loader/prompts.yaml b/rdagent/components/coder/data_science/raw_data_loader/prompts.yaml index c62f3055..ce494fe1 100644 --- a/rdagent/components/coder/data_science/raw_data_loader/prompts.yaml +++ b/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" %} diff --git a/rdagent/core/conf.py b/rdagent/core/conf.py index 417b0d24..74d3a19d 100644 --- a/rdagent/core/conf.py +++ b/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() diff --git a/rdagent/core/proposal.py b/rdagent/core/proposal.py index 84a25a91..137fc20c 100644 --- a/rdagent/core/proposal.py +++ b/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): diff --git a/rdagent/log/conf.py b/rdagent/log/conf.py index 72d0edf7..f8f81686 100644 --- a/rdagent/log/conf.py +++ b/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]] = {} diff --git a/rdagent/log/logger.py b/rdagent/log/logger.py index a8e9deca..a60bacb3 100644 --- a/rdagent/log/logger.py +++ b/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("") - # 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(".") diff --git a/rdagent/oai/backend/litellm.py b/rdagent/oai/backend/litellm.py index 0a65bed0..f2667c37 100644 --- a/rdagent/oai/backend/litellm.py +++ b/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: diff --git a/rdagent/scenarios/data_science/loop.py b/rdagent/scenarios/data_science/loop.py index d007be9c..f34df85f 100644 --- a/rdagent/scenarios/data_science/loop.py +++ b/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" diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/prompts_v2.yaml b/rdagent/scenarios/data_science/proposal/exp_gen/prompts_v2.yaml index 2a8aebb8..062be654 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/prompts_v2.yaml +++ b/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." } diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/prompts_v3.yaml b/rdagent/scenarios/data_science/proposal/exp_gen/prompts_v3.yaml index d02e039e..8f88ff8e 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/prompts_v3.yaml +++ b/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//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" %}/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 diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py b/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py index e6d6c553..41c7237f 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py +++ b/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( diff --git a/rdagent/scenarios/data_science/share.yaml b/rdagent/scenarios/data_science/share.yaml index 437951ee..8af03e51 100644 --- a/rdagent/scenarios/data_science/share.yaml +++ b/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: |- diff --git a/rdagent/utils/env.py b/rdagent/utils/env.py index dd4a8e1e..0a2c3fca 100644 --- a/rdagent/utils/env.py +++ b/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: diff --git a/rdagent/utils/workflow.py b/rdagent/utils/workflow.py deleted file mode 100644 index c246dbb1..00000000 --- a/rdagent/utils/workflow.py +++ /dev/null @@ -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 diff --git a/rdagent/utils/workflow/__init__.py b/rdagent/utils/workflow/__init__.py new file mode 100644 index 00000000..41c2ae3c --- /dev/null +++ b/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"] diff --git a/rdagent/utils/workflow/loop.py b/rdagent/utils/workflow/loop.py new file mode 100644 index 00000000..6859e70a --- /dev/null +++ b/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 = {} diff --git a/rdagent/utils/workflow/misc.py b/rdagent/utils/workflow/misc.py new file mode 100644 index 00000000..24b450ae --- /dev/null +++ b/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 diff --git a/rdagent/utils/workflow/tracking.py b/rdagent/utils/workflow/tracking.py new file mode 100644 index 00000000..1fcade86 --- /dev/null +++ b/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 diff --git a/test/utils/test_env.py b/test/utils/test_env.py index cca2dd01..92321ec3 100644 --- a/test/utils/test_env.py +++ b/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)