[Merged by Bors] - Optimize the use of model compiler

README.md

@@ -149,9 +149,9 @@ Assume builing with CUDA version 10.0.
    - `cuda-nvml-dev-10-0`
    - `libcudnn7=*+cuda10.0`
    - `libcudnn7-dev=*+cuda10.0`
-   - `libnvinfer6=*+cuda10.0`
+   - `libnvinfer7=*+cuda10.0`
    - `libnvinfer-dev=*+cuda10.0`
-   - `libnvonnxparsers6=*+cuda10.0`
+   - `libnvonnxparsers7=*+cuda10.0`
    - `libnvonnxparsers-dev=*+cuda10.0`
 2. Run the following command:
 

ci/azure-pipelines/install-tensorrt.yml

@@ -15,8 +15,8 @@ steps:
             cuda-nvml-dev-10-0 \
             'libcudnn7=*+cuda10.0' \
             'libcudnn7-dev=*+cuda10.0' \
-            'libnvinfer6=*+cuda10.0' \
+            'libnvinfer7=*+cuda10.0' \
             'libnvinfer-dev=*+cuda10.0' \
-            'libnvonnxparsers6=*+cuda10.0' \
+            'libnvonnxparsers7=*+cuda10.0' \
             'libnvonnxparsers-dev=*+cuda10.0'
     displayName: Install TensorRT libraries

model_compiler/setup.py

@@ -18,7 +18,8 @@
     'torch==1.1.0',
     'torchvision==0.3.0',
     'requests',
-    'tensorflow==1.14.0'
+    'tensorflow==1.14.0',
+    'jsonschema==3.1.1',
 ]
 
 _TEST_REQUIRES = [
@@ -39,4 +40,8 @@
     keywords='compile serving model',
     install_requires=_REQUIRED_PACKAGES,
     extras_require={'test': _TEST_REQUIRES},
+    package_data={
+        'model_compiler': ['*.json']
+    },
+
 )

model_compiler/src/model_compiler/__init__.py

@@ -1,2 +1,9 @@
 # Copyright 2019 ZTE corporation. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
+
+"""
+Model compiler package
+"""
+
+# flake8:noqa
+from .mo import compile_model

model_compiler/src/model_compiler/config.py

@@ -2,20 +2,20 @@
 # SPDX-License-Identifier: Apache-2.0
 
 """
-Message from proxy definition
+Request message definition
 """
 import json
 import os
 
+import jsonschema
+
 from .log_util import get_logger
 
 _LOGGER = get_logger(__name__)
 
 
 class _Config:
     def __init__(self, **kwargs):
-        if 'job_id' not in kwargs:
-            raise Exception('Not found job id in config, config: %s' % kwargs)
         self._config = kwargs
         _LOGGER.info('Dump config: %s', self._config)
 
@@ -61,38 +61,48 @@ def _env_to_list(name):
     max_batch_size = _getenv('max_batch_size', 1, int)
     checkpoint_path = _getenv('checkpoint_path')
     frozen_graph_path = _getenv('frozen_graph_path')
+    version = _getenv('version_number', None)
 
     return _Config(serving_type=serving_type, h5_path=h5_path, script_path=script_path, export_path=export_path,
                    onnx_path=onnx_path, input_layer_names=input_layer_names, output_layer_names=output_layer_names,
                    input_signatures=input_signatures, output_signatures=output_signatures, model_name=model_name,
                    max_batch_size=max_batch_size, job_id=job_id, callback=callback, checkpoint_path=checkpoint_path,
                    frozen_graph_path=frozen_graph_path, input_names=input_names, input_formats=input_formats,
-                   output_names=output_names)
+                   output_names=output_names, version=version)
 
 
-def _create_config_from_json(json_file):
+def _create_config_from_file(json_file):
     with open(json_file) as input_file:
         json_message = json.load(input_file)
     return _Config(**json_message)
 
 
