multiprocessing support

qlib/model/trainer.py

@@ -12,9 +12,11 @@ In ``DelayTrainer``, the first step is only to save some necessary info to model
 """
 
 import socket
+import time
 from typing import Callable, List
 
 from qlib.data.dataset import Dataset
+from qlib.log import get_module_logger
 from qlib.model.base import Model
 from qlib.utils import flatten_dict, get_cls_kwargs, init_instance_by_config
 from qlib.workflow import R
@@ -190,6 +192,8 @@ class TrainerR(Trainer):
         Returns:
             List[Recorder]: a list of Recorders
         """
+        if isinstance(tasks, dict):
+            tasks = [tasks]
         if len(tasks) == 0:
             return []
         if train_func is None:
@@ -213,6 +217,8 @@ class TrainerR(Trainer):
         Returns:
             List[Recorder]: the same list as the param.
         """
+        if isinstance(recs, Recorder):
+            recs = [recs]
         for rec in recs:
             rec.set_tags(**{self.STATUS_KEY: self.STATUS_END})
         return recs
@@ -250,6 +256,8 @@ class DelayTrainerR(TrainerR):
         Returns:
             List[Recorder]: a list of Recorders
         """
+        if isinstance(recs, Recorder):
+            recs = [recs]
         if end_train_func is None:
             end_train_func = self.end_train_func
         if experiment_name is None:
@@ -315,6 +323,8 @@ class TrainerRM(Trainer):
         Returns:
             List[Recorder]: a list of Recorders
         """
+        if isinstance(tasks, dict):
+            tasks = [tasks]
         if len(tasks) == 0:
             return []
         if train_func is None:
@@ -329,12 +339,24 @@ class TrainerRM(Trainer):
         run_task(
             train_func,
             task_pool,
+            query={"filter": {"$in": tasks}},  # only train these tasks
             experiment_name=experiment_name,
             before_status=before_status,
             after_status=after_status,
             **kwargs,
         )
 
+        # FIXME: reset to waiting automatically
+        for _id in _id_list:
+            is_prn = False
+            while tm.re_query(_id)["status"] == "running":
+                if not is_prn:
+                    get_module_logger("TrainerRM").warn(
+                        f"A task (_id: {_id}) is not being trained by this Trainer. Ignore this message if it is being trained by others."
+                    )
+                    is_prn = True
+                time.sleep(10)
+
         recs = []
         for _id in _id_list:
             rec = tm.re_query(_id)["res"]
@@ -352,10 +374,33 @@ class TrainerRM(Trainer):
         Returns:
             List[Recorder]: the same list as the param.
         """
+        if isinstance(recs, Recorder):
+            recs = [recs]
         for rec in recs:
             rec.set_tags(**{self.STATUS_KEY: self.STATUS_END})
         return recs
 
+    def worker(
+        self,
+        train_func: Callable = None,
+        experiment_name: str = None,
+    ):
+        """
+        The multiprocessing method for `train`. It can share a same task_pool with `train` and can run in other progress or other machines.
+
+        Args:
+            train_func (Callable): the training method which needs at least `task`s and `experiment_name`. None for the default training method.
+            experiment_name (str): the experiment name, None for use default name.
+        """
+        if train_func is None:
+            train_func = self.train_func
+        if experiment_name is None:
+            experiment_name = self.experiment_name
+        task_pool = self.task_pool
+        if task_pool is None:
+            task_pool = experiment_name
+        run_task(train_func, task_pool=task_pool, experiment_name=experiment_name)
+
 
 class DelayTrainerRM(TrainerRM):
     """
@@ -395,6 +440,8 @@ class DelayTrainerRM(TrainerRM):
         Returns:
             List[Recorder]: a list of Recorders
         """
+        if isinstance(tasks, dict):
+            tasks = [tasks]
         if len(tasks) == 0:
             return []
         return super().train(
@@ -410,8 +457,6 @@ class DelayTrainerRM(TrainerRM):
         Given a list of Recorder and return a list of trained Recorder.
         This class will finish real data loading and model fitting.
 
-        NOTE: This method will train all STATUS_PART_DONE tasks in the task pool, not only the ``recs``.
-
         Args:
             recs (list): a list of Recorder, the tasks have been saved to them.
             end_train_func (Callable, optional): the end_train method which need at least `recorder`s and `experiment_name`. Defaults to None for using self.end_train_func.
@@ -421,7 +466,8 @@ class DelayTrainerRM(TrainerRM):
         Returns:
             List[Recorder]: a list of Recorders
         """
-
+        if isinstance(recs, Recorder):
+            recs = [recs]
         if end_train_func is None:
             end_train_func = self.end_train_func
         if experiment_name is None:
@@ -441,6 +487,42 @@ class DelayTrainerRM(TrainerRM):
             before_status=TaskManager.STATUS_PART_DONE,
             **kwargs,
         )
+
+        # FIXME: reset to waiting automatically
+        tm = TaskManager(task_pool=task_pool)
+        for query_task in tm.query({"filter": {"$in": tasks}}):
+            _id = query_task["_id"]
+            is_prn = False
+            while tm.re_query(_id)["status"] == "running":
+                if not is_prn:
+                    get_module_logger("DelayTrainerRM").warn(
+                        f"A task (_id: {_id}) is not being trained by this Trainer. Ignore this message if it is being trained by others."
+                    )
+                    is_prn = True
+                time.sleep(10)
+
         for rec in recs:
             rec.set_tags(**{self.STATUS_KEY: self.STATUS_END})
         return recs
+
+    def worker(self, end_train_func=None, experiment_name: str = None):
+        """
+        The multiprocessing method for `end_train`. It can share a same task_pool with `end_train` and can run in other progress or other machines.
+
+        Args:
+            end_train_func (Callable, optional): the end_train method which need at least `recorder`s and `experiment_name`. Defaults to None for using self.end_train_func.
+            experiment_name (str): the experiment name, None for use default name.
+        """
+        if end_train_func is None:
+            end_train_func = self.end_train_func
+        if experiment_name is None:
+            experiment_name = self.experiment_name
+        task_pool = self.task_pool
+        if task_pool is None:
+            task_pool = experiment_name
+        run_task(
+            end_train_func,
+            task_pool=task_pool,
+            experiment_name=experiment_name,
+            before_status=TaskManager.STATUS_PART_DONE,
+        )