more clearly structure

examples/taskmanager/update_online_pred.py

@@ -1,8 +1,9 @@
 import qlib
 from qlib.model.trainer import task_train
-from qlib.workflow.task.update import ModelUpdater
+from qlib.workflow.task.online import RollingOnlineManager
 from qlib.config import REG_CN
 import fire
+from qlib.workflow import R
 
 data_handler_config = {
     "start_time": "2008-01-01",
@@ -50,33 +51,33 @@ task = {
     },
 }
 
-provider_uri = "~/.qlib/qlib_data/cn_data"  # target_dir
-
 
 def first_train(experiment_name="online_svr"):
 
-    qlib.init(provider_uri=provider_uri, region=REG_CN)
+    rom = RollingOnlineManager(experiment_name)
-    model_updater = ModelUpdater(experiment_name)
 
     rid = task_train(task_config=task, experiment_name=experiment_name)
-    model_updater.reset_online_model(rid)
+    
+    rom.reset_online_model(rid)
 
 
 def update_online_pred(experiment_name="online_svr"):
 
-    qlib.init(provider_uri=provider_uri, region=REG_CN)
+    rom = RollingOnlineManager(experiment_name)
-    model_updater = ModelUpdater(experiment_name)
 
     print("Here are the online models waiting for update:")
-    for rid, rec in model_updater.list_online_model().items():
+    for rid, rec in rom.list_online_model().items():
         print(rid)
 
-    model_updater.update_online_pred()
+    rom.update_online_pred()
 
 
 if __name__ == "__main__":
-    fire.Fire()
+    ## to train a model and set it to online model, use the command below
-    # to train a model and set it to online model, use the command below
     # python update_online_pred.py first_train
-    # to update online predictions once a day, use the command below
+    ## to update online predictions once a day, use the command below
     # python update_online_pred.py update_online_pred
+
+    provider_uri = "~/.qlib/qlib_data/cn_data"  # target_dir
+    qlib.init(provider_uri=provider_uri, region=REG_CN)
+    fire.Fire()

qlib/model/trainer.py

@@ -26,9 +26,11 @@ def task_train(task_config: dict, experiment_name: str) -> str:
     # model initiaiton
     model = init_instance_by_config(task_config["model"])
     dataset = init_instance_by_config(task_config["dataset"])
+    datahandler = dataset.handler
     
     # start exp
     with R.start(experiment_name=experiment_name):
+
         # train model
         R.log_params(**flatten_dict(task_config))
         model.fit(dataset)
@@ -36,6 +38,10 @@ def task_train(task_config: dict, experiment_name: str) -> str:
         R.save_objects(**{"params.pkl": model})
         R.save_objects(**{"task": task_config})  # keep the original format and datatype
 
+        artifact_uri = recorder.get_artifact_uri()[7:]  # delete "file://"
+        dataset.to_pickle(artifact_uri + "/dataset", exclude=["handler"])
+        datahandler.to_pickle(artifact_uri + "/datahandler")
+
         # generate records: prediction, backtest, and analysis
         records = task_config.get("record", [])
         if isinstance(records, dict):  # prevent only one dict
@@ -53,4 +59,5 @@ def task_train(task_config: dict, experiment_name: str) -> str:
                 record["kwargs"].update(rconf)
                 ar = init_instance_by_config(record)
                 ar.generate()
+
     return recorder.info["id"]

qlib/workflow/task/collect.py

@@ -16,48 +16,38 @@ class TaskCollector:
         self.exp = R.get_exp(experiment_name=experiment_name)
         self.logger = get_module_logger("TaskCollector")
 