-def create_config(from_source='env', json_file=None) -> object:
+def create_config(from_source='env', source=None) -> object:
     """
     Create message config, default from environment variables
-    :param from_source: from where parse message, e.g. 'env' or 'json'
-    :param json_file: json message
+    :param from_source: from where parse input message, e.g. 'env', 'file', 'object'
+    :param source: json message
     :return:
     """
+    if from_source not in ['env', 'file', 'object']:
+        raise IOError('Create config error, unsupported source {}'.format(from_source))
     if from_source == 'env':
         return _create_config_from_env()
+    elif from_source == 'json':
+        return _create_config_from_file(source)
     else:
-        return _create_config_from_json(json_file)
+        return create_config_from_obj(**source)
 
 
 def create_config_from_obj(obj) -> object:
     """
-    Create message config from a dictionary
+    Create message config from a dictionary which must match config_schema
     :param obj: dict
     :return:
     """
+    schema_path = os.path.join(os.path.dirname(__file__), 'config_schema.json')
+    with open(schema_path) as schema_file:
+        body_schema = json.load(schema_file)
+
+    jsonschema.validate(obj, body_schema)
     return _Config(**obj)

model_compiler/src/model_compiler/config_schema.json

@@ -0,0 +1,101 @@
+{
+  "$schema": "http://json-schema.org/draft-04/schema",
+  "properties": {
+    "serving_type": {
+      "type": "string",
+      "enum": [
+        "tf",
+        "tensorrt",
+        "openvino"
+      ],
+      "description": "target serving model type"
+    },
+    "model_name": {
+      "type": "string",
+      "description": "output model name"
+    },
+    "version": {
+      "type": "integer",
+      "minimum": 1,
+      "description": "version number of output model"
+    },
+    "max_batch_size": {
+      "type": "integer",
+      "minimum": 1,
+      "description": "maximum batch size of output serving model can support"
+    },
+    "input_model": {
+      "type": "string",
+      "description": "file path of a pre-trained model which can be Keras h5 model(*.h5), TensorFlow checkpoint or frozen graph(*.pb)"
+    },
+    "script_path": {
+      "type": "string",
+      "description": "file path of a model script which describes the model structure especially custom layers or operations , optional"
+    },
+    "input_layer_names": {
+      "type": "array",
+      "items": {
+        "type": "string"
+      },
+      "description": "required when input model is Keras h5 model"
+    },
+    "output_layer_names": {
+      "type": "array",
+      "items": {
+        "type": "string"
+      },
+      "description": "required when input model is Keras h5 model"
+    },
+    "input_formats": {
+      "type": "array",
+      "items": {
+        "enum": [
+          "channels_last",
+          "channels_first",
+          null
+        ]
+      },
+      "description": "required when input is image"
+    },
+    "input_names": {
+      "type": "array",
+      "items": {
+        "type": "string"
+      },
+      "description": "input tensor names, required when input model is TensorFlow checkpoint or frozen graph"
+    },
+    "output_names": {
+      "type": "array",
+      "items": {
+        "type": "string"
+      },
+      "description": "output tensor names, required when input model is TensorFlow checkpoint or frozen graph"
+    },
+    "input_signatures": {
+      "type": "array",
+      "items": {
+        "type": "string"
+      },
+      "description": "alias of inputs, required when serving_type is tf"
+    },
+    "output_signatures": {
+      "type": "array",
+      "items": {
+        "type": "string"
+      },
+      "description": "alias of outputs, required when serving_type is tf"
+    },
+    "export_path": {
+      "type": "string",
+      "description": "path of output model"
+    }
+  },
+  "required": [
+    "serving_type",
+    "model_name",
+    "max_batch_size",
+    "input_model",
+    "export_path"
+  ],
+  "additionalProperties": true
+}

model_compiler/src/model_compiler/mo.py

@@ -0,0 +1,19 @@
+# Copyright 2019 ZTE corporation. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Entry for model compiler.
+"""
+
+from .config import create_config_from_obj
+from .runtime import create_compiler
+
+
+def compile_model(request):
+    """
+    Compile serving model
+    :param request: dict, must match config_schema.json
+    :return:
+    """
+    compiler = create_compiler(create_config_from_obj(request))
+    return compiler.compile()