[IOTDB-5728] Implement config parser & model/dataset factory on MLNode

mlnode/iotdb/mlnode/exception.py

@@ -29,3 +29,13 @@ def __init__(self, node_url: str):
 class ModelNotExistError(_BaseError):
     def __init__(self, file_path: str):
         self.message = "Model path: ({}) not exists".format(file_path)
+
+
+class ModelNotSupportedError(_BaseError):
+    def __init__(self, model_name: str):
+        self.message = "Model: ({}) not supported".format(model_name)
+
+
+class MissingConfigError(_BaseError):
+    def __init__(self, config_name: str):
+        self.message = "Missing config: ({})".format(config_name)

mlnode/iotdb/mlnode/parser.py

@@ -0,0 +1,125 @@
+import argparse
+from iotdb.mlnode.exception import ModelNotSupportedError, MissingConfigError
+
+
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+
+class ConfigParser(object):
+    def __init__(self):
+        self.parser = argparse.ArgumentParser()
+        self.argument_list = []
+        self.required_list = []
+
+    def parse_configs(self, config):
+        conf_dict = vars(self.parser.parse_args([]))
+        args = []
+        for k in self.required_list:
+            if k not in config.keys():
+                raise MissingConfigError(k)
+        for k, v in config.items():
+            if k in conf_dict.keys():
+                args.append("--{}".format(k))
+                args.append(eval(v) if type(conf_dict[k]) is list and type(v) is not list else v)
+        return vars(self.parser.parse_args(args))
+
+
+class DataConfigParser(ConfigParser):
+    def __init__(self):
+        super().__init__()
+        self.parser.add_argument('--source_type', type=str)
+        self.parser.add_argument('--filename', type=str)
+        self.parser.add_argument('--dataset_type', type=str, default='window')
+        self.parser.add_argument('--time_embed', type=str, default='h')
+        self.parser.add_argument('--input_len', type=int, default=96)
+        self.parser.add_argument('--pred_len', type=int, default=96)
+        self.parser.add_argument('--input_vars', type=int, default=7)
+        self.parser.add_argument('--output_vars', type=int, default=7)
+        self.argument_list = ['source_type', 'filename', 'dataset_type', 'time_embed',
+                              'input_len', 'pred_len', 'input_vars', 'output_vars']
+        self.required_list = ['source_type', 'filename', 'input_vars', 'output_vars']
+
+
+class TaskConfigParser(ConfigParser):
+    def __init__(self):
+        super().__init__()
+        self.parser.add_argument('--model_id', type=str)
+        self.parser.add_argument('--tuning', type=bool, default=False)
+        self.parser.add_argument('--task_type', type=str, default='m')
+        self.parser.add_argument('--task_class', type=str, default='forecast_training_task')
+        self.parser.add_argument('--input_len', type=int, default=96)
+        self.parser.add_argument('--pred_len', type=int, default=96)
+        self.parser.add_argument('--input_vars', type=int, default=7)
+        self.parser.add_argument('--output_vars', type=int, default=7)
+        self.parser.add_argument('--learning_rate', type=float, default=0.0001)
+        self.parser.add_argument('--batch_size', type=int, default=32)
+        self.parser.add_argument('--num_workers', type=int, default=0)
+        self.parser.add_argument('--epochs', type=int, default=10)
+        self.parser.add_argument('--use_gpu', type=bool, default=False)
+        self.parser.add_argument('--gpu', type=int, default=0)
+        self.parser.add_argument('--use_multi_gpu', type=bool, default=False)
+        self.parser.add_argument('--devices', type=list, default=[0])
+        self.parser.add_argument('--metric_names', type=list, default=['MSE', 'MAE'])
+        self.argument_list = ['model_id', 'tuning', 'task_type', 'task_class', 'input_len', 'pred_len',
+                              'input_vars', 'output_vars', 'learning_rate', 'batch_size', 'num_workers',
+                              'epochs', 'use_gpu', 'gpu', 'use_multi_gpu', 'devices', 'metric_names']
+        self.required_list = ['model_id', 'input_vars', 'output_vars']
+
+
+class DLinearConfigParser(ConfigParser):
+    def __init__(self):
+        super().__init__()
+        self.parser.add_argument('--model_name', type=str, default='dlinear')
+        self.parser.add_argument('--input_len', type=int, default=96)
+        self.parser.add_argument('--pred_len', type=int, default=96)
+        self.parser.add_argument('--input_vars', type=int, default=7)
+        self.parser.add_argument('--output_vars', type=int, default=7)
+        self.parser.add_argument('--task_type', type=str, default='m')
+        self.parser.add_argument('--kernel_size', type=int, default=25)
+        self.parser.add_argument('--block_type', type=str, default='g')
+        self.parser.add_argument('--d_model', type=int, default=128)
+        self.parser.add_argument('--inner_layers', type=int, default=4)
+        self.parser.add_argument('--outer_layers', type=int, default=4)
+        self.argument_list = ['model_name', 'input_len', 'pred_len', 'input_vars', 'output_vars',
+                              'task_type', 'kernel_size', 'block_type', 'd_model', 'inner_layers', 'outer_layers']
+        self.required_list = ['model_name', 'input_vars', 'output_vars']
+
+
+class ConfigManager(object):
+    def __init__(self, config):
+        self.config = config
+        if 'model_name' not in config:
+            raise MissingConfigError('model_name')
+        self.data_conf_parser = DataConfigParser()
+        self.task_conf_parser = TaskConfigParser()
+        self.model_conf_parser = get_model_config_parser(config['model_name'])
+
+    def get_configs(self):
+        return self.data_conf_parser.parse_configs(self.config), \
+            self.task_conf_parser.parse_configs(self.config), \
+            self.model_conf_parser.parse_configs(self.config)
+
+
+def get_model_config_parser(model_name):
+    parser = None
+    if model_name == 'dlinear':
+        parser = DLinearConfigParser()
+    else:
+        raise ModelNotSupportedError(model_name)
+    return parser

