From 4f484e009ad685e681f8c5442cc2410f63f00531 Mon Sep 17 00:00:00 2001 From: you-n-g Date: Tue, 23 Jul 2024 16:37:41 +0800 Subject: [PATCH] Workflow Support Loading and saving sessions (#98) * Successfully logging the trace * Start debugging & Add policy file * Support loading sessions * Add docs * Add tqdm --- .gitignore | 2 +- docs/policy.rst | 24 +++++++ rdagent/app/qlib_rd_loop/model_w_sc.py | 87 +++++++++++++++++++++++ rdagent/log/logger.py | 6 ++ rdagent/log/storage.py | 45 ++++++++++-- rdagent/utils/workflow.py | 97 ++++++++++++++++++++++++++ 6 files changed, 254 insertions(+), 7 deletions(-) create mode 100644 docs/policy.rst create mode 100644 rdagent/app/qlib_rd_loop/model_w_sc.py create mode 100644 rdagent/utils/workflow.py diff --git a/.gitignore b/.gitignore index 2e5dd32d..296b3105 100644 --- a/.gitignore +++ b/.gitignore @@ -64,7 +64,7 @@ coverage.xml # Django stuff: *.log -^log +log/ local_settings.py db.sqlite3 db.sqlite3-journal diff --git a/docs/policy.rst b/docs/policy.rst new file mode 100644 index 00000000..da22d5f7 --- /dev/null +++ b/docs/policy.rst @@ -0,0 +1,24 @@ +====== +Policy +====== + +This project welcomes contributions and suggestions. Most contributions require you to agree to a +Contributor License Agreement (CLA) declaring that you have the right to, and actually do, grant us +the rights to use your contribution. For details, visit https://cla.opensource.microsoft.com. + +When you submit a pull request, a CLA bot will automatically determine whether you need to provide +a CLA and decorate the PR appropriately (e.g., status check, comment). Simply follow the instructions +provided by the bot. You will only need to do this once across all repos using our CLA. + +This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/). +For more information see the [Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/) or +contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with any additional questions or comments. + +Trademarks +========== + +This project may contain trademarks or logos for projects, products, or services. Authorized use of Microsoft +trademarks or logos is subject to and must follow +[Microsoft's Trademark & Brand Guidelines](https://www.microsoft.com/en-us/legal/intellectualproperty/trademarks/usage/general). +Use of Microsoft trademarks or logos in modified versions of this project must not cause confusion or imply Microsoft sponsorship. +Any use of third-party trademarks or logos are subject to those third-party's policies. diff --git a/rdagent/app/qlib_rd_loop/model_w_sc.py b/rdagent/app/qlib_rd_loop/model_w_sc.py new file mode 100644 index 00000000..aa6c7b31 --- /dev/null +++ b/rdagent/app/qlib_rd_loop/model_w_sc.py @@ -0,0 +1,87 @@ +""" +Model workflow with session control +It is from `rdagent/app/qlib_rd_loop/model.py` and try to replace `rdagent/app/qlib_rd_loop/RDAgent.py` +""" + +import fire +from typing import Any +from rdagent.app.qlib_rd_loop.conf import PROP_SETTING +from rdagent.core.developer import Developer +from rdagent.core.exception import ModelEmptyError +from rdagent.core.proposal import ( + Hypothesis2Experiment, + HypothesisExperiment2Feedback, + HypothesisGen, + Trace, +) +from rdagent.core.scenario import Scenario +from rdagent.core.utils import import_class +from rdagent.log import rdagent_logger as logger + +from rdagent.utils.workflow import LoopMeta, LoopBase + +class ModelLoop(LoopBase, metaclass=LoopMeta): + # TODO: supporting customized loop control like catching `ModelEmptyError` + + def __init__(self): + scen: Scenario = import_class(PROP_SETTING.model_scen)() + + self.hypothesis_gen: HypothesisGen = import_class(PROP_SETTING.model_hypothesis_gen)(scen) + + self.hypothesis2experiment: Hypothesis2Experiment = import_class(PROP_SETTING.model_hypothesis2experiment)() + + self.qlib_model_coder: Developer = import_class(PROP_SETTING.model_coder)(scen) + self.qlib_model_runner: Developer = import_class(PROP_SETTING.model_runner)(scen) + + self.qlib_model_summarizer: HypothesisExperiment2Feedback = import_class(PROP_SETTING.model_summarizer)(scen) + self.trace = Trace(scen=scen) + super().__init__() + + def propose(self, prev_out: dict[str, Any]): + with logger.tag("r"): # research + hypothesis = self.hypothesis_gen.gen(self.trace) + logger.log_object(hypothesis, tag="hypothesis generation") + return hypothesis + + def exp_gen(self, prev_out: dict[str, Any]): + with logger.tag("r"): # research + exp = self.hypothesis2experiment.convert(prev_out["propose"], self.trace) + logger.log_object(exp.sub_tasks, tag="experiment generation") + return exp + + def coding(self, prev_out: dict[str, Any]): + with logger.tag("d"): # develop + exp = self.qlib_model_coder.develop(prev_out["exp_gen"]) + logger.log_object(exp.sub_workspace_list, tag="model coder result") + return exp + + def running(self, prev_out: dict[str, Any]): + with logger.tag("ef"): # evaluate and feedback + exp = self.qlib_model_runner.develop(prev_out["coding"]) + logger.log_object(exp, tag="model runner result") + return exp + + def feedback(self, prev_out: dict[str, Any]): + feedback = self.qlib_model_summarizer.generate_feedback(prev_out["running"], prev_out["propose"], self.trace) + logger.log_object(feedback, tag="feedback") + self.trace.hist.append((prev_out["propose"],prev_out["running"] , feedback)) + + +def main(path=None): + """ + You can continue running session by + + .. code-block:: python + + dotenv run -- python rdagent/app/qlib_rd_loop/model_w_sc.py $LOG_PATH/__session__/1/0_propose + + """ + if path is None: + model_loop = ModelLoop() + else: + model_loop = ModelLoop.load(path) + model_loop.run() + + +if __name__ == "__main__": + fire.Fire(main) diff --git a/rdagent/log/logger.py b/rdagent/log/logger.py index 2c7c81ea..e69ed8b4 100644 --- a/rdagent/log/logger.py +++ b/rdagent/log/logger.py @@ -61,6 +61,10 @@ class RDAgentLog(SingletonBaseClass): self.main_pid = os.getpid() + def set_trace_path(self, log_trace_path): + self.log_trace_path = Path(log_trace_path) + self.storage = FileStorage(log_trace_path) + @contextmanager def tag(self, tag: str): if tag.strip() == "": @@ -91,12 +95,14 @@ class RDAgentLog(SingletonBaseClass): return pid_chain def file_format(self, record, raw: bool = False): + # FIXME: the formmat is tightly coupled with the message reading in storage. record["message"] = LogColors.remove_ansi_codes(record["message"]) if raw: return "{message}" return "{time:YYYY-MM-DD HH:mm:ss.SSS} | {level: <8} | {name}:{function}:{line} - {message}\n" 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/log/storage.py b/rdagent/log/storage.py index d07baa58..208ed9ec 100644 --- a/rdagent/log/storage.py +++ b/rdagent/log/storage.py @@ -56,12 +56,13 @@ class FileStorage(Storage): f.write(obj) return path + log_pattern = re.compile( + r"(?P\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{3}) \| " + r"(?PDEBUG|INFO|WARNING|ERROR|CRITICAL) *\| " + r"(?P.+:.+:\d+) - " + ) + def iter_msg(self, watch: bool = False) -> Generator[Message, None, None]: - log_pattern = re.compile( - r"(?P\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{3}) \| " - r"(?PDEBUG|INFO|WARNING|ERROR|CRITICAL) *\| " - r"(?P.+:.+:\d+) - " - ) msg_l = [] for file in self.path.glob("**/*.log"): tag = '.'.join(str(file.relative_to(self.path)).replace("/", ".").split(".")[:-3]) @@ -70,7 +71,7 @@ class FileStorage(Storage): with file.open("r") as f: content = f.read() - matches, next_matches = log_pattern.finditer(content), log_pattern.finditer(content) + matches, next_matches = self.log_pattern.finditer(content), self.log_pattern.finditer(content) next_match = next(next_matches, None) # NOTE: the content will be the text between `match` and `next_match` for match in matches: @@ -107,3 +108,35 @@ class FileStorage(Storage): msg_l.sort(key=lambda x: x.timestamp) for m in msg_l: yield m + + def truncate(self, time: datetime) -> None: + # any message later than `time` will be removed + for file in self.path.glob("**/*.log"): + + with file.open("r") as f: + content = f.read() + + new_content = "" + + matches, next_matches = self.log_pattern.finditer(content), self.log_pattern.finditer(content) + + next_match = next(next_matches, None) + for match in matches: + next_match = next(next_matches, None) + timestamp_str = match.group("timestamp") + timestamp = datetime.strptime(timestamp_str, "%Y-%m-%d %H:%M:%S.%f").replace(tzinfo=timezone.utc) + + log_start = match.start() + log_end = next_match.start() if next_match else len(content) + msg = content[match.end():log_end].strip() + + if timestamp > time: + if "Logging object in" in msg: + absolute_p = msg.split("Logging object in ")[1] + p = Path(absolute_p) + p.unlink() + continue + + new_content += content[log_start:log_end] + with file.open("w") as f: + f.write(new_content) diff --git a/rdagent/utils/workflow.py b/rdagent/utils/workflow.py new file mode 100644 index 00000000..ee37d989 --- /dev/null +++ b/rdagent/utils/workflow.py @@ -0,0 +1,97 @@ +""" +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) + +""" +from pathlib import Path +import pickle +from tqdm.auto import tqdm + + +from collections import defaultdict +from dataclasses import dataclass +import datetime +from typing import Callable +from rdagent.log import rdagent_logger as logger + + +class LoopMeta(type): + + def __new__(cls, clsname, bases, attrs): + + # move custommized steps into steps + steps = [] + for name in attrs.keys(): + if not name.startswith("__"): + steps.append(name) + attrs["steps"] = steps + + return super().__new__(cls, clsname, bases, attrs) + + +@dataclass +class LoopTrace: + start: datetime.datetime # the start time of the trace + end: datetime.datetime # the end time of the trace + # TODO: more information about the trace + + +class LoopBase: + steps: list[Callable] # a list of steps to work on + loop_trace: dict[int, list[LoopTrace]] + + def __init__(self): + self.loop_idx = 0 # current loop index + self.step_idx = 0 # the index of next step to be run + self.loop_prev_out = {} # 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__" + + def run(self): + with tqdm(total=len(self.steps), desc="Workflow Progress", unit="step") as pbar: + while True: + li, si = self.loop_idx, self.step_idx + + start = datetime.datetime.now(datetime.timezone.utc) + + name = self.steps[si] + func = getattr(self, name) + self.loop_prev_out[name] = func(self.loop_prev_out) + + end = datetime.datetime.now(datetime.timezone.utc) + + self.loop_trace[li].append(LoopTrace(start, end)) + + # Update tqdm progress bar + pbar.set_postfix(loop_index=li, step_index=si, step_name=name) + pbar.update(1) + + # 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 + 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): + 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: str | Path): + path = Path(path) + with path.open("rb") as f: + session = pickle.load(f) + logger.set_trace_path(session.session_folder.parent) + + max_loop = max(session.loop_trace.keys()) + logger.storage.truncate(time=session.loop_trace[max_loop][-1].end) + return session