Merge branch 'online_srv' of github.com:you-n-g/qlib into online_srv

docs/advanced/task_managment.rst

@@ -0,0 +1,67 @@
+.. _task_managment:
+
+=================================
+Task Management
+=================================
+.. currentmodule:: qlib
+
+
+Introduction
+=============
+
+The `Workflow <../component/introduction.html>`_ part introduce how to run research workflow in a loosely-coupled way. But it can only execute one ``task`` when you use ``qrun``. To automatically generate and execute different tasks, Task Management module provide a whole process including `Task Generating`_, `Task Storing`_, `Task Running`_ and `Task Collecting`_. 
+With this module, users can run their ``task`` automatically at different periods, in different losses or even by different models.
+
+An example of the entire process is shown `here <>`_.
+
+Task Generating
+===============
+A ``task`` consists of `Model`, `Dataset`, `Record` or anything added by users. 
+The specific task template can be viewed in 
+`Task Section <../component/workflow.html#task-section>`_.
+Even though the task template is fixed, Users can use ``TaskGen`` to generate different ``task`` by task template.
+
+Here is the base class of TaskGen:
+
+.. autoclass:: qlib.workflow.task.gen.TaskGen
+    :members:
+
+``Qlib`` provider a class `RollingGen<https://github.com/microsoft/qlib/tree/main/qlib/workflow/task/gen.py>`_ to generate a list of ``task`` of dataset in different date segments.
+This allows users to verify the effect of data from different periods on the model in one experiment.
+
+Task Storing
+===============
+In order to achieve higher efficiency and the possibility of cluster operation, ``Task Manager`` will store all tasks in `MongoDB <https://www.mongodb.com/>`_.
+Users **MUST** finished the configuration of `MongoDB <https://www.mongodb.com/>`_ when using this module.
+
+Users need to provide the url and database of ``task`` storing like this.
+
+    .. code-block:: python
+
+        from qlib.config import C
+        C["mongo"] = {
+            "task_url" : "mongodb://localhost:27017/", # maybe you need to change it to your url
+            "task_db_name" : "rolling_db" # you can custom database name
+        }
+
+The CRUD methods of ``task`` can be found in TaskManager. More methods can be seen in the `Github<https://github.com/microsoft/qlib/tree/main/qlib/workflow/task/manage.py>`_.
+
+.. autoclass:: qlib.workflow.task.manage.TaskManager
+    :members:
+
+Task Running
+===============
+After generating and storing those ``task``, it's time to run the ``task`` in the *WAITING* status.
+``qlib`` provide a method to run those ``task`` in task pool, however users can also customize how tasks are executed.
+An easy way to get the ``task_func`` is using ``qlib.model.trainer.task_train`` directly.
+It will run the whole workflow defined by ``task``, which includes *Model*, *Dataset*, *Record*.
+
+.. autofunction:: qlib.workflow.task.manage.run_task
+
+Task Collecting
+===============
+To see the results of ``task`` after running, ``Qlib`` provide a task collector to collect the tasks by filter condition (optional).
+The collector will return a dict of filtered key (users defined by task config) and value (predict scores from ``pred.pkl``).
+
+.. autoclass:: qlib.workflow.task.collect.TaskCollector
+    :members:

examples/taskmanager/task_manager_rolling.ipynb

