From 982b43e4fd46b75c46e9b230bc8abde33f3d06b9 Mon Sep 17 00:00:00 2001 From: Roland Minrui <114476598+RolandMinrui@users.noreply.github.com> Date: Fri, 21 Feb 2025 16:46:58 +0800 Subject: [PATCH] feat: add output_path to load function of LoopBase (#628) * feat: add output_path to load from checkpoint function * add default value to output_path * sort import * sort imports --------- Co-authored-by: Xu --- rdagent/app/data_science/loop.py | 6 ++++-- rdagent/utils/workflow.py | 12 ++++++++++-- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/rdagent/app/data_science/loop.py b/rdagent/app/data_science/loop.py index 30393960..6cef32f5 100644 --- a/rdagent/app/data_science/loop.py +++ b/rdagent/app/data_science/loop.py @@ -142,13 +142,15 @@ class DataScienceRDLoop(RDLoop): logger.log_object(self.trace.sota_experiment(), tag="SOTA experiment") -def main(path=None, step_n=None, loop_n=None, competition="bms-molecular-translation"): +def main(path=None, output_path=None, step_n=None, loop_n=None, competition="bms-molecular-translation"): """ Parameters ---------- path : path like `$LOG_PATH/__session__/1/0_propose`. It indicates that we restore the state that after finish the step 0 in loop1 + output_path : + path like `$LOG_PATH`. It indicates that where we want to save our session and log information. step_n : How many steps to run; if None, it will run forever until error or KeyboardInterrupt loop_n : @@ -179,7 +181,7 @@ def main(path=None, step_n=None, loop_n=None, competition="bms-molecular-transla if path is None: kaggle_loop = DataScienceRDLoop(DS_RD_SETTING) else: - kaggle_loop = DataScienceRDLoop.load(path) + kaggle_loop = DataScienceRDLoop.load(path, output_path) kaggle_loop.run(step_n=step_n, loop_n=loop_n) diff --git a/rdagent/utils/workflow.py b/rdagent/utils/workflow.py index eb0f8be6..d187b966 100644 --- a/rdagent/utils/workflow.py +++ b/rdagent/utils/workflow.py @@ -14,7 +14,7 @@ import time from collections import defaultdict from dataclasses import dataclass from pathlib import Path -from typing import Any, Callable, TypeVar, cast +from typing import Any, Callable, Optional, TypeVar, Union, cast from tqdm.auto import tqdm @@ -161,10 +161,18 @@ class LoopBase: pickle.dump(self, f) @classmethod - def load(cls, path: str | Path) -> "LoopBase": + def load(cls, path: Union[str, Path], output_path: Optional[Union[str, Path]] = None) -> "LoopBase": path = Path(path) with path.open("rb") as f: session = cast(LoopBase, pickle.load(f)) + + # set session folder + if output_path: + output_path = Path(output_path) + output_path.mkdir(parents=True, exist_ok=True) + session.session_folder = output_path / "__session__" + + # set trace path logger.set_trace_path(session.session_folder.parent) max_loop = max(session.loop_trace.keys())