-    def list_recorders(self, rec_filter_func=None, task_filter_func=None, only_finished=True, only_have_task=False):
+    def list_recorders(self, rec_filter_func=None):
-        """
+        """"""
-        Return a dict of {rid: Recorder} by recorder filter and task filter. It is not necessary to use those filter.
+        recs = self.exp.list_recorders()
-        If you don't train with "task_train", then there is no "task"(a file in mlruns/artifacts) which includes the task config.
+        recs_flt = {}
-        If there is a "task", then it will become rec.task which can be get simply.
+        for rid, rec in recs.items():
+            if rec_filter_func is None or rec_filter_func(rec):
+                recs_flt[rid] = rec
+
+        return recs_flt
+
+    def get_recorder_by_id(self, recorder_id):
+        return self.exp.get_recorder(recorder_id, create=False)
+
+    def list_recorders_by_task(self, task_filter_func):
+        """[summary]
 
         Parameters
         ----------
-        rec_filter_func : Callable[[Recorder], bool], optional
+        task_filter_func : [type], optional
-            judge whether you need this recorder, by default None
+            [description], by default None
-        task_filter_func : Callable[[dict], bool], optional
-            judge whether you need this task, by default None
-        only_finished : bool, optional
-            whether always use finished recorder, by default True
-        only_have_task : bool, optional
-            whether it is necessary to get the task config
-
-        Returns
-        -------
-        dict
-            a dict of {rid: Recorder}
-
         """
-        recs = self.exp.list_recorders()
-        recs_flt = {}
-        if task_filter_func is not None:
-            only_have_task = True
-        for rid, rec in recs.items():
-            if (only_finished and rec.status == rec.STATUS_FI) or only_finished == False:
-                if rec_filter_func is None or rec_filter_func(rec):
-                    task = None
-                    try:
-                        task = rec.load_object("task")
-                    except OSError:
-                        pass
-                    if task is None and only_have_task:
-                        continue
-                    if task_filter_func is None or task_filter_func(task):
-                        rec.task = task
-                        recs_flt[rid] = rec
 
-        return recs_flt
+        def rec_filter_func(recorder):
+            try:
+                task = recorder.load_object("task")
+            except OSError:
+                raise OSError(
+                    f"Can't find task in {recorder.info['id']}, have you trained with model.trainer.task_train?"
+                )
+            return task_filter_func(task)
+
+        return self.list_recorders(rec_filter_func)
 
     def collect_predictions(
         self,

qlib/workflow/task/online.py

@@ -0,0 +1,124 @@
+from typing import Union, List
+from qlib import get_module_logger
+from qlib.workflow import R
+from qlib.model.trainer import task_train
+from qlib.workflow.recorder import Recorder
+from qlib.workflow.task.collect import TaskCollector
+from qlib.workflow.task.update import ModelUpdater
+
+
+class OnlineManagement:
+    def __init__(self, experiment_name):
+        pass
+
+    def update_online_pred(self, recorder: Union[str, Recorder]):
+        """update the predictions of online models
+
+        Parameters
+        ----------
+        recorder : Union[str, Recorder]
+            the id or the instance of Recorder
+
+        """
+        raise NotImplementedError(f"Please implement the `update_pred` method.")
+
+    def prepare_new_models(self, tasks: List[dict]):
+        """prepare(train) new models
+
+        Parameters
+        ----------
+        tasks : List[dict]
+            a list of tasks
+
+        """
+        raise NotImplementedError(f"Please implement the `prepare_new_models` method.")
+
+    def reset_online_model(self, recorders: List[Union[str, Recorder]]):
+        """reset online model
+
+        Parameters
+        ----------
+        recorders : List[Union[str, Recorder]]
+            a list of the recorder id or the instance
+
+        """
+        raise NotImplementedError(f"Please implement the `reset_online_model` method.")
+
+
+class RollingOnlineManager(OnlineManagement):
+
+    ONLINE_TAG = "online_model"
+    ONLINE_TAG_TRUE = "True"
+    ONLINE_TAG_FALSE = "False"
+
+    def __init__(self, experiment_name: str) -> None:
+        """ModelUpdater needs experiment name to find the records
+
+        Parameters
+        ----------
+        experiment_name : str
+            experiment name string
+        """
+        super(RollingOnlineManager, self).__init__(experiment_name)
+        self.logger = get_module_logger("RollingOnlineManager")
+        self.exp_name = experiment_name
+        self.tc = TaskCollector(experiment_name)
+
+    def set_online_model(self, recorder: Union[str, Recorder]):
+        """online model will be identified at the tags of the record
+
+        Parameters
+        ----------
+        recorder: Union[str,Recorder]
+            the id of a Recorder or the Recorder instance
+        """
+        if isinstance(recorder, str):
+            recorder = self.tc.get_recorder_by_id(recorder_id=recorder)
+        recorder.set_tags(**{self.ONLINE_TAG: self.ONLINE_TAG_TRUE})
+
+    def cancel_online_model(self, recorder: Union[str, Recorder]):
+        if isinstance(recorder, str):
+            recorder = self.tc.get_recorder_by_id(recorder_id=recorder)
+        recorder.set_tags(**{self.ONLINE_TAG: self.ONLINE_TAG_FALSE})
+
+    def cancel_all_online_model(self):
+        recs = self.tc.list_recorders()
+        for rid, rec in recs.items():
+            self.cancel_online_model(rec)
+
+    def reset_online_model(self, recorders: Union[str, List[Union[str, Recorder]]]):
+        """cancel all online model and reset the given model to online model
+
+        Parameters
+        ----------
+        recorders: List[Union[str,Recorder]]
+            the list of the id of a Recorder or the Recorder instance
+        """
+        self.cancel_all_online_model()
+        if isinstance(recorders, str):
+            recorders = [recorders]
+        for rec_or_rid in recorders:
+            self.set_online_model(rec_or_rid)
+
+    def online_filter(self, recorder):
+        tags = recorder.list_tags()
+        if tags.get(self.ONLINE_TAG, self.ONLINE_TAG_FALSE) == self.ONLINE_TAG_TRUE:
+            return True
+        return False
+
+    def list_online_model(self):
+        """list the record of online model
+
+        Returns
+        -------
+        dict
+            {rid : recorder of the online model}
+        """
+
+        return self.tc.list_recorders(rec_filter_func=self.online_filter)
+
+    def update_online_pred(self):
+        """update all online model predictions to the latest day in Calendar."""
+        mu = ModelUpdater(self.exp_name)
+        cnt = mu.update_all_pred(self.online_filter)
+        self.logger.info(f"Finish updating {cnt} online model predictions of {self.exp_name}.")

qlib/workflow/task/update.py

@@ -1,9 +1,7 @@
 from typing import Union, List
 from qlib.workflow import R
-from tqdm.auto import tqdm
 from qlib.data import D
 import pandas as pd
-from qlib.utils import init_instance_by_config
 from qlib import get_module_logger
 from qlib.workflow import R
 from qlib.model.trainer import task_train
@@ -11,15 +9,11 @@ from qlib.workflow.recorder import Recorder
 from qlib.workflow.task.collect import TaskCollector
 
 
-class ModelUpdater(TaskCollector):
+class ModelUpdater:
     """
