Merge pull request #1 from WenWeiTHU/mlnode/forecast

Mlnode/forecast

Author: hangzhou188

mlnode/iotdb/mlnode/handler.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 #
-from iotdb.mlnode.algorithm.factory import create_forecast_model
+
 from iotdb.mlnode.constant import TSStatusCode
 from iotdb.mlnode.data_access.factory import create_forecast_dataset
 from iotdb.mlnode.parser import parse_training_request
@@ -51,13 +51,8 @@ def createTrainingTask(self, req: TCreateTrainingTaskReq):
             model_config['input_vars'] = data_config['input_vars']
             model_config['output_vars'] = data_config['output_vars']
 
-            # create model & check model config legitimacy
-            model, model_config = create_forecast_model(**model_config)
-
-            model_config['input_vars'] = data_config['input_vars']
-            model_config['output_vars'] = data_config['output_vars']
             # create task & check task config legitimacy
-            task = self.__task_manager.create_training_task(dataset, model, model_config, task_config)
+            task = self.__task_manager.create_training_task(dataset, data_config, model_config, task_config)
 
             return get_status(TSStatusCode.SUCCESS_STATUS)
         except Exception as e:

mlnode/iotdb/mlnode/process/manager.py

@@ -18,13 +18,10 @@
 
 import multiprocessing as mp
 
-from typing import Dict
-
-from torch import nn
+from typing import Dict, Union
 from torch.utils.data import Dataset
-
 from iotdb.mlnode.log import logger
-from iotdb.mlnode.process.task import ForecastingTrainingTask
+from iotdb.mlnode.process.task import ForecastingSingleTrainingTask, ForecastingTuningTrainingTask
 
 
 class TaskManager(object):
@@ -43,14 +40,14 @@ def __init__(self, pool_size: int):
 
     def create_training_task(self,
                              dataset: Dataset,
-                             model: nn.Module,
+                             data_configs: Dict,
                              model_configs: Dict,
-                             task_configs: Dict) -> ForecastingTrainingTask:
+                             task_configs: Dict):
         """
 
         Args:
             dataset: a torch dataset to be used for training
-            model: torch.nn.Module
+            data_configs: dict of data configurations
             model_configs: dict of model configurations
             task_configs: dict of task configurations
 
@@ -59,16 +56,27 @@ def create_training_task(self,
         """
         model_id = task_configs['model_id']
         self.__pid_info[model_id] = self.__shared_resource_manager.dict()
-        task = ForecastingTrainingTask(
-            task_configs,
-            model_configs,
-            model,
-            dataset,
-            self.__pid_info
-        )
+        if task_configs['tuning']:
+            task = ForecastingTuningTrainingTask(
+                task_configs,
+                model_configs,
+                self.__pid_info,
+                data_configs,
+                dataset,
+                model_id,
+            )
+        else:
+            task = ForecastingSingleTrainingTask(
+                task_configs,
+                model_configs,
+                self.__pid_info,
+                data_configs,
+                dataset,
+                model_id,
+            )
         return task
 
-    def submit_training_task(self, task: ForecastingTrainingTask) -> None:
+    def submit_training_task(self, task: Union[ForecastingTuningTrainingTask, ForecastingSingleTrainingTask]) -> None:
         if task is not None:
             self.__training_process_pool.apply_async(task, args=())
             logger.info(f'Task: ({task.model_id}) - Training process submitted successfully')

mlnode/iotdb/mlnode/process/task.py

@@ -18,42 +18,49 @@
 import multiprocessing
 import os
 from abc import abstractmethod
-from typing import Dict
+from typing import Dict, Tuple
 
 import optuna
 from torch import nn
 from torch.utils.data import Dataset
+from multiprocessing import Pipe
 
 from iotdb.mlnode.log import logger
 from iotdb.mlnode.process.trial import ForecastingTrainingTrial
 from iotdb.mlnode.algorithm.factory import create_forecast_model
-from iotdb.mlnode.client import client_manager
+from iotdb.mlnode.client import client_manager, ConfigNodeClient
 from iotdb.mlnode.config import descriptor
 from iotdb.thrift.common.ttypes import TrainingState
 
 