@@ -0,0 +1,176 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": true
+   },
+   "outputs": [],
+   "source": [
+    "import qlib\n",
+    "from qlib.config import REG_CN\n",
+    "from qlib.workflow.task.gen import RollingGen, task_generator\n",
+    "from qlib.workflow.task.manage import TaskManager\n",
+    "from qlib.config import C\n",
+    "\n",
+    "data_handler_template = {\n",
+    "    \"start_time\": \"2008-01-01\",\n",
+    "    \"end_time\": \"2020-08-01\",\n",
+    "    \"fit_start_time\": \"2008-01-01\",\n",
+    "    \"fit_end_time\": \"2014-12-31\",\n",
+    "    \"instruments\": 'csi100',\n",
+    "}\n",
+    "\n",
+    "dataset_template = {\n",
+    "        \"class\": \"DatasetH\",\n",
+    "        \"module_path\": \"qlib.data.dataset\",\n",
+    "        \"kwargs\": {\n",
+    "            \"handler\": {\n",
+    "                \"class\": \"Alpha158\",\n",
+    "                \"module_path\": \"qlib.contrib.data.handler\",\n",
+    "                \"kwargs\": data_handler_template,\n",
+    "            },\n",
+    "            \"segments\": {\n",
+    "                \"train\": (\"2008-01-01\", \"2014-12-31\"),\n",
+    "                \"valid\": (\"2015-01-01\", \"2016-12-31\"),\n",
+    "                \"test\": (\"2017-01-01\", \"2020-08-01\"),\n",
+    "            },\n",
+    "        },\n",
+    "    }\n",
+    "\n",
+    "record_template = [\n",
+    "    {\n",
+    "        \"class\": \"SignalRecord\",\n",
+    "        \"module_path\": \"qlib.workflow.record_temp\",\n",
+    "    },\n",
+    "    {\n",
+    "        \"class\": \"SigAnaRecord\",\n",
+    "        \"module_path\": \"qlib.workflow.record_temp\",\n",
+    "    }\n",
+    "]\n",
+    "\n",
+    "# use lgb\n",
+    "lgb_task_template = {\n",
+    "    \"model\": {\n",
+    "        \"class\": \"LGBModel\",\n",
+    "        \"module_path\": \"qlib.contrib.model.gbdt\",\n",
+    "    },\n",
+    "    \"dataset\": dataset_template,\n",
+    "    \"record\": record_template,\n",
+    "}\n",
+    "\n",
+    "# use xgboost\n",
+    "xgboost_task_template = {\n",
+    "    \"model\": {\n",
+    "        \"class\": \"XGBModel\",\n",
+    "        \"module_path\": \"qlib.contrib.model.xgboost\",\n",
+    "    },\n",
+    "    \"dataset\": dataset_template,\n",
+    "    \"record\": record_template,\n",
+    "}\n",
+    "\n",
+    "provider_uri = \"~/.qlib/qlib_data/cn_data\"  # target_dir\n",
+    "qlib.init(provider_uri=provider_uri, region=REG_CN)\n",
+    "\n",
+    "C[\"mongo\"] = {\n",
+    "    \"task_url\" : \"mongodb://localhost:27017/\", # maybe you need to change it to your url\n",
+    "    \"task_db_name\" : \"rolling_db\"\n",
+    "}\n",
+    "\n",
+    "exp_name = 'rolling_exp' # experiment name, will be used as the experiment in MLflow\n",
+    "task_pool = 'rolling_task' # task pool name, will be used as the document in MongoDB"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "tasks = task_generator(\n",
+    "    xgboost_task_template, # default task name\n",
+    "    RollingGen(step=550,rtype=RollingGen.ROLL_SD), # generate different date segment\n",
+    "    task_lgb=lgb_task_template # use \"task_lgb\" as the task name\n",
+    ")\n",
+    "# Uncomment next two lines to see the generated tasks\n",
+    "# from pprint import pprint\n",
+    "# pprint(tasks)\n",
+    "tm = TaskManager(task_pool=task_pool)\n",
+    "tm.create_task(tasks) # all tasks will be saved to MongoDB"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from qlib.workflow.task.manage import run_task\n",
+    "from qlib.workflow.task.collect import TaskCollector\n",
+    "from qlib.model.trainer import task_train\n",
+    "\n",
+    "run_task(task_train, task_pool, experiment_name=exp_name) # all tasks will be trained using \"task_train\" method"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "def get_task_key(task):\n",
+    "    task_key = task[\"task_key\"]\n",
+    "    rolling_end_timestamp = task[\"dataset\"][\"kwargs\"][\"segments\"][\"test\"][1]\n",
+    "    return task_key, rolling_end_timestamp.strftime('%Y-%m-%d')\n",
+    "\n",
+    "def my_filter(task):\n",
+    "    # only choose the results of \"task_lgb\" and test segment end in 2019 from all tasks\n",
+    "    task_key, rolling_end = get_task_key(task)\n",
+    "    if task_key==\"task_lgb\" and rolling_end.startswith('2019'):\n",
+    "        return True\n",
+    "    return False\n",
+    "\n",
+    "# name tasks by \"get_task_key\" and filter tasks by \"my_filter\"\n",
+    "pred_rolling = TaskCollector.collect_predictions(exp_name, get_task_key, my_filter) \n",
+    "pred_rolling"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "3.6.5-final"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}

examples/taskmanager/task_manager_rolling.py

