diff --git a/tests/test_all_pipeline.py b/tests/test_all_pipeline.py index 4c20405fa..13420335b 100644 --- a/tests/test_all_pipeline.py +++ b/tests/test_all_pipeline.py @@ -36,7 +36,7 @@ port_analysis_config = { } -def train(): +def train(uri_path: str = None): """train model Returns @@ -55,7 +55,7 @@ def train(): print(R) # start exp - with R.start(experiment_name="workflow"): + with R.start(experiment_name="workflow", uri=uri_path): R.log_params(**flatten_dict(CSI300_GBDT_TASK)) model.fit(dataset) @@ -79,7 +79,7 @@ def train(): return pred_score, {"ic": ic, "ric": ric}, rid -def train_with_sigana(): +def train_with_sigana(uri_path: str = None): """train model followed by SigAnaRecord Returns @@ -91,9 +91,8 @@ def train_with_sigana(): """ model = init_instance_by_config(CSI300_GBDT_TASK["model"]) dataset = init_instance_by_config(CSI300_GBDT_TASK["dataset"]) - # start exp - with R.start(experiment_name="workflow_with_sigana"): + with R.start(experiment_name="workflow_with_sigana", uri=uri_path): R.log_params(**flatten_dict(CSI300_GBDT_TASK)) model.fit(dataset) @@ -130,7 +129,7 @@ def fake_experiment(): return default_uri == default_uri_to_check, current_uri == current_uri_to_check, current_uri -def backtest_analysis(pred, rid): +def backtest_analysis(pred, rid, uri_path: str = None): """backtest and analysis Parameters @@ -139,6 +138,8 @@ def backtest_analysis(pred, rid): predict scores rid : str the id of the recorder to be used in this function + uri_path: str + mlflow uri path Returns ------- @@ -146,7 +147,8 @@ def backtest_analysis(pred, rid): the analysis result """ - recorder = R.get_recorder(experiment_name="workflow", recorder_id=rid) + with R.start(experiment_name="workflow", recorder_id=rid, uri=uri_path): + recorder = R.get_recorder(experiment_name="workflow", recorder_id=rid) # backtest par = PortAnaRecord(recorder, port_analysis_config) par.generate() @@ -160,24 +162,24 @@ class TestAllFlow(TestAutoData): REPORT_NORMAL = None POSITIONS = None RID = None + URI_PATH = "file:" + str(Path(__file__).parent.joinpath("test_all_flow_mlruns").resolve()) @classmethod def tearDownClass(cls) -> None: - shutil.rmtree(str(Path(C["exp_manager"]["kwargs"]["uri"].strip("file:")).resolve())) + shutil.rmtree(cls.URI_PATH.lstrip("file:")) def test_0_train_with_sigana(self): - TestAllFlow.PRED_SCORE, ic_ric, uri_path = train_with_sigana() + TestAllFlow.PRED_SCORE, ic_ric, uri_path = train_with_sigana(self.URI_PATH) self.assertGreaterEqual(ic_ric["ic"].all(), 0, "train failed") self.assertGreaterEqual(ic_ric["ric"].all(), 0, "train failed") - shutil.rmtree(str(Path(uri_path.strip("file:")).resolve())) def test_1_train(self): - TestAllFlow.PRED_SCORE, ic_ric, TestAllFlow.RID = train() + TestAllFlow.PRED_SCORE, ic_ric, TestAllFlow.RID = train(self.URI_PATH) self.assertGreaterEqual(ic_ric["ic"].all(), 0, "train failed") self.assertGreaterEqual(ic_ric["ric"].all(), 0, "train failed") def test_2_backtest(self): - analyze_df = backtest_analysis(TestAllFlow.PRED_SCORE, TestAllFlow.RID) + analyze_df = backtest_analysis(TestAllFlow.PRED_SCORE, TestAllFlow.RID, self.URI_PATH) self.assertGreaterEqual( analyze_df.loc(axis=0)["excess_return_with_cost", "annualized_return"].values[0], 0.10, diff --git a/tests/test_contrib_workflow.py b/tests/test_contrib_workflow.py index 9b1edbd4e..2f5feb78e 100644 --- a/tests/test_contrib_workflow.py +++ b/tests/test_contrib_workflow.py @@ -12,10 +12,10 @@ from qlib.tests import TestAutoData from qlib.tests.config import CSI300_GBDT_TASK -def train_multiseg(): +def train_multiseg(uri_path: str = None): model = init_instance_by_config(CSI300_GBDT_TASK["model"]) dataset = init_instance_by_config(CSI300_GBDT_TASK["dataset"]) - with R.start(experiment_name="workflow"): + with R.start(experiment_name="workflow", uri=uri_path): R.log_params(**flatten_dict(CSI300_GBDT_TASK)) model.fit(dataset) recorder = R.get_recorder() @@ -25,10 +25,10 @@ def train_multiseg(): return uri -def train_mse(): +def train_mse(uri_path: str = None): model = init_instance_by_config(CSI300_GBDT_TASK["model"]) dataset = init_instance_by_config(CSI300_GBDT_TASK["dataset"]) - with R.start(experiment_name="workflow"): + with R.start(experiment_name="workflow", uri=uri_path): R.log_params(**flatten_dict(CSI300_GBDT_TASK)) model.fit(dataset) recorder = R.get_recorder() @@ -39,13 +39,17 @@ def train_mse(): class TestAllFlow(TestAutoData): + URI_PATH = "file:" + str(Path(__file__).parent.joinpath("test_contrib_mlruns").resolve()) + + @classmethod + def tearDownClass(cls) -> None: + shutil.rmtree(cls.URI_PATH.lstrip("file:")) + def test_0_multiseg(self): - uri_path = train_multiseg() - shutil.rmtree(str(Path(uri_path.strip("file:")).resolve())) + uri_path = train_multiseg(self.URI_PATH) def test_1_mse(self): - uri_path = train_mse() - shutil.rmtree(str(Path(uri_path.strip("file:")).resolve())) + uri_path = train_mse(self.URI_PATH) def suite():