opt local trainer (better mem releasing) (#1116)

* opt local trainer (better mem releasing)

* Update setup.py

* Update data.py

* fix CI

qlib/model/trainer.py

@@ -15,13 +15,22 @@ import socket
 from typing import Callable, List
 
 from tqdm.auto import tqdm
+
+from qlib.config import C
 from qlib.data.dataset import Dataset
+from qlib.data.dataset.weight import Reweighter
+from qlib.log import get_module_logger
 from qlib.model.base import Model
-from qlib.utils import flatten_dict, init_instance_by_config, auto_filter_kwargs, fill_placeholder
+from qlib.utils import (
+    auto_filter_kwargs,
+    fill_placeholder,
+    flatten_dict,
+    init_instance_by_config,
+)
+from qlib.utils.paral import call_in_subproc
 from qlib.workflow import R
 from qlib.workflow.recorder import Recorder
 from qlib.workflow.task.manage import TaskManager, run_task
-from qlib.data.dataset.weight import Reweighter
 
 
 def _log_task_info(task_config: dict):
@@ -210,17 +219,19 @@ class TrainerR(Trainer):
     STATUS_BEGIN = "begin_task_train"
     STATUS_END = "end_task_train"
 
-    def __init__(self, experiment_name: str = None, train_func: Callable = task_train):
+    def __init__(self, experiment_name: str = None, train_func: Callable = task_train, call_in_subproc: bool = False):
         """
         Init TrainerR.
 
         Args:
             experiment_name (str, optional): the default name of experiment.
             train_func (Callable, optional): default training method. Defaults to `task_train`.
+            call_in_subproc (bool): call the process in subprocess to force memory release
         """
         super().__init__()
         self.experiment_name = experiment_name
         self.train_func = train_func
+        self._call_in_subproc = call_in_subproc
 
     def train(self, tasks: list, train_func: Callable = None, experiment_name: str = None, **kwargs) -> List[Recorder]:
         """
@@ -245,6 +256,9 @@ class TrainerR(Trainer):
             experiment_name = self.experiment_name
         recs = []
         for task in tqdm(tasks, desc="train tasks"):
+            if self._call_in_subproc:
+                get_module_logger("TrainerR").info("running models in sub process (for forcing release memroy).")
+                train_func = call_in_subproc(train_func, C)
             rec = train_func(task, experiment_name, **kwargs)
             rec.set_tags(**{self.STATUS_KEY: self.STATUS_BEGIN})
             recs.append(rec)