@@ -0,0 +1,94 @@
+import qlib
+from qlib.config import REG_CN
+from qlib.workflow.task.gen import RollingGen, task_generator
+from qlib.workflow.task.manage import TaskManager
+from qlib.config import C
+
+data_handler_config = {
+    "start_time": "2008-01-01",
+    "end_time": "2020-08-01",
+    "fit_start_time": "2008-01-01",
+    "fit_end_time": "2014-12-31",
+    "instruments": "csi100",
+}
+
+dataset_config = {
+    "class": "DatasetH",
+    "module_path": "qlib.data.dataset",
+    "kwargs": {
+        "handler": {"class": "Alpha158", "module_path": "qlib.contrib.data.handler", "kwargs": data_handler_config,},
+        "segments": {
+            "train": ("2008-01-01", "2014-12-31"),
+            "valid": ("2015-01-01", "2016-12-31"),
+            "test": ("2017-01-01", "2020-08-01"),
+        },
+    },
+}
+
+record_config = [
+    {"class": "SignalRecord", "module_path": "qlib.workflow.record_temp",},
+    {"class": "SigAnaRecord", "module_path": "qlib.workflow.record_temp",},
+]
+
+# use lgb
+task_lgb_config = {
+    "model": {"class": "LGBModel", "module_path": "qlib.contrib.model.gbdt",},
+    "dataset": dataset_config,
+    "record": record_config,
+}
+
+# use xgboost
+task_xgboost_config = {
+    "model": {"class": "XGBModel", "module_path": "qlib.contrib.model.xgboost",},
+    "dataset": dataset_config,
+    "record": record_config,
+}
+
+provider_uri = "~/.qlib/qlib_data/cn_data"  # target_dir
+qlib.init(provider_uri=provider_uri, region=REG_CN)
+
+C["mongo"] = {
+    "task_url": "mongodb://localhost:27017/",  # maybe you need to change it to your url
+    "task_db_name": "rolling_db",
+}
+
+exp_name = "rolling_exp"  # experiment name, will be used as the experiment in MLflow
+task_pool = "rolling_task"  # task pool name, will be used as the document in MongoDB
+
+tasks = task_generator(
+    task_xgboost_config,  # default task name
+    RollingGen(step=550, rtype=RollingGen.ROLL_SD),  # generate different date segment
+    task_lgb=task_lgb_config,  # use "task_lgb" as the task name
+)
+
+# Uncomment next two lines to see the generated tasks
+# from pprint import pprint
+# pprint(tasks)
+
+tm = TaskManager(task_pool=task_pool)
+tm.create_task(tasks)  # all tasks will be saved to MongoDB
+
+from qlib.workflow.task.manage import run_task
+from qlib.workflow.task.collect import TaskCollector
+from qlib.model.trainer import task_train
+
+run_task(task_train, task_pool, experiment_name=exp_name)  # all tasks will be trained using "task_train" method
+
+
+def get_task_key(task_config):
+    task_key = task_config["task_key"]
+    rolling_end_timestamp = task_config["dataset"]["kwargs"]["segments"]["test"][1]
+    return task_key, rolling_end_timestamp.strftime("%Y-%m-%d")
+
+
+def my_filter(task_config):
+    # only choose the results of "task_lgb" and test in 2019 from all tasks
+    task_key, rolling_end = get_task_key(task_config)
+    if task_key == "task_lgb" and rolling_end.startswith("2019"):
+        return True
+    return False
+
+
+# name tasks by "get_task_key" and filter tasks by "my_filter"
+pred_rolling = TaskCollector.collect_predictions(exp_name, get_task_key, my_filter)
+pred_rolling

examples/taskmanager/update_online_pred.py

@@ -0,0 +1,77 @@
+import qlib
+from qlib.model.trainer import task_train
+from qlib.workflow.task.update import ModelUpdater
+from qlib.config import REG_CN
+import fire
+
+data_handler_config = {
+        "start_time": "2008-01-01",
+        "end_time": "2020-08-01",
+        "fit_start_time": "2008-01-01",
+        "fit_end_time": "2014-12-31",
+        "instruments": "csi100",
+    }
+
+task = {
+    "model": {
+        "class": "LGBModel",
+        "module_path": "qlib.contrib.model.gbdt",
+        "kwargs": {
+            "loss": "mse",
+            "colsample_bytree": 0.8879,
+            "learning_rate": 0.0421,
+            "subsample": 0.8789,
+            "lambda_l1": 205.6999,
+            "lambda_l2": 580.9768,
+            "max_depth": 8,
+            "num_leaves": 210,
+            "num_threads": 20,
+        },
+    },
+    "dataset": {
+        "class": "DatasetH",
+        "module_path": "qlib.data.dataset",
+        "kwargs": {
+            "handler": {
+                "class": "Alpha158",
+                "module_path": "qlib.contrib.data.handler",
+                "kwargs": data_handler_config,
+            },
+            "segments": {
+                "train": ("2008-01-01", "2014-12-31"),
+                "valid": ("2015-01-01", "2016-12-31"),
+                "test": ("2017-01-01", "2020-08-01"),
+            },
+        },
+    },
+    "record": {"class": "SignalRecord", "module_path": "qlib.workflow.record_temp",},
+}
+
+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)
+    model_updater = ModelUpdater(experiment_name)
+
+    rid = task_train(task_config=task, experiment_name=experiment_name)
+    model_updater.reset_online_model(rid)
+
+def update_online_pred(experiment_name="online_svr"):
+    
+    qlib.init(provider_uri=provider_uri, region=REG_CN)
+    model_updater = ModelUpdater(experiment_name)
+
+    print("Here are the online models waiting for update:")
+    for rid, rec in model_updater.list_online_model().items():
+        print(rid)
+
+    model_updater.update_online_pred()
+    
+if __name__ == '__main__':
+    fire.Fire()
+    # 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
+    # python update_online_pred.py update_online_pred

