Workflow Support Loading and saving sessions (#98)

* Successfully logging the trace

* Start debugging & Add policy file

* Support loading sessions

* Add docs

* Add tqdm

.gitignore

@@ -64,7 +64,7 @@ coverage.xml
 
 # Django stuff:
 *.log
-^log
+log/
 local_settings.py
 db.sqlite3
 db.sqlite3-journal

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.

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)

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(".")
 

rdagent/log/storage.py

@@ -56,12 +56,13 @@ class FileStorage(Storage):
                 f.write(obj)
             return path
 
-    def iter_msg(self, watch: bool = False) -> Generator[Message, None, None]:
     log_pattern = re.compile(
         r"(?P<timestamp>\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{3}) \| "
         r"(?P<level>DEBUG|INFO|WARNING|ERROR|CRITICAL) *\| "
         r"(?P<caller>.+:.+:\d+) - "
     )
+
+    def iter_msg(self, watch: bool = False) -> Generator[Message, None, None]:
         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)

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