-class TrainingTrialObjective:
+class ForestingTrainingObjective:
     """
     A class which serve as a function, should accept trial as args
     and return the optimization objective.
     Optuna will try to minimize the objective.
     """
 
-    def __init__(self, trial_configs: Dict, model_configs: Dict, dataset: Dataset, pid_info: Dict):
+    def __init__(
+            self,
+            trial_configs: Dict,
+            model_configs: Dict,
+            dataset: Dataset,
+            # pid_info: Dict
+    ):
         self.trial_configs = trial_configs
         self.model_configs = model_configs
         self.dataset = dataset
-        self.pid_info = pid_info
+        # self.pid_info = pid_info
 
     def __call__(self, trial: optuna.Trial):
         # TODO: decide which parameters to tune
         trial_configs = self.trial_configs
         trial_configs['learning_rate'] = trial.suggest_float("lr", 1e-7, 1e-1, log=True)
         trial_configs['trial_id'] = 'tid_' + str(trial._trial_id)
         # TODO: check args
-        model, model_cfg = create_forecast_model(**self.model_configs)
-        self.pid_info[self.trial_configs['model_id']][trial._trial_id] = os.getpid()
-        _trial = ForecastingTrainingTrial(trial_configs, model, self.model_configs, self.dataset)
+        model, model_configs = create_forecast_model(**self.model_configs)
+        # self.pid_info[self.trial_configs['model_id']][trial._trial_id] = os.getpid()
+        _trial = ForecastingTrainingTrial(trial_configs, model, model_configs, self.dataset)
         loss = _trial.start()
         return loss
 
@@ -65,71 +72,163 @@ class _BasicTask(object):
     """
 
     def __init__(
-        self,
-        task_configs: Dict,
-        model_configs: Dict,
-        model: nn.Module,
-        dataset: Dataset,
-        pid_info: Dict
+            self,
+            task_configs: Dict,
+            model_configs: Dict,
+            pid_info: Dict
     ):
         """
         Args:
             task_configs:
             model_configs:
-            model:
-            dataset:
             pid_info:
         """
         self.pid_info = pid_info
         self.task_configs = task_configs
         self.model_configs = model_configs
-        self.model = model
+
+    @abstractmethod
+    def __call__(self):
+        raise NotImplementedError
+
+
+class _BasicTrainingTask(_BasicTask):
+    def __init__(
+            self,
+            task_configs: Dict,
+            model_configs: Dict,
+            pid_info: Dict,
+            data_configs: Dict,
+            dataset: Dataset,
+    ):
+        """
+        Args:
+            task_configs:
+            model_configs:
+            pid_info:
+            data_configs:
+            dataset:
+        """
+        super().__init__(task_configs, model_configs, pid_info)
+        self.data_configs = data_configs
         self.dataset = dataset
+        self.confignode_client = client_manager.borrow_config_node_client()
+
+    @abstractmethod
+    def __call__(self):
+        raise NotImplementedError
+
+
+class _BasicInferenceTask(_BasicTask):
+    def __int__(
+            self,
+            task_configs: Dict,
+            model_configs: Dict,
+            pid_info: Dict,
+            data: Tuple,
+    ):
+        super().__init__(task_configs, model_configs, pid_info)
+        self.data = data
+        self.model, self.model_configs = create_forecast_model(**self.model_configs)
 
     @abstractmethod
     def __call__(self):
         raise NotImplementedError
 
+    @abstractmethod
+    def data_align(self):
+        raise NotImplementedError
+
 
-class ForecastingTrainingTask(_BasicTask):
-    def __init__(self, task_configs: Dict, model_configs: Dict, model: nn.Module, dataset: Dataset,
-                 pid_info: Dict):
+class ForecastingSingleTrainingTask(_BasicTrainingTask):
+    def __init__(
+            self,
+            task_configs: Dict,
+            model_configs: Dict,
+            pid_info: Dict,
+            data_configs: Dict,
+            dataset: Dataset,
+            model_id: str,
+    ):
         """
         Args:
             task_configs: dict of task configurations
             model_configs: dict of model configurations