qlib/model/trainer.py

@@ -6,7 +6,7 @@ from qlib.workflow import R
 from qlib.workflow.record_temp import SignalRecord
 
 
-def task_train(task_config: dict, experiment_name):
+def task_train(task_config: dict, experiment_name: str) -> str:
     """
     task based training
 
@@ -14,6 +14,13 @@ def task_train(task_config: dict, experiment_name):
     ----------
     task_config : dict
         A dict describes a task setting.
+    experiment_name: str
+        The name of experiment
+
+    Returns
+    ----------
+    rid : str
+        The id of the recorder of this task
     """
 
     # model initiaiton
@@ -27,16 +34,23 @@ def task_train(task_config: dict, experiment_name):
         model.fit(dataset)
         recorder = R.get_recorder()
         R.save_objects(**{"params.pkl": model})
+        R.save_objects(**{"task.pkl": task_config})  # keep the original format and datatype
 
         # generate records: prediction, backtest, and analysis
-        for record in task_config["record"]:
+        records = task_config.get("record", [])
+        if isinstance(records, dict):  # prevent only one dict
+            records = [records]
+        for record in records:
             if record["class"] == SignalRecord.__name__:
                 srconf = {"model": model, "dataset": dataset, "recorder": recorder}
+                record.setdefault("kwargs", {})
                 record["kwargs"].update(srconf)
                 sr = init_instance_by_config(record)
                 sr.generate()
             else:
                 rconf = {"recorder": recorder}
+                record.setdefault("kwargs", {})
                 record["kwargs"].update(rconf)
                 ar = init_instance_by_config(record)
                 ar.generate()
+    return recorder.info["id"]

qlib/workflow/task/collect.py

@@ -1,20 +1,77 @@
 from qlib.workflow import R
 import pandas as pd
 from typing import Union
-from tqdm.auto import tqdm
+from qlib import get_module_logger
 
 
-class RollingEnsemble:
+class TaskCollector:
+    """
+    Collect the record results of the finished tasks with key and filter
+    """
+
+    @staticmethod
+    def collect_predictions(
+        experiment_name: str,
+        get_key_func,
+        filter_func=None,
+    ):
+        """
+
+        Parameters
+        ----------
+        experiment_name : str
+        get_key_func : function(task: dict) -> Union[Number, str, tuple]
+            get the key of a task when collect it
+        filter_func : function(task: dict) -> bool
+            to judge a task will be collected or not
+
+        Returns
+        -------
+
+        """
+        exp = R.get_exp(experiment_name=experiment_name)
+        # filter records
+        recs = exp.list_recorders()
+
+        recs_flt = {}
+        for rid, rec in recs.items():
+            params = rec.load_object("task.pkl")
+            if rec.status == rec.STATUS_FI:
+                if filter_func is None or filter_func(params):
+                    rec.params = params
+                    recs_flt[rid] = rec
+
+        # group
+        recs_group = {}
+        for _, rec in recs_flt.items():
+            params = rec.params
+            group_key = get_key_func(params)
+            recs_group.setdefault(group_key, []).append(rec)
+
+        # reduce group
+        reduce_group = {}
+        for k, rec_l in recs_group.items():
+            pred_l = []
+            for rec in rec_l:
+                pred_l.append(rec.load_object("pred.pkl").iloc[:, 0])
+            pred = pd.concat(pred_l).sort_index()
+            reduce_group[k] = pred
+
+        get_module_logger("TaskCollector").info(f"Collect {len(reduce_group)} predictions in {experiment_name}")
+        return reduce_group
+
+
+class RollingCollector:
     """
     Rolling Models Ensemble based on (R)ecord
 
     This shares nothing with Ensemble
     """
 
-    # TODO: 这边还可以加加速
+    # TODO: speed up this class
     def __init__(self, get_key_func, flt_func=None):
-        self.get_key_func = get_key_func
+        self.get_key_func = get_key_func  # get the key of a task based on task config
-        self.flt_func = flt_func
+        self.flt_func = flt_func  # determine whether a task can be retained based on task config
 
     def __call__(self, exp_name) -> Union[pd.Series, dict]:
         # TODO;
@@ -26,8 +83,7 @@ class RollingEnsemble:
 
         recs_flt = {}
         for rid, rec in tqdm(recs.items(), desc="Loading data"):
-            # rec = exp.get_recorder(recorder_id=rid)
+            params = rec.load_object("task.pkl")
-            params = rec.load_object("param")
             if rec.status == rec.STATUS_FI:
                 if self.flt_func is None or self.flt_func(params):
                     rec.params = params

qlib/workflow/task/gen.py

@@ -9,23 +9,87 @@ import typing
 from .utils import TimeAdjuster
 
 