-    The model updater to re-train model or update predictions
+    The model updater to update model results in new data.
     """
 
-    ONLINE_TAG = "online_model"
-    ONLINE_TAG_TRUE = "True"
-    ONLINE_TAG_FALSE = "False"
-
     def __init__(self, experiment_name: str) -> None:
         """ModelUpdater needs experiment name to find the records
 
@@ -29,42 +23,35 @@ class ModelUpdater(TaskCollector):
             experiment name string
         """
         self.exp_name = experiment_name
-        self.exp = R.get_exp(experiment_name=experiment_name)
         self.logger = get_module_logger("ModelUpdater")
+        self.tc = TaskCollector(experiment_name)
 
-    def set_online_model(self, recorder: Union[str, Recorder]):
+    def _reload_dataset(self, recorder, start_time, end_time):
-        """online model will be identified at the tags of the record
+        """reload dataset from pickle file
 
         Parameters
         ----------
-        recorder: Union[str,Recorder]
+        recorder : Recorder
-            the id of a Recorder or the Recorder instance
+            the instance of the Recorder
+        start_time : Timestamp
+            the start time you want to load
+        end_time : Timestamp
+            the end time you want to load
+
+        Returns
+        -------
+        Dataset
+            the instance of Dataset
         """
-        if isinstance(recorder, str):
+        segments = {"test": (start_time, end_time)}
-            recorder = self.exp.get_recorder(recorder_id=recorder)
-        recorder.set_tags(**{ModelUpdater.ONLINE_TAG: ModelUpdater.ONLINE_TAG_TRUE})
 