-            model: nn.Module
-            dataset: training dataset
             pid_info: a map shared between processes, can be used to find the pid with model_id and trial_id
+            data_configs: dict of data configurations
+            dataset: training dataset
         """
-        super(ForecastingTrainingTask, self).__init__(task_configs, model_configs, model, dataset, pid_info)
-        self.model_id = self.task_configs['model_id']
-        self.tuning = self.task_configs['tuning']
-        self.confignode_client = client_manager.borrow_config_node_client()
-
-        if self.tuning:
-            self.study = optuna.create_study(direction='minimize')
-        else:
-            self.default_trial_id = 'tid_0'
-            self.task_configs['trial_id'] = self.default_trial_id
-            self.trial = ForecastingTrainingTrial(self.task_configs, self.model, self.model_configs, self.dataset)
-            self.pid_info[self.model_id][self.default_trial_id] = os.getpid()
+        super().__init__(task_configs, model_configs, pid_info, data_configs, dataset)
+        self.model_id = model_id
+        self.default_trial_id = 'tid_0'
+        self.task_configs['trial_id'] = self.default_trial_id
+        model, model_configs = create_forecast_model(**model_configs)
+        self.trial = ForecastingTrainingTrial(task_configs, model, model_configs, dataset)
+        self.pid_info[self.model_id][self.default_trial_id] = os.getpid()
 
     def __call__(self):
         try:
-            if self.tuning:
-                self.study.optimize(TrainingTrialObjective(
-                    self.task_configs,
-                    self.model_configs,
-                    self.dataset,
-                    self.pid_info),
-                    n_trials=descriptor.get_config().get_mn_tuning_trial_num(),
-                    n_jobs=descriptor.get_config().get_mn_tuning_trial_concurrency())
-                best_trial_id = 'tid_' + str(self.study.best_trial._trial_id)
-                self.confignode_client.update_model_state(self.model_id, TrainingState.FINISHED, best_trial_id)
-            else:
-                self.trial.start()
-                self.confignode_client.update_model_state(self.model_id, TrainingState.FINISHED, self.default_trial_id)
+            self.trial.start()
+            self.confignode_client.update_model_state(self.model_id, TrainingState.FINISHED, self.default_trial_id)
         except Exception as e:
             logger.warn(e)
             raise e
+
+
+class ForecastingTuningTrainingTask(_BasicTrainingTask):
+    def __init__(
+            self,
+            task_configs: Dict,
+            model_configs: Dict,
+            pid_info: Dict,
+            data_configs: Dict,
+            dataset: Dataset,
+            model_id: str,
+    ):
+        """
+        Args:
+            task_configs: dict of task configurations
+            model_configs: dict of model configurations
+            pid_info: a map shared between processes, can be used to find the pid with model_id and trial_id
+            data_configs: dict of data configurations
+            dataset: training dataset
+        """
+        super().__init__(task_configs, model_configs, pid_info, data_configs, dataset)
+        self.model_id = model_id
+        self.study = optuna.create_study(direction='minimize')
+
+    def __call__(self):
+        # try:
+        self.study.optimize(ForestingTrainingObjective(
+            self.task_configs,
+            self.model_configs,
+            self.dataset),
+            n_trials=descriptor.get_config().get_mn_tuning_trial_num(),
+            n_jobs=descriptor.get_config().get_mn_tuning_trial_concurrency())
+        best_trial_id = 'tid_' + str(self.study.best_trial._trial_id)
+        self.confignode_client.update_model_state(self.model_id, TrainingState.FINISHED, best_trial_id)
+        #
+        # except Exception as e:
+        #     logger.warn(e)
+        #     raise e
+
+
+class ForecastingInferenceTask(_BasicInferenceTask):
+    def __int__(
+            self,
+            task_configs: Dict,
+            model_configs: Dict,
+            pid_info: Dict,
+            data: Tuple,
+            pipe: Pipe,
+    ):
+        super().__init__(task_configs, model_configs, pid_info, data)
+
+    def __call__(self):
+        pass
+
+    def data_align(self):
+        pass
+
+    def generate_future_mark(self):
+        pass