-class TaskGen(metaclass=abc.ABCMeta):
+def task_generator(*args, **kwargs) -> list:
-    @abc.abstractmethod
-    def __call__(self, *args, **kwargs) -> typing.List[dict]:
     """
-        generate
+    Accept the dict of task config and the TaskGen to generate different tasks.
+    There is no limit to the number and position of input.
+    The key of input will add to task config.
+
+    for example:
+        There are 3 task_config(a,b,c) and 2 TaskGen(A,B). A will double the task_config and B will triple.
+        task_generator(a_key=a, b_key=b, c_key=c, A, B) will finally generate 3*2*3 = 18 task_config.
 
     Parameters
     ----------
-        args, kwargs:
+    args : dict or TaskGen
-            The info for generating tasks
+    kwargs : dict or TaskGen
-            Example 1):
+
-                input: a specific task template
+    Returns
+    -------
+    gen_task_list : list
+        a list of task config after generating
+    """
+    tasks_list = []
+    gen_list = []
+
+    tmp_id = 1
+    for task in args:
+        if isinstance(task, dict):
+            task["task_key"] = tmp_id
+            tmp_id += 1
+            tasks_list.append(task)
+        elif isinstance(task, TaskGen):
+            gen_list.append(task)
+        else:
+            raise NotImplementedError(f"{type(task)} is not supported in task_generator")
+
+    for key, task in kwargs.items():
+        if isinstance(task, dict):
+            task["task_key"] = key
+            tasks_list.append(task)
+        elif isinstance(task, TaskGen):
+            gen_list.append(task)
+        else:
+            raise NotImplementedError(f"{type(task)} is not supported in task_generator")
+
+    # generate gen_task_list
+    gen_task_list = []
+    for gen in gen_list:
+        new_task_list = []
+        for task in tasks_list:
+            new_task_list.extend(gen.generate(task))
+        gen_task_list = new_task_list
+    return gen_task_list
+
+
+class TaskGen(metaclass=abc.ABCMeta):
+    """
+    the base class for generate different tasks
+
+    Example 1:
+
+        input: a specific task template and rolling steps
+
         output: rolling version of the tasks
-            Example 2):
+
-                input: a specific task template
+    Example 2:
+
+        input: a specific task template and losses list
+
         output: a set of tasks with different losses
 
+    """
+
+    @abc.abstractmethod
+    def generate(self, task: dict) -> typing.List[dict]:
+        """
+        generate different tasks based on a task template
+
+        Parameters
+        ----------
+        task: dict
+            a task template
+
         Returns
         -------
         typing.List[dict]:
@@ -35,9 +99,8 @@ class TaskGen(metaclass=abc.ABCMeta):
 
 
 class RollingGen(TaskGen):
-
+    ROLL_EX = TimeAdjuster.SHIFT_EX  # fixed start date, expanding end date
-    ROLL_EX = TimeAdjuster.SHIFT_EX
+    ROLL_SD = TimeAdjuster.SHIFT_SD  # fixed segments size, slide it from start date
-    ROLL_SD = TimeAdjuster.SHIFT_SD
 
     def __init__(self, step: int = 40, rtype: str = ROLL_EX):
         """
@@ -48,16 +111,17 @@ class RollingGen(TaskGen):
         step : int
             step to rolling
         rtype : str
-            rolling type (expanding, rolling)
+            rolling type (expanding, sliding)
         """
         self.step = step
         self.rtype = rtype
-        self.ta = TimeAdjuster(future=True)  # 为了保证test最后的日期不是None, 所以这边要改一改
+        # TODO: Ask pengrong to update future date in dataset
+        self.ta = TimeAdjuster(future=True)
 
         self.test_key = "test"
         self.train_key = "train"
 
-    def __call__(self, task: dict):
+    def generate(self, task: dict):
         """
         Converting the task into a rolling task
 
@@ -101,22 +165,23 @@ class RollingGen(TaskGen):
             # calculate segments
             if prev_seg is None:
                 # First rolling
-                # 1) prepare the end porint
+                # 1) prepare the end point
                 segments = copy.deepcopy(self.ta.align_seg(t["dataset"]["kwargs"]["segments"]))
-                test_end = self.ta.max() if segments[self.test_key][1] is None else segments[self.test_key][1]
+                test_end = self.ta.last_date() if segments[self.test_key][1] is None else segments[self.test_key][1]
-                # 2) and the init test segments
+                # 2) and init test segments
                 test_start_idx = self.ta.align_idx(segments[self.test_key][0])
                 segments[self.test_key] = (self.ta.get(test_start_idx), self.ta.get(test_start_idx + self.step - 1))
             else:
                 segments = {}
                 try:
                     for k, seg in prev_seg.items():
-                        # 决定怎么shift
+                        # decide how to shift
+                        # expanding only for train data, the segments size of test data and valid data won't change
                         if k == self.train_key and self.rtype == self.ROLL_EX:
                             rtype = self.ta.SHIFT_EX
                         else:
                             rtype = self.ta.SHIFT_SD