-    def cancel_online_model(self, recorder: Union[str, Recorder]):
+        dataset = recorder.load_object("dataset")
-        if isinstance(recorder, str):
+        datahandler = recorder.load_object("datahandler")
-            recorder = self.exp.get_recorder(recorder_id=recorder)
-        recorder.set_tags(**{ModelUpdater.ONLINE_TAG: ModelUpdater.ONLINE_TAG_FALSE})
 
-    def cancel_all_online_model(self):
+        datahandler.conf_data(**{"start_time": start_time, "end_time": end_time})
-        recs = self.exp.list_recorders()
+        dataset.setup_data(handler=datahandler, segments=segments)
-        for rid, rec in recs.items():
+        datahandler.init(datahandler.IT_LS)
-            self.cancel_online_model(rec)
+        return dataset
-
-    def reset_online_model(self, recorders: List[Union[str, Recorder]]):
-        """cancel all online model and reset the given model to online model
-
-        Parameters
-        ----------
-        recorders: List[Union[str,Recorder]]
-            the list of the id of a Recorder or the Recorder instance
-        """
-        self.cancel_all_online_model()
-        for rec_or_rid in recorders:
-            self.set_online_model(rec_or_rid)
 
     def update_pred(self, recorder: Union[str, Recorder]):
         """update predictions to the latest day in Calendar based on rid
@@ -75,10 +62,9 @@ class ModelUpdater(TaskCollector):
             the id of a Recorder or the Recorder instance
         """
         if isinstance(recorder, str):
-            recorder = self.exp.get_recorder(recorder_id=recorder)
+            recorder = self.tc.get_recorder_by_id(recorder_id=recorder)
         old_pred = recorder.load_object("pred.pkl")
         last_end = old_pred.index.get_level_values("datetime").max()
-        task_config = recorder.load_object("task")  # recorder.task
 
         # updated to the latest trading day
         cal = D.calendar(start_time=last_end + pd.Timedelta(days=1), end_time=None)
@@ -90,10 +76,8 @@ class ModelUpdater(TaskCollector):
             return
 
         start_time, end_time = cal[0], cal[-1]
-        task_config["dataset"]["kwargs"]["segments"]["test"] = (start_time, end_time)
-        task_config["dataset"]["kwargs"]["handler"]["kwargs"]["end_time"] = end_time
 
-        dataset = init_instance_by_config(task_config["dataset"])
+        dataset = self._reload_dataset(recorder, start_time, end_time)
 
         model = recorder.load_object("params.pkl")
         new_pred = model.predict(dataset)
@@ -131,29 +115,7 @@ class ModelUpdater(TaskCollector):
             the count of updated record
 
         """
-        recs = self.list_recorders(rec_filter_func=rec_filter_func, only_have_task=True)
+        recs = self.tc.list_recorders(rec_filter_func=rec_filter_func)
         for rid, rec in recs.items():
             self.update_pred(rec)
         return len(recs)
-
-    def online_filter(self, recorder):
-        tags = recorder.list_tags()
-        if tags.get(ModelUpdater.ONLINE_TAG, ModelUpdater.ONLINE_TAG_FALSE) == ModelUpdater.ONLINE_TAG_TRUE:
-            return True
-        return False
-
-    def update_online_pred(self):
-        """update all online model predictions to the latest day in Calendar."""
-        cnt = self.update_all_pred(self.online_filter)
-        self.logger.info(f"Finish updating {cnt} online model predictions of {self.exp_name}.")
-
-    def list_online_model(self):
-        """list the record of online model
-
-        Returns
-        -------
-        dict
-            {rid : recorder of the online model}
-        """
-
-        return self.list_recorders(rec_filter_func=self.online_filter)