mlnode/iotdb/mlnode/util.py

@@ -15,20 +15,21 @@
 # specific language governing permissions and limitations
 # under the License.
 #
+
+
 from iotdb.mlnode.exception import BadNodeUrlError
-from iotdb.mlnode.log import logger
 from iotdb.thrift.common.ttypes import TEndPoint
+from iotdb.thrift.mlnode.ttypes import TCreateTrainingTaskReq
+from iotdb.mlnode.log import logger
+from iotdb.mlnode.parser import ConfigManager
 
 
 def parse_endpoint_url(endpoint_url: str) -> TEndPoint:
     """ Parse TEndPoint from a given endpoint url.
-
     Args:
         endpoint_url: an endpoint url, format: ip:port
-
     Returns:
         TEndPoint
-
     Raises:
         BadNodeUrlError
     """
@@ -45,3 +46,21 @@ def parse_endpoint_url(endpoint_url: str) -> TEndPoint:
     except ValueError as e:
         logger.warning("Illegal endpoint url format: {} ({})".format(endpoint_url, e))
         raise BadNodeUrlError(endpoint_url)
+
+
+# TODO: may have many bug
+def parse_training_request(req: TCreateTrainingTaskReq):
+    """
+    Parse TCreateTrainingTaskReq with given yaml template
+    Args:
+        req: TCreateTrainingTaskReq
+    Returns:
+        data_conf: configurations related to data
+        model_conf: configurations related to model
+        task_conf: configurations related to task
+    """
+    config = req.modelConfigs
+    config.update(model_id=req.modelId)
+    config.update(tuning=req.isAuto)
+    config_manager = ConfigManager(config)
+    return config_manager.get_configs()

mlnode/test/test_parse_training_request.py

@@ -0,0 +1,57 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+from iotdb.mlnode.util import parse_training_request
+from iotdb.thrift.mlnode.ttypes import TCreateTrainingTaskReq
+
+
+def test_parse_training_request():
+    modelId = 'mid_etth1_dlinear_default'
+    isAuto = False
+    modelConfigs = {
+        'task_class': 'forecast_training_task',
+        'source_type': 'thrift',
+        'dataset_type': 'window',
+        'filename': 'ETTh1.csv',
+        'time_embed': 'h',
+        'input_len': 96,
+        'pred_len': 96,
+        'model_name': 'dlinear',
+        'input_vars': 7,
+        'output_vars': 7,
+        'task_type': 'm',
+        'kernel_size': 25,
+        'learning_rate': 1e-3,
+        'batch_size': 32,
+        'num_workers': 0,
+        'epochs': 10,
+        'metric_names': ['MSE', 'MAE']
+    }
+    req = TCreateTrainingTaskReq(
+        modelId=str(modelId),
+        isAuto=isAuto,
+        modelConfigs={k: str(v) for k, v in modelConfigs.items()},
+    )
+    data_conf, model_conf, task_conf = parse_training_request(req)
+    for config in modelConfigs:
+        if config in data_conf:
+            assert data_conf[config] == modelConfigs[config]
+        if config in model_conf:
+            assert model_conf[config] == modelConfigs[config]
+        if config in task_conf:
+            assert task_conf[config] == modelConfigs[config]