-                        # 整段数据做shift
+                        # shift the segments data
                         segments[k] = self.ta.shift(seg, step=self.step, rtype=rtype)
                     if segments[self.test_key][0] > test_end:
                         break
@@ -125,6 +190,7 @@ class RollingGen(TaskGen):
                     # No more rolling
                     break
 
+            # update segments of this task
             t["dataset"]["kwargs"]["segments"] = copy.deepcopy(segments)
             prev_seg = segments
             res.append(t)

qlib/workflow/task/manage.py

@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 """
-A task consists of 2 parts
+A task consists of 3 parts
 - tasks description: the desc will define the task
 - tasks status: the status of the task
 - tasks result information : A user can get the task with the task description and task result.
@@ -26,6 +26,9 @@ from qlib import auto_init
 class TaskManager:
     """TaskManager
     here is the what will a task looks like
+
+    .. code-block:: python
+
         {
             'def': pickle serialized task definition.  using pickle will make it easier
             'filter': json-like data. This is for filtering the tasks.
@@ -36,12 +39,9 @@ class TaskManager:
     The tasks manager assume that you will only update the tasks you fetched.
     The mongo fetch one and update will make it date updating secure.
 
-    Usage Examples from the CLI.
+    .. note::
-    python -m blocks.tasks.__init__ task_stat --task_pool meta_task_rule
 
-
+        assumption: the data in MongoDB was encoded and the data out of MongoDB was decoded
-    NOTE:
-    - 假设： 存储在db里面的都是encode过的， 拿出来的都是decode过的
     """
 
     STATUS_WAITING = "waiting"
@@ -52,6 +52,14 @@ class TaskManager:
     ENCODE_FIELDS_PREFIX = ["def", "res"]
 
     def __init__(self, task_pool=None):
+        """
+        init Task Manager, remember to make the statement of MongoDB url and database name firstly.
+
+        Parameters
+        ----------
+        task_pool: str
+            the name of Collection in MongoDB
+        """
         self.mdb = get_mongodb()
         self.task_pool = task_pool
 
@@ -85,7 +93,7 @@ class TaskManager:
         return {k: str(v) for k, v in flt.items()}
 
     def replace_task(self, task, new_task, task_pool=None):
-        # 这里的假设是从接口拿出来的都是decode过的，在接口内部的都是 encode过的
+        # assume that the data out of interface was decoded and the data in interface was encoded
         new_task = self._encode_task(new_task)
         task_pool = self._get_task_pool(task_pool)
         query = {"_id": ObjectId(task["_id"])}
@@ -104,6 +112,19 @@ class TaskManager:
             task_pool.insert_one(task)
 
     def insert_task_def(self, task_def, task_pool=None):
+        """
+        insert a task to task_pool
+
+        Parameters
+        ----------
+        task_def: dict
+        task_pool: str
+            the name of Collection in MongoDB
+
+        Returns
+        -------
+
+        """
         task_pool = self._get_task_pool(task_pool)
         task = self._encode_task(
             {
@@ -115,6 +136,23 @@ class TaskManager:
         self.insert_task(task, task_pool)
 
     def create_task(self, task_def_l, task_pool=None, dry_run=False, print_nt=False):
+        """
+        if the tasks in task_def_l is new, then insert new tasks into the task_pool
+
+        Parameters
+        ----------
+        task_def_l: list
+            a list of task
+        task_pool: str
+            the name of task_pool (collection name of MongoDB)
+        dry_run: bool
+            if insert those new tasks to task pool
+        print_nt: bool
+            if print new task
+        Returns
+        -------
+
+        """
         task_pool = self._get_task_pool(task_pool)
         new_tasks = []
         for t in task_def_l:
@@ -145,7 +183,7 @@ class TaskManager:
         task = task_pool.find_one_and_update(
             query, {"$set": {"status": self.STATUS_RUNNING}}, sort=[("priority", pymongo.DESCENDING)]
         )
-        # 这里我的 priority 必须是 高数优先级更高，因为 null会被在 ASCENDING时被排在最前面
+        # null will be at the top after sorting when using ASCENDING, so the larger the number higher, the higher the priority
         if task is None:
             return None
         task["status"] = self.STATUS_RUNNING
@@ -153,6 +191,20 @@ class TaskManager:
 
     @contextmanager
     def safe_fetch_task(self, query={}, task_pool=None):
+        """
+        fetch task from task_pool using query with contextmanager
+
+        Parameters
+        ----------
+        query: dict
+            the dict of query
+        task_pool: str
+            the name of Collection in MongoDB
+
+        Returns
+        -------
+
+        """
         task = self.fetch_task(query=query, task_pool=task_pool)
         try:
             yield task
@@ -171,12 +223,20 @@ class TaskManager:
                 yield task
 
     def query(self, query={}, decode=True, task_pool=None):
-        """query
+        """
         This function may raise exception `pymongo.errors.CursorNotFound: cursor id not found` if it takes too long to iterate the generator
 
-        :param query:
+        Parameters
-        :param decode:
+        ----------
-        :param task_pool:
+        query: dict
+            the dict of query
+        decode: bool
+        task_pool: str
+            the name of Collection in MongoDB
+
+        Returns
+        -------
+
         """
         query = query.copy()
         if "_id" in query:
@@ -200,6 +260,20 @@ class TaskManager:
         task_pool.update_one({"_id": task["_id"]}, update_dict)
 
     def remove(self, query={}, task_pool=None):
+        """
+        remove the task using query
+
+        Parameters
+        ----------
+        query: dict
+            the dict of query
+        task_pool: str
+            the name of Collection in MongoDB
+
+        Returns
+        -------
+
+        """
         query = query.copy()
         task_pool = self._get_task_pool(task_pool)
         if "_id" in query:
@@ -254,15 +328,15 @@ class TaskManager:
 
 
 def run_task(task_func, task_pool, force_release=False, *args, **kwargs):
-    """run_task.
+    """
-    While task pool is not empty, use task_func to fetch and run tasks in task_pool
+    While task pool is not empty (has WAITING tasks), use task_func to fetch and run tasks in task_pool
 
     Parameters
     ----------
     task_func : def (task_def, *args, **kwargs) -> <res which will be committed>
         the function to run the task
-    task_pool :
+    task_pool : str
-        The name of the task pool
+        the name of the task pool (Collection in MongoDB)
     force_release :
         will the program force to release the resource
     args :

qlib/workflow/task/update.py

@@ -0,0 +1,154 @@
+from typing import Union
+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
+
+
+class ModelUpdater:
+    """
+    The model updater to re-train model or update predictions
+    """
+
+    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
+        """
+        self.exp_name = experiment_name
+        self.exp = R.get_exp(experiment_name=experiment_name)
+        self.logger = get_module_logger("ModelUpdater")
+
+    def set_online_model(self, rid: str):
+        """online model will be identified at the tags of the record
+
+        Parameters
+        ----------
+        rid : str
+            the id of a record
+        """
+        rec = self.exp.get_recorder(recorder_id=rid)
+        rec.set_tags(**{self.ONLINE_TAG: self.ONLINE_TAG_TRUE})
+
+    def cancel_online_model(self, rid: str):
+        rec = self.exp.get_recorder(recorder_id=rid)
+        rec.set_tags(**{self.ONLINE_TAG: self.ONLINE_TAG_FALSE})
+
+    def cancel_all_online_model(self):
+        recs = self.exp.list_recorders()
+        for rid, rec in recs.items():
+            self.cancel_online_model(rid)
+
+    def reset_online_model(self, rids: Union[str, list]):
+        """cancel all online model and reset the given model to online model
+
+        Parameters
+        ----------
+        rids : Union[str, list]
+            the name of a record or the list of the name of records
+        """
+        self.cancel_all_online_model()
+        if isinstance(rids, str):
+            rids = [rids]
+        for rid in rids:
+            self.set_online_model(rid)
+
+    def update_pred(self, rid: str):
+        """update predictions to the latest day in Calendar based on rid
+
+        Parameters
+        ----------
+        rid : str
+            the id of the record
+        """
+        rec = self.exp.get_recorder(recorder_id=rid)
+        old_pred = rec.load_object("pred.pkl")
+        last_end = old_pred.index.get_level_values("datetime").max()
+        task_config = rec.load_object("task.pkl")
+
+        # updated to the latest trading day
+        cal = D.calendar(start_time=last_end + pd.Timedelta(days=1), end_time=None)
+
+        if len(cal) == 0:
+            self.logger.info(f"All prediction in {rid} of {self.exp_name} are latest. No need to update.")
+            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"])
+
+        model = rec.load_object("params.pkl")
+        new_pred = model.predict(dataset)
+
+        cb_pred = pd.concat([old_pred, new_pred.to_frame("score")], axis=0)
+        cb_pred = cb_pred.sort_index()
+
+        rec.save_objects(**{"pred.pkl": cb_pred})
+
+        self.logger.info(f"Finish updating new {new_pred.shape[0]} predictions in {rid} of {self.exp_name}.")
+
+    def update_all_pred(self, filter_func=None):
+        """update all predictions in this experiment after filter.
+
+            An example of filter function:
+
+            .. code-block:: python
+
+                def record_filter(record):
+                    task_config = record.load_object("task.pkl")
+                    if task_config["model"]["class"]=="LGBModel":
+                        return True
+                    return False
+
+        Parameters
+        ----------
+        filter_func : function, optional
+            the filter function to decide whether this record will be updated, by default None
+
+        Returns
+        ----------
+        cnt: int
+            the count of updated record
+
+        """
+        cnt = 0
+        recs = self.exp.list_recorders()
+        for rid, rec in recs.items():
+            if rec.status == rec.STATUS_FI:
+                if filter_func != None and filter_func(rec) == False:
+                    # records that should be filtered out
+                    continue
+                self.update_pred(rid)
+                cnt += 1
+        return cnt
+
+    def online_filter(self, record):
+        tags = record.list_tags()
+        if tags[self.ONLINE_TAG] == self.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):
+        recs = self.exp.list_recorders()
+        online_rec = {}
+        for rid, rec in recs.items():
+            if self.online_filter(rec):
+                online_rec[rid] = rec
+        return online_rec

qlib/workflow/task/utils.py

@@ -6,9 +6,20 @@ from qlib.data import D
 from qlib.config import C
 from qlib.log import get_module_logger
 from pymongo import MongoClient
+from typing import Union
 
 
 def get_mongodb():
+    """
+
+    get database in MongoDB, which means you need to declare the address and the name of database.
+    for example:
+        C["mongo"] = {
+            "task_url" : "mongodb://localhost:27017/",
+            "task_db_name" : "rolling_db"
+        }
+
+    """
     try:
         cfg = C["mongo"]
     except KeyError:
@@ -20,7 +31,9 @@ def get_mongodb():
 
 
 class TimeAdjuster:
-    """找到合适的日期，然后adjust date"""
+    """
+    find appropriate date and adjust date.
+    """
 
     def __init__(self, future=False):
         self.cals = D.calendar(future=future)
@@ -40,11 +53,30 @@ class TimeAdjuster:
 
     def max(self):
         """
-        Return return the max calendar date
+        (Deprecated)
+        Return the max calendar datetime
         """
         return max(self.cals)
 
+    def last_date(self) -> pd.Timestamp:
+        """
+        Return the last datetime in the calendar
+        """
+        return self.cals[-1]
+
     def align_idx(self, time_point, tp_type="start"):
+        """
+        align the index of time_point in the calendar
+
+        Parameters
+        ----------
+        time_point
+        tp_type : str
+
+        Returns
+        -------
+        index : int
+        """
         time_point = pd.Timestamp(time_point)
         if tp_type == "start":
             idx = bisect.bisect_left(self.cals, time_point)
@@ -56,18 +88,36 @@ class TimeAdjuster:
 
     def align_time(self, time_point, tp_type="start"):
         """
-        Align a timepoint to calendar  weekdays
+        Align time_point to trade date of calendar
 
         Parameters
         ----------
-        time_point :
+        time_point
             Time point
         tp_type : str
             time point type (`"start"`, `"end"`)
         """
         return self.cals[self.align_idx(time_point, tp_type=tp_type)]
 
-    def align_seg(self, segment):
+    def align_seg(self, segment: Union[dict, tuple]):
+        """
+        align the given date to trade date
+
+        for example:
+            input: {'train': ('2008-01-01', '2014-12-31'), 'valid': ('2015-01-01', '2016-12-31'), 'test': ('2017-01-01', '2020-08-01')}
+
+            output: {'train': (Timestamp('2008-01-02 00:00:00'), Timestamp('2014-12-31 00:00:00')),
+                    'valid': (Timestamp('2015-01-05 00:00:00'), Timestamp('2016-12-30 00:00:00')),
+                    'test': (Timestamp('2017-01-03 00:00:00'), Timestamp('2020-07-31 00:00:00'))}
+
+        Parameters
+        ----------
+        segment
+
+        Returns
+        -------
+        the start and end trade date (pd.Timestamp) between the given start and end date.
+        """
         if isinstance(segment, dict):
             return {k: self.align_seg(seg) for k, seg in segment.items()}
         elif isinstance(segment, tuple):
@@ -75,17 +125,18 @@ class TimeAdjuster:
         else:
             raise NotImplementedError(f"This type of input is not supported")
 
-    def truncate(self, segment, test_start, days: int):
+    def truncate(self, segment: tuple, test_start, days: int):
         """
         truncate the segment based on the test_start date
 
         Parameters
         ----------
-        segment :
+        segment : tuple
             time segment
+        test_start
         days : int
             The trading days to be truncated
-            大部分情况是因为这个时间段的数据(一般是特征)会用到 `days` 天的数据
+            the data in this segment may need 'days' data
         """
         test_idx = self.align_idx(test_start)
         if isinstance(segment, tuple):
@@ -101,9 +152,9 @@ class TimeAdjuster:
     SHIFT_SD = "sliding"
     SHIFT_EX = "expanding"
 
-    def shift(self, seg, step: int, rtype=SHIFT_SD):
+    def shift(self, seg: tuple, step: int, rtype=SHIFT_SD):
         """
-        shift the datatiem of segment
+        shift the datatime of segment
 
         Parameters
         ----------