Add MultiSegRecord in contrib.workflow and decouple its tests from test_all_pipeline

tests/test_all_pipeline.py

@@ -6,24 +6,11 @@ import shutil
 import unittest
 from pathlib import Path
 
-import numpy as np
-import pandas as pd
-
 import qlib
-from qlib.config import REG_CN, C
-from qlib.utils import drop_nan_by_y_index
-from qlib.contrib.model.gbdt import LGBModel
-from qlib.contrib.data.handler import Alpha158
-from qlib.contrib.strategy.strategy import TopkDropoutStrategy
-from qlib.contrib.evaluate import (
-    backtest as normal_backtest,
-    risk_analysis,
-)
-from qlib.contrib.workflow.record_temp import SignalMseRecord
-from qlib.utils import exists_qlib_data, init_instance_by_config, flatten_dict
+from qlib.config import C
+from qlib.utils import init_instance_by_config, flatten_dict
 from qlib.workflow import R
 from qlib.workflow.record_temp import SignalRecord, SigAnaRecord, PortAnaRecord
-from qlib.tests.data import GetData
 from qlib.tests import TestAutoData
 
 
@@ -166,8 +153,6 @@ def train_with_sigana():
         ric = sar.load(sar.get_path("ric.pkl"))
         pred_score = sar.load("pred.pkl")
 
-        smr = SignalMseRecord(recorder)
-        smr.generate()
         uri_path = R.get_uri()
     return pred_score, {"ic": ic, "ric": ric}, uri_path
 
@@ -256,8 +241,10 @@ class TestAllFlow(TestAutoData):
 
 def suite():
     _suite = unittest.TestSuite()
-    _suite.addTest(TestAllFlow("test_0_train"))
-    _suite.addTest(TestAllFlow("test_1_backtest"))
+    _suite.addTest(TestAllFlow("test_0_train_with_sigana"))
+    _suite.addTest(TestAllFlow("test_1_train"))
+    _suite.addTest(TestAllFlow("test_2_backtest"))
+    _suite.addTest(TestAllFlow("test_3_expmanager"))
     return _suite
 
 

tests/test_contrib_workflow.py

@@ -0,0 +1,97 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import sys
+import shutil
+import unittest
+from pathlib import Path
+
+import qlib
+from qlib.config import C
+from qlib.contrib.workflow import MultiSegRecord, SignalMseRecord
+from qlib.utils import init_instance_by_config, flatten_dict
+from qlib.workflow import R
+from qlib.tests import TestAutoData
+
+
+market = "csi300"
+benchmark = "SH000300"
+
+###################################
+# train model
+###################################
+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": market,
+}
+
+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"),
+            },
+        },
+    },
+}
+
+
+def test_multiseg():
+    model = init_instance_by_config(task["model"])
+    dataset = init_instance_by_config(task["dataset"])
+    with R.start(experiment_name="workflow"):
+        R.log_params(**flatten_dict(task))
+        model.fit(dataset)
+
+        # prediction
+        recorder = R.get_recorder()
+        sr = MultiSegRecord(model, dataset, recorder)
+        sr.generate(dict(valid="valid", test="test"))
+
+        uri = R.get_uri()
+
+    return uri
+
+
+class TestAllFlow(TestAutoData):
+    def test_0_multiseg(self):
+        uri_path = test_multiseg()
+        shutil.rmtree(str(Path(uri_path.strip("file:")).resolve()))
+
+
+def suite():
+    _suite = unittest.TestSuite()
+    _suite.addTest(TestAllFlow("test_0_multiseg"))
+    return _suite
+
+
+if __name__ == "__main__":
+    runner = unittest.TextTestRunner()
+    runner.run(suite())