NGTH_STDEV,
                                         sequence_name=name)
                    last_choices[client_idx] = "valid-valid"
                else:
                    sequence_valid(client_metadata,
                                   rng,
                                   trial,
                                   model_name,
                                   dtype,
                                   SEQUENCE_LENGTH_MEAN,
                                   SEQUENCE_LENGTH_STDEV,
                                   sequence_name=name)
                    last_choices[client_idx] = "valid"

    except Exception as ex:
        _thread_exceptions_mutex.acquire()
        try:
            _thread_exceptions.append(traceback.format_exc())
        finally:
            _thread_exceptions_mutex.release()

    # We need to explicitly close each client so that streams get
    # cleaned up and closed correctly, otherwise the application
    # can hang when exiting.
    for c, i in client_metadata_list:
        print("thread {} closing client {}".format(name, i))
        c.close()

    print("Exiting thread {}".format(name))


def check_status(model_name):
    client = grpcclient.InferenceServerClient("localhost:8001",
                                              verbose=FLAGS.verbose)
    stats = client.get_inference_statistics(model_name)
    print(stats)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-v',
                        '--verbose',
                        action="store_true",
                        required=False,
                        default=False,
                        help='Enable verbose output')
    parser.add_argument('-r',
                        '--random-seed',
                        type=int,
                        required=False,
                        help='Random seed.')
    parser.add_argument('-t',
                        '--concurrency',
                        type=int,
                        required=False,
                        default=8,
                        help='Request concurrency. Default is 8.')
    parser.add_argument(
        '-i',
        '--iterations',
        type=int,
        required=False,
        default=200,
        help='Number of iterations of stress test to run. Default is 200.')
    FLAGS = parser.parse_args()

    # Initialize the random seed. For reproducibility each thread
    # maintains its own RNG which is initialized based on this seed.
    randseed = 0
    if FLAGS.random_seed != None:
        randseed = FLAGS.random_seed
    else:
        randseed = int(time.time())
    np.random.seed(randseed)

    print("random seed = {}".format(randseed))
    print("concurrency = {}".format(FLAGS.concurrency))
    print("iterations = {}".format(FLAGS.iterations))

    trial = "custom"
    dtype = get_datatype(trial)
    model_name = tu.get_sequence_model_name(trial, dtype)

    threads = []
    for idx, thd in enumerate(range(FLAGS.concurrency)):
        thread_name = "thread_{}".format(idx)

        # Create the seed for the thread. Since these are created in
        # reproducible order off of the initial seed we will get
        # reproducible results when given the same seed.
        seed = np.random.randint(2**32)

        # Each thread is reserved a block of correlation IDs or size
        # CORRELATION_ID_BLOCK_SIZE
        correlation_id_base = 1 + (idx * CORRELATION_ID_BLOCK_SIZE)

        threads.append(
            threading.Thread(target=stress_thread,
                             args=(thread_name, seed, FLAGS.iterations,
                                   correlation_id_base, trial, model_name,
                                   dtype)))

    for t in threads:
        t.start()
    for t in threads:
        t.join()

    check_status(model_name)

    _thread_exceptions_mutex.acquire()
    try:
        if len(_thread_exceptions) > 0:
            for ex in _thread_exceptions:
                print("*********\n{}".format(ex))
            sys.exit(1)
    finally:
        _thread_exceptions_mutex.release()

    print("Exiting stress test")
    sys.exit(0)


# qa/L0_server_status/server_status_test.py

# Copyright 2018-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../common")

import numpy as np
import os
import unittest
import infer_util as iu
import test_util as tu
import tritongrpcclient as grpcclient
import tritonhttpclient as httpclient
from tritonclientutils import *


class ServerMetadataTest(tu.TestResultCollector):

    def test_basic(self):
        try:
            for pair in [("localhost:8000", "http"),
                         ("localhost:8001", "grpc")]:
                model_name = "graphdef_int32_int8_int8"
                extensions = [
                    'classification', 'sequence', 'model_repository',
                    'schedule_policy', 'model_configuration',
                    'system_shared_memory', 'cuda_shared_memory',
                    'binary_tensor_data', 'statistics'
                ]
                if pair[1] == "http":
                    triton_client = httpclient.InferenceServerClient(
                        url=pair[0], verbose=True)
                else:
                    triton_client = grpcclient.InferenceServerClient(
                        url=pair[0], verbose=True)

                self.assertTrue(triton_client.is_server_live())
                self.assertTrue(triton_client.is_server_ready())
                server_metadata = triton_client.get_server_metadata()
                model_metadata = triton_client.get_model_metadata(model_name)

                if pair[1] == "http":
                    self.assertEqual(os.environ["TRITON_SERVER_VERSION"],
                                     server_metadata['version'])
                    self.assertEqual("triton", server_metadata['name'])
                    for ext in extensions:
                        self.assertIn(ext, server_metadata['extensions'])

                    self.assertEqual(model_name, model_metadata['name'])
                else:
                    self.assertEqual(os.environ["TRITON_SERVER_VERSION"],
                                     server_metadata.version)
                    self.assertEqual("triton", server_metadata.name)
                    for ext in extensions:
                        self.assertIn(ext, server_metadata.extensions)

                    self.assertEqual(model_name, model_metadata.name)
        except InferenceServerException as ex:
            self.assertTrue(False, "unexpected error {}".format(ex))

    def test_unknown_model(self):
        try:
            for pair in [("localhost:8000", "http"),
                         ("localhost:8001", "grpc")]:
                model_name = "foo"
                if pair[1] == "http":
                    triton_client = httpclient.InferenceServerClient(
                        url=pair[0], verbose=True)
                else:
                    triton_client = grpcclient.InferenceServerClient(
                        url=pair[0], verbose=True)

                self.assertTrue(triton_client.is_server_live())
                self.assertTrue(triton_client.is_server_ready())
                server_metadata = triton_client.get_server_metadata()
                if pair[1] == "http":
                    self.assertEqual(os.environ["TRITON_SERVER_VERSION"],
                                     server_metadata['version'])
                    self.assertEqual("triton", server_metadata['name'])
                else:
                    self.assertEqual(os.environ["TRITON_SERVER_VERSION"],
                                     server_metadata.version)
                    self.assertEqual("triton", server_metadata.name)

                model_metadata = triton_client.get_model_metadata(model_name)
                self.assertTrue(False, "expected unknown model failure")
        except InferenceServerException as ex:
            self.assertTrue(ex.message().startswith(
                "Request for unknown model: 'foo' is not found"))

    def test_unknown_model_version(self):
        try:
            for pair in [("localhost:8000", "http"),
                         ("localhost:8001", "grpc")]:
                model_name = "graphdef_int32_int8_int8"
                if pair[1] == "http":
                    triton_client = httpclient.InferenceServerClient(
                        url=pair[0], verbose=True)
                else:
                    triton_client = grpcclient.InferenceServerClient(
                        url=pair[0], verbose=True)

                self.assertTrue(triton_client.is_server_live())
                self.assertTrue(triton_client.is_server_ready())

                model_metadata = triton_client.get_model_metadata(
                    model_name, model_version="99")
                self.assertTrue(False, "expected unknown model version failure")
        except InferenceServerException as ex:
            self.assertTrue(ex.message().startswith(
                "Request for unknown model: 'graphdef_int32_int8_int8' version 99 is not found"
            ))

    def test_model_latest_infer(self):
        input_size = 16
        tensor_shape = (1, input_size)
        platform_name = {
            'graphdef': 'tensorflow_graphdef',
            'onnx': 'onnxruntime_onnx'
        }

        # There are 3 versions of *_int32_int32_int32 and all
        # should be available.
        for platform in ('graphdef', 'onnx'):
            model_name = platform + "_int32_int32_int32"

            # Initially there should be no version stats..
            try:
                for pair in [("localhost:8000", "http"),
                             ("localhost:8001", "grpc")]:
                    if pair[1] == "http":
                        triton_client = httpclient.InferenceServerClient(
                            url=pair[0], verbose=True)
                    else:
                        triton_client = grpcclient.InferenceServerClient(
                            url=pair[0], verbose=True)

                    self.assertTrue(triton_client.is_server_live())
                    self.assertTrue(triton_client.is_server_ready())
                    model_metadata = triton_client.get_model_metadata(
                        model_name)
                    # verify all versions are reported when no model version is specified
                    if pair[1] == "http":
                        self.assertEqual(model_name, model_metadata['name'])
                        self.assertEqual(len(model_metadata['versions']), 3)
                        for v in (1, 2, 3):
                            self.assertIn(str(v), model_metadata['versions'])
                    else:
                        self.assertEqual(model_name, model_metadata.name)
                        self.assertEqual(len(model_metadata.versions), 3)
                        for v in (1, 2, 3):
                            self.assertIn(str(v), model_metadata.versions)

                    # verify contents of model metadata
                    if pair[1] == "http":
                        model_platform = model_metadata['platform']
                        model_inputs = model_metadata['inputs']
                        model_outputs = model_metadata['outputs']
                    else:
                        model_platform = model_metadata.platform
                        model_inputs = model_metadata.inputs
                        model_outputs = model_metadata.outputs

                    self.assertEqual(platform_name[platform], model_platform)
                    self.assertEqual(len(model_inputs), 2)
                    self.assertEqual(len(model_outputs), 2)

                    for model_input in model_inputs:
                        if pair[1] == "http":
                            input_dtype = model_input['datatype']
                            input_shape = model_input['shape']
                            input_name = model_input['name']
                        else:
                            input_dtype = model_input.datatype
                            input_shape = model_input.shape
                            input_name = model_input.name
                        self.assertIn(input_name, ["INPUT0", "INPUT1"])
                        self.assertEqual(input_dtype, "INT32")
                        self.assertEqual(input_shape, [-1, 16])

                    for model_output in model_outputs:
                        if pair[1] == "http":
                            output_dtype = model_output['datatype']
                            output_shape = model_output['shape']
                            output_name = model_output['name']
                        else:
                            output_dtype = model_output.datatype
                            output_shape = model_output.shape
                            output_name = model_output.name
                        self.assertIn(output_name, ["OUTPUT0", "OUTPUT1"])
                        self.assertEqual(output_dtype, "INT32")
                        self.assertEqual(output_shape, [-1, 16])

            except InferenceServerException as ex:
                self.assertTrue(False, "unexpected error {}".format(ex))

            # Infer using latest version (which is 3)...
            iu.infer_exact(self,
                           platform,
                           tensor_shape,
                           1,
                           np.int32,
                           np.int32,
                           np.int32,
                           model_version=None,
                           swap=True)

            try:
                for pair in [("localhost:8000", "http"),
                             ("localhost:8001", "grpc")]:
                    if pair[1] == "http":
                        triton_client = httpclient.InferenceServerClient(
                            url=pair[0], verbose=True)
                    else:
                        triton_client = grpcclient.InferenceServerClient(
                            url=pair[0], verbose=True)

                    self.assertTrue(triton_client.is_server_live())
                    self.assertTrue(triton_client.is_server_ready())
                    for v in (1, 2, 3):
                        self.assertTrue(
                            triton_client.is_model_ready(model_name,
                                                         model_version=str(v)))

                    # Only version 3 should have infer stats
                    infer_stats = triton_client.get_inference_statistics(
                        model_name)
                    if pair[1] == "http":
                        stats = infer_stats['model_stats']
                    else:
                        stats = infer_stats.model_stats
                    self.assertEqual(
                        len(stats), 3,
                        "expected 3 infer stats for model " + model_name)
                    for s in stats:
                        if pair[1] == "http":
                            v = s['version']
                            stat = s['inference_stats']
                        else:
                            v = s.version
                            stat = s.inference_stats

                        if v == "3":
                            if pair[1] == "http":
                                self.assertTrue(stat['success']['count'], 3)
                            else:
                                self.assertTrue(stat.success.count, 3)
                        else:
                            if pair[1] == "http":
                                self.assertEqual(
                                    stat['success']['count'], 0,
                                    "unexpected infer success counts for version "
                                    + str(v) + " of model " + model_name)
                            else:
                                self.assertEqual(
                                    stat.success.count, 0,
                                    "unexpected infer success counts for version "
                                    + str(v) + " of model " + model_name)

            except InferenceServerException as ex:
                self.assertTrue(False, "unexpected error {}".format(ex))

    def test_model_specific_infer(self):
        input_size = 16

        # There are 3 versions of *_float32_float32_float32 but only
        # versions 1 and 3 should be available.
        for platform in ('graphdef', 'onnx', 'plan'):
            tensor_shape = (1, input_size)
            model_name = platform + "_float32_float32_float32"

            # Initially there should be no version status...
            try:
                for pair in [("localhost:8000", "http"),
                             ("localhost:8001", "grpc")]:
                    if pair[1] == "http":
                        triton_client = httpclient.InferenceServerClient(
                            url=pair[0], verbose=True)
                    else:
                        triton_client = grpcclient.InferenceServerClient(
                            url=pair[0], verbose=True)

                    self.assertTrue(triton_client.is_server_live())
                    self.assertTrue(triton_client.is_server_ready())
                    self.assertTrue(
                        triton_client.is_model_ready(model_name,
                                                     model_version="1"))
                    self.assertFalse(
                        triton_client.is_model_ready(model_name,
                                                     model_version="2"))
                    self.assertTrue(
                        triton_client.is_model_ready(model_name,
                                                     model_version="3"))
            except InferenceServerException as ex:
                self.assertTrue(False, "unexpected error {}".format(ex))

            # Infer using version 1...
            iu.infer_exact(self,
                           platform,
                           tensor_shape,
                           1,
                           np.float32,
                           np.float32,
                           np.float32,
                           model_version=1,
                           swap=False)

            try:
                for pair in [("localhost:8000", "http"),
                             ("localhost:8001", "grpc")]:
                    if pair[1] == "http":
                        triton_client = httpclient.InferenceServerClient(
                            url=pair[0], verbose=True)
                    else:
                        triton_client = grpcclient.InferenceServerClient(
                            url=pair[0], verbose=True)

                    self.assertTrue(triton_client.is_server_live())
                    self.assertTrue(triton_client.is_server_ready())
                    self.assertTrue(
                        triton_client.is_model_ready(model_name,
                                                     model_version="1"))
                    self.assertFalse(
                        triton_client.is_model_ready(model_name,
                                                     model_version="2"))
                    self.assertTrue(
                        triton_client.is_model_ready(model_name,
                                                     model_version="3"))

                    # Only version 1 should have infer stats
                    infer_stats = triton_client.get_inference_statistics(
                        model_name, model_version='1')
                    if pair[1] == "http":
                        self.assertEqual(
                            len(infer_stats['model_stats']), 1,
                            "expected 1 infer stats for version 1"
                            " of model " + model_name)
                        stats = infer_stats['model_stats'][0]['inference_stats']
                        self.assertTrue(stats['success']['count'], 3)
                    else:
                        self.assertEqual(
                            len(infer_stats.model_stats), 1,
                            "expected 1 infer stats for version 1"
                            " of model " + model_name)
                        stats = infer_stats.model_stats[0].inference_stats
                        self.assertTrue(stats.success.count, 3)
                    infer_stats = triton_client.get_inference_statistics(
                        model_name, model_version='3')
                    if pair[1] == "http":
                        stats = infer_stats['model_stats'][0]['inference_stats']
                        self.assertEqual(
                            stats['success']['count'], 0,
                            "unexpected infer stats for version 3"
                            " of model " + model_name)
                    else:
                        stats = infer_stats.model_stats[0].inference_stats
                        self.assertEqual(
                            stats.success.count, 0,
                            "unexpected infer stats for version 3"
                            " of model " + model_name)

            except InferenceServerException as ex:
                self.assertTrue(False, "unexpected error {}".format(ex))


class ModelMetadataTest(tu.TestResultCollector):
    '''
    These tests must be run after the ServerMetadataTest. See test.sh
    file for correct test running.
    '''

    def test_model_versions_deleted(self):
        # Originally There were 3 versions of *_int32_int32_int32 and
        # version 3 was executed once. Version 2 and 3 models were
        # deleted from the model repository so now only expect version 1 to
        # be ready and show stats.
        for platform in ('graphdef', 'onnx'):
            model_name = platform + "_int32_int32_int32"

            try:
                for pair in [("localhost:8000", "http"),
                             ("localhost:8001", "grpc")]:
                    if pair[1] == "http":
                        triton_client = httpclient.InferenceServerClient(
                            url=pair[0], verbose=True)
                    else:
                        triton_client = grpcclient.InferenceServerClient(
                            url=pair[0], verbose=True)

                    self.assertTrue(triton_client.is_server_live())
                    self.assertTrue(triton_client.is_server_ready())
                    model_metadata = triton_client.get_model_metadata(
                        model_name)
                    if pair[1] == "http":
                        self.assertEqual(model_name, model_metadata['name'])
                        self.assertEqual(len(model_metadata['versions']), 1)
                        self.assertEqual("1", model_metadata['versions'][0])
                    else:
                        self.assertEqual(model_name, model_metadata.name)
                        self.assertEqual(len(model_metadata.versions), 1)
                        self.assertEqual("1", model_metadata.versions[0])

                    # Only version 3 should have infer stats, only 1 is ready
                    for v in (1, 2, 3):
                        if v == 1:
                            self.assertTrue(
                                triton_client.is_model_ready(
                                    model_name, model_version=str(v)))
                            infer_stats = triton_client.get_inference_statistics(
                                model_name, model_version=str(v))
                            if pair[1] == "http":
                                self.assertEqual(
                                    len(infer_stats['model_stats']), 1,
                                    "expected 1 infer stats for version " +
                                    str(v) + " of model " + model_name)
                                stats = infer_stats['model_stats'][0][
                                    'inference_stats']
                                self.assertEqual(stats['success']['count'], 0)
                            else:
                                self.assertEqual(
                                    len(infer_stats.model_stats), 1,
                                    "expected 1 infer stats for version " +
                                    str(v) + " of model " + model_name)
                                stats = infer_stats.model_stats[
                                    0].inference_stats
                                self.assertEqual(stats.success.count, 0)

                        else:
                            self.assertFalse(
                                triton_client.is_model_ready(
                                    model_name, model_version=str(v)))

            except InferenceServerException as ex:
                self.assertTrue(False, "unexpected error {}".format(ex))

    def test_model_versions_added(self):
        # Originally There was version 1 of *_float16_float32_float32.
        # Version 7 was added so now expect just version 7 to be ready
        # and provide infer stats.
        for platform in ('graphdef',):
            model_name = platform + "_float16_float32_float32"

            try:
                for pair in [("localhost:8000", "http"),
                             ("localhost:8001", "grpc")]:
                    if pair[1] == "http":
                        triton_client = httpclient.InferenceServerClient(
                            url=pair[0], verbose=True)
                    else:
                        triton_client = grpcclient.InferenceServerClient(
                            url=pair[0], verbose=True)

                    self.assertTrue(triton_client.is_server_live())
                    self.assertTrue(triton_client.is_server_ready())
                    model_metadata = triton_client.get_model_metadata(
                        model_name)
                    if pair[1] == "http":
                        self.assertEqual(
                            model_name, model_metadata['name'],
                            "expected status for model " + model_name)
                        self.assertEqual(
                            len(model_metadata['versions']), 1,
                            "expected status for 1 versions for model " +
                            model_name)
                        self.assertEqual("7", model_metadata['versions'][0])
                    else:
                        self.assertEqual(
                            model_name, model_metadata.name,
                            "expected status for model " + model_name)
                        self.assertEqual(
                            len(model_metadata.versions), 1,
                            "expected status for 1 versions for model " +
                            model_name)
                        self.assertEqual("7", model_metadata.versions[0])

                    # Only version 7 should be ready and show infer stat.
                    for v in (1, 7):
                        if v == 7:
                            self.assertTrue(
                                triton_client.is_model_ready(
                                    model_name, model_version=str(v)))
                            infer_stats = triton_client.get_inference_statistics(
                                model_name, model_version=str(v))
                            if pair[1] == "http":
                                stats = infer_stats['model_stats'][0][
                                    'inference_stats']
                                self.assertEqual(
                                    stats['success']['count'], 0,
                                    "unexpected infer stats for version " +
                                    str(v) + " of model " + model_name)
                            else:
                                stats = infer_stats.model_stats[
                                    0].inference_stats
                                self.assertEqual(
                                    stats.success.count, 0,
                                    "unexpected infer stats for version " +
                                    str(v) + " of model " + model_name)

                        else:
                            self.assertFalse(
                                triton_client.is_model_ready(
                                    model_name, model_version=str(v)))
                            try:
                                infer_stats = triton_client.get_inference_statistics(
                                    model_name, model_version=str(v))
                                self.assertTrue(
                                    False,
                                    "unexpected infer stats for the model that is not ready"
                                )
                            except InferenceServerException as ex:
                                self.assertIn(
                                    "requested model version is not available for model",
                                    str(ex))

            except InferenceServerException as ex:
                self.assertTrue(False, "unexpected error {}".format(ex))

    def test_infer_stats_no_model_version(self):
        # Originally There were 3 versions of *_int32_int32_int32 and
        # version 3 was executed once. Version 2 and 3 models were
        # deleted from the model repository so now only expect version 1 to
        # be ready and show infer stats.
        for platform in ('graphdef', 'onnx'):
            model_name = platform + "_int32_int32_int32"

            try:
                for pair in [("localhost:8000", "http"),
                             ("localhost:8001", "grpc")]:
                    if pair[1] == "http":
                        triton_client = httpclient.InferenceServerClient(
                            url=pair[0], verbose=True)
                    else:
                        triton_client = grpcclient.InferenceServerClient(
                            url=pair[0], verbose=True)

                    self.assertTrue(triton_client.is_server_live())
                    self.assertTrue(triton_client.is_server_ready())
                    model_metadata = triton_client.get_model_metadata(
                        model_name)
                    if pair[1] == "http":
                        self.assertEqual(model_name, model_metadata['name'])
                        self.assertEqual(len(model_metadata['versions']), 1)
                        self.assertEqual("1", model_metadata['versions'][0])
                    else:
                        self.assertEqual(model_name, model_metadata.name)
                        self.assertEqual(len(model_metadata.versions), 1)
                        self.assertEqual("1", model_metadata.versions[0])

                    # Only version 3 should have infer stats, only 1 is ready
                    for v in (1, 2, 3):
                        if v == 1:
                            self.assertTrue(
                                triton_client.is_model_ready(
                                    model_name, model_version=str(v)))
                        else:
                            self.assertFalse(
                                triton_client.is_model_ready(
                                    model_name, model_version=str(v)))

                    infer_stats = triton_client.get_inference_statistics(
                        model_name)
                    if pair[1] == "http":
                        stats = infer_stats['model_stats']
                    else:
                        stats = infer_stats.model_stats
                    self.assertEqual(
                        len(stats), 1,
                        "expected 1 infer stats for model " + model_name)

                    if pair[1] == "http":
                        version = stats[0]['version']
                        stat = stats[0]['inference_stats']
                    else:
                        version = stats[0].version
                        stat = stats[0].inference_stats

                    if version != "1":
                        self.assertTrue(
                            False,
                            "expected version 1 for infer stat, got " + version)
                    else:
                        if pair[1] == "http":
                            self.assertEqual(
                                stat['success']['count'], 0,
                                "unexpected infer stats for version " +
                                str(version) + " of model " + model_name)
                        else:
                            self.assertEqual(
                                stat.success.count, 0,
                                "unexpected infer stats for version " +
                                str(version) + " of model " + model_name)

            except InferenceServerException as ex:
                self.assertTrue(False, "unexpected error {}".format(ex))

    def test_infer_stats_no_model(self):
        # Test get_inference_statistics when no model/model_version is passed.
        try:
            for pair in [("localhost:8000", "http"),
                         ("localhost:8001", "grpc")]:
                if pair[1] == "http":
                    triton_client = httpclient.InferenceServerClient(
                        url=pair[0], verbose=True)
                else:
                    triton_client = grpcclient.InferenceServerClient(
                        url=pair[0], verbose=True)

                self.assertTrue(triton_client.is_server_live())
                self.assertTrue(triton_client.is_server_ready())

                # Returns infer stats for ALL models + ready versions
                infer_stats = triton_client.get_inference_statistics()
                if pair[1] == "http":
                    stats = infer_stats['model_stats']
                else:
                    stats = infer_stats.model_stats
                self.assertEqual(
                    len(stats), 219,
                    "expected 219 infer stats for all ready versions of all model"
                )

        except InferenceServerException as ex:
            self.assertTrue(False, "unexpected error {}".format(ex))


if __name__ == '__main__':
    unittest.main()


# qa/L0_shared_memory/shared_memory_test.py

# Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../common")

import numpy as np
import unittest
import os
import test_util as tu
import tritonclient.grpc as grpcclient
import tritonclient.http as httpclient
import tritonclient.utils.shared_memory as shm
from tritonclient import utils


class SharedMemoryTest(tu.TestResultCollector):

    def test_invalid_create_shm(self):
        # Raises error since tried to create invalid system shared memory region
        try:
            shm_op0_handle = shm.create_shared_memory_region(
                "dummy_data", "/dummy_data", -1)
            shm.destroy_shared_memory_region(shm_op0_handle)
        except Exception as ex:
            self.assertTrue(str(ex) == "unable to initialize the size")

    def test_valid_create_set_register(self):
        # Create a valid system shared memory region, fill data in it and register
        if _protocol == "http":
            triton_client = httpclient.InferenceServerClient(_url, verbose=True)
        else:
            triton_client = grpcclient.InferenceServerClient(_url, verbose=True)
        shm_op0_handle = shm.create_shared_memory_region(
            "dummy_data", "/dummy_data", 8)
        shm.set_shared_memory_region(shm_op0_handle,
                                     [np.array([1, 2], dtype=np.float32)])
        triton_client.register_system_shared_memory("dummy_data", "/dummy_data",
                                                    8)
        shm_status = triton_client.get_system_shared_memory_status()
        if _protocol == "http":
            self.assertTrue(len(shm_status) == 1)
        else:
            self.assertTrue(len(shm_status.regions) == 1)
        shm.destroy_shared_memory_region(shm_op0_handle)

    def test_unregister_before_register(self):
        # Create a valid system shared memory region and unregister before register
        if _protocol == "http":
            triton_client = httpclient.InferenceServerClient(_url, verbose=True)
        else:
            triton_client = grpcclient.InferenceServerClient(_url, verbose=True)
        shm_op0_handle = shm.create_shared_memory_region(
            "dummy_data", "/dummy_data", 8)
        triton_client.unregister_system_shared_memory("dummy_data")
        shm_status = triton_client.get_system_shared_memory_status()
        if _protocol == "http":
            self.assertTrue(len(shm_status) == 0)
        else:
            self.assertTrue(len(shm_status.regions) == 0)
        shm.destroy_shared_memory_region(shm_op0_handle)

    def test_unregister_after_register(self):
        # Create a valid system shared memory region and unregister after register
        if _protocol == "http":
            triton_client = httpclient.InferenceServerClient(_url, verbose=True)
        else:
            triton_client = grpcclient.InferenceServerClient(_url, verbose=True)
        shm_op0_handle = shm.create_shared_memory_region(
            "dummy_data", "/dummy_data", 8)
        triton_client.register_system_shared_memory("dummy_data", "/dummy_data",
                                                    8)
        triton_client.unregister_system_shared_memory("dummy_data")
        shm_status = triton_client.get_system_shared_memory_status()
        if _protocol == "http":
            self.assertTrue(len(shm_status) == 0)
        else:
            self.assertTrue(len(shm_status.regions) == 0)
        shm.destroy_shared_memory_region(shm_op0_handle)

    def test_reregister_after_register(self):
        # Create a valid system shared memory region and unregister after register
        if _protocol == "http":
            triton_client = httpclient.InferenceServerClient(_url, verbose=True)
        else:
            triton_client = grpcclient.InferenceServerClient(_url, verbose=True)
        shm_op0_handle = shm.create_shared_memory_region(
            "dummy_data", "/dummy_data", 8)
        triton_client.register_system_shared_memory("dummy_data", "/dummy_data",
                                                    8)
        try:
            triton_client.register_system_shared_memory("dummy_data",
                                                        "/dummy_data", 8)
        except Exception as ex:
            self.assertTrue(
                "shared memory region 'dummy_data' already in manager" in str(
                    ex))
        shm_status = triton_client.get_system_shared_memory_status()
        if _protocol == "http":
            self.assertTrue(len(shm_status) == 1)
        else:
            self.assertTrue(len(shm_status.regions) == 1)
        shm.destroy_shared_memory_region(shm_op0_handle)

    def _configure_sever(self):
        shm_ip0_handle = shm.create_shared_memory_region(
            "input0_data", "/input0_data", 64)
        shm_ip1_handle = shm.create_shared_memory_region(
            "input1_data", "/input1_data", 64)
        shm_op0_handle = shm.create_shared_memory_region(
            "output0_data", "/output0_data", 64)
        shm_op1_handle = shm.create_shared_memory_region(
            "output1_data", "/output1_data", 64)
        input0_data = np.arange(start=0, stop=16, dtype=np.int32)
        input1_data = np.ones(shape=16, dtype=np.int32)
        shm.set_shared_memory_region(shm_ip0_handle, [input0_data])
        shm.set_shared_memory_region(shm_ip1_handle, [input1_data])
        if _protocol == "http":
            triton_client = httpclient.InferenceServerClient(_url, verbose=True)
        else:
            triton_client = grpcclient.InferenceServerClient(_url, verbose=True)
        triton_client.register_system_shared_memory("input0_data",
                                                    "/input0_data", 64)
        triton_client.register_system_shared_memory("input1_data",
                                                    "/input1_data", 64)
        triton_client.register_system_shared_memory("output0_data",
                                                    "/output0_data", 64)
        triton_client.register_system_shared_memory("output1_data",
                                                    "/output1_data", 64)
        return [shm_ip0_handle, shm_ip1_handle, shm_op0_handle, shm_op1_handle]

    def _cleanup_server(self, shm_handles):
        for shm_handle in shm_handles:
            shm.destroy_shared_memory_region(shm_handle)

    def _basic_inference(self,
                         shm_ip0_handle,
                         shm_ip1_handle,
                         shm_op0_handle,
                         shm_op1_handle,
                         error_msg,
                         big_shm_name="",
                         big_shm_size=64):
        input0_data = np.arange(start=0, stop=16, dtype=np.int32)
        input1_data = np.ones(shape=16, dtype=np.int32)
        inputs = []
        outputs = []
        if _protocol == "http":
            triton_client = httpclient.InferenceServerClient(_url, verbose=True)
            inputs.append(httpclient.InferInput("INPUT0", [1, 16], "INT32"))
            inputs.append(httpclient.InferInput("INPUT1", [1, 16], "INT32"))
            outputs.append(
                httpclient.InferRequestedOutput('OUTPUT0', binary_data=True))
            outputs.append(
                httpclient.InferRequestedOutput('OUTPUT1', binary_data=False))
        else:
            triton_client = grpcclient.InferenceServerClient(_url, verbose=True)
            inputs.append(grpcclient.InferInput("INPUT0", [1, 16], "INT32"))
            inputs.append(grpcclient.InferInput("INPUT1", [1, 16], "INT32"))
            outputs.append(grpcclient.InferRequestedOutput('OUTPUT0'))
            outputs.append(grpcclient.InferRequestedOutput('OUTPUT1'))

        inputs[0].set_shared_memory("input0_data", 64)

        if type(shm_ip1_handle) == np.array:
            inputs[1].set_data_from_numpy(input0_data, binary_data=True)
        elif big_shm_name != "":
            inputs[1].set_shared_memory(big_shm_name, big_shm_size)
        else:
            inputs[1].set_shared_memory("input1_data", 64)

        outputs[0].set_shared_memory("output0_data", 64)
        outputs[1].set_shared_memory("output1_data", 64)

        try:
            results = triton_client.infer("simple",
                                          inputs,
                                          model_version="",
                                          outputs=outputs)
            output = results.get_output('OUTPUT0')
            if _protocol == "http":
                output_datatype = output['datatype']
                output_shape = output['shape']
            else:
                output_datatype = output.datatype
                output_shape = output.shape
            output_dtype = utils.triton_to_np_dtype(output_datatype)
            output_data = shm.get_contents_as_numpy(shm_op0_handle,
                                                    output_dtype, output_shape)
            self.assertTrue(
                (output_data[0] == (input0_data + input1_data)).all(),
                "Model output does not match expected output")
        except Exception as ex:
            error_msg.append(str(ex))

    def test_unregister_after_inference(self):
        # Unregister after inference
        error_msg = []
        shm_handles = self._configure_sever()
        self._basic_inference(shm_handles[0], shm_handles[1], shm_handles[2],
                              shm_handles[3], error_msg)
        if len(error_msg) > 0:
            raise Exception(str(error_msg))
        if _protocol == "http":
            triton_client = httpclient.InferenceServerClient(_url, verbose=True)
        else:
            triton_client = grpcclient.InferenceServerClient(_url, verbose=True)
        triton_client.unregister_system_shared_memory("output0_data")
        shm_status = triton_client.get_system_shared_memory_status()
        if _protocol == "http":
            self.assertTrue(len(shm_status) == 3)
        else:
            self.assertTrue(len(shm_status.regions) == 3)
        self._cleanup_server(shm_handles)

    def test_register_after_inference(self):
        # Register after inference
        error_msg = []
        shm_handles = self._configure_sever()
        if _protocol == "http":
            triton_client = httpclient.InferenceServerClient(_url, verbose=True)
        else:
            triton_client = grpcclient.InferenceServerClient(_url, verbose=True)
        self._basic_inference(shm_handles[0], shm_handles[1], shm_handles[2],
                              shm_handles[3], error_msg)
        if len(error_msg) > 0:
            raise Exception(str(error_msg))
        shm_ip2_handle = shm.create_shared_memory_region(
            "input2_data", "/input2_data", 64)
        triton_client.register_system_shared_memory("input2_data",
                                                    "/input2_data", 64)
        shm_status = triton_client.get_system_shared_memory_status()
        if _protocol == "http":
            self.assertTrue(len(shm_status) == 5)
        else:
            self.assertTrue(len(shm_status.regions) == 5)
        shm_handles.append(shm_ip2_handle)
        self._cleanup_server(shm_handles)

    def test_too_big_shm(self):
        # Shared memory input region larger than needed - Throws error
        error_msg = []
        shm_handles = self._configure_sever()
        shm_ip2_handle = shm.create_shared_memory_region(
            "input2_data", "/input2_data", 128)
        if _protocol == "http":
            triton_client = httpclient.InferenceServerClient(_url, verbose=True)
        else:
            triton_client = grpcclient.InferenceServerClient(_url, verbose=True)
        triton_client.register_system_shared_memory("input2_data",
                                                    "/input2_data", 128)
        self._basic_inference(shm_handles[0], shm_ip2_handle, shm_handles[2],
                              shm_handles[3], error_msg, "input2_data", 128)
        if len(error_msg) > 0:
            self.assertTrue(
                "unexpected total byte size 128 for input 'INPUT1', expecting 64"
                in error_msg[-1])
        shm_handles.append(shm_ip2_handle)
        self._cleanup_server(shm_handles)

    def test_mixed_raw_shm(self):
        # Mix of shared memory and RAW inputs
        error_msg = []
        shm_handles = self._configure_sever()
        input1_data = np.ones(shape=16, dtype=np.int32)
        self._basic_inference(shm_handles[0], [input1_data], shm_handles[2],
                              shm_handles[3], error_msg)
        if len(error_msg) > 0:
            raise Exception(error_msg[-1])
        self._cleanup_server(shm_handles)

    def test_unregisterall(self):
        # Unregister all shared memory blocks
        shm_handles = self._configure_sever()
        if _protocol == "http":
            triton_client = httpclient.InferenceServerClient(_url, verbose=True)
        else:
            triton_client = grpcclient.InferenceServerClient(_url, verbose=True)
        status_before = triton_client.get_system_shared_memory_status()
        if _protocol == "http":
            self.assertTrue(len(status_before) == 4)
        else:
            self.assertTrue(len(status_before.regions) == 4)
        triton_client.unregister_system_shared_memory()
        status_after = triton_client.get_system_shared_memory_status()
        if _protocol == "http":
            self.assertTrue(len(status_after) == 0)
        else:
            self.assertTrue(len(status_after.regions) == 0)
        self._cleanup_server(shm_handles)


if __name__ == '__main__':
    _protocol = os.environ.get('CLIENT_TYPE', "http")
    if _protocol == "http":
        _url = "localhost:8000"
    else:
        _url = "localhost:8001"
    unittest.main()


# qa/L0_simple_ensemble/ensemble_test.py

# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../common")
sys.path.append("../clients")

import logging

import os
import unittest
import numpy as np
import infer_util as iu
import test_util as tu
import tritonhttpclient


class EnsembleTest(tu.TestResultCollector):

    def _get_infer_count_per_version(self, model_name):
        triton_client = tritonhttpclient.InferenceServerClient("localhost:8000",
                                                               verbose=True)
        stats = triton_client.get_inference_statistics(model_name)
        self.assertEqual(len(stats["model_stats"]), 2)
        infer_count = [0, 0]
        for model_stat in stats["model_stats"]:
            self.assertEqual(model_stat["name"], model_name,
                             "expected stats for model " + model_name)
            model_version = model_stat['version']
            if model_version == "1":
                infer_count[0] = model_stat["inference_stats"]["success"][
                    "count"]
            elif model_version == "2":
                infer_count[1] = model_stat["inference_stats"]["success"][
                    "count"]
            else:
                self.assertTrue(
                    False, "unexpected version {} for model {}".format(
                        model_version, model_name))
        return infer_count

    def test_ensemble_add_sub(self):
        for bs in (1, 8):
            iu.infer_exact(self, "ensemble_add_sub", (bs, 16), bs, np.int32,
                           np.int32, np.int32)

        infer_count = self._get_infer_count_per_version("simple")
        # The two 'simple' versions should have the same infer count
        if (infer_count[0] != infer_count[1]):
            self.assertTrue(
                False,
                "unexpeced different infer count for different 'simple' versions"
            )

    def test_ensemble_add_sub_one_output(self):
        for bs in (1, 8):
            iu.infer_exact(self,
                           "ensemble_add_sub", (bs, 16),
                           bs,
                           np.int32,
                           np.int32,
                           np.int32,
                           outputs=("OUTPUT0",))

        infer_count = self._get_infer_count_per_version("simple")
        # Only 'simple' version 2 should have non-zero infer count
        # as it is in charge of producing OUTPUT0
        if (infer_count[0] != 0):
            self.assertTrue(
                False, "unexpeced non-zero infer count for 'simple' version 1")
        elif (infer_count[1] == 0):
            self.assertTrue(
                False, "unexpeced zero infer count for 'simple' version 2")


if __name__ == '__main__':
    logging.basicConfig(stream=sys.stderr)
    unittest.main()


# qa/L0_storage_swiftstack/infer_test.py

# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../common")

from builtins import range
from future.utils import iteritems
import unittest
import numpy as np
import infer_util as iu
import test_util as tu
import os


class InferTest(tu.TestResultCollector):

    def _full_exact(self, input_dtype, output0_dtype, output1_dtype,
                    output0_raw, output1_raw, swap):

        def _infer_exact_helper(tester,
                                pf,
                                tensor_shape,
                                batch_size,
                                input_dtype,
                                output0_dtype,
                                output1_dtype,
                                output0_raw=True,
                                output1_raw=True,
                                model_version=None,
                                swap=False,
                                outputs=("OUTPUT0", "OUTPUT1"),
                                use_http=True,
                                use_grpc=True,
                                skip_request_id_check=False,
                                use_streaming=True,
                                correlation_id=0):
            for bs in (1, batch_size):
                iu.infer_exact(tester,
                               pf, (bs,) + tensor_shape,
                               bs,
                               input_dtype,
                               output0_dtype,
                               output1_dtype,
                               output0_raw=output0_raw,
                               output1_raw=output1_raw,
                               model_version=model_version,
                               swap=swap,
                               outputs=outputs,
                               use_http=use_http,
                               use_grpc=use_grpc,
                               skip_request_id_check=skip_request_id_check,
                               use_streaming=use_streaming,
                               correlation_id=correlation_id)

        input_size = 16

        if tu.validate_for_tf_model(input_dtype, output0_dtype, output1_dtype,
                                    (input_size,), (input_size,),
                                    (input_size,)):
            for pf in ["graphdef", "savedmodel"]:
                _infer_exact_helper(self,
                                    pf, (input_size,),
                                    8,
                                    input_dtype,
                                    output0_dtype,
                                    output1_dtype,
                                    output0_raw=output0_raw,
                                    output1_raw=output1_raw,
                                    swap=swap)

        if tu.validate_for_trt_model(input_dtype, output0_dtype, output1_dtype,
                                     (input_size, 1, 1), (input_size, 1, 1),
                                     (input_size, 1, 1)):
            if input_dtype == np.int8:
                _infer_exact_helper(self,
                                    'plan', (input_size, 1, 1),
                                    8,
                                    input_dtype,
                                    output0_dtype,
                                    output1_dtype,
                                    output0_raw=output0_raw,
                                    output1_raw=output1_raw,
                                    swap=swap)
            else:
                _infer_exact_helper(self,
                                    'plan', (input_size,),
                                    8,
                                    input_dtype,
                                    output0_dtype,
                                    output1_dtype,
                                    output0_raw=output0_raw,
                                    output1_raw=output1_raw,
                                    swap=swap)

        if tu.validate_for_onnx_model(input_dtype, output0_dtype, output1_dtype,
                                      (input_size,), (input_size,),
                                      (input_size,)):
            _infer_exact_helper(self,
                                'onnx', (input_size,),
                                8,
                                input_dtype,
                                output0_dtype,
                                output1_dtype,
                                output0_raw=output0_raw,
                                output1_raw=output1_raw,
                                swap=swap)

        if tu.validate_for_libtorch_model(input_dtype, output0_dtype,
                                          output1_dtype, (input_size,),
                                          (input_size,), (input_size,)):
            _infer_exact_helper(self,
                                'libtorch', (input_size,),
                                8,
                                input_dtype,
                                output0_dtype,
                                output1_dtype,
                                output0_raw=output0_raw,
                                output1_raw=output1_raw,
                                swap=swap)

    def test_raw_fff(self):
        self._full_exact(np.float32,
                         np.float32,
                         np.float32,
                         output0_raw=True,
                         output1_raw=True,
                         swap=True)

    def test_class_fff(self):
        self._full_exact(np.float32,
                         np.float32,
                         np.float32,
                         output0_raw=False,
                         output1_raw=False,
                         swap=True)


if __name__ == '__main__':
    unittest.main()


# qa/L0_string_io/string_client_test.py

#!/usr/bin/env python
# Copyright 2019-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append('../common')

import argparse
import numpy as np
import os
from builtins import range
import tritonclient.http as tritonhttpclient
import tritonclient.grpc as tritongrpcclient
import tritonclient.utils as tritonutils
import unittest
import test_util as tu


class ClientStringTest(tu.TestResultCollector):

    def _test_infer_unicode(self, model_name, client, input_):
        # Send inference request to the inference server. Get results for
        # both output tensors.
        inputs = []
        outputs = []
        inputs.append(client[1].InferInput('INPUT0', input_.shape, "BYTES"))

        if client[1] == tritonhttpclient:
            inputs[0].set_data_from_numpy(input_, client[3])
        else:
            inputs[0].set_data_from_numpy(input_)

        if client[1] == tritonhttpclient:
            outputs.append(client[1].InferRequestedOutput(
                'OUTPUT0', binary_data=client[2]))
        else:
            outputs.append(client[1].InferRequestedOutput('OUTPUT0'))

        results = client[0].infer(model_name=model_name,
                                  inputs=inputs,
                                  outputs=outputs)

        out0 = results.as_numpy('OUTPUT0')
        # We expect there to be 1 results (with batch-size 1). Verify
        # that all 8 result elements are the same as the input.
        self.assertTrue(np.array_equal(input_, out0))
        return out0

    def _test_infer_non_unicode(self,
                                model_name,
                                client,
                                input_,
                                binary_data=True):
        # Send inference request to the inference server. Get results for
        # both output tensors.
        inputs = []
        outputs = []
        inputs.append(client[1].InferInput('INPUT0', input_.shape, "BYTES"))

        if client[1] == tritonhttpclient:
            inputs[0].set_data_from_numpy(input_, client[3])
        else:
            inputs[0].set_data_from_numpy(input_)

        if client[1] == tritonhttpclient:
            outputs.append(client[1].InferRequestedOutput(
                'OUTPUT0', binary_data=client[2]))
        else:
            outputs.append(client[1].InferRequestedOutput('OUTPUT0'))

        results = client[0].infer(model_name=model_name,
                                  inputs=inputs,
                                  outputs=outputs)

        out0 = results.as_numpy('OUTPUT0')
        # We expect there to be 1 results (with batch-size 1). Verify
        # that all 8 result elements are the same as the input.
        if client[2]:
            self.assertTrue(np.array_equal(input_.astype(np.bytes_), out0))
        else:
            self.assertTrue(
                np.array_equal(input_.astype(np.bytes_),
                               out0.astype(np.bytes_)))
        return out0

    def _test_unicode_bytes_dtype(self, client, model_name, dtype='|S78'):
        # Create the data for the input tensor. Initialize the tensor to 8
        # byte strings. (dtype of np.bytes_)
        # Sample string that should no longer cause failure
        in0 = np.array([
            [
                b'\nF\n\'\n\x01a\x12"\x1a \n\x1e\xfa\x03\x94\x01\x0f\xd7\x02\xf1\x05\xdf\x01\x82\x03\xb5\x05\xc1\x07\xba\x06\xff\x06\xc7\x07L\xf5\x03\xe2\x07\xa9\x03\n\x0c\n\x01b\x12\x07\x1a\x05\n\x03\x89\xcc=\n\r\n\x01c\x12\x08\x12\x06\n\x04\xdf\\\xcb\xbf'
            ],
            [
                b'\n:\n\x1a\n\x01a\x12\x15\x1a\x13\n\x11*\xe3\x05\xc5\x06\xda\x07\xcb\x06~\xb1\x05\xb3\x01\xa9\x02\x15\n\r\n\x01b\x12\x08\x1a\x06\n\x04\xf6\xa2\xc5\x01\n\r\n\x01c\x12\x08\x12\x06\n\x04\xbb[\n\xbf'
            ],
            [
                b'\nL\n-\n\x01a\x12(\x1a&\n$\x87\x07\xce\x01\xe7\x06\xee\x04\xe1\x03\xf1\x03\xd7\x07\xbe\x02\xb8\x05\xe0\x05\xe4\x01\x88\x06\xb6\x03\xb9\x05\x83\x06\xf8\x04\xe2\x04\xf4\x06\n\x0c\n\x01b\x12\x07\x1a\x05\n\x03\x89\xcc=\n\r\n\x01c\x12\x08\x12\x06\n\x04\xbc\x99+@'
            ],
            [
                b'\n2\n\x12\n\x01a\x12\r\x1a\x0b\n\t\x99\x02\xde\x04\x9f\x04\xc5\x053\n\r\n\x01b\x12\x08\x1a\x06\n\x04\xf6\xa2\xc5\x01\n\r\n\x01c\x12\x08\x12\x06\n\x04\x12\x07\x83\xbe'
            ],
            [
                b'\nJ\n\r\n\x01b\x12\x08\x1a\x06\n\x04\x9b\x94\xad\x04\n\r\n\x01c\x12\x08\x12\x06\n\x04\xc3\x8a\x08\xbf\n*\n\x01a\x12%\x1a#\n!\x9c\x02\xb2\x02\xcd\x02\x9d\x07\x8d\x01\xb6\x05a\xf1\x01\xf0\x05\xdb\x02\xac\x04\xbd\x05\xe0\x04\xd2\x06\xaf\x02\xa8\x01\x8b\x04'
            ],
            [
                b'\n3\n\x13\n\x01a\x12\x0e\x1a\x0c\n\n<\xe2\x05\x8a\x01\xb3\x07?\xfd\x01\n\r\n\x01b\x12\x08\x1a\x06\n\x04\xf6\xa2\xc5\x01\n\r\n\x01c\x12\x08\x12\x06\n\x04\x1b\x931\xbf\x00\x00'
            ],
            [
                b'\n&\n\x07\n\x01a\x12\x02\x1a\x00\n\x0c\n\x01b\x12\x07\x1a\x05\n\x03\x89\xcc=\n\r\n\x01c\x12\x08\x12\x06\n\x04{\xbc\x0e>\x00\x00\x00'
            ],
            [
                b'\nF\n\'\n\x01a\x12"\x1a \n\x1e\x97\x01\x93\x02\x9e\x01\xac\x06\xff\x01\xd8\x05\xe1\x07\xd8\x04g]\x9a\x05\xff\x06\xde\x07\x8f\x04\x97\x04\xda\x03\n\x0c\n\x01b\x12\x07\x1a\x05\n\x03\x9a\xb7I\n\r\n\x01c\x12\x08\x12\x06\n\x04\xfb\x87\x83\xbf'
            ]
        ],
                       dtype=dtype).flatten()
        self._test_infer_unicode(model_name, client, in0)

    def _test_str_dtype(self, client, model_name, dtype=np.object_):
        in0_bytes = np.array([str(i) for i in range(10000, 10008)], dtype=dtype)
        self._test_infer_non_unicode(model_name, client, in0_bytes)

        in0_bytes = np.array([i for i in range(10000, 10008)], dtype=dtype)
        self._test_infer_non_unicode(model_name, client, in0_bytes)

    def _test_bytes(self, model_name):
        dtypes = [np.object_, np.bytes_]

        # This clients will fail for binary_data=False when the binary input
        # is not UTF-8 encodable. They should work for other cases however.
        binary_false_clients = [
            (tritonhttpclient.InferenceServerClient("localhost:8000",
                                                    verbose=True),
             tritonhttpclient, True, False),
            (tritonhttpclient.InferenceServerClient("localhost:8000",
                                                    verbose=True),
             tritonhttpclient, False, False),
            (tritonhttpclient.InferenceServerClient("localhost:8000",
                                                    verbose=True),
             tritonhttpclient, False, True),
        ]

        # These clients work for every data type
        other_clients = [
            (tritongrpcclient.InferenceServerClient("localhost:8001",
                                                    verbose=True),
             tritongrpcclient, False),
            (tritonhttpclient.InferenceServerClient("localhost:8000",
                                                    verbose=True),
             tritonhttpclient, True, True),
        ]

        for client in other_clients + binary_false_clients:
            self._test_str_dtype(client, model_name)
            for dtype in dtypes:
                self._test_str_dtype(client, model_name, dtype)

        for client in other_clients:
            self._test_unicode_bytes_dtype(client, model_name)
            for dtype in dtypes:
                self._test_unicode_bytes_dtype(client, model_name, dtype)

        for client in binary_false_clients:
            with self.assertRaises(tritonutils.InferenceServerException):
                self._test_unicode_bytes_dtype(client, model_name)
            for dtype in dtypes:
                with self.assertRaises(tritonutils.InferenceServerException):
                    self._test_unicode_bytes_dtype(client, model_name, dtype)

    def test_tf_unicode_bytes(self):
        self._test_bytes("graphdef_nobatch_zero_1_object")
        self._test_bytes("string_identity")


if __name__ == '__main__':
    unittest.main()


# qa/L0_tf_gpu_io/tf_gpu_io_test.py

# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../common")

import unittest
import numpy as np
import os
import infer_util as iu
import test_util as tu
import tritonhttpclient as httpclient

TENSOR_SIZE = 16384


class TfGpuIoTest(tu.TestResultCollector):

    def _test_helper(self,
                     model_name,
                     shape,
                     override_input_names=[],
                     override_output_names=[],
                     batching_enabled=False):
        try:
            bs = 1
            if batching_enabled:
                shape = [[
                    bs,
                ] + shape]
            iu.infer_zero(
                self,
                'graphdef',
                bs,
                np.float32,
                shape,
                shape,
                override_model_name=model_name,
                override_input_names=override_input_names,
                override_output_names=override_output_names,
            )

        except Exception as ex:
            self.assertTrue(False, "unexpected error {}".format(ex))

    def test_sig_tag0(self):
        self._test_helper("sig_tag0", [16],
                          override_input_names=["INPUT"],
                          override_output_names=["OUTPUT"])

    def test_graphdef_zero_1_float32_def(self):
        self._test_helper("graphdef_zero_1_float32_def", [TENSOR_SIZE],
                          batching_enabled=True)

    def test_graphdef_zero_1_float32_gpu(self):
        self._test_helper("graphdef_zero_1_float32_gpu", [TENSOR_SIZE],
                          batching_enabled=True)

    def test_savedmodel_zero_1_float32_def(self):
        self._test_helper("savedmodel_zero_1_float32_def", [TENSOR_SIZE],
                          batching_enabled=True)

    def test_savedmodel_zero_1_float32_gpu(self):
        self._test_helper("savedmodel_zero_1_float32_gpu", [TENSOR_SIZE],
                          batching_enabled=True)


if __name__ == '__main__':
    unittest.main()


# qa/L0_tf_parameters/tf_parameter_test.py

# Copyright 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys
sys.path.append('../common')

import test_util as tu
import tritonclient.http as tritonhttpclient
import tritonclient.utils
import numpy as np
import unittest


class TFParameterTest(tu.TestResultCollector):

    def setUp(self):
        self._client = tritonhttpclient.InferenceServerClient("localhost:8000",
                                                              verbose=True)

    def _infer_helper(self):
        # The model has a single variable which is added to the input.  Since the
        # variable is initialized to zero the input and output must match.
        model_name = 'graphdef_variable'
        input = np.array([10], dtype=np.int32)

        inputs = []
        inputs.append(tritonhttpclient.InferInput('INPUT', input.shape,
                                                  'INT32'))
        inputs[-1].set_data_from_numpy(input)

        outputs = []
        outputs.append(tritonhttpclient.InferRequestedOutput('OUTPUT'))

        results = self._client.infer(model_name=model_name,
                                     inputs=inputs,
                                     outputs=outputs)
        output = results.as_numpy('OUTPUT')
        np.testing.assert_array_equal(output, input)

    def test_tf_variable(self):
        self._infer_helper()

    def test_tf_variable_error(self):
        with self.assertRaises(
                tritonclient.utils.InferenceServerException) as e:
            self._infer_helper()
        self.assertIn("Attempting to use uninitialized value VARIABLE",
                      e.exception.message())


if __name__ == '__main__':
    unittest.main()


# qa/L0_tf_tag_sigdef/tf_tag_sigdef_test.py

# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../common")

from builtins import range
from future.utils import iteritems
import unittest
import numpy as np
import os
import test_util as tu
import tritonhttpclient as httpclient
from tritonclientutils import InferenceServerException


class TagSigdefTest(tu.TestResultCollector):
    base_model_name = "sig_tag"
    base_tag = "serve"
    test_tag = "testTag"
    base_sig_def = "serving_default"
    test_sig_def = "testSigDef"
    dims = 16

    def _test_helper(self, modelVersion, tag, sig_def):
        shape = [self.dims]
        model_name = self.base_model_name + str(modelVersion)
        # The multiplier is defined during model creation. See server/qa/common/gen_tag_sigdef.py
        # for details
        multiplier = modelVersion + 1
        output_name = "OUTPUT"
        triton_client = httpclient.InferenceServerClient("localhost:8000",
                                                         verbose=True)
        inputs = []
        outputs = []
        inputs.append(httpclient.InferInput('INPUT', shape, "FP32"))
        input_data = np.ones(shape=shape).astype(np.float32)
        inputs[0].set_data_from_numpy(input_data, binary_data=True)

        outputs.append(
            httpclient.InferRequestedOutput(output_name, binary_data=True))
        results = triton_client.infer(model_name, inputs, outputs=outputs)
        output_data = results.as_numpy(output_name)
        test_output = input_data * multiplier
        self.assertTrue(np.isclose(output_data, test_output).all())

    def test_default(self):
        self._test_helper(0, self.base_tag, self.base_sig_def)

    def test_sig_def(self):
        self._test_helper(1, self.base_tag, self.test_sig_def)

    def test_tag(self):
        self._test_helper(2, self.test_tag, self.base_sig_def)

    def test_tag_sig_def(self):
        self._test_helper(3, self.test_tag, self.test_sig_def)


if __name__ == '__main__':
    unittest.main()


# qa/L0_tf_unknown_rank/tf_unknown_rank_test.py

# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../common")

import unittest
import numpy as np
import test_util as tu
import tritonhttpclient
from tritonclientutils import *


class UnknownRankTest(tu.TestResultCollector):
    # helper function to generate requests to the server
    def infer_unknown(self, model_name, tensor_shape):
        print("About to run the test")
        input_data = np.random.random_sample(tensor_shape).astype(np.float32)
        client = tritonhttpclient.InferenceServerClient('localhost:8000')
        inputs = [
            tritonhttpclient.InferInput("INPUT", input_data.shape,
                                        np_to_triton_dtype(input_data.dtype))
        ]
        inputs[0].set_data_from_numpy(input_data)
        results = client.infer(model_name, inputs)
        self.assertTrue(np.array_equal(results.as_numpy('OUTPUT'), input_data))

    def test_success(self):
        model_name = "unknown_rank_success"
        tensor_shape = (1)
        try:
            self.infer_unknown(model_name, tensor_shape)
        except InferenceServerException as ex:
            self.assertTrue(False, "unexpected error {}".format(ex))

    def test_wrong_input(self):
        model_name = "unknown_rank_wrong_output"
        tensor_shape = (1, 2)
        try:
            self.infer_unknown(model_name, tensor_shape)
            self.fail(
                "Found success when expected failure with model given " \
                "wrong input tensor [1,2] for input [-1,1]."
            )
        except InferenceServerException as ex:
            self.assertIn(
                "unexpected shape for input \'INPUT\' for model " \
                "\'unknown_rank_wrong_output\'. Expected [1], got [1,2]",
                ex.message())


if __name__ == '__main__':
    unittest.main()


# qa/L0_tftrt_optimization/tftrt_optimization_test.py

# Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../common")

import unittest
import numpy as np
import test_util as tu
import tritonhttpclient as httpclient
from tritonclientutils import InferenceServerException


class TFTRTOptimizationTest(tu.TestResultCollector):

    def setUp(self):
        self.input0_ = np.arange(start=0, stop=16,
                                 dtype=np.float32).reshape(1, 16)
        self.input1_ = np.ones(shape=16, dtype=np.float32).reshape(1, 16)
        self.expected_output0_ = self.input0_ + self.input1_
        self.expected_output1_ = self.input0_ - self.input1_

    def _addsub_infer(self, model_name):
        triton_client = httpclient.InferenceServerClient("localhost:8000",
                                                         verbose=True)

        inputs = []
        outputs = []
        inputs.append(httpclient.InferInput('INPUT0', [1, 16], "FP32"))
        inputs.append(httpclient.InferInput('INPUT1', [1, 16], "FP32"))

        # Initialize the data
        inputs[0].set_data_from_numpy(self.input0_, binary_data=True)
        inputs[1].set_data_from_numpy(self.input1_, binary_data=False)

        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT0', binary_data=True))
        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT1', binary_data=True))

        results = triton_client.infer(model_name, inputs, outputs=outputs)

        output0_data = results.as_numpy('OUTPUT0')
        output1_data = results.as_numpy('OUTPUT1')

        self.assertTrue(np.array_equal(self.expected_output0_, output0_data),
                        "incorrect sum")
        self.assertTrue(np.array_equal(self.expected_output1_, output1_data),
                        "incorrect difference")

    def test_graphdef(self):
        self._addsub_infer("graphdef_float32_float32_float32_trt")
        self._addsub_infer("graphdef_float32_float32_float32_param")

    def test_savedmodel(self):
        self._addsub_infer("savedmodel_float32_float32_float32_trt")
        self._addsub_infer("savedmodel_float32_float32_float32_param")


if __name__ == '__main__':
    unittest.main()


# qa/L0_trace/trace_endpoint_test.py

#!/usr/bin/python

# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../common")

import numpy as np
import sys
import unittest
import tritonclient.http as httpclient
import tritonclient.grpc as grpcclient
import json
from google.protobuf import json_format
import test_util as tu


# Similar set up as dynamic batcher tests
class TraceEndpointTest(tu.TestResultCollector):

    def tearDown(self):
        # Clear all trace settings to initial state.
        # Note that the tearDown function uses HTTP client so the pass/fail
        # of the HTTP trace setting test cases should be checked to make sure
        # tearDown() is properly executed and not affecting start state of
        # other test cases
        clear_settings = {
            "trace_file": None,
            "trace_level": None,
            "trace_rate": None,
            "trace_count": None,
            "log_frequency": None
        }
        triton_client = httpclient.InferenceServerClient("localhost:8000")
        triton_client.update_trace_settings(model_name="simple",
                                            settings=clear_settings)
        triton_client.update_trace_settings(model_name=None,
                                            settings=clear_settings)

    def check_server_initial_state(self):
        # Helper function to make sure the trace setting is properly
        # initialized / reset before actually running the test case.
        # Note that this function uses HTTP client so the pass/fail of
        # the HTTP trace setting test cases should be checked to make sure
        # the initial state is checked properly before running other test cases.
        initial_settings = {
            "trace_file": "global_unittest.log",
            "trace_level": ["TIMESTAMPS"],
            "trace_rate": "1",
            "trace_count": "-1",
            "log_frequency": "0"
        }
        triton_client = httpclient.InferenceServerClient("localhost:8000")
        self.assertEqual(initial_settings,
                         triton_client.get_trace_settings(model_name="simple"))
        self.assertEqual(initial_settings, triton_client.get_trace_settings())

    def test_http_get_settings(self):
        # Model trace settings will be the same as global trace settings since
        # no update has been made.
        initial_settings = {
            "trace_file": "global_unittest.log",
            "trace_level": ["TIMESTAMPS"],
            "trace_rate": "1",
            "trace_count": "-1",
            "log_frequency": "0"
        }
        triton_client = httpclient.InferenceServerClient("localhost:8000")
        self.assertEqual(initial_settings,
                         triton_client.get_trace_settings(model_name="simple"),
                         "Unexpected initial model trace settings")
        self.assertEqual(initial_settings, triton_client.get_trace_settings(),
                         "Unexpected initial global settings")

    def test_grpc_get_settings(self):
        # Model trace settings will be the same as global trace settings since
        # no update has been made.
        initial_settings = grpcclient.service_pb2.TraceSettingResponse()
        json_format.Parse(
            json.dumps({
                "settings": {
                    "trace_file": {
                        "value": ["global_unittest.log"]
                    },
                    "trace_level": {
                        "value": ["TIMESTAMPS"]
                    },
                    "trace_rate": {
                        "value": ["1"]
                    },
                    "trace_count": {
                        "value": ["-1"]
                    },
                    "log_frequency": {
                        "value": ["0"]
                    },
                }
            }), initial_settings)

        triton_client = grpcclient.InferenceServerClient("localhost:8001")
        self.assertEqual(initial_settings,
                         triton_client.get_trace_settings(model_name="simple"),
                         "Unexpected initial model trace settings")
        self.assertEqual(initial_settings, triton_client.get_trace_settings(),
                         "Unexpected initial global settings")

    def test_http_update_settings(self):
        # Update model and global trace settings in order,
        # and expect the global trace settings will only reflect to
        # the model setting fields that haven't been specified.
        self.check_server_initial_state()

        expected_first_model_settings = {
            "trace_file": "model.log",
            "trace_level": ["TIMESTAMPS"],
            "trace_rate": "1",
            "trace_count": "-1",
            "log_frequency": "0"
        }
        expected_second_model_settings = {
            "trace_file": "model.log",
            "trace_level": ["TIMESTAMPS", "TENSORS"],
            "trace_rate": "1",
            "trace_count": "-1",
            "log_frequency": "0"
        }
        expected_global_settings = {
            "trace_file": "another.log",
            "trace_level": ["TIMESTAMPS", "TENSORS"],
            "trace_rate": "1",
            "trace_count": "-1",
            "log_frequency": "0"
        }

        model_update_settings = {"trace_file": "model.log"}
        global_update_settings = {
            "trace_file": "another.log",
            "trace_level": ["TIMESTAMPS", "TENSORS"]
        }

        triton_client = httpclient.InferenceServerClient("localhost:8000")
        self.assertEqual(
            expected_first_model_settings,
            triton_client.update_trace_settings(model_name="simple",
                                                settings=model_update_settings),
            "Unexpected updated model trace settings")
        # Note that 'trace_level' may be mismatch due to the order of
        # the levels listed, currently we assume the order is the same
        # for simplicity. But the order shouldn't be enforced and this checking
        # needs to be improved when this kind of failure is reported
        self.assertEqual(
            expected_global_settings,
            triton_client.update_trace_settings(
                settings=global_update_settings),
            "Unexpected updated global settings")
        self.assertEqual(expected_second_model_settings,
                         triton_client.get_trace_settings(model_name="simple"),
                         "Unexpected model trace settings after global update")

    def test_grpc_update_settings(self):
        # Update model and global trace settings in order,
        # and expect the global trace settings will only reflect to
        # the model setting fields that haven't been specified.
        self.check_server_initial_state()

        expected_first_model_settings = grpcclient.service_pb2.TraceSettingResponse(
        )
        json_format.Parse(
            json.dumps({
                "settings": {
                    "trace_file": {
                        "value": ["model.log"]
                    },
                    "trace_level": {
                        "value": ["TIMESTAMPS"]
                    },
                    "trace_rate": {
                        "value": ["1"]
                    },
                    "trace_count": {
                        "value": ["-1"]
                    },
                    "log_frequency": {
                        "value": ["0"]
                    },
                }
            }), expected_first_model_settings)

        expected_second_model_settings = grpcclient.service_pb2.TraceSettingResponse(
        )
        json_format.Parse(
            json.dumps({
                "settings": {
                    "trace_file": {
                        "value": ["model.log"]
                    },
                    "trace_level": {
                        "value": ["TIMESTAMPS", "TENSORS"]
                    },
                    "trace_rate": {
                        "value": ["1"]
                    },
                    "trace_count": {
                        "value": ["-1"]
                    },
                    "log_frequency": {
                        "value": ["0"]
                    },
                }
            }), expected_second_model_settings)

        expected_global_settings = grpcclient.service_pb2.TraceSettingResponse()
        json_format.Parse(
            json.dumps({
                "settings": {
                    "trace_file": {
                        "value": ["another.log"]
                    },
                    "trace_level": {
                        "value": ["TIMESTAMPS", "TENSORS"]
                    },
                    "trace_rate": {
                        "value": ["1"]
                    },
                    "trace_count": {
                        "value": ["-1"]
                    },
                    "log_frequency": {
                        "value": ["0"]
                    },
                }
            }), expected_global_settings)

        model_update_settings = {"trace_file": "model.log"}
        global_update_settings = {
            "trace_file": "another.log",
            "trace_level": ["TIMESTAMPS", "TENSORS"]
        }

        triton_client = grpcclient.InferenceServerClient("localhost:8001")
        self.assertEqual(
            expected_first_model_settings,
            triton_client.update_trace_settings(model_name="simple",
                                                settings=model_update_settings),
            "Unexpected updated model trace settings")
        # Note that 'trace_level' may be mismatch due to the order of
        # the levels listed, currently we assume the order is the same
        # for simplicity. But the order shouldn't be enforced and this checking
        # needs to be improved when this kind of failure is reported
        self.assertEqual(
            expected_global_settings,
            triton_client.update_trace_settings(
                settings=global_update_settings),
            "Unexpected updated global settings")
        self.assertEqual(expected_second_model_settings,
                         triton_client.get_trace_settings(model_name="simple"),
                         "Unexpected model trace settings after global update")

    def test_http_clear_settings(self):
        # Clear global and model trace settings in order,
        # and expect the default / global trace settings are
        # propagated properly.
        self.check_server_initial_state()

        # First set up the model / global trace setting that:
        # model 'simple' has 'trace_rate' and 'log_frequency' specified
        # global has 'trace_level', 'trace_count' and 'trace_rate' specified
        triton_client = httpclient.InferenceServerClient("localhost:8000")
        triton_client.update_trace_settings(model_name="simple",
                                            settings={
                                                "trace_rate": "12",
                                                "log_frequency": "34"
                                            })
        triton_client.update_trace_settings(settings={
            "trace_rate": "56",
            "trace_count": "78",
            "trace_level": ["OFF"]
        })

        expected_global_settings = {
            "trace_file": "global_unittest.log",
            "trace_level": ["OFF"],
            "trace_rate": "1",
            "trace_count": "-1",
            "log_frequency": "0"
        }
        expected_first_model_settings = {
            "trace_file": "global_unittest.log",
            "trace_level": ["OFF"],
            "trace_rate": "12",
            "trace_count": "-1",
            "log_frequency": "34"
        }
        expected_second_model_settings = {
            "trace_file": "global_unittest.log",
            "trace_level": ["OFF"],
            "trace_rate": "1",
            "trace_count": "-1",
            "log_frequency": "34"
        }
        global_clear_settings = {"trace_rate": None, "trace_count": None}
        model_clear_settings = {"trace_rate": None, "trace_level": None}

        # Clear global
        self.assertEqual(
            expected_global_settings,
            triton_client.update_trace_settings(settings=global_clear_settings),
            "Unexpected cleared global trace settings")
        self.assertEqual(expected_first_model_settings,
                         triton_client.get_trace_settings(model_name="simple"),
                         "Unexpected model trace settings after global clear")
        self.assertEqual(
            expected_second_model_settings,
            triton_client.update_trace_settings(model_name="simple",
                                                settings=model_clear_settings),
            "Unexpected model trace settings after model clear")
        self.assertEqual(expected_global_settings,
                         triton_client.get_trace_settings(),
                         "Unexpected global trace settings after model clear")

    def test_grpc_clear_settings(self):
        # Clear global and model trace settings in order,
        # and expect the default / global trace settings are
        # propagated properly.
        self.check_server_initial_state()

        # First set up the model / global trace setting that:
        # model 'simple' has 'trace_rate' and 'log_frequency' specified
        # global has 'trace_level', 'trace_count' and 'trace_rate' specified
        triton_client = grpcclient.InferenceServerClient("localhost:8001")
        triton_client.update_trace_settings(model_name="simple",
                                            settings={
                                                "trace_rate": "12",
                                                "log_frequency": "34"
                                            })
        triton_client.update_trace_settings(settings={
            "trace_rate": "56",
            "trace_count": "78",
            "trace_level": ["OFF"]
        })

        expected_global_settings = grpcclient.service_pb2.TraceSettingResponse()
        json_format.Parse(
            json.dumps({
                "settings": {
                    "trace_file": {
                        "value": ["global_unittest.log"]
                    },
                    "trace_level": {
                        "value": ["OFF"]
                    },
                    "trace_rate": {
                        "value": ["1"]
                    },
                    "trace_count": {
                        "value": ["-1"]
                    },
                    "log_frequency": {
                        "value": ["0"]
                    },
                }
            }), expected_global_settings)
        expected_first_model_settings = grpcclient.service_pb2.TraceSettingResponse(
        )
        json_format.Parse(
            json.dumps({
                "settings": {
                    "trace_file": {
                        "value": ["global_unittest.log"]
                    },
                    "trace_level": {
                        "value": ["OFF"]
                    },
                    "trace_rate": {
                        "value": ["12"]
                    },
                    "trace_count": {
                        "value": ["-1"]
                    },
                    "log_frequency": {
                        "value": ["34"]
                    },
                }
            }), expected_first_model_settings)
        expected_second_model_settings = grpcclient.service_pb2.TraceSettingResponse(
        )
        json_format.Parse(
            json.dumps({
                "settings": {
                    "trace_file": {
                        "value": ["global_unittest.log"]
                    },
                    "trace_level": {
                        "value": ["OFF"]
                    },
                    "trace_rate": {
                        "value": ["1"]
                    },
                    "trace_count": {
                        "value": ["-1"]
                    },
                    "log_frequency": {
                        "value": ["34"]
                    },
                }
            }), expected_second_model_settings)

        global_clear_settings = {"trace_rate": None, "trace_count": None}
        model_clear_settings = {"trace_rate": None, "trace_level": None}

        # Clear global
        self.assertEqual(
            expected_global_settings,
            triton_client.update_trace_settings(settings=global_clear_settings),
            "Unexpected cleared global trace settings")
        self.assertEqual(expected_first_model_settings,
                         triton_client.get_trace_settings(model_name="simple"),
                         "Unexpected model trace settings after global clear")
        self.assertEqual(
            expected_second_model_settings,
            triton_client.update_trace_settings(model_name="simple",
                                                settings=model_clear_settings),
            "Unexpected model trace settings after model clear")
        self.assertEqual(expected_global_settings,
                         triton_client.get_trace_settings(),
                         "Unexpected global trace settings after model clear")


if __name__ == '__main__':
    unittest.main()


# qa/L0_trt_data_dependent_shape/trt_data_dependent_shape_test.py

# Copyright 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../common")

import unittest
import numpy as np
import test_util as tu
import tritonclient.http as client
from tritonclient.utils import InferenceServerException

class TrtDataDependentShapeTest(tu.TestResultCollector):
    def setUp(self):
        self.triton_client = client.InferenceServerClient(
            "localhost:8000", verbose=True)

    def test_fixed(self):
        model_name = "plan_nobatch_nonzero_fixed"
        input_np = np.arange(16, dtype=np.int32).reshape((4, 4))
        expected_output_np = np.nonzero(input_np)

        inputs = []
        inputs.append(client.InferInput('INPUT', [4, 4], "INT32"))
        inputs[-1].set_data_from_numpy(input_np)

        results = self.triton_client.infer(model_name=model_name,
                                           inputs=inputs)
        # Validate the results by comparing with precomputed values.
        output_np = results.as_numpy('OUTPUT')
        self.assertTrue(
            np.array_equal(output_np, expected_output_np),
            "OUTPUT expected: {}, got {}".format(expected_output_np,
                                                  output_np))

    def test_dynamic(self):
        model_name = "plan_nobatch_nonzero_dynamic"
        input_data = []
        for i in range(20*16):
            input_data.append(i if (i % 2) == 0 else 0)
        input_np = np.array(input_data, dtype=np.int32).reshape((20, 16))
        expected_output_np = np.nonzero(input_np)

        inputs = []
        inputs.append(client.InferInput('INPUT', [20, 16], "INT32"))
        inputs[-1].set_data_from_numpy(input_np)

        results = self.triton_client.infer(model_name=model_name,
                                           inputs=inputs)
        # Validate the results by comparing with precomputed values.
        output_np = results.as_numpy('OUTPUT')
        self.assertTrue(
            np.array_equal(output_np, expected_output_np),
            "OUTPUT expected: {}, got {}".format(expected_output_np,
                                                  output_np))


if __name__ == '__main__':
    unittest.main()


# qa/L0_trt_dla/dla_test.py

#!/usr/bin/env python
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../common")

import unittest
import numpy as np
from PIL import Image
import test_util as tu

import tritonclient.http as httpclient
from tritonclient.utils import InferenceServerException


class InferTest(tu.TestResultCollector):

    def _preprocess(self, img, dtype):
        """
        Pre-process an image to meet the size and type
        requirements specified by the parameters.
        """

        sample_img = img.convert('RGB')
        resized_img = sample_img.resize((224, 224), Image.BILINEAR)
        resized = np.array(resized_img)

        typed = resized.astype(dtype)
        scaled = typed - np.asarray((123, 117, 104), dtype=dtype)
        ordered = np.transpose(scaled, (2, 0, 1))

        return ordered

    def test_resnet50(self):
        try:
            triton_client = httpclient.InferenceServerClient(
                url="localhost:8000")
        except Exception as e:
            print("channel creation failed: " + str(e))
            sys.exit(1)

        image_filename = "../images/vulture.jpeg"
        model_name = "resnet50_plan"
        batch_size = 32

        img = Image.open(image_filename)
        image_data = self._preprocess(img, np.int8)
        image_data = np.expand_dims(image_data, axis=0)

        batched_image_data = image_data
        for i in range(1, batch_size):
            batched_image_data = np.concatenate(
                (batched_image_data, image_data), axis=0)

        inputs = [
            httpclient.InferInput('input_tensor_0', [batch_size, 3, 224, 224],
                                  'INT8')
        ]
        inputs[0].set_data_from_numpy(batched_image_data, binary_data=True)

        outputs = [
            httpclient.InferRequestedOutput('topk_layer_output_index',
                                            binary_data=True)
        ]

        results = triton_client.infer(model_name, inputs, outputs=outputs)

        output_data = results.as_numpy('topk_layer_output_index')
        print(output_data)

        # Validate the results by comparing with precomputed values.
        # VULTURE class corresponds with index 23
        EXPECTED_CLASS_INDEX = 23
        for i in range(batch_size):
            self.assertEqual(output_data[i][0][0], EXPECTED_CLASS_INDEX)


if __name__ == '__main__':
    unittest.main()


# qa/L0_trt_dynamic_shape/trt_dynamic_shape_test.py

# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../common")

from builtins import range
from future.utils import iteritems
import os
import shutil
import time
import unittest
import numpy as np
import infer_util as iu
import test_util as tu
import tritonhttpclient
from tritonclientutils import InferenceServerException


class TrtDynamicShapeTest(tu.TestResultCollector):

    def setUp(self):
        self.dtype_ = np.float32
        self.model_name_ = 'plan'

    def test_load_specific_optimization_profile(self):
        # Only OP 5 should be available, which only allow batch size 8
        tensor_shape = (1,)
        try:
            iu.infer_exact(self, self.model_name_, (1,) + tensor_shape, 1,
                           self.dtype_, self.dtype_, self.dtype_)
        except InferenceServerException as ex:
            self.assertTrue(
                "model expected the shape of dimension 0 to be between 6 and 8 but received 1"
                in ex.message())

        try:
            iu.infer_exact(self, self.model_name_, (8,) + tensor_shape, 8,
                           self.dtype_, self.dtype_, self.dtype_)
        except InferenceServerException as ex:
            self.assertTrue(False, "unexpected error {}".format(ex))

    def test_load_default_optimization_profile(self):
        # Only default OP (OP 0) has max tensor shape 33
        tensor_shape = (33,)

        try:
            iu.infer_exact(self, self.model_name_, (8,) + tensor_shape, 8,
                           self.dtype_, self.dtype_, self.dtype_)
        except InferenceServerException as ex:
            self.assertTrue(False, "unexpected error {}".format(ex))

        over_tensor_shape = (34,)
        try:
            iu.infer_exact(self, self.model_name_, (8,) + over_tensor_shape, 8,
                           self.dtype_, self.dtype_, self.dtype_)
        except InferenceServerException as ex:
            self.assertTrue(
                "model expected the shape of dimension 1 to be between 1 and 33 but received 34"
                in ex.message())

    def test_select_optimization_profile(self):
        # Different profile has different optimized input shape
        batch_size = 4
        tensor_shape = (16,)
        try:
            iu.infer_exact(self, self.model_name_, (batch_size,) + tensor_shape,
                           batch_size, self.dtype_, self.dtype_, self.dtype_)
        except InferenceServerException as ex:
            self.assertTrue(False, "unexpected error {}".format(ex))

    def test_load_wrong_optimization_profile(self):
        client = tritonhttpclient.InferenceServerClient("localhost:8000")
        model_name = tu.get_model_name(self.model_name_, self.dtype_,
                                       self.dtype_, self.dtype_)
        model_status = client.is_model_ready(model_name, "1")
        self.assertFalse(model_status, "expected model to be not ready")


if __name__ == '__main__':
    unittest.main()


# qa/L0_trt_plugin/trt_plugin_test.py

# Copyright (c) 2018-2020, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../common")

from builtins import range
from future.utils import iteritems
import unittest
import numpy as np
import os
import test_util as tu

import tritonclient.http as httpclient
from tritonclient.utils import InferenceServerException

# By default, find tritonserver on "localhost", but can be overridden
# with TRITONSERVER_IPADDR envvar
_tritonserver_ipaddr = os.environ.get('TRITONSERVER_IPADDR', 'localhost')


class PluginModelTest(tu.TestResultCollector):

    def _full_exact(self, model_name, plugin_name, shape):
        print(f"{_tritonserver_ipaddr}:8000")
        triton_client = httpclient.InferenceServerClient(
            f"{_tritonserver_ipaddr}:8000")

        inputs = []
        outputs = []
        inputs.append(httpclient.InferInput('INPUT0', list(shape), "FP32"))

        input0_data = np.ones(shape=shape).astype(np.float32)
        inputs[0].set_data_from_numpy(input0_data, binary_data=True)

        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT0', binary_data=True))

        results = triton_client.infer(model_name + '_' + plugin_name,
                                      inputs,
                                      outputs=outputs)

        output0_data = results.as_numpy('OUTPUT0')

        # Verify values of Clip, GELU, and Normalize
        if plugin_name == 'CustomClipPlugin':
            # Clip data to minimum of .1, maximum of .5
            test_output = np.clip(input0_data, 0.1, 0.5)
            self.assertTrue(np.isclose(output0_data, test_output).all())
        elif plugin_name == 'CustomGeluPluginDynamic':
            # Add bias
            input0_data += 1
            # Calculate Gelu activation
            test_output = (input0_data *
                           0.5) * (1 + np.tanh((0.797885 * input0_data) +
                                               (0.035677 * (input0_data**3))))
            self.assertTrue(np.isclose(output0_data, test_output).all())
        elif plugin_name == 'Normalize_TRT':
            # L2 norm is sqrt(sum([1]*16)))
            test_output = input0_data / np.sqrt(sum([1] * 16))
            self.assertTrue(np.isclose(output0_data, test_output).all())
        else:
            self.assertTrue(False, "Unexpected plugin: " + plugin_name)

    def test_raw_fff_clip(self):
        for bs in (1, 8):
            self._full_exact('plan_float32_float32_float32', 'CustomClipPlugin',
                             (bs, 16))

    def test_raw_fff_gelu(self):
        self._full_exact('plan_nobatch_float32_float32_float32',
                         'CustomGeluPluginDynamic', (16, 1, 1))

    def test_raw_fff_norm(self):
        # model that supports batching
        for bs in (1, 8):
            self._full_exact('plan_float32_float32_float32', 'Normalize_TRT',
                             (bs, 16, 16, 16))


if __name__ == '__main__':
    unittest.main()


# qa/L0_trt_reformat_free/trt_reformat_free_test.py

# Copyright 2020-2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../common")

from builtins import range
from future.utils import iteritems
import os
import shutil
import time
import unittest
import numpy as np
import infer_util as iu
import test_util as tu
import tritonhttpclient
import tritonclient.utils.shared_memory as shm
from tritonclientutils import InferenceServerException


def div_up(a, b):
    return (a + b - 1) // b


def reformat(format, tensor_np):
    if format == "CHW2":
        factor = 2
    if format == "CHW32":
        factor = 32
    shape = list(tensor_np.shape) + [factor]
    shape[-4] = div_up(shape[-4], factor)
    reformatted_tensor_np = np.empty(shape, tensor_np.dtype)
    if len(tensor_np.shape) == 3:
        batch = [(tensor_np, reformatted_tensor_np)]
    elif len(tensor_np.shape) == 4:
        batch = [(tensor_np[idx], reformatted_tensor_np[idx])
                 for idx in range(tensor_np.shape[0])]
    else:
        raise ValueError(
            "Unexpected numpy shape {} for testing reformat-free input".format(
                tensor_np.shape))
    for (tensor, reformatted_tensor) in batch:
        for c in range(tensor.shape[0]):
            for h in range(tensor.shape[1]):
                for w in range(tensor.shape[2]):
                    reformatted_tensor[c //
                                       factor][h][w][c %
                                                     factor] = tensor[c][h][w]
    return reformatted_tensor_np


class TrtReformatFreeTest(tu.TestResultCollector):

    def add_reformat_free_data_as_shared_memory(self, name, tensor, tensor_np):
        byte_size = tensor_np.size * tensor_np.dtype.itemsize
        self.shm_handles.append(
            shm.create_shared_memory_region(name, name, byte_size))
        # Put data values into shared memory
        shm.set_shared_memory_region(self.shm_handles[-1], [tensor_np])
        # Register shared memory with Triton Server
        self.triton_client.register_system_shared_memory(name, name, byte_size)
        # Set the parameters to use data from shared memory
        tensor.set_shared_memory(name, byte_size)

    def setUp(self):
        self.shm_handles = []
        self.triton_client = tritonhttpclient.InferenceServerClient(
            "localhost:8000", verbose=True)

    def tearDown(self):
        self.triton_client.unregister_system_shared_memory()
        for handle in self.shm_handles:
            shm.destroy_shared_memory_region(handle)

    def test_nobatch_chw2_input(self):
        model_name = "plan_nobatch_CHW2_LINEAR_float16_float16_float16"
        input_np = np.arange(26, dtype=np.float16).reshape((13, 2, 1))
        expected_output0_np = input_np + input_np
        expected_output1_np = input_np - input_np
        reformatted_input_np = reformat("CHW2", input_np)

        # Use shared memory to bypass the shape check in client library, because
        # for non-linear format tensor, the data buffer is padded and thus the
        # data byte size may not match what is calculated from tensor shape
        inputs = []
        inputs.append(tritonhttpclient.InferInput('INPUT0', [13, 2, 1], "FP16"))
        self.add_reformat_free_data_as_shared_memory("input0", inputs[-1],
                                                     reformatted_input_np)
        inputs.append(tritonhttpclient.InferInput('INPUT1', [13, 2, 1], "FP16"))
        self.add_reformat_free_data_as_shared_memory("input1", inputs[-1],
                                                     reformatted_input_np)

        outputs = []
        outputs.append(
            tritonhttpclient.InferRequestedOutput('OUTPUT0', binary_data=True))
        outputs.append(
            tritonhttpclient.InferRequestedOutput('OUTPUT1', binary_data=True))

        results = self.triton_client.infer(model_name=model_name,
                                           inputs=inputs,
                                           outputs=outputs)
        # Validate the results by comparing with precomputed values.
        output0_np = results.as_numpy('OUTPUT0')
        output1_np = results.as_numpy('OUTPUT1')
        self.assertTrue(
            np.array_equal(output0_np, expected_output0_np),
            "OUTPUT0 expected: {}, got {}".format(expected_output0_np,
                                                  output0_np))
        self.assertTrue(
            np.array_equal(output1_np, expected_output1_np),
            "OUTPUT0 expected: {}, got {}".format(expected_output1_np,
                                                  output1_np))

    def test_chw2_input(self):
        model_name = "plan_CHW2_LINEAR_float16_float16_float16"
        for bs in [1, 8]:
            input_np = np.arange(26 * bs, dtype=np.float16).reshape(
                (bs, 13, 2, 1))
            expected_output0_np = input_np + input_np
            expected_output1_np = input_np - input_np
            reformatted_input_np = reformat("CHW2", input_np)

            # Use shared memory to bypass the shape check in client library,
            # because for non-linear format tensor, the data buffer is padded
            # and thus the data byte size may not match what is calculated from
            # tensor shape
            inputs = []
            inputs.append(
                tritonhttpclient.InferInput('INPUT0', [bs, 13, 2, 1], "FP16"))
            self.add_reformat_free_data_as_shared_memory(
                "input0" + str(bs), inputs[-1], reformatted_input_np)
            inputs.append(
                tritonhttpclient.InferInput('INPUT1', [bs, 13, 2, 1], "FP16"))
            self.add_reformat_free_data_as_shared_memory(
                "input1" + str(bs), inputs[-1], reformatted_input_np)

            outputs = []
            outputs.append(
                tritonhttpclient.InferRequestedOutput('OUTPUT0',
                                                      binary_data=True))
            outputs.append(
                tritonhttpclient.InferRequestedOutput('OUTPUT1',
                                                      binary_data=True))

            results = self.triton_client.infer(model_name=model_name,
                                               inputs=inputs,
                                               outputs=outputs)
            # Validate the results by comparing with precomputed values.
            output0_np = results.as_numpy('OUTPUT0')
            output1_np = results.as_numpy('OUTPUT1')
            self.assertTrue(
                np.array_equal(output0_np, expected_output0_np),
                "OUTPUT0 expected: {}, got {}".format(expected_output0_np,
                                                      output0_np))
            self.assertTrue(
                np.array_equal(output1_np, expected_output1_np),
                "OUTPUT0 expected: {}, got {}".format(expected_output1_np,
                                                      output1_np))

    def test_nobatch_chw32_input(self):
        model_name = "plan_nobatch_CHW32_LINEAR_float32_float32_float32"
        input_np = np.arange(26, dtype=np.float32).reshape((13, 2, 1))
        expected_output0_np = input_np + input_np
        expected_output1_np = input_np - input_np
        reformatted_input_np = reformat("CHW32", input_np)

        # Use shared memory to bypass the shape check in client library, because
        # for non-linear format tensor, the data buffer is padded and thus the
        # data byte size may not match what is calculated from tensor shape
        inputs = []
        inputs.append(tritonhttpclient.InferInput('INPUT0', [13, 2, 1], "FP32"))
        self.add_reformat_free_data_as_shared_memory("input0", inputs[-1],
                                                     reformatted_input_np)
        inputs.append(tritonhttpclient.InferInput('INPUT1', [13, 2, 1], "FP32"))
        self.add_reformat_free_data_as_shared_memory("input1", inputs[-1],
                                                     reformatted_input_np)

        outputs = []
        outputs.append(
            tritonhttpclient.InferRequestedOutput('OUTPUT0', binary_data=True))
        outputs.append(
            tritonhttpclient.InferRequestedOutput('OUTPUT1', binary_data=True))

        results = self.triton_client.infer(model_name=model_name,
                                           inputs=inputs,
                                           outputs=outputs)
        # Validate the results by comparing with precomputed values.
        output0_np = results.as_numpy('OUTPUT0')
        output1_np = results.as_numpy('OUTPUT1')
        self.assertTrue(
            np.array_equal(output0_np, expected_output0_np),
            "OUTPUT0 expected: {}, got {}".format(expected_output0_np,
                                                  output0_np))
        self.assertTrue(
            np.array_equal(output1_np, expected_output1_np),
            "OUTPUT0 expected: {}, got {}".format(expected_output1_np,
                                                  output1_np))

    def test_chw32_input(self):
        model_name = "plan_CHW32_LINEAR_float32_float32_float32"
        for bs in [1, 8]:
            input_np = np.arange(26 * bs, dtype=np.float32).reshape(
                (bs, 13, 2, 1))
            expected_output0_np = input_np + input_np
            expected_output1_np = input_np - input_np
            reformatted_input_np = reformat("CHW32", input_np)

            # Use shared memory to bypass the shape check in client library,
            # because for non-linear format tensor, the data buffer is padded
            # and thus the data byte size may not match what is calculated from
            # tensor shape
            inputs = []
            inputs.append(
                tritonhttpclient.InferInput('INPUT0', [bs, 13, 2, 1], "FP32"))
            self.add_reformat_free_data_as_shared_memory(
                "input0" + str(bs), inputs[-1], reformatted_input_np)
            inputs.append(
                tritonhttpclient.InferInput('INPUT1', [bs, 13, 2, 1], "FP32"))
            self.add_reformat_free_data_as_shared_memory(
                "input1" + str(bs), inputs[-1], reformatted_input_np)

            outputs = []
            outputs.append(
                tritonhttpclient.InferRequestedOutput('OUTPUT0',
                                                      binary_data=True))
            outputs.append(
                tritonhttpclient.InferRequestedOutput('OUTPUT1',
                                                      binary_data=True))

            results = self.triton_client.infer(model_name=model_name,
                                               inputs=inputs,
                                               outputs=outputs)
            # Validate the results by comparing with precomputed values.
            output0_np = results.as_numpy('OUTPUT0')
            output1_np = results.as_numpy('OUTPUT1')
            self.assertTrue(
                np.array_equal(output0_np, expected_output0_np),
                "OUTPUT0 expected: {}, got {}".format(expected_output0_np,
                                                      output0_np))
            self.assertTrue(
                np.array_equal(output1_np, expected_output1_np),
                "OUTPUT0 expected: {}, got {}".format(expected_output1_np,
                                                      output1_np))


if __name__ == '__main__':
    unittest.main()


# qa/L0_trt_shape_tensors/trt_shape_tensor_test.py

# Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../common")

from builtins import range
from future.utils import iteritems
import os
import unittest
import time
import threading
import traceback
import numpy as np
import infer_util as iu
import test_util as tu
import sequence_util as su

import tritongrpcclient as grpcclient

TEST_SYSTEM_SHARED_MEMORY = bool(
    int(os.environ.get('TEST_SYSTEM_SHARED_MEMORY', 0)))

_model_instances = 1
_max_queue_delay_ms = 10000
_max_sequence_idle_ms = 5000

_deferred_exceptions_lock = threading.Lock()
_deferred_exceptions = []


class InferShapeTensorTest(tu.TestResultCollector):

    def setUp(self):
        # The helper client for setup will be GRPC for simplicity.
        self.triton_client_ = grpcclient.InferenceServerClient("localhost:8001")
        global _deferred_exceptions
        _deferred_exceptions = []

    def tearDown(self):
        self.triton_client_.unregister_system_shared_memory()
        self.triton_client_.unregister_cuda_shared_memory()
        super().tearDown()

    def add_deferred_exception(self, ex):
        global _deferred_exceptions
        with _deferred_exceptions_lock:
            _deferred_exceptions.append(ex)

    def check_deferred_exception(self):
        # Just raise one of the exceptions...
        with _deferred_exceptions_lock:
            if len(_deferred_exceptions) > 0:
                raise _deferred_exceptions[0]

    def check_response(self,
                       bs,
                       thresholds,
                       shape_values,
                       dummy_input_shapes,
                       shm_region_names=None,
                       precreated_shm_regions=None,
                       shm_suffix=""):
        try:
            # Add batch size to shape as full shape is expected
            for i in range(len(dummy_input_shapes)):
                dummy_input_shapes[i] = [
                    bs,
                ] + dummy_input_shapes[i]
            start_ms = int(round(time.time() * 1000))

            iu.infer_shape_tensor(
                self,
                'plan',
                np.float32,
                shape_values,
                dummy_input_shapes,
                use_grpc=False,
                use_streaming=False,
                shm_suffix=shm_suffix,
                use_system_shared_memory=TEST_SYSTEM_SHARED_MEMORY,
                batch_size=bs)

            end_ms = int(round(time.time() * 1000))

            lt_ms = thresholds[0]
            gt_ms = thresholds[1]
            if lt_ms is not None:
                self.assertTrue(
                    (end_ms - start_ms) < lt_ms,
                    "expected less than " + str(lt_ms) +
                    "ms response time, got " + str(end_ms - start_ms) + " ms")
            if gt_ms is not None:
                self.assertTrue(
                    (end_ms - start_ms) > gt_ms,
                    "expected greater than " + str(gt_ms) +
                    "ms response time, got " + str(end_ms - start_ms) + " ms")
        except Exception as ex:
            self.add_deferred_exception(ex)

    def check_setup(self, model_name):
        # Make sure test.sh set up the correct batcher settings
        config = self.triton_client_.get_model_config(model_name).config
        bconfig = config.dynamic_batching
        self.assertTrue(2 in bconfig.preferred_batch_size)
        self.assertTrue(6 in bconfig.preferred_batch_size)
        self.assertEqual(bconfig.max_queue_delay_microseconds,
                         _max_queue_delay_ms * 1000)  # 10 secs

    def check_status(self, model_name, batch_exec, exec_cnt, infer_cnt):
        stats = self.triton_client_.get_inference_statistics(model_name, "1")
        self.assertEqual(len(stats.model_stats), 1, "expect 1 model stats")
        self.assertEqual(stats.model_stats[0].name, model_name,
                         "expect model stats for model {}".format(model_name))
        self.assertEqual(
            stats.model_stats[0].version, "1",
            "expect model stats for model {} version 1".format(model_name))

        if batch_exec is not None:
            batch_stats = stats.model_stats[0].batch_stats
            print(batch_stats)
            self.assertEqual(
                len(batch_stats), len(batch_exec),
                "expected {} different batch-sizes, got {}".format(
                    len(batch_exec), len(batch_stats)))

            for batch_stat in batch_stats:
                bs = batch_stat.batch_size
                bc = batch_stat.compute_infer.count
                self.assertTrue(
                    bs in batch_exec,
                    "did not find expected batch-size {}".format(bs))
                # Get count from one of the stats
                self.assertEqual(
                    bc, batch_exec[bs],
                    "expected model-execution-count {} for batch size {}, got {}"
                    .format(batch_exec[bs], bs, bc))

        actual_exec_cnt = stats.model_stats[0].execution_count
        self.assertEqual(
            actual_exec_cnt, exec_cnt,
            "expected model-exec-count {}, got {}".format(
                exec_cnt, actual_exec_cnt))

        actual_infer_cnt = stats.model_stats[0].inference_count
        self.assertEqual(
            actual_infer_cnt, infer_cnt,
            "expected model-inference-count {}, got {}".format(
                infer_cnt, actual_infer_cnt))

        actual_infer_cnt = stats.model_stats[0].inference_count
        self.assertEqual(
            actual_infer_cnt, infer_cnt,
            "expected model-inference-count {}, got {}".format(
                infer_cnt, actual_infer_cnt))

    def test_static_batch(self):
        iu.infer_shape_tensor(
            self,
            'plan',
            np.float32, [[32, 32]], [[8, 4, 4]],
            use_system_shared_memory=TEST_SYSTEM_SHARED_MEMORY,
            batch_size=8)
        iu.infer_shape_tensor(
            self,
            'plan',
            np.float32, [[4, 4]], [[8, 32, 32]],
            use_system_shared_memory=TEST_SYSTEM_SHARED_MEMORY,
            batch_size=8)
        iu.infer_shape_tensor(
            self,
            'plan',
            np.float32, [[4, 4]], [[8, 4, 4]],
            use_system_shared_memory=TEST_SYSTEM_SHARED_MEMORY,
            batch_size=8)

    def test_nobatch(self):
        iu.infer_shape_tensor(
            self,
            'plan_nobatch',
            np.float32, [[32, 32]], [[4, 4]],
            use_system_shared_memory=TEST_SYSTEM_SHARED_MEMORY)
        iu.infer_shape_tensor(
            self,
            'plan_nobatch',
            np.float32, [[4, 4]], [[32, 32]],
            use_system_shared_memory=TEST_SYSTEM_SHARED_MEMORY)
        iu.infer_shape_tensor(
            self,
            'plan_nobatch',
            np.float32, [[4, 4]], [[4, 4]],
            use_system_shared_memory=TEST_SYSTEM_SHARED_MEMORY)

    def test_wrong_shape_values(self):
        over_shape_values = [[32, 33]]
        try:
            iu.infer_shape_tensor(
                self,
                'plan',
                np.float32,
                over_shape_values, [[8, 4, 4]],
                use_system_shared_memory=TEST_SYSTEM_SHARED_MEMORY,
                batch_size=8)
        # InferenceServerException will be raised from different namespace,
        # use dynamic type characteristic to catch both ex
        except Exception as ex:
            self.assertTrue(
                "The shape value at index 2 is expected to be in range from 1 to 32, Got: 33"
                in ex.message())

    # Dynamic Batcher tests
    def test_dynamic_different_shape_values(self):
        # Send two requests with sum of static batch sizes ==
        # preferred size, but with different shape values. This
        # should cause the requests to not be batched. The first
        # response will come back immediately and the second
        # delayed by the max batch queue delay
        try:
            model_name = tu.get_zero_model_name("plan", 1, np.float32)
            self.check_setup(model_name)
            self.assertFalse("TRITONSERVER_DELAY_SCHEDULER" in os.environ)

            threads = []
            threads.append(
                threading.Thread(target=self.check_response,
                                 args=(3, (6000, None)),
                                 kwargs={
                                     'shape_values': [[2, 2]],
                                     'dummy_input_shapes': [[16, 16]],
                                     'shm_suffix': '{}'.format(len(threads))
                                 }))
            threads.append(
                threading.Thread(target=self.check_response,
                                 args=(3, (_max_queue_delay_ms * 1.5,
                                           _max_queue_delay_ms)),
                                 kwargs={
                                     'shape_values': [[4, 4]],
                                     'dummy_input_shapes': [[16, 16]],
                                     'shm_suffix': '{}'.format(len(threads))
                                 }))
            threads[0].start()
            time.sleep(1)
            threads[1].start()
            for t in threads:
                t.join()
            self.check_deferred_exception()
            self.check_status(model_name, {3: 2}, 2, 6)
        except Exception as ex:
            self.assertTrue(False, "unexpected error {}".format(ex))

    def test_dynamic_identical_shape_values(self):
        # Send two requests with sum of static batch sizes ==
        # preferred size, but with identical shape values. This
        # should cause the requests to get batched. Both
        # responses should come back immediately.
        try:
            model_name = tu.get_zero_model_name("plan", 1, np.float32)
            self.check_setup(model_name)
            self.assertFalse("TRITONSERVER_DELAY_SCHEDULER" in os.environ)

            threads = []
            threads.append(
                threading.Thread(target=self.check_response,
                                 args=(4, (6000, None)),
                                 kwargs={
                                     'shape_values': [[4, 4]],
                                     'dummy_input_shapes': [[16, 16]],
                                     'shm_suffix': '{}'.format(len(threads))
                                 }))
            threads.append(
                threading.Thread(target=self.check_response,
                                 args=(2, (6000, None)),
                                 kwargs={
                                     'shape_values': [[4, 4]],
                                     'dummy_input_shapes': [[16, 16]],
                                     'shm_suffix': '{}'.format(len(threads))
                                 }))
            threads[0].start()
            time.sleep(1)
            threads[1].start()
            for t in threads:
                t.join()
            self.check_deferred_exception()
            self.check_status(model_name, {6: 1}, 1, 6)
        except Exception as ex:
            self.assertTrue(False, "unexpected error {}".format(ex))


class SequenceBatcherShapeTensorTest(su.SequenceBatcherTestUtil):

    def get_expected_result(self, expected_result, value, flag_str=None):
        # Adjust the expected_result for models
        expected_result = value
        if (flag_str is not None) and ("start" in flag_str):
            expected_result += 1
        return expected_result

    def test_sequence_identical_shape_values(self):
        # Test model instances together are configured with
        # total-batch-size 4. Send four equal-length sequences
        # with identical shape values in parallel and make sure
        # they get completely batched into batch-size 4
        # inferences.
        self.clear_deferred_exceptions()
        dtype = np.float32
        try:
            model_name = tu.get_sequence_model_name("plan", dtype)
            self.check_setup(model_name)

            # Need scheduler to wait for queue to contain all
            # inferences for both sequences.
            self.assertTrue("TRITONSERVER_DELAY_SCHEDULER" in os.environ)
            self.assertEqual(int(os.environ["TRITONSERVER_DELAY_SCHEDULER"]),
                             12)
            self.assertTrue(
                "TRITONSERVER_BACKLOG_DELAY_SCHEDULER" in os.environ)
            self.assertEqual(
                int(os.environ["TRITONSERVER_BACKLOG_DELAY_SCHEDULER"]), 0)
            precreated_shm0_handles = self.precreate_register_shape_tensor_regions(
                ((2, 1), (4, 2), (8, 3)), dtype, 0)
            precreated_shm1_handles = self.precreate_register_shape_tensor_regions(
                ((2, 11), (4, 12), (8, 13)), dtype, 1)
            precreated_shm2_handles = self.precreate_register_shape_tensor_regions(
                ((2, 111), (4, 112), (8, 113)), dtype, 2)
            precreated_shm3_handles = self.precreate_register_shape_tensor_regions(
                ((2, 1111), (4, 1112), (8, 1113)), dtype, 3)
            threads = []
            threads.append(
                threading.Thread(
                    target=self.check_sequence_shape_tensor_io,
                    args=(
                        model_name,
                        dtype,
                        1001,
                        (None, None),
                        # (flag_str, shape_value, value, pre_delay_ms)
                        (("start", 2, 1, None), (None, 4, 2, None), ("end", 8,
                                                                     3, None)),
                        self.get_expected_result(6, 3, "end"),
                        precreated_shm0_handles),
                    kwargs={'sequence_name': "{}".format(self._testMethodName)
                           }))
            threads.append(
                threading.Thread(
                    target=self.check_sequence_shape_tensor_io,
                    args=(
                        model_name,
                        dtype,
                        1002,
                        (None, None),
                        # (flag_str, shape_value, value, pre_delay_ms)
                        (("start", 2, 11, None), (None, 4, 12, None),
                         ("end", 8, 13, None)),
                        self.get_expected_result(36, 13, "end"),
                        precreated_shm1_handles),
                    kwargs={'sequence_name': "{}".format(self._testMethodName)
                           }))
            threads.append(
                threading.Thread(
                    target=self.check_sequence_shape_tensor_io,
                    args=(
                        model_name,
                        dtype,
                        1003,
                        (None, None),
                        # (flag_str, shape_value, value, pre_delay_ms)
                        (("start", 2, 111, None), (None, 4, 112, None),
                         ("end", 8, 113, None)),
                        self.get_expected_result(336, 113, "end"),
                        precreated_shm2_handles),
                    kwargs={'sequence_name': "{}".format(self._testMethodName)
                           }))
            threads.append(
                threading.Thread(
                    target=self.check_sequence_shape_tensor_io,
                    args=(
                        model_name,
                        dtype,
                        1004,
                        (None, None),
                        # (flag_str, shape_value, value, pre_delay_ms)
                        (("start", 2, 1111, None), (None, 4, 1112, None),
                         ("end", 8, 1113, None)),
                        self.get_expected_result(3336, 1113, "end"),
                        precreated_shm3_handles),
                    kwargs={'sequence_name': "{}".format(self._testMethodName)
                           }))

            for t in threads:
                t.start()
            for t in threads:
                t.join()
            self.check_deferred_exception()
            self.check_status(model_name, {4: 3}, 3, 12)
        except Exception as ex:
            self.assertTrue(False, "unexpected error {}".format(ex))
        finally:
            if TEST_SYSTEM_SHARED_MEMORY:
                self.cleanup_shm_regions(precreated_shm0_handles)
                self.cleanup_shm_regions(precreated_shm1_handles)
                self.cleanup_shm_regions(precreated_shm2_handles)
                self.cleanup_shm_regions(precreated_shm3_handles)

    def test_sequence_different_shape_values(self):
        # Test model instances together are configured with
        # total-batch-size 4. Send four equal-length sequences with
        # different shape values in 2 sequences and 2 sequences that
        # share the same shape value. Make sure that the 2 sequences
        # with same shapes batch together but other two sequences do
        # not.
        self.clear_deferred_exceptions()
        dtype = np.float32

        precreated_shm0_handles = self.precreate_register_shape_tensor_regions(
            ((1, 1), (1, 2), (1, 3)), dtype, 0)
        precreated_shm1_handles = self.precreate_register_shape_tensor_regions(
            ((32, 11), (32, 12), (32, 13)), dtype, 1)
        precreated_shm2_handles = self.precreate_register_shape_tensor_regions(
            ((16, 111), (16, 112), (16, 113)), dtype, 2)
        precreated_shm3_handles = self.precreate_register_shape_tensor_regions(
            ((1, 1111), (1, 1112), (1, 1113)), dtype, 3)
        try:
            model_name = tu.get_sequence_model_name("plan", dtype)
            self.check_setup(model_name)

            # Need scheduler to wait for queue to contain all
            # inferences for both sequences.
            self.assertTrue("TRITONSERVER_DELAY_SCHEDULER" in os.environ)
            self.assertEqual(int(os.environ["TRITONSERVER_DELAY_SCHEDULER"]),
                             12)
            self.assertTrue(
                "TRITONSERVER_BACKLOG_DELAY_SCHEDULER" in os.environ)
            self.assertEqual(
                int(os.environ["TRITONSERVER_BACKLOG_DELAY_SCHEDULER"]), 0)

            threads = []
            threads.append(
                threading.Thread(
                    target=self.check_sequence_shape_tensor_io,
                    args=(
                        model_name,
                        dtype,
                        1001,
                        (None, None),
                        # (flag_str, shape_value, value, pre_delay_ms)
                        (("start", 1, 1, None), (None, 1, 2, None), ("end", 1,
                                                                     3, None)),
                        self.get_expected_result(6, 3, "end"),
                        precreated_shm0_handles),
                    kwargs={'sequence_name': "{}".format(self._testMethodName)
                           }))
            threads.append(
                threading.Thread(
                    target=self.check_sequence_shape_tensor_io,
                    args=(
                        model_name,
                        dtype,
                        1002,
                        (None, None),
                        # (flag_str, shape_value, value, pre_delay_ms)
                        (("start", 32, 11, None), (None, 32, 12, None),
                         ("end", 32, 13, None)),
                        self.get_expected_result(36, 13, "end"),
                        precreated_shm1_handles),
                    kwargs={'sequence_name': "{}".format(self._testMethodName)
                           }))
            threads.append(
                threading.Thread(
                    target=self.check_sequence_shape_tensor_io,
                    args=(
                        model_name,
                        dtype,
                        1003,
                        (None, None),
                        # (flag_str, shape_value, value, pre_delay_ms)
                        (("start", 16, 111, None), (None, 16, 112, None),
                         ("end", 16, 113, None)),
                        self.get_expected_result(336, 113, "end"),
                        precreated_shm2_handles),
                    kwargs={'sequence_name': "{}".format(self._testMethodName)
                           }))
            threads.append(
                threading.Thread(
                    target=self.check_sequence_shape_tensor_io,
                    args=(
                        model_name,
                        dtype,
                        1004,
                        (None, None),
                        # (flag_str, shape_value, value, pre_delay_ms)
                        (("start", 1, 1111, None), (None, 1, 1112, None),
                         ("end", 1, 1113, None)),
                        self.get_expected_result(3336, 1113, "end"),
                        precreated_shm3_handles),
                    kwargs={'sequence_name': "{}".format(self._testMethodName)
                           }))

            for t in threads:
                t.start()
                time.sleep(1)
            for t in threads:
                t.join()

            self.check_deferred_exception()
            self.check_status(model_name, {4: 3, 3: 6}, 9, 12)
        except Exception as ex:
            self.assertTrue(False, "unexpected error {}".format(ex))
        finally:
            if TEST_SYSTEM_SHARED_MEMORY:
                self.cleanup_shm_regions(precreated_shm0_handles)
                self.cleanup_shm_regions(precreated_shm1_handles)
                self.cleanup_shm_regions(precreated_shm2_handles)
                self.cleanup_shm_regions(precreated_shm3_handles)


class DynaSequenceBatcherTest(su.SequenceBatcherTestUtil):

    def get_expected_result(self,
                            expected_result,
                            corrid,
                            value,
                            flag_str=None):
        expected_result = value
        if flag_str is not None:
            if "start" in flag_str:
                expected_result += 1
            if "end" in flag_str:
                expected_result += corrid
        return expected_result

    def _multi_sequence_different_shape_impl(self, sleep_secs):
        self.clear_deferred_exceptions()
        dtype = np.float32

        precreated_shm0_handles = self.precreate_register_dynaseq_shape_tensor_regions(
            ((1, 1), (12, 2), (2, 3)), dtype, 0)
        precreated_shm1_handles = self.precreate_register_dynaseq_shape_tensor_regions(
            ((3, 11), (4, 12), (5, 13)), dtype, 1)
        precreated_shm2_handles = self.precreate_register_dynaseq_shape_tensor_regions(
            ((6, 111), (7, 112), (8, 113)), dtype, 2)
        precreated_shm3_handles = self.precreate_register_dynaseq_shape_tensor_regions(
            ((9, 1111), (10, 1112), (11, 1113)), dtype, 3)

        try:
            model_name = tu.get_dyna_sequence_model_name("plan", dtype)
            self.check_setup(model_name)
            self.assertFalse("TRITONSERVER_DELAY_SCHEDULER" in os.environ)
            self.assertFalse(
                "TRITONSERVER_BACKLOG_DELAY_SCHEDULER" in os.environ)

            corrids = [1001, 1002, 1003, 1004]
            threads = []
            threads.append(
                threading.Thread(
                    target=self.check_sequence_shape_tensor_io,
                    args=(
                        model_name,
                        dtype,
                        corrids[0],
                        (None, None),
                        # (flag_str, shape_value, value, pre_delay_ms)
                        (("start", 1, 1, None), (None, 12, 2, None), ("end", 2,
                                                                      3, None)),
                        self.get_expected_result(4 + corrids[0], corrids[0], 3,
                                                 "end"),
                        precreated_shm0_handles),
                    kwargs={
                        'sequence_name':
                            "{}_{}".format(self._testMethodName, corrids[0]),
                        'using_dynamic_batcher':
                            True
                    }))
            threads.append(
                threading.Thread(
                    target=self.check_sequence_shape_tensor_io,
                    args=(
                        model_name,
                        dtype,
                        corrids[1],
                        (None, None),
                        # (flag_str, shape_value, value, pre_delay_ms)
                        (("start", 3, 11, None), (None, 4, 12, None),
                         ("end", 5, 13, None)),
                        self.get_expected_result(36 + corrids[1], corrids[1],
                                                 13, "end"),
                        precreated_shm1_handles),
                    kwargs={
                        'sequence_name':
                            "{}_{}".format(self._testMethodName, corrids[1]),
                        'using_dynamic_batcher':
                            True
                    }))
            threads.append(
                threading.Thread(
                    target=self.check_sequence_shape_tensor_io,
                    args=(
                        model_name,
                        dtype,
                        corrids[2],
                        (None, None),
                        # (flag_str, shape_value, value, pre_delay_ms)
                        (("start", 6, 111, None), (None, 7, 112, None),
                         ("end", 8, 113, None)),
                        self.get_expected_result(336 + corrids[2], corrids[2],
                                                 113, "end"),
                        precreated_shm2_handles),
                    kwargs={
                        'sequence_name':
                            "{}_{}".format(self._testMethodName, corrids[2]),
                        'using_dynamic_batcher':
                            True
                    }))
            threads.append(
                threading.Thread(
                    target=self.check_sequence_shape_tensor_io,
                    args=(
                        model_name,
                        dtype,
                        corrids[3],
                        (None, None),
                        # (flag_str, shape_value, value, pre_delay_ms)
                        (("start", 9, 1111, None), (None, 10, 1112, None),
                         ("end", 11, 1113, None)),
                        self.get_expected_result(3336 + corrids[3], corrids[3],
                                                 1113, "end"),
                        precreated_shm3_handles),
                    kwargs={
                        'sequence_name':
                            "{}_{}".format(self._testMethodName, corrids[3]),
                        'using_dynamic_batcher':
                            True
                    }))

            for t in threads:
                t.start()
                if sleep_secs > 0:
                    time.sleep(sleep_secs)
            for t in threads:
                t.join()
            self.check_deferred_exception()
            self.check_status(model_name, {1: 12}, 12, 12)
        except Exception as ex:
            self.assertTrue(False, "unexpected error {}".format(ex))
        finally:
            if TEST_SYSTEM_SHARED_MEMORY:
                self.cleanup_shm_regions(precreated_shm0_handles)
                self.cleanup_shm_regions(precreated_shm1_handles)
                self.cleanup_shm_regions(precreated_shm2_handles)
                self.cleanup_shm_regions(precreated_shm3_handles)

    def _multi_sequence_identical_shape_impl(self, sleep_secs):
        self.clear_deferred_exceptions()
        dtype = np.float32

        precreated_shm0_handles = self.precreate_register_dynaseq_shape_tensor_regions(
            ((2, 1), (4, 2), (8, 3)), dtype, 0)
        precreated_shm1_handles = self.precreate_register_dynaseq_shape_tensor_regions(
            ((2, 11), (4, 12), (8, 13)), dtype, 1)
        precreated_shm2_handles = self.precreate_register_dynaseq_shape_tensor_regions(
            ((2, 111), (4, 112), (8, 113)), dtype, 2)
        precreated_shm3_handles = self.precreate_register_dynaseq_shape_tensor_regions(
            ((2, 1111), (4, 1112), (8, 1113)), dtype, 3)

        try:
            model_name = tu.get_dyna_sequence_model_name("plan", dtype)

            self.check_setup(model_name)
            self.assertFalse("TRITONSERVER_DELAY_SCHEDULER" in os.environ)
            self.assertFalse(
                "TRITONSERVER_BACKLOG_DELAY_SCHEDULER" in os.environ)

            corrids = [1001, 1002, 1003, 1004]
            threads = []
            threads.append(
                threading.Thread(
                    target=self.check_sequence_shape_tensor_io,
                    args=(
                        model_name,
                        dtype,
                        corrids[0],
                        (None, None),
                        # (flag_str, shape_value, value, pre_delay_ms)
                        (("start", 2, 1, None), (None, 4, 2, None), ("end", 8,
                                                                     3, None)),
                        self.get_expected_result(4 + corrids[0], corrids[0], 3,
                                                 "end"),
                        precreated_shm0_handles),
                    kwargs={
                        'sequence_name':
                            "{}_{}".format(self._testMethodName, corrids[0]),
                        'using_dynamic_batcher':
                            True
                    }))
            threads.append(
                threading.Thread(
                    target=self.check_sequence_shape_tensor_io,
                    args=(
                        model_name,
                        dtype,
                        corrids[1],
                        (None, None),
                        # (flag_str, shape_value, value, pre_delay_ms)
                        (("start", 2, 11, None), (None, 4, 12, None),
                         ("end", 8, 13, None)),
                        self.get_expected_result(36 + corrids[1], corrids[1],
                                                 13, "end"),
                        precreated_shm1_handles),
                    kwargs={
                        'sequence_name':
                            "{}_{}".format(self._testMethodName, corrids[1]),
                        'using_dynamic_batcher':
                            True
                    }))
            threads.append(
                threading.Thread(
                    target=self.check_sequence_shape_tensor_io,
                    args=(
                        model_name,
                        dtype,
                        corrids[2],
                        (None, None),
                        # (flag_str, shape_value, value, pre_delay_ms)
                        (("start", 2, 111, None), (None, 4, 112, None),
                         ("end", 8, 113, None)),
                        self.get_expected_result(336 + corrids[2], corrids[2],
                                                 113, "end"),
                        precreated_shm2_handles),
                    kwargs={
                        'sequence_name':
                            "{}_{}".format(self._testMethodName, corrids[2]),
                        'using_dynamic_batcher':
                            True
                    }))
            threads.append(
                threading.Thread(
                    target=self.check_sequence_shape_tensor_io,
                    args=(
                        model_name,
                        dtype,
                        corrids[3],
                        (None, None),
                        # (flag_str, shape_value, value, pre_delay_ms)
                        (("start", 2, 1111, None), (None, 4, 1112, None),
                         ("end", 8, 1113, None)),
                        self.get_expected_result(3336 + corrids[3], corrids[3],
                                                 1113, "end"),
                        precreated_shm3_handles),
                    kwargs={
                        'sequence_name':
                            "{}_{}".format(self._testMethodName, corrids[3]),
                        'using_dynamic_batcher':
                            True
                    }))

            for t in threads:
                t.start()
                if sleep_secs > 0:
                    time.sleep(sleep_secs)
            for t in threads:
                t.join()
            self.check_deferred_exception()
            self.check_status(model_name, {4: 3}, 3, 12)
        except Exception as ex:
            self.assertTrue(False, "unexpected error {}".format(ex))
        finally:
            if TEST_SYSTEM_SHARED_MEMORY:
                self.cleanup_shm_regions(precreated_shm0_handles)
                self.cleanup_shm_regions(precreated_shm1_handles)
                self.cleanup_shm_regions(precreated_shm2_handles)
                self.cleanup_shm_regions(precreated_shm3_handles)

    def test_dynaseq_identical_shape_values_series(self):
        # Send four sequences with identical shape values in series
        # and make sure they get completely batched into batch-size
        # 4 inferences.
        self._multi_sequence_identical_shape_impl(1)

    def test_dynaseq_identical_shape_values_parallel(self):
        # Send four sequences with identical shape values in parallel
        # and make sure they get completely batched into batch-size
        # 4 inferences.
        self._multi_sequence_identical_shape_impl(0)

    def test_dynaseq_different_shape_values_series(self):
        # Send four sequences with different shape values in series
        # and make sure they don't get batched together.
        self._multi_sequence_different_shape_impl(1)

    def test_dynaseq_different_shape_values_parallel(self):
        # Send four sequences with different shape values in parallel
        # and make sure they don't get batched together.
        self._multi_sequence_different_shape_impl(0)


if __name__ == '__main__':
    unittest.main()


# qa/L0_vertex_ai/vertex_ai_test.py

#!/usr/bin/python
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../common")

import os
import shutil
import time
import unittest
import numpy as np
import infer_util as iu
import test_util as tu
import tritonclient.http as httpclient

import argparse
import csv
import json
import os
import requests
import socket
import sys


class VertexAiTest(tu.TestResultCollector):

    def setUp(self):
        port = os.getenv('AIP_HTTP_PORT', '8080')
        predict_endpoint = os.getenv('AIP_PREDICT_ROUTE', '/predict')
        self.model_ = os.getenv('TEST_EXPLICIT_MODEL_NAME', 'addsub')
        self.url_ = "http://localhost:{}{}".format(port, predict_endpoint)
        self.input_data_ = [
            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15
        ]
        self.expected_output0_data_ = [x * 2 for x in self.input_data_]
        self.expected_output1_data_ = [0 for x in self.input_data_]

    def test_predict(self):
        inputs = []
        outputs = []
        inputs.append(httpclient.InferInput('INPUT0', [1, 16], "INT32"))
        inputs.append(httpclient.InferInput('INPUT1', [1, 16], "INT32"))

        # Initialize the data
        input_data = np.array(self.input_data_, dtype=np.int32)
        input_data = np.expand_dims(input_data, axis=0)
        inputs[0].set_data_from_numpy(input_data, binary_data=False)
        inputs[1].set_data_from_numpy(input_data, binary_data=False)

        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT0', binary_data=False))
        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT1', binary_data=False))
        request_body, _ = httpclient.InferenceServerClient.generate_request_body(
            inputs, outputs=outputs)

        headers = {'Content-Type': 'application/json'}
        r = requests.post(self.url_, data=request_body, headers=headers)
        r.raise_for_status()

        result = httpclient.InferenceServerClient.parse_response_body(
            r._content)

        output0_data = result.as_numpy('OUTPUT0')
        output1_data = result.as_numpy('OUTPUT1')
        for i in range(16):
            self.assertEqual(output0_data[0][i], self.expected_output0_data_[i])
            self.assertEqual(output1_data[0][i], self.expected_output1_data_[i])

    def test_predict_specified_model(self):
        inputs = []
        outputs = []
        inputs.append(httpclient.InferInput('INPUT0', [1, 16], "INT32"))
        inputs.append(httpclient.InferInput('INPUT1', [1, 16], "INT32"))

        # Initialize the data
        input_data = np.array(self.input_data_, dtype=np.int32)
        input_data = np.expand_dims(input_data, axis=0)
        inputs[0].set_data_from_numpy(input_data, binary_data=False)
        inputs[1].set_data_from_numpy(input_data, binary_data=False)

        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT0', binary_data=False))
        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT1', binary_data=False))
        request_body, _ = httpclient.InferenceServerClient.generate_request_body(
            inputs, outputs=outputs)

        headers = {
            'Content-Type':
                'application/json',
            "X-Vertex-Ai-Triton-Redirect":
                "v2/models/{}/infer".format(self.model_)
        }
        r = requests.post(self.url_, data=request_body, headers=headers)
        r.raise_for_status()

        result = httpclient.InferenceServerClient.parse_response_body(
            r._content)

        output0_data = result.as_numpy('OUTPUT0')
        output1_data = result.as_numpy('OUTPUT1')
        if self.model_ == "addsub":
            expected_output0_data = [x * 2 for x in self.input_data_]
            expected_output1_data = [0 for x in self.input_data_]
        else:
            expected_output0_data = [0 for x in self.input_data_]
            expected_output1_data = [x * 2 for x in self.input_data_]
        for i in range(16):
            self.assertEqual(output0_data[0][i], expected_output0_data[i])
            self.assertEqual(output1_data[0][i], expected_output1_data[i])

    def test_predict_request_binary(self):
        inputs = []
        outputs = []
        inputs.append(httpclient.InferInput('INPUT0', [1, 16], "INT32"))
        inputs.append(httpclient.InferInput('INPUT1', [1, 16], "INT32"))

        # Initialize the data
        input_data = np.array(self.input_data_, dtype=np.int32)
        input_data = np.expand_dims(input_data, axis=0)
        inputs[0].set_data_from_numpy(input_data, binary_data=True)
        inputs[1].set_data_from_numpy(input_data, binary_data=False)

        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT0', binary_data=False))
        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT1', binary_data=False))
        request_body, header_length = httpclient.InferenceServerClient.generate_request_body(
            inputs, outputs=outputs)

        headers = {
            'Content-Type':
                'application/vnd.vertex-ai-triton.binary+json;json-header-size={}'
                .format(header_length)
        }
        r = requests.post(self.url_, data=request_body, headers=headers)
        r.raise_for_status()

        result = httpclient.InferenceServerClient.parse_response_body(
            r._content)
        output0_data = result.as_numpy('OUTPUT0')
        output1_data = result.as_numpy('OUTPUT1')
        for i in range(16):
            self.assertEqual(output0_data[0][i], self.expected_output0_data_[i])
            self.assertEqual(output1_data[0][i], self.expected_output1_data_[i])

    def test_predict_response_binary(self):
        inputs = []
        outputs = []
        inputs.append(httpclient.InferInput('INPUT0', [1, 16], "INT32"))
        inputs.append(httpclient.InferInput('INPUT1', [1, 16], "INT32"))

        # Initialize the data
        input_data = np.array(self.input_data_, dtype=np.int32)
        input_data = np.expand_dims(input_data, axis=0)
        inputs[0].set_data_from_numpy(input_data, binary_data=False)
        inputs[1].set_data_from_numpy(input_data, binary_data=False)

        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT0', binary_data=True))
        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT1', binary_data=False))
        request_body, _ = httpclient.InferenceServerClient.generate_request_body(
            inputs, outputs=outputs)

        headers = {'Content-Type': 'application/json'}
        r = requests.post(self.url_, data=request_body, headers=headers)
        r.raise_for_status()

        header_length_str = r.headers['Inference-Header-Content-Length']
        result = httpclient.InferenceServerClient.parse_response_body(
            r._content, header_length=int(header_length_str))

        output0_data = result.as_numpy('OUTPUT0')
        output1_data = result.as_numpy('OUTPUT1')
        for i in range(16):
            self.assertEqual(output0_data[0][i], self.expected_output0_data_[i])
            self.assertEqual(output1_data[0][i], self.expected_output1_data_[i])

    def test_malformed_binary_header(self):
        inputs = []
        outputs = []
        inputs.append(httpclient.InferInput('INPUT0', [1, 16], "INT32"))
        inputs.append(httpclient.InferInput('INPUT1', [1, 16], "INT32"))

        # Initialize the data
        input_data = np.array(self.input_data_, dtype=np.int32)
        input_data = np.expand_dims(input_data, axis=0)
        inputs[0].set_data_from_numpy(input_data, binary_data=True)
        inputs[1].set_data_from_numpy(input_data, binary_data=False)

        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT0', binary_data=False))
        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT1', binary_data=False))
        request_body, header_length = httpclient.InferenceServerClient.generate_request_body(
            inputs, outputs=outputs)

        headers = {
            'Content-Type':
                'additional-string/application/vnd.vertex-ai-triton.binary+json;json-header-size={}'
                .format(header_length)
        }
        r = requests.post(self.url_, data=request_body, headers=headers)
        self.assertEqual(
            400, r.status_code,
            "Expected error code {} returned for the request; got: {}".format(
                400, r.status_code))

    def test_malformed_binary_header_not_number(self):
        inputs = []
        outputs = []
        inputs.append(httpclient.InferInput('INPUT0', [1, 16], "INT32"))
        inputs.append(httpclient.InferInput('INPUT1', [1, 16], "INT32"))

        # Initialize the data
        input_data = np.array(self.input_data_, dtype=np.int32)
        input_data = np.expand_dims(input_data, axis=0)
        inputs[0].set_data_from_numpy(input_data, binary_data=True)
        inputs[1].set_data_from_numpy(input_data, binary_data=False)

        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT0', binary_data=False))
        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT1', binary_data=False))
        request_body, header_length = httpclient.InferenceServerClient.generate_request_body(
            inputs, outputs=outputs)

        headers = {
            'Content-Type':
                'application/vnd.vertex-ai-triton.binary+json;json-header-size=additional-string{}'
                .format(header_length)
        }
        r = requests.post(self.url_, data=request_body, headers=headers)
        self.assertEqual(
            400, r.status_code,
            "Expected error code {} returned for the request; got: {}".format(
                400, r.status_code))

    def test_malformed_binary_header_negative_number(self):
        inputs = []
        outputs = []
        inputs.append(httpclient.InferInput('INPUT0', [1, 16], "INT32"))
        inputs.append(httpclient.InferInput('INPUT1', [1, 16], "INT32"))

        # Initialize the data
        input_data = np.array(self.input_data_, dtype=np.int32)
        input_data = np.expand_dims(input_data, axis=0)
        inputs[0].set_data_from_numpy(input_data, binary_data=True)
        inputs[1].set_data_from_numpy(input_data, binary_data=False)

        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT0', binary_data=False))
        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT1', binary_data=False))
        request_body, header_length = httpclient.InferenceServerClient.generate_request_body(
            inputs, outputs=outputs)

        headers = {
            'Content-Type':
                'application/vnd.vertex-ai-triton.binary+json;json-header-size=-123'
        }
        r = requests.post(self.url_, data=request_body, headers=headers)
        self.assertEqual(
            400, r.status_code,
            "Expected error code {} returned for the request; got: {}".format(
                400, r.status_code))

    def test_malformed_binary_header_large_number(self):
        inputs = []
        outputs = []
        inputs.append(httpclient.InferInput('INPUT0', [1, 16], "INT32"))
        inputs.append(httpclient.InferInput('INPUT1', [1, 16], "INT32"))

        # Initialize the data
        input_data = np.array(self.input_data_, dtype=np.int32)
        input_data = np.expand_dims(input_data, axis=0)
        inputs[0].set_data_from_numpy(input_data, binary_data=True)
        inputs[1].set_data_from_numpy(input_data, binary_data=False)

        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT0', binary_data=False))
        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT1', binary_data=False))
        request_body, header_length = httpclient.InferenceServerClient.generate_request_body(
            inputs, outputs=outputs)

        headers = {
            'Content-Type':
                'application/vnd.vertex-ai-triton.binary+json;json-header-size=12345'
        }
        r = requests.post(self.url_, data=request_body, headers=headers)
        self.assertEqual(
            400, r.status_code,
            "Expected error code {} returned for the request; got: {}".format(
                400, r.status_code))


if __name__ == '__main__':
    unittest.main()


# qa/common/check_copyright.py

#!/usr/bin/env python3

# Copyright 2018-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
import os
import re
import pathlib

FLAGS = None
SKIP_EXTS = ('jpeg', 'jpg', 'pgm', 'png', 'log', 'serverlog', 'preprocessed',
             'jmx', 'gz', 'json', 'pdf', 'so', 'onnx', 'svg')
REPO_PATH_FROM_THIS_FILE = '../..'
SKIP_PATHS = (
    'build', 'deploy/gke-marketplace-app/.gitignore',
    'deploy/gke-marketplace-app/server-deployer/chart/.helmignore',
    'deploy/gcp/.helmignore', 'deploy/aws/.helmignore',
    'deploy/fleetcommand/.helmignore', 'docs/.gitignore',
    'docs/_static/.gitattributes', 'docs/examples/model_repository',
    'docs/examples/jetson', 'docker', 'qa/common/cuda_op_kernel.cu.cc.patch',
    'qa/ensemble_models/mix_platform_float32_float32_float32/output0_labels.txt',
    'qa/ensemble_models/mix_type_int32_float32_float32/output0_labels.txt',
    'qa/ensemble_models/mix_ensemble_int32_float32_float32/output0_labels.txt',
    'qa/ensemble_models/wrong_label_int32_float32_float32/output0_labels.txt',
    'qa/ensemble_models/label_override_int32_float32_float32/output0_labels.txt',
    'qa/L0_model_config/noautofill_platform',
    'qa/L0_model_config/autofill_noplatform',
    'qa/L0_model_config/autofill_noplatform_success',
    'qa/L0_model_config/special_cases',
    'qa/L0_model_config/cli_messages/cli_override/expected',
    'qa/L0_model_config/cli_messages/cli_deprecation/expected',
    'qa/L0_perf_nomodel/baseline', 'qa/L0_perf_nomodel/legacy_baseline',
    'qa/L0_warmup/raw_mug_data', 'qa/L0_java_resnet/expected_output_data',
    'qa/L0_trt_dla_jetson/trt_dla_model_store',
    'qa/openvino_models/dynamic_batch', 'qa/openvino_models/fixed_batch',
    'TRITON_VERSION')

COPYRIGHT_YEAR_RE = 'Copyright( \\(c\\))? 20[1-9][0-9](-(20)?[1-9][0-9])?(,((20[2-9][0-9](-(20)?[2-9][0-9])?)|([2-9][0-9](-[2-9][0-9])?)))*,? NVIDIA CORPORATION( & AFFILIATES)?. All rights reserved.'

COPYRIGHT = '''

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions
are met:
 * Redistributions of source code must retain the above copyright
   notice, this list of conditions and the following disclaimer.
 * Redistributions in binary form must reproduce the above copyright
   notice, this list of conditions and the following disclaimer in the
   documentation and/or other materials provided with the distribution.
 * Neither the name of NVIDIA CORPORATION nor the names of its
   contributors may be used to endorse or promote products derived
   from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
'''

repo_abs_path = pathlib.Path(__file__).parent.joinpath(
    REPO_PATH_FROM_THIS_FILE).resolve()

copyright_year_re = re.compile(COPYRIGHT_YEAR_RE)


def visit(path):
    if FLAGS.verbose:
        print("visiting " + path)

    for skip in SKIP_EXTS:
        if path.endswith('.' + skip):
            if FLAGS.verbose:
                print("skipping due to extension: " + path)
            return True

    for skip in SKIP_PATHS:
        if str(pathlib.Path(path).resolve()).startswith(
                str(repo_abs_path.joinpath(skip).resolve())):
            if FLAGS.verbose:
                print("skipping due to path prefix: " + path)
            return True

    with open(path, 'r') as f:
        first_line = True
        line = None
        try:
            for fline in f:
                line = fline

                # Skip any '#!', '..', '<!--', '\*' or '{{/*' lines at the
                # start of the file
                if first_line:
                    first_line = False
                    if (fline.startswith("#!") or fline.startswith("..") or
                            fline.startswith("<!--") or
                            fline.startswith("/*") or
                            fline.startswith("{{/*")):
                        continue
                # Skip empty lines...
                if len(fline.strip()) != 0:
                    break
        except UnicodeDecodeError as ex:
            # If we get this exception on the first line then assume a
            # non-text file.
            if not first_line:
                raise ex
            if FLAGS.verbose:
                print("skipping binary file: " + path)
            return True

        if line is None:
            if FLAGS.verbose:
                print("skipping empty file: " + path)
            return True

        line = line.strip()

        # The next line must be the copyright line with a single year
        # or a year range. It is optionally allowed to have '# ' or
        # '// ' prefix.
        prefix = ""
        if line.startswith('# '):
            prefix = '# '
        elif line.startswith('// '):
            prefix = '// '
        elif not line.startswith(COPYRIGHT_YEAR_RE[0]):
            print(
                "incorrect prefix for copyright line, allowed prefixes '# ' or '// ', for "
                + path + ": " + line)
            return False

        # Check if the copyright year line matches the regex
        # and see if the year(s) are reasonable
        years = []

        copyright_row = line[len(prefix):]
        if copyright_year_re.match(copyright_row):
            for year in copyright_row.split("(c) " if "(c) " in
                                            copyright_row else "Copyright "
                                           )[1].split(" NVIDIA ")[0].split(","):
                if len(year) == 4:  # 2021
                    years.append(int(year))
                elif len(year) == 2:  # 21
                    years.append(int(year) + 2000)
                elif len(year) == 9:  # 2021-2022
                    years.append(int(year[0:4]))
                    years.append(int(year[5:9]))
                elif len(year) == 7:  # 2021-22
                    years.append(int(year[0:4]))
                    years.append(int(year[5:7]) + 2000)
                elif len(year) == 5:  # 21-23
                    years.append(int(year[0:2]) + 2000)
                    years.append(int(year[3:5]) + 2000)
        else:
            print("copyright year is not recognized for " + path + ": " + line)
            return False

        if years[0] > FLAGS.year:
            print("copyright start year greater than current year for " + path +
                  ": " + line)
            return False
        if years[-1] > FLAGS.year:
            print("copyright end year greater than current year for " + path +
                  ": " + line)
            return False
        for i in range(1, len(years)):
            if years[i - 1] >= years[i]:
                print("copyright years are not increasing for " + path + ": " +
                      line)
                return False

        # Subsequent lines must match the copyright body.
        copyright_body = [
            l.rstrip() for i, l in enumerate(COPYRIGHT.splitlines()) if i > 0
        ]
        copyright_idx = 0
        for line in f:
            if copyright_idx >= len(copyright_body):
                break

            if len(prefix) == 0:
                line = line.rstrip()
            else:
                line = line.strip()

            if len(copyright_body[copyright_idx]) == 0:
                expected = prefix.strip()
            else:
                expected = (prefix + copyright_body[copyright_idx])
            if line != expected:
                print("incorrect copyright body for " + path)
                print("  expected: '" + expected + "'")
                print("       got: '" + line + "'")
                return False
            copyright_idx += 1

        if copyright_idx != len(copyright_body):
            print("missing " + str(len(copyright_body) - copyright_idx) +
                  " lines of the copyright body")
            return False

    if FLAGS.verbose:
        print("copyright correct for " + path)
    return True


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-v',
                        '--verbose',
                        action="store_true",
                        required=False,
                        default=False,
                        help='Enable verbose output')
    parser.add_argument('-y',
                        '--year',
                        type=int,
                        required=True,
                        help='Copyright year')
    parser.add_argument('paths',
                        type=str,
                        nargs='*',
                        default=None,
                        help='Directories or files to check')
    FLAGS = parser.parse_args()

    if FLAGS.paths is None or len(FLAGS.paths) == 0:
        parser.print_help()
        exit(1)

    ret = True
    for path in FLAGS.paths:
        if not os.path.isdir(path):
            if not visit(path):
                ret = False
        else:
            for root, dirs, files in os.walk(path):
                for name in files:
                    if not visit(os.path.join(root, name)):
                        ret = False

    exit(0 if ret else 1)


# qa/common/check_massif_log.py

# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys
import re
from collections import defaultdict
import math


def parse_massif_out(filename):
    """
    Extract the allocation data from the massif output file, and compile
    it into a dictionary.    
    
    """
    # Read the file
    with open(filename, 'r') as f:
        contents = f.read()
        snapshots = re.findall('snapshot=(.*?)heap_tree',
                               contents,
                               flags=re.DOTALL)

    # Create snapshot dictionary
    summary = defaultdict(list)

    for snapshot in snapshots:
        # Split the record and ignore first two columns
        columns = snapshot.split()[2:]

        # Put columns and values into dictionary
        for col in columns:
            k, v = col.split('=')
            summary[k].append(int(v))

    # Return dict
    return summary


def is_unbounded_growth(summary, max_allowed_alloc, start_from_middle):
    """
    Check whether the heap allocations is increasing     
    
    """
    totals = summary['mem_heap_B']

    if len(totals) < 5:
        print("Error: Not enough snapshots")
        return False

    # Measure difference between mean and maximum memory usage
    processed_snapshot = totals[len(totals) //
                                2:] if start_from_middle else totals
    processed_snapshot.sort(reverse=True)
    # Remove 5% of the max value which will be treated as outlier
    num_max_min_dropout = math.ceil(0.05 * len(processed_snapshot))
    start = num_max_min_dropout
    end = len(processed_snapshot) - num_max_min_dropout
    mem_heap_avg = sum(processed_snapshot[start:end]) / len(
        processed_snapshot[start:end])
    mem_heap_max = max(processed_snapshot[start:end])

    # Compute change in allocation rate
    memory_allocation_delta_mb = (mem_heap_max - mem_heap_avg) / 1e6

    print("Change in memory allocation: %f MB, MAX ALLOWED: %f MB" %
          (memory_allocation_delta_mb, max_allowed_alloc))

    return (memory_allocation_delta_mb > max_allowed_alloc)


if __name__ == '__main__':
    # FIXME turn to proper argument handling
    summary = parse_massif_out(sys.argv[1])
    max_allowed_alloc = float(sys.argv[2])
    start_from_middle = ((len(sys.argv) == 4) and
                         (sys.argv[3] == "--start-from-middle"))
    if is_unbounded_growth(summary, max_allowed_alloc, start_from_middle):
        sys.exit(1)
    else:
        sys.exit(0)


# qa/common/check_valgrind_log.py

# Copyright 2021-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys
import argparse

# Check the valgrind logs for memory leaks, ignoring known memory leaks
#   * cnmem https://github.com/NVIDIA/cnmem/issues/12
#   * Tensorflow::NewSession
#   * dl-open leak could be due to https://bugs.kde.org/show_bug.cgi?id=358980
#   * dlerror leak in tensorflow::HadoopFileSystem::HadoopFileSystem()
#     -> tensorflow::LibHDFS::LoadAndBind()::{lambda(char const*, void**)#1}::operator()(char const*, void**)
#     -> tensorflow::internal::LoadLibrary
#     -> dlerror

LEAK_WHITE_LIST = [
    'cnmem', 'tensorflow::NewSession', 'dl-init', 'dl-open', 'dlerror',
    'libtorch'
]


def check_valgrind_log(log_file):
    """
    Counts the definite leaks reported
    by valgrind, matches them against
    the whitelist.

    Parameters
    ----------
    log_file: str
        The path to the log file
    
    Returns
    -------
    list of str
        a list of the leak records as strings
    """

    with open(args.input_log_file, 'r') as f:
        logs = f.read()

    # Find the pid and start and end of definite leak reports
    pid_token_end = logs.find('==', logs.find('==') + 1) + 2
    pid_token = logs[:pid_token_end]
    leaks_start = logs.find('are definitely lost')
    first_leak_line = logs.rfind('\n', 0, leaks_start)
    if leaks_start == -1 or first_leak_line == -1:
        # No leaks in log
        return []
    end_of_leaks = logs.find(f"{pid_token} LEAK SUMMARY:")
    if end_of_leaks == -1:
        print(
            f"\n***\n*** Test Failed for {log_file}: Malformed Valgrind log.\n***"
        )
        sys.exit(1)
    leak_records_section = logs[first_leak_line + 1:end_of_leaks]

    # Each leak record is separated by a line containing '==<pid>== \n'
    record_separator = f"{pid_token} \n"
    leak_records = leak_records_section.split(record_separator)

    # Check each leak against whitelist
    filtered_leak_records = []
    for leak in leak_records:
        for token in LEAK_WHITE_LIST:
            if not leak or leak.find(token) != -1:
                break
        else:
            filtered_leak_records.append(leak)

    return filtered_leak_records


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '-f',
        '--input-log-file',
        type=str,
        required=True,
        help="The name of the file containing the valgrind logs.")
    args = parser.parse_args()

    leak_records = check_valgrind_log(log_file=args.input_log_file)
    if leak_records:
        for leak in leak_records:
            print(leak)
        print(
            f"\n***\n*** Test Failed: {len(leak_records)} leaks detected.\n***")
        sys.exit(1)
    sys.exit(0)


# qa/common/gen_ensemble_model_utils.py

# Copyright (c) 2019-2022, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import os
import test_util as tu
import numpy as np

BASIC_ENSEMBLE_TYPES = ["simple", "sequence", "fan"]

np_dtype_string = np.dtype(object)


def np_to_model_dtype(np_dtype):
    if np_dtype == bool:
        return "TYPE_BOOL"
    elif np_dtype == np.int8:
        return "TYPE_INT8"
    elif np_dtype == np.int16:
        return "TYPE_INT16"
    elif np_dtype == np.int32:
        return "TYPE_INT32"
    elif np_dtype == np.int64:
        return "TYPE_INT64"
    elif np_dtype == np.uint8:
        return "TYPE_UINT8"
    elif np_dtype == np.uint16:
        return "TYPE_UINT16"
    elif np_dtype == np.float16:
        return "TYPE_FP16"
    elif np_dtype == np.float32:
        return "TYPE_FP32"
    elif np_dtype == np.float64:
        return "TYPE_FP64"
    elif np_dtype == np_dtype_string:
        return "TYPE_STRING"
    return None


def fixed_to_variable_size(shape):
    return [-1] * len(shape)


def platform_types_and_validation():
    res = [("graphdef", tu.validate_for_tf_model),
           ("savedmodel", tu.validate_for_tf_model),
           ("plan", tu.validate_for_trt_model),
           ("onnx", tu.validate_for_onnx_model),
           ("libtorch", tu.validate_for_libtorch_model)]
    return res


class AddSubEnsembleSchedule:
    """
    Helper class to generate ensemble schedule that behaves the same as
    addsub model given an ensemble type
    """

    def __init__(self, ensemble_type):
        if ensemble_type == "fan":
            self._get_schedule = AddSubEnsembleSchedule._get_fan_ensemble_schedule
        elif ensemble_type == "sequence":
            self._get_schedule = AddSubEnsembleSchedule._get_sequence_ensemble_schedule
        else:
            self._get_schedule = AddSubEnsembleSchedule._get_simple_ensemble_schedule

    def get_schedule(self, base_model_name, input_shape, output0_shape,
                     output1_shape, input_model_dtype, output0_model_dtype,
                     output1_model_dtype):
        return self._get_schedule(base_model_name, input_shape, output0_shape,
                                  output1_shape, input_model_dtype,
                                  output0_model_dtype, output1_model_dtype)

    @classmethod
    def _get_simple_ensemble_schedule(cls, base_model_name, input_shape,
                                      output0_shape, output1_shape, input_dtype,
                                      output0_dtype, output1_dtype):
        # libtorch model uses other naming convention for outputs
        output_index_delimiter = "__" if "libtorch" in base_model_name else ""
        # ensemble input -> addsub -> ensemble output
        schedule = '''
ensemble_scheduling {{
  step [
    {{
      model_name: "{}"
      model_version: -1
      input_map {{
        key: "INPUT0"
        value: "INPUT0"
      }}
      input_map {{
        key: "INPUT1"
        value: "INPUT1"
      }}
      output_map {{
        key: "OUTPUT{delimiter}0"
        value: "OUTPUT0"
      }}
      output_map {{
        key: "OUTPUT{delimiter}1"
        value: "OUTPUT1"
      }}
    }}
  ]
}}
'''.format(base_model_name, delimiter=output_index_delimiter)
        return schedule

    @classmethod
    def _get_sequence_ensemble_schedule(cls, base_model_name, input_shape,
                                        output0_shape, output1_shape,
                                        input_dtype, output0_dtype,
                                        output1_dtype):
        # libtorch model uses other naming convention for outputs
        output_index_delimiter = "__" if "libtorch" in base_model_name else ""
        # ensemble input -> nop -> addsub -> ensemble output
        nop_input_shape = fixed_to_variable_size(input_shape)
        schedule = '''
ensemble_scheduling {{
  step [
    {{
      model_name: "nop_{}_{}"
      model_version: -1
      input_map {{
        key: "INPUT0"
        value: "INPUT0"
      }}
      input_map {{
        key: "INPUT1"
        value: "INPUT1"
      }}
      output_map {{
        key: "OUTPUT0"
        value: "same_input0"
      }}
      output_map {{
        key: "OUTPUT1"
        value: "same_input1"
      }}
    }},
    {{
      model_name: "{}"
      model_version: -1
      input_map {{
        key: "INPUT0"
        value: "same_input0"
      }}
      input_map {{
        key: "INPUT1"
        value: "same_input1"
      }}
      output_map {{
        key: "OUTPUT{delimiter}0"
        value: "OUTPUT0"
      }}
      output_map {{
        key: "OUTPUT{delimiter}1"
        value: "OUTPUT1"
      }}
    }}
  ]
}}
'''.format(input_dtype,
           tu.shape_to_dims_str(nop_input_shape),
           base_model_name,
           delimiter=output_index_delimiter)
        return schedule

    @classmethod
    def _get_fan_ensemble_schedule(cls, base_model_name, input_shape,
                                   output0_shape, output1_shape, input_dtype,
                                   output0_dtype, output1_dtype):
        # libtorch model uses other naming convention for outputs
        output_index_delimiter = "__" if "libtorch" in base_model_name else ""

        # ensemble input -> nop -> addsub ->
        # nop (fan out, one output send to one nop) -> ensemble output (fan in)
        nop_input_shape = fixed_to_variable_size(input_shape)
        nop_output0_shape = fixed_to_variable_size(output0_shape)
        nop_output1_shape = fixed_to_variable_size(output1_shape)
        schedule = '''
ensemble_scheduling {{
  step [
    {{
      model_name: "nop_{}_{}"
      model_version: -1
      input_map {{
        key: "INPUT0"
        value: "INPUT0"
      }}
      input_map {{
        key: "INPUT1"
        value: "INPUT1"
      }}
      output_map {{
        key: "OUTPUT0"
        value: "same_input0"
      }}
      output_map {{
        key: "OUTPUT1"
        value: "same_input1"
      }}
    }},
    {{
      model_name: "{}"
      model_version: -1
      input_map {{
        key: "INPUT0"
        value: "same_input0"
      }}
      input_map {{
        key: "INPUT1"
        value: "same_input1"
      }}
      output_map {{
        key: "OUTPUT{delimiter}0"
        value: "same_output0"
      }}
      output_map {{
        key: "OUTPUT{delimiter}1"
        value: "same_output1"
      }}
    }},
    {{
      model_name: "nop_{}_{}"
      model_version: -1
      input_map {{
        key: "INPUT0"
        value: "same_output0"
      }}
      input_map {{
        key: "INPUT1"
        value: "same_output0"
      }}
      output_map {{
        key: "OUTPUT0"
        value: "OUTPUT0"
      }}
    }},
    {{
      model_name: "nop_{}_{}"
      model_version: -1
      input_map {{
        key: "INPUT0"
        value: "same_output1"
      }}
      input_map {{
        key: "INPUT1"
        value: "same_output1"
      }}
      output_map {{
        key: "OUTPUT1"
        value: "OUTPUT1"
      }}
    }}
  ]
}}
'''.format(input_dtype,
           tu.shape_to_dims_str(nop_input_shape),
           base_model_name,
           output0_dtype,
           tu.shape_to_dims_str(nop_output0_shape),
           output1_dtype,
           tu.shape_to_dims_str(nop_output1_shape),
           delimiter=output_index_delimiter)
        return schedule


class IdentityEnsembleSchedule:
    """
    Helper class to generate ensemble schedule that behaves the same as
    identity model given an ensemble type
    """

    def __init__(self, ensemble_type, ensemble_test_type="zero"):
        self._test_type = ensemble_test_type
        if ensemble_type == "fan":
            self._get_schedule = IdentityEnsembleSchedule._get_fan_ensemble_schedule
        elif ensemble_type == "sequence":
            self._get_schedule = IdentityEnsembleSchedule._get_sequence_ensemble_schedule
        else:
            self._get_schedule = IdentityEnsembleSchedule._get_simple_ensemble_schedule

    def get_schedule(self, dtype, input_shapes, input_model_shapes,
                     output_shapes, output_model_shapes):
        return self._get_schedule(dtype, input_shapes, input_model_shapes,
                                  output_shapes, output_model_shapes,
                                  self._test_type)

    @classmethod
    def _get_simple_ensemble_schedule(cls, dtype, input_shapes,
                                      input_model_shapes, output_shapes,
                                      output_model_shapes, test_type):
        # ensemble reshaped input -> nop with reshaped tensor shape -> ensemble
        # reshaped output (actual ensemble input/output is not visible in schedule)
        steps = []
        for idx in range(len(input_shapes)):
            steps.append('''
    {{
      model_name: "nop_{}_{}"
      model_version: -1
      input_map {{
        key: "INPUT0"
        value: "INPUT{}"
      }}
      input_map {{
        key: "INPUT1"
        value: "INPUT{}"
      }}
      output_map {{
        key: "OUTPUT0"
        value: "OUTPUT{}"
      }}
    }}
'''.format(np_to_model_dtype(dtype),
            tu.shape_to_dims_str(input_model_shapes[idx]), idx, idx, idx))

        schedule = '''
ensemble_scheduling {{
  step [
{}
  ]
}}
'''.format(",".join(steps))

        return schedule

    @classmethod
    def _get_sequence_ensemble_schedule(cls, dtype, input_shapes,
                                        input_model_shapes, output_shapes,
                                        output_model_shapes, test_type):
        in_str = "tunnel_in_" if test_type == "reshape" else ""
        out_str = "tunnel_out_" if test_type == "reshape" else ""
        # ensemble reshaped input -> nop with another input only reshape ->
        # nop with output only reshape -> ensemble reshaped output
        steps = []
        for idx in range(len(input_shapes)):
            steps.append('''
    {{
      model_name: "nop_{in_str}{type}_{shape}"
      model_version: -1
      input_map {{
        key: "INPUT0"
        value: "INPUT{idx}"
      }}
      input_map {{
        key: "INPUT1"
        value: "INPUT{idx}"
      }}
      output_map {{
        key: "OUTPUT0"
        value: "temp_{idx}"
      }}
    }},
    {{
      model_name: "nop_{out_str}{type}_{shape}"
      model_version: -1
      input_map {{
        key: "INPUT0"
        value: "temp_{idx}"
      }}
      input_map {{
        key: "INPUT1"
        value: "temp_{idx}"
      }}
      output_map {{
        key: "OUTPUT0"
        value: "OUTPUT{idx}"
      }}
    }}
'''.format(type=np_to_model_dtype(dtype),
            in_str=in_str,
            out_str=out_str,
            idx=idx,
            shape=tu.shape_to_dims_str(input_model_shapes[idx])))

        schedule = '''
ensemble_scheduling {{
  step [
{}
  ]
}}
'''.format(",".join(steps))

        return schedule

    @classmethod
    def _get_fan_ensemble_schedule(cls, dtype, input_shapes, input_model_shapes,
                                   output_shapes, output_model_shapes,
                                   test_type):
        # Note that the simple and sequence test already test "fan" in some
        # degree, because there is no direct match from nop input/output
        # like what is in addsub-like ensemble.
        #
        # ensemble reshaped input -> nop with another input only reshape ->
        # nop with variable size -> nop with output only reshape ->
        # ensemble reshaped output
        in_str = ""
        out_str = ""
        intermediate_shapes = input_model_shapes
        if test_type == "reshape":
            in_str = "tunnel_in_"
            out_str = "tunnel_out_"
            intermediate_shapes = [[-1]] * len(input_model_shapes)
        steps = []
        for idx in range(len(input_shapes)):
            steps.append('''
    {{
      model_name: "nop_{in_str}{type}_{shape}"
      model_version: -1
      input_map {{
        key: "INPUT0"
        value: "INPUT{idx}"
      }}
      input_map {{
        key: "INPUT1"
        value: "INPUT{idx}"
      }}
      output_map {{
        key: "OUTPUT0"
        value: "temp_in_{idx}"
      }}
    }},
    {{
      model_name: "nop_{type}_{intermediate_shape}"
      model_version: -1
      input_map {{
        key: "INPUT0"
        value: "temp_in_{idx}"
      }}
      input_map {{
        key: "INPUT1"
        value: "temp_in_{idx}"
      }}
      output_map {{
        key: "OUTPUT0"
        value: "temp_out_{idx}"
      }}
    }},
    {{
      model_name: "nop_{out_str}{type}_{shape}"
      model_version: -1
      input_map {{
        key: "INPUT0"
        value: "temp_out_{idx}"
      }}
      input_map {{
        key: "INPUT1"
        value: "temp_out_{idx}"
      }}
      output_map {{
        key: "OUTPUT0"
        value: "OUTPUT{idx}"
      }}
    }}
'''.format(type=np_to_model_dtype(dtype),
            in_str=in_str,
            out_str=out_str,
            intermediate_shape=tu.shape_to_dims_str(intermediate_shapes[idx]),
            idx=idx,
            shape=tu.shape_to_dims_str(input_model_shapes[idx])))

        schedule = '''
ensemble_scheduling {{
  step [
{}
  ]
}}
'''.format(",".join(steps))

        return schedule


class SequenceEnsembleSchedule:
    """
    Helper class to generate ensemble schedule that behaves the same as
    sequence model given an ensemble type
    """

    def __init__(self, ensemble_type):
        if ensemble_type == "fan":
            self._get_schedule = SequenceEnsembleSchedule._get_fan_ensemble_schedule
        elif ensemble_type == "sequence":
            self._get_schedule = SequenceEnsembleSchedule._get_sequence_ensemble_schedule
        else:
            self._get_schedule = SequenceEnsembleSchedule._get_simple_ensemble_schedule

    def get_schedule(self, base_model_name, shape, model_dtype):
        return self._get_schedule(base_model_name, shape, model_dtype)

    @classmethod
    def _get_simple_ensemble_schedule(cls, base_model_name, shape, model_dtype):
        # libtorch model uses other naming convention
        index_suffix = "__0" if "libtorch" in base_model_name else ""
        # ensemble input -> sequence -> ensemble output
        schedule = '''
ensemble_scheduling {{
  step [
    {{
      model_name: "{}"
      model_version: -1
      input_map {{
        key: "INPUT{index}"
        value: "INPUT"
      }}
      output_map {{
        key: "OUTPUT{index}"
        value: "OUTPUT"
      }}
    }}
  ]
}}
'''.format(base_model_name, index=index_suffix)
        return schedule

    @classmethod
    def _get_sequence_ensemble_schedule(cls, base_model_name, shape,
                                        model_dtype):
        # nop cannot handle STRING data type, fall back to simple
        if model_dtype == "TYPE_STRING":
            return SequenceEnsembleSchedule._get_simple_ensemble_schedule(
                base_model_name, shape, model_dtype)

        # libtorch model uses other naming convention
        index_suffix = "__0" if "libtorch" in base_model_name else ""
        # ensemble input -> nop -> sequence -> ensemble output
        nop_input_shape = fixed_to_variable_size(shape)
        schedule = '''
ensemble_scheduling {{
  step [
    {{
      model_name: "nop_{}_{}"
      model_version: -1
      input_map {{
        key: "INPUT0"
        value: "INPUT"
      }}
      input_map {{
        key: "INPUT1"
        value: "INPUT"
      }}
      output_map {{
        key: "OUTPUT0"
        value: "same_input"
      }}
    }},
    {{
      model_name: "{}"
      model_version: -1
      input_map {{
        key: "INPUT{index}"
        value: "same_input"
      }}
      output_map {{
        key: "OUTPUT{index}"
        value: "OUTPUT"
      }}
    }}
  ]
}}
'''.format(model_dtype,
           tu.shape_to_dims_str(nop_input_shape),
           base_model_name,
           index=index_suffix)
        return schedule

    @classmethod
    def _get_fan_ensemble_schedule(cls, base_model_name, shape, model_dtype):
        # nop cannot handle STRING data type, fall back to simple
        if model_dtype == "TYPE_STRING":
            return SequenceEnsembleSchedule._get_simple_ensemble_schedule(
                base_model_name, shape, model_dtype)

        # libtorch model uses other naming convention
        index_suffix = "__0" if "libtorch" in base_model_name else ""
        # Not a "fan" due to configuration of base sequence model
        # ensemble input -> nop -> sequence -> nop -> ensemble output
        nop_shape = fixed_to_variable_size(shape)
        schedule = '''
ensemble_scheduling {{
  step [
    {{
      model_name: "nop_{}_{}"
      model_version: -1
      input_map {{
        key: "INPUT0"
        value: "INPUT"
      }}
      input_map {{
        key: "INPUT1"
        value: "INPUT"
      }}
      output_map {{
        key: "OUTPUT0"
        value: "same_input"
      }}
    }},
    {{
      model_name: "{}"
      model_version: -1
      input_map {{
        key: "INPUT{index}"
        value: "same_input"
      }}
      output_map {{
        key: "OUTPUT{index}"
        value: "same_output"
      }}
    }},
    {{
      model_name: "nop_{}_{}"
      model_version: -1
      input_map {{
        key: "INPUT0"
        value: "same_output"
      }}
      input_map {{
        key: "INPUT1"
        value: "same_output"
      }}
      output_map {{
        key: "OUTPUT0"
        value: "OUTPUT"
      }}
    }}
  ]
}}
'''.format(model_dtype,
           tu.shape_to_dims_str(nop_shape),
           base_model_name,
           model_dtype,
           tu.shape_to_dims_str(nop_shape),
           index=index_suffix)
        return schedule


def create_ensemble_modelfile(base_model,
                              models_dir,
                              max_batch,
                              model_version,
                              input_shape,
                              output0_shape,
                              output1_shape,
                              input_dtype,
                              output0_dtype,
                              output1_dtype,
                              swap=False):

    # No actual model file in ensemble model

    # Use a different model name for the non-batching variant
    for ensemble_type in BASIC_ENSEMBLE_TYPES:
        ensemble_model_name = "{}_{}{}".format(
            ensemble_type, base_model, "_nobatch" if max_batch == 0 else "")
        model_name = tu.get_model_name(ensemble_model_name, input_dtype,
                                       output0_dtype, output1_dtype)
        model_version_dir = models_dir + "/" + model_name + "/" + str(
            model_version)

        try:
            os.makedirs(model_version_dir)
        except OSError as ex:
            pass  # ignore existing dir


def create_ensemble_modelconfig(base_model, models_dir, max_batch,
                                model_version, input_shape, output0_shape,
                                output1_shape, input_dtype, output0_dtype,
                                output1_dtype, output0_label_cnt,
                                version_policy):

    # No validation as long as the base model supports the type and shape

    input_model_dtype = np_to_model_dtype(input_dtype)
    output0_model_dtype = np_to_model_dtype(output0_dtype)
    output1_model_dtype = np_to_model_dtype(output1_dtype)

    for ensemble_type in BASIC_ENSEMBLE_TYPES:
        # Only in "fan" that ensemble output is not directly from addsub. In
        # other case, ensemble should still generate proper labell without
        # label file
        labels = None if ensemble_type != "fan" else "output0_labels.txt"

        # Use a different model name for the non-batching variant
        ensemble_model_name = "{}_{}{}".format(
            ensemble_type, base_model, "_nobatch" if max_batch == 0 else "")
        model_name = tu.get_model_name(ensemble_model_name, input_dtype,
                                       output0_dtype, output1_dtype)
        base_model_name = tu.get_model_name(
            "{}{}".format(base_model, "_nobatch" if max_batch == 0 else ""),
            input_dtype, output0_dtype, output1_dtype)

        ensemble_schedule = AddSubEnsembleSchedule(ensemble_type).get_schedule(
            base_model_name, input_shape, output0_shape, output1_shape,
            input_model_dtype, output0_model_dtype, output1_model_dtype)

        config_dir = models_dir + "/" + model_name
        config = create_general_modelconfig(model_name,
                                            "ensemble",
                                            max_batch,
                                            repeat(input_dtype, 2),
                                            repeat(input_shape, 2),
                                            repeat(None, 2),
                                            [output0_dtype, output1_dtype],
                                            [output0_shape, output1_shape],
                                            repeat(None, 2), [labels, None],
                                            version_policy=version_policy,
                                            force_tensor_number_suffix=True)
        config += ensemble_schedule

        try:
            os.makedirs(config_dir)
        except OSError as ex:
            pass  # ignore existing dir

        with open(config_dir + "/config.pbtxt", "w") as cfile:
            cfile.write(config)

        if labels is not None:
            with open(config_dir + "/output0_labels.txt", "w") as lfile:
                for l in range(output0_label_cnt):
                    lfile.write("label" + str(l) + "\n")


def create_identity_ensemble_modelfile(ensemble_test_type, models_dir,
                                       model_version, max_batch, dtype,
                                       input_shapes, output_shapes):
    io_cnt = len(input_shapes)

    # Use a different model name for the non-batching variant
    for ensemble_type in BASIC_ENSEMBLE_TYPES:
        ensemble_prefix = "{}_{}".format(ensemble_type, ensemble_test_type)
        model_name = tu.get_zero_model_name(
            ensemble_prefix + ("_nobatch" if max_batch == 0 else ""), io_cnt,
            dtype)
        model_version_dir = models_dir + "/" + model_name + "/" + str(
            model_version)

        try:
            os.makedirs(model_version_dir)
        except OSError as ex:
            pass  # ignore existing dir


def create_identity_ensemble_modelconfig(ensemble_test_type, models_dir,
                                         model_version, max_batch, dtype,
                                         input_shapes, input_model_shapes,
                                         output_shapes, output_model_shapes):
    io_cnt = len(input_shapes)

    for ensemble_type in BASIC_ENSEMBLE_TYPES:
        # Use a different model name for the non-batching variant
        ensemble_prefix = "{}_{}".format(ensemble_type, ensemble_test_type)
        model_name = tu.get_zero_model_name(
            ensemble_prefix + ("_nobatch" if max_batch == 0 else ""), io_cnt,
            dtype)

        ensemble_schedule = IdentityEnsembleSchedule(
            ensemble_type,
            ensemble_test_type).get_schedule(dtype, input_shapes,
                                             input_model_shapes, output_shapes,
                                             output_model_shapes)

        config_dir = models_dir + "/" + model_name
        config = create_general_modelconfig(model_name,
                                            "ensemble",
                                            max_batch,
                                            repeat(dtype, io_cnt),
                                            input_shapes,
                                            input_model_shapes,
                                            repeat(dtype, io_cnt),
                                            output_shapes,
                                            output_model_shapes,
                                            repeat(None, io_cnt),
                                            force_tensor_number_suffix=True)
        config += ensemble_schedule

        try:
            os.makedirs(config_dir)
        except OSError as ex:
            pass  # ignore existing dir

        with open(config_dir + "/config.pbtxt", "w") as cfile:
            cfile.write(config)


def create_sequence_ensemble_modelfile(base_model, models_dir, max_batch,
                                       model_version, shape, dtype):

    # No actual model file in ensemble model

    # Use a different model name for the non-batching variant
    for ensemble_type in BASIC_ENSEMBLE_TYPES:
        ensemble_model_name = "{}_{}{}".format(
            ensemble_type, base_model, "_nobatch" if max_batch == 0 else "")
        model_name = tu.get_sequence_model_name(ensemble_model_name, dtype)
        model_version_dir = models_dir + "/" + model_name + "/" + str(
            model_version)

        try:
            os.makedirs(model_version_dir)
        except OSError as ex:
            pass  # ignore existing dir


def create_sequence_ensemble_modelconfig(base_model, models_dir, max_batch,
                                         model_version, shape, dtype):

    # No validation as long as the base model supports the type and shape

    model_dtype = np_to_model_dtype(dtype)

    for ensemble_type in BASIC_ENSEMBLE_TYPES:
        # Use a different model name for the non-batching variant
        ensemble_model_name = "{}_{}{}".format(
            ensemble_type, base_model, "_nobatch" if max_batch == 0 else "")
        model_name = tu.get_sequence_model_name(ensemble_model_name, dtype)
        base_model_name = tu.get_sequence_model_name(
            "{}{}".format(base_model, "_nobatch" if max_batch == 0 else ""),
            dtype)

        ensemble_schedule = SequenceEnsembleSchedule(
            ensemble_type).get_schedule(base_model_name, shape, model_dtype)

        config_dir = models_dir + "/" + model_name
        config = create_general_modelconfig(model_name, "ensemble", max_batch,
                                            [dtype], [shape], [None], [dtype],
                                            [shape], [None], [None])
        config += ensemble_schedule

        try:
            os.makedirs(config_dir)
        except OSError as ex:
            pass  # ignore existing dir

        with open(config_dir + "/config.pbtxt", "w") as cfile:
            cfile.write(config)


def create_nop_modelconfig(models_dir,
                           tensor_shape,
                           tensor_dtype,
                           tensor_model_shape=None):
    model_name = "nop_{}_{}".format(dtype_str(tensor_dtype),
                                    tu.shape_to_dims_str(tensor_shape))
    # Make [] to [1].
    # Note that this doesn't affect the naming ("nop_{}_" instead of "nop_{}_1")
    if len(tensor_shape) == 0:
        tensor_shape = [1]

    config_dir = models_dir + "/" + model_name
    config = create_general_modelconfig(
        model_name,
        "",
        1024,
        repeat(tensor_dtype, 2),
        repeat(tensor_shape, 2),
        repeat(tensor_model_shape, 2),
        repeat(tensor_dtype, 2),
        repeat(tensor_shape, 2),
        repeat(tensor_model_shape, 2),
        repeat(None, 2),
        backend="identity",
        instance_group_str="instance_group [ { kind: KIND_CPU } ]")

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_nop_tunnel_modelconfig(models_dir, tensor_shape, tensor_dtype):
    # Must be fixed size
    in_model_name = "nop_tunnel_in_{}_{}".format(
        dtype_str(tensor_dtype), tu.shape_to_dims_str(tensor_shape))
    out_model_name = "nop_tunnel_out_{}_{}".format(
        dtype_str(tensor_dtype), tu.shape_to_dims_str(tensor_shape))
    # Make [] to [1].
    # Note that this doesn't affect the naming ("nop_{}_" instead of "nop_{}_1")
    if len(tensor_shape) == 0:
        tensor_shape = [1]
    internal_shape = 1
    for dim in tensor_shape:
        if dim < 0:
            raise Exception(
                "Must specify fixed size input / output for nop tunnel")
        internal_shape *= dim

    # Tunnel in nop (reshape to one dimension)
    config_dir = models_dir + "/" + in_model_name
    config = create_general_modelconfig(
        in_model_name,
        "",
        1024,
        repeat(tensor_dtype, 2),
        repeat(tensor_shape, 2),
        repeat([internal_shape], 2),
        repeat(tensor_dtype, 2),
        repeat([internal_shape], 2),
        repeat(None, 2),
        repeat(None, 2),
        backend="identity",
        instance_group_str="instance_group [ { kind: KIND_CPU } ]")

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)

    # Tunnel out nop (reshape back to original shape)
    config_dir = models_dir + "/" + out_model_name
    config = create_general_modelconfig(
        out_model_name,
        "",
        1024,
        repeat(tensor_dtype, 2),
        repeat([internal_shape], 2),
        repeat(tensor_shape, 2),
        repeat(tensor_dtype, 2),
        repeat(tensor_shape, 2),
        repeat(None, 2),
        repeat(None, 2),
        backend="identity",
        instance_group_str="instance_group [ { kind: KIND_CPU } ]")

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_general_modelconfig(model_name,
                               platform,
                               max_batch,
                               input_dtypes,
                               input_shapes,
                               input_model_shapes,
                               output_dtypes,
                               output_shapes,
                               output_model_shapes,
                               label_filenames,
                               backend=None,
                               version_policy=None,
                               default_model_filename=None,
                               instance_group_str="",
                               force_tensor_number_suffix=False):
    assert len(input_dtypes) == len(input_shapes)
    assert len(input_model_shapes) == len(input_shapes)
    assert len(output_dtypes) == len(output_shapes)
    assert len(output_model_shapes) == len(output_shapes)
    assert len(label_filenames) == len(output_shapes)

    # Unpack version policy
    version_policy_str = "{ latest { num_versions: 1 }}"
    if version_policy is not None:
        type, val = version_policy
        if type == 'latest':
            version_policy_str = "{{ latest {{ num_versions: {} }}}}".format(
                val)
        elif type == 'specific':
            version_policy_str = "{{ specific {{ versions: {} }}}}".format(val)
        else:
            version_policy_str = "{ all { }}"

    default_model_filename_str = ""
    if default_model_filename is not None:
        default_model_filename_str = 'default_model_filename: "{}"'.format(
            default_model_filename)

    # If backend is specified use backend instead of platform
    if backend is not None:
        key = "backend"
        val = backend
    else:
        key = "platform"
        val = platform

    config = '''
name: "{}"
{}: "{}"
max_batch_size: {}
version_policy: {}
{}
{}
'''.format(model_name, key, val, max_batch, version_policy_str,
           default_model_filename_str, instance_group_str)

    for idx in range(len(input_dtypes)):
        idx_str = ""
        if len(input_dtypes) != 1 or force_tensor_number_suffix:
            idx_str = str(idx)
        config += '''
input [
  {{
    name: "INPUT{}"
    data_type: {}
    dims: [ {} ]
    {}
  }}
]'''.format(idx_str, dtype_str(input_dtypes[idx]),
            tu.shape_to_dims_str(input_shapes[idx]),
            reshape_str(input_shapes[idx], input_model_shapes[idx]))

    for idx in range(len(output_dtypes)):
        idx_str = ""
        if len(input_dtypes) != 1 or force_tensor_number_suffix:
            idx_str = str(idx)
        config += '''
output [
  {{
    name: "OUTPUT{}"
    data_type: {}
    dims: [ {} ]
    {}
    {}
  }}
]'''.format(idx_str, dtype_str(output_dtypes[idx]),
            tu.shape_to_dims_str(output_shapes[idx]),
            reshape_str(output_shapes[idx], output_model_shapes[idx]),
            label_str(label_filenames[idx]))
    return config


def repeat(obj, cnt):
    return [obj] * cnt


def dtype_str(dtype):
    return dtype if isinstance(dtype, str) else np_to_model_dtype(dtype)


def reshape_str(shape, model_shape):
    if model_shape is None or shape == model_shape:
        return ""
    return "reshape: {{ shape: [ {} ] }}".format(
        tu.shape_to_dims_str(model_shape))


def label_str(label):
    if label is None:
        return ""
    return 'label_filename: "{}"'.format(label)


# qa/common/gen_qa_custom_ops_models.py

# Copyright 2019-2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
import os

FLAGS = None


def create_zeroout_modelfile(create_savedmodel, models_dir, model_version):
    # Load the zero-out custom operator
    _zero_out_module = tf.load_op_library(os.path.join(FLAGS.zero_out_lib_path))
    zero_out = _zero_out_module.zero_out

    # Create the model that uses custom operator.
    tf.reset_default_graph()
    zin = tf.placeholder(tf.int32, [
        None,
    ], "to_zero")
    zout = zero_out(zin, name="zeroed")

    model_name = "savedmodel_zeroout" if create_savedmodel else "graphdef_zeroout"
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    if create_savedmodel:
        with tf.Session() as sess:
            input_name = "to_zero"
            output_name = "zeroed"
            input_tensor = tf.get_default_graph().get_tensor_by_name(
                input_name + ":0")
            output_tensor = tf.get_default_graph().get_tensor_by_name(
                output_name + ":0")
            input_dict = dict()
            output_dict = dict()
            input_dict[input_name] = input_tensor
            output_dict[output_name] = output_tensor
            tf.saved_model.simple_save(sess,
                                       model_version_dir + "/model.savedmodel",
                                       inputs=input_dict,
                                       outputs=output_dict)
    else:
        with tf.Session() as sess:
            graph_io.write_graph(sess.graph.as_graph_def(),
                                 model_version_dir,
                                 "model.graphdef",
                                 as_text=False)


def create_zeroout_modelconfig(create_savedmodel, models_dir, model_version):
    model_name = "savedmodel_zeroout" if create_savedmodel else "graphdef_zeroout"
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "{}"
max_batch_size: 0
input [
  {{
    name: "to_zero"
    data_type: TYPE_INT32
    dims: [ -1 ]
  }}
]
output [
  {{
    name: "zeroed"
    data_type: TYPE_INT32
    dims: [ -1 ]
  }}
]
'''.format(
        model_name,
        "tensorflow_savedmodel" if create_savedmodel else "tensorflow_graphdef")

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_cudaop_modelfile(create_savedmodel, models_dir, model_version):
    # Load the add_one custom operator
    _cuda_op_module = tf.load_op_library(os.path.join(FLAGS.cuda_op_lib_path))
    add_one = _cuda_op_module.add_one

    # Create the model that uses custom operator.
    tf.reset_default_graph()
    zin = tf.placeholder(tf.int32, [
        None,
    ], "in")
    zout = add_one(zin, name="out")

    model_name = "savedmodel_cudaop" if create_savedmodel else "graphdef_cudaop"
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    if create_savedmodel:
        with tf.Session() as sess:
            input_name = "in"
            output_name = "out"
            input_tensor = tf.get_default_graph().get_tensor_by_name(
                input_name + ":0")
            output_tensor = tf.get_default_graph().get_tensor_by_name(
                output_name + ":0")
            input_dict = dict()
            output_dict = dict()
            input_dict[input_name] = input_tensor
            output_dict[output_name] = output_tensor
            tf.saved_model.simple_save(sess,
                                       model_version_dir + "/model.savedmodel",
                                       inputs=input_dict,
                                       outputs=output_dict)
    else:
        with tf.Session() as sess:
            graph_io.write_graph(sess.graph.as_graph_def(),
                                 model_version_dir,
                                 "model.graphdef",
                                 as_text=False)


def create_cudaop_modelconfig(create_savedmodel, models_dir, model_version):
    model_name = "savedmodel_cudaop" if create_savedmodel else "graphdef_cudaop"
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "{}"
max_batch_size: 0
input [
  {{
    name: "in"
    data_type: TYPE_INT32
    dims: [ -1 ]
  }}
]
output [
  {{
    name: "out"
    data_type: TYPE_INT32
    dims: [ -1 ]
  }}
]
'''.format(
        model_name,
        "tensorflow_savedmodel" if create_savedmodel else "tensorflow_graphdef")

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_busyop_modelfile(create_savedmodel, models_dir, model_version):
    # Load the busy_loop custom operator
    _busy_op_module = tf.load_op_library(os.path.join(FLAGS.busy_op_lib_path))
    busy_loop = _busy_op_module.busy_loop

    # Create the model that uses custom operator.
    tf.reset_default_graph()
    zin = tf.placeholder(tf.int32, [
        None,
    ], "in")
    zout = busy_loop(zin, name="out")

    model_name = "savedmodel_busyop" if create_savedmodel else "graphdef_busyop"
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    if create_savedmodel:
        with tf.Session() as sess:
            input_name = "in"
            output_name = "out"
            input_tensor = tf.get_default_graph().get_tensor_by_name(
                input_name + ":0")
            output_tensor = tf.get_default_graph().get_tensor_by_name(
                output_name + ":0")
            input_dict = dict()
            output_dict = dict()
            input_dict[input_name] = input_tensor
            output_dict[output_name] = output_tensor
            tf.saved_model.simple_save(sess,
                                       model_version_dir + "/model.savedmodel",
                                       inputs=input_dict,
                                       outputs=output_dict)
    else:
        with tf.Session() as sess:
            graph_io.write_graph(sess.graph.as_graph_def(),
                                 model_version_dir,
                                 "model.graphdef",
                                 as_text=False)


def create_busyop_modelconfig(create_savedmodel, models_dir, model_version):
    model_name = "savedmodel_busyop" if create_savedmodel else "graphdef_busyop"
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "{}"
max_batch_size: 0
input [
  {{
    name: "in"
    data_type: TYPE_INT32
    dims: [ -1 ]
  }}
]
output [
  {{
    name: "out"
    data_type: TYPE_INT32
    dims: [ -1 ]
  }}
]
'''.format(
        model_name,
        "tensorflow_savedmodel" if create_savedmodel else "tensorflow_graphdef")

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_moduloop_modelfile(models_dir, model_version):
    model_name = "libtorch_modulo"

    op_source = """
    #include <torch/script.h>
    torch::Tensor custom_modulo(torch::Tensor input1, torch::Tensor input2) {
      torch::Tensor output = torch::fmod(input1, input2);
      return output.clone();
    }
    static auto registry =
      torch::RegisterOperators("my_ops::custom_modulo", &custom_modulo);
    """

    torch.utils.cpp_extension.load_inline(
        name="custom_modulo",
        cpp_sources=op_source,
        is_python_module=False,
        verbose=True,
    )

    class ModuloCustomNet(nn.Module):

        def __init__(self):
            super(ModuloCustomNet, self).__init__()

        def forward(self, input0, input1):
            return torch.ops.my_ops.custom_modulo(input0, input1)

    moduloCustomModel = ModuloCustomNet()
    example_input0 = torch.arange(1, 11, dtype=torch.float32)
    example_input1 = torch.tensor([2] * 10, dtype=torch.float32)
    traced = torch.jit.trace(moduloCustomModel,
                             (example_input0, example_input1))

    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    traced.save(model_version_dir + "/model.pt")


def create_moduloop_modelconfig(models_dir, model_version):
    model_name = "libtorch_modulo"
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "pytorch_libtorch"
max_batch_size: 0
input [
  {{
    name: "INPUT__0"
    data_type: TYPE_FP32
    dims: [ 10 ]
  }},
  {{
    name: "INPUT__1"
    data_type: TYPE_FP32
    dims: [ 10 ]
  }}
]
output [
  {{
    name: "OUTPUT__0"
    data_type: TYPE_FP32
    dims: [ 10 ]
  }}
]
'''.format(model_name)

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


# Use Torchvision ops
def create_visionop_modelfile(models_dir, model_version):
    model_name = "libtorch_visionop"

    class CustomVisionNet(nn.Module):

        def __init__(self):
            super(CustomVisionNet, self).__init__()

        def forward(self, input, boxes):
            return torchvision.ops.roi_align(input, boxes, [5, 5], 1.0, -1,
                                             False)

    visionCustomModel = CustomVisionNet()
    visionCustomModel.eval()
    scripted = torch.jit.script(visionCustomModel)

    model_version_dir = models_dir + "/" + \
        model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    scripted.save(model_version_dir + "/model.pt")


def create_visionop_modelconfig(models_dir, model_version):
    model_name = "libtorch_visionop"
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "pytorch_libtorch"
max_batch_size: 0
input [
  {{
    name: "INPUT__0"
    data_type: TYPE_FP32
    dims: [ 1, 3, 10, 10 ]
  }},
  {{
    name: "INPUT__1"
    data_type: TYPE_FP32
    dims: [1, 5]
  }}
]
output [
  {{
    name: "OUTPUT__0"
    data_type: TYPE_FP32
    dims: [1, 3, 5, 5]
  }}
]
'''.format(model_name)

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_zero_out_models(models_dir):
    model_version = 1

    if FLAGS.graphdef:
        create_zeroout_modelconfig(False, models_dir, model_version)
        create_zeroout_modelfile(False, models_dir, model_version)

    if FLAGS.savedmodel:
        create_zeroout_modelconfig(True, models_dir, model_version)
        create_zeroout_modelfile(True, models_dir, model_version)


def create_cuda_op_models(models_dir):
    model_version = 1

    if FLAGS.graphdef:
        create_cudaop_modelconfig(False, models_dir, model_version)
        create_cudaop_modelfile(False, models_dir, model_version)

    if FLAGS.savedmodel:
        create_cudaop_modelconfig(True, models_dir, model_version)
        create_cudaop_modelfile(True, models_dir, model_version)


def create_busy_op_models(models_dir):
    model_version = 1

    if FLAGS.graphdef:
        create_busyop_modelconfig(False, models_dir, model_version)
        create_busyop_modelfile(False, models_dir, model_version)

    if FLAGS.savedmodel:
        create_busyop_modelconfig(True, models_dir, model_version)
        create_busyop_modelfile(True, models_dir, model_version)


def create_modulo_op_models(models_dir):
    model_version = 1

    if FLAGS.libtorch:
        create_moduloop_modelconfig(models_dir, model_version)
        create_moduloop_modelfile(models_dir, model_version)


def create_vision_op_models(models_dir):
    model_version = 1

    if FLAGS.libtorch:
        create_visionop_modelconfig(models_dir, model_version)
        create_visionop_modelfile(models_dir, model_version)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--models_dir',
                        type=str,
                        required=True,
                        help='Top-level model directory')
    parser.add_argument('--zero_out_lib_path',
                        type=str,
                        required=False,
                        default="./libzeroout.so",
                        help='Fullpath to libzeroout.so')
    parser.add_argument('--cuda_op_lib_path',
                        type=str,
                        required=False,
                        default="./libcudaop.so",
                        help='Fullpath to libcudaop.so')
    parser.add_argument('--busy_op_lib_path',
                        type=str,
                        required=False,
                        default="./libbusyop.so",
                        help='Fullpath to libbusyop.so')
    parser.add_argument('--graphdef',
                        required=False,
                        action='store_true',
                        help='Generate GraphDef models')
    parser.add_argument('--savedmodel',
                        required=False,
                        action='store_true',
                        help='Generate SavedModel models')
    parser.add_argument('--libtorch',
                        required=False,
                        action='store_true',
                        help='Generate Pytorch LibTorch models')
    FLAGS, unparsed = parser.parse_known_args()

    if FLAGS.graphdef or FLAGS.savedmodel:
        # Use Tensorflow 2 as default. Need to disable the v2 behavior for
        # model generation scripts.
        import tensorflow.compat.v1 as tf
        tf.disable_v2_behavior()
        from tensorflow.python.framework import graph_io
        create_zero_out_models(FLAGS.models_dir)
        create_cuda_op_models(FLAGS.models_dir)
        create_busy_op_models(FLAGS.models_dir)

    if FLAGS.libtorch:
        import torch
        from torch import nn
        import torchvision
        import torch.utils.cpp_extension
        create_modulo_op_models(FLAGS.models_dir)
        create_vision_op_models(FLAGS.models_dir)


# qa/common/gen_qa_dyna_sequence_implicit_models.py

# Copyright 2019-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
import os
import numpy as np

FLAGS = None
np_dtype_string = np.dtype(object)


def np_to_onnx_dtype(np_dtype):
    if np_dtype == bool:
        return onnx.TensorProto.BOOL
    elif np_dtype == np.int8:
        return onnx.TensorProto.INT8
    elif np_dtype == np.int16:
        return onnx.TensorProto.INT16
    elif np_dtype == np.int32:
        return onnx.TensorProto.INT32
    elif np_dtype == np.int64:
        return onnx.TensorProto.INT64
    elif np_dtype == np.uint8:
        return onnx.TensorProto.UINT8
    elif np_dtype == np.uint16:
        return onnx.TensorProto.UINT16
    elif np_dtype == np.float16:
        return onnx.TensorProto.FLOAT16
    elif np_dtype == np.float32:
        return onnx.TensorProto.FLOAT
    elif np_dtype == np.float64:
        return onnx.TensorProto.DOUBLE
    elif np_dtype == np_dtype_string:
        return onnx.TensorProto.STRING


def np_to_model_dtype(np_dtype):
    if np_dtype == bool:
        return "TYPE_BOOL"
    elif np_dtype == np.int8:
        return "TYPE_INT8"
    elif np_dtype == np.int16:
        return "TYPE_INT16"
    elif np_dtype == np.int32:
        return "TYPE_INT32"
    elif np_dtype == np.int64:
        return "TYPE_INT64"
    elif np_dtype == np.uint8:
        return "TYPE_UINT8"
    elif np_dtype == np.uint16:
        return "TYPE_UINT16"
    elif np_dtype == np.float16:
        return "TYPE_FP16"
    elif np_dtype == np.float32:
        return "TYPE_FP32"
    elif np_dtype == np.float64:
        return "TYPE_FP64"
    elif np_dtype == np_dtype_string:
        return "TYPE_STRING"
    return None


def np_to_trt_dtype(np_dtype):
    if np_dtype == bool:
        return trt.bool
    elif np_dtype == np.int8:
        return trt.int8
    elif np_dtype == np.int32:
        return trt.int32
    elif np_dtype == np.float16:
        return trt.float16
    elif np_dtype == np.float32:
        return trt.float32
    return None


def create_onnx_modelfile(models_dir, model_version, max_batch, dtype, shape):

    if not tu.validate_for_onnx_model(dtype, dtype, dtype, shape, shape, shape):
        return

    model_name = tu.get_dyna_sequence_model_name(
        "onnx_nobatch" if max_batch == 0 else "onnx", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    # Create the model. For now don't implement a proper accumulator
    # just return 0 if not-ready and 'INPUT'+'START'*('END'*'CORRID')
    # otherwise...  the tests know to expect this.
    onnx_dtype = np_to_onnx_dtype(dtype)
    onnx_input_shape, idx = tu.shape_to_onnx_shape(shape, 0)
    onnx_output_shape, idx = tu.shape_to_onnx_shape(shape, idx)

    # If the input is a string then use int32 for operation and just
    # cast to/from string for input and output.
    onnx_control_dtype = onnx_dtype
    if onnx_dtype == onnx.TensorProto.STRING:
        onnx_control_dtype = onnx.TensorProto.INT32

    # If input dtype is bool, then use bool type for control and
    # int32 type for input/output
    if onnx_dtype == onnx.TensorProto.BOOL:
        onnx_dtype = onnx.TensorProto.INT32

    batch_dim = [] if max_batch == 0 else [None]

    onnx_input = onnx.helper.make_tensor_value_info(
        "INPUT", onnx_dtype, batch_dim + onnx_input_shape)
    onnx_input_state = onnx.helper.make_tensor_value_info(
        "INPUT_STATE", onnx_dtype, batch_dim + onnx_input_shape)
    onnx_start = onnx.helper.make_tensor_value_info("START", onnx_control_dtype,
                                                    batch_dim + [1])
    onnx_ready = onnx.helper.make_tensor_value_info("READY", onnx_control_dtype,
                                                    batch_dim + [1])
    onnx_corrid = onnx.helper.make_tensor_value_info("CORRID",
                                                     onnx.TensorProto.UINT64,
                                                     batch_dim + [1])
    onnx_end = onnx.helper.make_tensor_value_info("END", onnx_control_dtype,
                                                  batch_dim + [1])
    onnx_output = onnx.helper.make_tensor_value_info(
        "OUTPUT", onnx_dtype, batch_dim + onnx_output_shape)
    onnx_output_state = onnx.helper.make_tensor_value_info(
        "OUTPUT_STATE", onnx_dtype, batch_dim + onnx_output_shape)

    internal_input = onnx.helper.make_node("Identity", ["INPUT"], ["_INPUT"])
    internal_input_state = onnx.helper.make_node("Identity", ["INPUT_STATE"],
                                                 ["_INPUT_STATE"])
    # cast int8, int16 input to higer precision int as Onnx Add/Sub operator doesn't support those type
    # Also casting String data type to int32
    if ((onnx_dtype == onnx.TensorProto.INT8) or
        (onnx_dtype == onnx.TensorProto.INT16) or
        (onnx_dtype == onnx.TensorProto.STRING)):
        internal_input = onnx.helper.make_node("Cast", ["INPUT"], ["_INPUT"],
                                               to=onnx.TensorProto.INT32)
        internal_input_state = onnx.helper.make_node("Cast", ["INPUT_STATE"],
                                                     ["_INPUT_STATE"],
                                                     to=onnx.TensorProto.INT32)

    # Convert boolean value to int32 value
    if onnx_control_dtype == onnx.TensorProto.BOOL:
        internal_input1 = onnx.helper.make_node("Cast", ["START"], ["_START"],
                                                to=onnx.TensorProto.INT32)
        internal_input2 = onnx.helper.make_node("Cast", ["READY"], ["_READY"],
                                                to=onnx.TensorProto.INT32)
        not_start_cast = onnx.helper.make_node("Not", ["START"],
                                               ["_NOT_START_CAST"])
        not_start = onnx.helper.make_node("Cast", ["_NOT_START_CAST"],
                                          ["_NOT_START"],
                                          to=onnx.TensorProto.INT32)
        not_ready_cast = onnx.helper.make_node("Not", ["START"],
                                               ["_NOT_READY_CAST"])
        not_ready = onnx.helper.make_node("Cast", ["_NOT_READY_CAST"],
                                          ["_NOT_READY"],
                                          to=onnx.TensorProto.INT32)

        input_state_cond = onnx.helper.make_node("And",
                                                 ["READY", "_NOT_START_CAST"],
                                                 ["input_state_cond"])
        input_state_cond_cast = onnx.helper.make_node("Cast",
                                                      ["input_state_cond"],
                                                      ["input_state_cond_cast"],
                                                      to=onnx.TensorProto.INT32)
        mul_state = onnx.helper.make_node(
            "Mul", ["_INPUT_STATE", "input_state_cond_cast"], ["mul_state"])
        add = onnx.helper.make_node("Add", ["_INPUT", "mul_state"], ["CAST"])

    else:
        start_cast = onnx.helper.make_node("Cast", ["START"], ["_START_CAST"],
                                           to=onnx.TensorProto.BOOL)
        not_start_cast = onnx.helper.make_node("Not", ["_START_CAST"],
                                               ["_NOT_START_CAST"])
        not_start = onnx.helper.make_node("Cast", ["_NOT_START_CAST"],
                                          ["_NOT_START"],
                                          to=onnx.TensorProto.INT32)

        ready_cast = onnx.helper.make_node("Cast", ["READY"], ["_READY_CAST"],
                                           to=onnx.TensorProto.BOOL)
        not_ready_cast = onnx.helper.make_node("Not", ["_READY_CAST"],
                                               ["_NOT_READY_CAST"])
        not_ready = onnx.helper.make_node("Cast", ["_NOT_READY_CAST"],
                                          ["_NOT_READY"],
                                          to=onnx.TensorProto.INT32)

        # Take advantage of knowledge that the READY false value is 0 and true is 1
        input_state_cond = onnx.helper.make_node(
            "And", ["_NOT_START_CAST", "_READY_CAST"], ["input_state_cond"])
        input_state_cond_cast = onnx.helper.make_node("Cast",
                                                      ["input_state_cond"],
                                                      ["input_state_cond_cast"],
                                                      to=onnx.TensorProto.INT32)
        mul_state = onnx.helper.make_node(
            "Mul", ["_INPUT_STATE", "input_state_cond_cast"], ["mul_state"])
        add = onnx.helper.make_node("Add", ["_INPUT", "mul_state"], ["CAST"])

    cast = onnx.helper.make_node("Cast", ["CAST"], ["OUTPUT"], to=onnx_dtype)
    cast_output_state = onnx.helper.make_node("Cast", ["CAST"],
                                              ["OUTPUT_STATE"],
                                              to=onnx_dtype)

    # Avoid cast from float16 to float16
    # (bug in Onnx Runtime, cast from float16 to float16 will become cast from float16 to float32)
    if onnx_dtype == onnx.TensorProto.FLOAT16:
        cast = onnx.helper.make_node("Identity", ["CAST"], ["OUTPUT"])
        cast_output_state = onnx.helper.make_node("Identity", ["CAST"],
                                                  ["OUTPUT_STATE"])

    if onnx_control_dtype == onnx.TensorProto.BOOL:
        onnx_nodes = [
            internal_input, internal_input_state, internal_input1,
            internal_input2, not_start_cast, not_start, not_ready_cast,
            not_ready, input_state_cond, input_state_cond_cast, mul_state, add,
            cast, cast_output_state
        ]
    else:
        onnx_nodes = [
            internal_input, internal_input_state, start_cast, not_start_cast,
            not_start, ready_cast, not_ready_cast, not_ready, input_state_cond,
            input_state_cond_cast, mul_state, add, cast, cast_output_state
        ]

    onnx_inputs = [
        onnx_end, onnx_corrid, onnx_input_state, onnx_input, onnx_start,
        onnx_ready
    ]
    onnx_outputs = [onnx_output, onnx_output_state]
    graph_proto = onnx.helper.make_graph(onnx_nodes, model_name, onnx_inputs,
                                         onnx_outputs)

    if FLAGS.onnx_opset > 0:
        model_opset = onnx.helper.make_operatorsetid("", FLAGS.onnx_opset)
        model_def = onnx.helper.make_model(graph_proto,
                                           producer_name="triton",
                                           opset_imports=[model_opset])
    else:
        model_def = onnx.helper.make_model(graph_proto, producer_name="triton")

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    onnx.save(model_def, model_version_dir + "/model.onnx")


def create_onnx_modelconfig(models_dir, model_version, max_batch, dtype, shape):

    if not tu.validate_for_onnx_model(dtype, dtype, dtype, shape, shape, shape):
        return

    model_name = tu.get_dyna_sequence_model_name(
        "onnx_nobatch" if max_batch == 0 else "onnx", dtype)
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "onnxruntime_onnx"
max_batch_size: {}
sequence_batching {{
  max_sequence_idle_microseconds: 5000000
  {}
  control_input [
    {{
      name: "START"
      control [
        {{
          kind: CONTROL_SEQUENCE_START
          {type}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "END"
      control [
        {{
          kind: CONTROL_SEQUENCE_END
          {type}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "READY"
      control [
        {{
          kind: CONTROL_SEQUENCE_READY
          {type}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "CORRID"
      control [
        {{
          kind: CONTROL_SEQUENCE_CORRID
          data_type: TYPE_UINT64
        }}
      ]
    }}
  ]
  state [
    {{
      input_name: "INPUT_STATE"
      output_name: "OUTPUT_STATE"
      data_type: {dtype}
      dims: {dims}
    }} 
  ]
}}
input [
  {{
    name: "INPUT"
    data_type: {dtype}
    dims: [ {dims} ]
  }}
]
output [
  {{
    name: "OUTPUT"
    data_type: {dtype}
    dims: [ {dims} ]
  }}
]
instance_group [
  {{
    kind: KIND_CPU
  }}
]
'''.format(
        model_name,
        max_batch,
        "oldest { max_candidate_sequences: 6\npreferred_batch_size: [ 4 ]\nmax_queue_delay_microseconds: 0\n}"
        if max_batch > 0 else "",
        dtype=np_to_model_dtype(dtype),
        dims=tu.shape_to_dims_str(shape),
        type="fp32" if dtype == np.float32 else "int32")

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_plan_fixed_modelfile(models_dir, model_version, max_batch, dtype,
                                shape):
    trt_dtype = np_to_trt_dtype(dtype)
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network()
    in0 = network.add_input("INPUT", trt_dtype, shape)
    in_state0 = network.add_input("INPUT_STATE", trt_dtype, shape)
    start0 = network.add_input("START", trt_dtype, [1 for i in shape])
    network.add_input("END", trt_dtype, [1 for i in shape])
    ready0 = network.add_input("READY", trt_dtype, [1 for i in shape])
    network.add_input("CORRID", trt.int32, [1 for i in shape])

    constant_1_data = trt.Weights(np.ones([1 for i in shape], dtype=dtype))
    constant_1 = network.add_constant([1 for i in shape], constant_1_data)
    not_start = network.add_elementwise(constant_1.get_output(0), start0,
                                        trt.ElementWiseOperation.SUB)
    not_start.set_output_type(0, trt_dtype)

    input_state_cond_temp = network.add_elementwise(
        ready0, not_start.get_output(0), trt.ElementWiseOperation.SUM)
    constant_2 = network.add_elementwise(constant_1.get_output(0),
                                         constant_1.get_output(0),
                                         trt.ElementWiseOperation.SUM)
    input_state_cond = network.add_elementwise(
        input_state_cond_temp.get_output(0), constant_2.get_output(0),
        trt.ElementWiseOperation.FLOOR_DIV)
    internal_state = network.add_elementwise(in_state0,
                                             input_state_cond.get_output(0),
                                             trt.ElementWiseOperation.PROD)
    out0 = network.add_elementwise(internal_state.get_output(0), in0,
                                   trt.ElementWiseOperation.SUM)
    out0_state = network.add_elementwise(internal_state.get_output(0), in0,
                                         trt.ElementWiseOperation.SUM)

    out0.get_output(0).name = "OUTPUT"
    network.mark_output(out0.get_output(0))

    out0_state.get_output(0).name = "OUTPUT_STATE"
    network.mark_output(out0_state.get_output(0))

    config = builder.create_builder_config()
    config.max_workspace_size = 1 << 20
    builder.max_batch_size = max(1, max_batch)
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine
    del network

    model_name = tu.get_dyna_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_fixed_rf_modelfile(models_dir, model_version, max_batch, dtype,
                                   shape):
    trt_dtype = np_to_trt_dtype(dtype)
    trt_memory_format = trt.TensorFormat.LINEAR

    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network()
    in0 = network.add_input("INPUT", trt_dtype, shape)
    in_state0 = network.add_input("INPUT_STATE", trt_dtype, shape)
    start0 = network.add_input("START", trt_dtype, [1 for i in shape])
    network.add_input("END", trt_dtype, [1 for i in shape])
    ready0 = network.add_input("READY", trt_dtype, [1 for i in shape])
    network.add_input("CORRID", trt.int32, [1 for i in shape])

    constant_1_data = trt.Weights(np.ones([1 for i in shape], dtype=dtype))
    constant_1 = network.add_constant([1 for i in shape], constant_1_data)
    not_start = network.add_elementwise(constant_1.get_output(0), start0,
                                        trt.ElementWiseOperation.SUB)
    not_start.set_output_type(0, trt_dtype)

    input_state_cond_temp = network.add_elementwise(
        ready0, not_start.get_output(0), trt.ElementWiseOperation.SUM)
    constant_2 = network.add_elementwise(constant_1.get_output(0),
                                         constant_1.get_output(0),
                                         trt.ElementWiseOperation.SUM)
    input_state_cond = network.add_elementwise(
        input_state_cond_temp.get_output(0), constant_2.get_output(0),
        trt.ElementWiseOperation.FLOOR_DIV)
    internal_state = network.add_elementwise(in_state0,
                                             input_state_cond.get_output(0),
                                             trt.ElementWiseOperation.PROD)
    out0 = network.add_elementwise(internal_state.get_output(0), in0,
                                   trt.ElementWiseOperation.SUM)
    out0_state = network.add_elementwise(internal_state.get_output(0), in0,
                                         trt.ElementWiseOperation.SUM)

    out0.get_output(0).name = "OUTPUT"
    network.mark_output(out0.get_output(0))
    out0.get_output(0).dtype = trt_dtype

    out0_state.get_output(0).name = "OUTPUT_STATE"
    network.mark_output(out0_state.get_output(0))
    out0_state.get_output(0).dtype = trt_dtype

    in0.allowed_formats = 1 << int(trt_memory_format)
    in_state0.allowed_formats = 1 << int(trt_memory_format)
    start0.allowed_formats = 1 << int(trt_memory_format)
    ready0.allowed_formats = 1 << int(trt_memory_format)
    out0.get_output(0).allowed_formats = 1 << int(trt_memory_format)
    out0_state.get_output(0).allowed_formats = 1 << int(trt_memory_format)

    if (trt_dtype == trt.int8):
        in0.dynamic_range = (-128.0, 127.0)
        in_state0.dynamic_range = (-128.0, 127.0)
        out0.dynamic_range = (-128.0, 127.0)
        out0_state.dynamic_range = (-128.0, 127.0)
        start0.dynamic_range = (-128.0, 127.0)
        ready0.dynamic_range = (-128.0, 127.0)

    flags = 1 << int(trt.BuilderFlag.STRICT_TYPES)

    if (trt_dtype == trt.int8):
        flags |= 1 << int(trt.BuilderFlag.INT8)
    elif (trt_dtype == trt.float16):
        flags |= 1 << int(trt.BuilderFlag.FP16)

    config = builder.create_builder_config()
    config.flags = flags
    config.max_workspace_size = 1 << 20
    builder.max_batch_size = max(1, max_batch)
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    model_name = tu.get_dyna_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_modelfile(models_dir, model_version, max_batch, dtype, shape):

    if not tu.validate_for_trt_model(dtype, dtype, dtype, shape, shape, shape):
        return

    if dtype != np.float32:
        create_plan_fixed_rf_modelfile(models_dir, model_version, max_batch,
                                       dtype, shape)
    else:
        create_plan_fixed_modelfile(models_dir, model_version, max_batch, dtype,
                                    shape)


def create_plan_modelconfig(models_dir, model_version, max_batch, dtype, shape):

    if not tu.validate_for_trt_model(dtype, dtype, dtype, shape, shape, shape):
        return

    model_name = tu.get_dyna_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "tensorrt_plan"
max_batch_size: {}
sequence_batching {{
  max_sequence_idle_microseconds: 5000000
  {}
  control_input [
    {{
      name: "START"
      control [
        {{
          kind: CONTROL_SEQUENCE_START
          {type}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "END"
      control [
        {{
          kind: CONTROL_SEQUENCE_END
          {type}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "READY"
      control [
        {{
          kind: CONTROL_SEQUENCE_READY
          {type}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "CORRID"
      control [
        {{
          kind: CONTROL_SEQUENCE_CORRID
          data_type: TYPE_INT32
        }}
      ]
    }}
  ]
  state [
    {{
      input_name: "INPUT_STATE"
      output_name: "OUTPUT_STATE"
      data_type: {dtype}
      dims: {dims}
    }} 
  ]
}}
input [
  {{
    name: "INPUT"
    data_type: {dtype}
    dims: [ {dims} ]
  }}
]
output [
  {{
    name: "OUTPUT"
    data_type: {dtype}
    dims: [ {dims} ]
  }}
]
instance_group [
  {{
    kind: KIND_GPU
  }}
]
'''.format(
        model_name,
        max_batch,
        "oldest { max_candidate_sequences: 6\npreferred_batch_size: [ 4 ]\nmax_queue_delay_microseconds: 0\n}"
        if max_batch > 0 else "",
        dtype=np_to_model_dtype(dtype),
        dims=tu.shape_to_dims_str(shape),
        type="fp32" if dtype == np.float32 else "int32")

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_models(models_dir, dtype, shape, no_batch=True):
    model_version = 1

    if FLAGS.onnx:
        create_onnx_modelconfig(models_dir, model_version, 8, dtype, shape)
        create_onnx_modelfile(models_dir, model_version, 8, dtype, shape)
        if no_batch:
            create_onnx_modelconfig(models_dir, model_version, 0, dtype, shape)
            create_onnx_modelfile(models_dir, model_version, 0, dtype, shape)

    if FLAGS.tensorrt:
        if dtype == bool:
            return

        create_plan_modelconfig(models_dir, model_version, 8, dtype, shape)
        create_plan_modelfile(models_dir, model_version, 8, dtype, shape)
        if no_batch:
            create_plan_modelconfig(models_dir, model_version, 0, dtype, shape)
            create_plan_modelfile(models_dir, model_version, 0, dtype, shape)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--models_dir',
                        type=str,
                        required=True,
                        help='Top-level model directory')
    parser.add_argument('--graphdef',
                        required=False,
                        action='store_true',
                        help='Generate GraphDef models')
    parser.add_argument('--savedmodel',
                        required=False,
                        action='store_true',
                        help='Generate SavedModel models')
    parser.add_argument('--tensorrt',
                        required=False,
                        action='store_true',
                        help='Generate TensorRT PLAN models')
    parser.add_argument(
        '--tensorrt-shape-io',
        required=False,
        action='store_true',
        help='Generate TensorRT PLAN models w/ shape tensor i/o')
    parser.add_argument('--onnx',
                        required=False,
                        action='store_true',
                        help='Generate Onnx models')
    parser.add_argument(
        '--onnx_opset',
        type=int,
        required=False,
        default=0,
        help='Opset used for Onnx models. Default is to use ONNXRT default')
    parser.add_argument('--libtorch',
                        required=False,
                        action='store_true',
                        help='Generate Pytorch LibTorch models')
    parser.add_argument('--openvino',
                        required=False,
                        action='store_true',
                        help='Generate OpenVino models')
    parser.add_argument('--variable',
                        required=False,
                        action='store_true',
                        help='Used variable-shape tensors for input/output')
    FLAGS, unparsed = parser.parse_known_args()

    if FLAGS.onnx:
        import onnx

    if FLAGS.tensorrt:
        import tensorrt as trt

    import test_util as tu

    # Tests with models that accept fixed-shape input/output tensors
    if not FLAGS.variable:
        create_models(FLAGS.models_dir, np.int32, [
            1,
        ])

    # Tests with models that accept variable-shape input/output tensors
    if FLAGS.variable:
        create_models(FLAGS.models_dir, np.int32, [
            -1,
        ], False)


# qa/common/gen_qa_dyna_sequence_models.py

# Copyright 2019-2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
import os
import numpy as np

FLAGS = None
np_dtype_string = np.dtype(object)


def np_to_model_dtype(np_dtype):
    if np_dtype == bool:
        return "TYPE_BOOL"
    elif np_dtype == np.int8:
        return "TYPE_INT8"
    elif np_dtype == np.int16:
        return "TYPE_INT16"
    elif np_dtype == np.int32:
        return "TYPE_INT32"
    elif np_dtype == np.int64:
        return "TYPE_INT64"
    elif np_dtype == np.uint8:
        return "TYPE_UINT8"
    elif np_dtype == np.uint16:
        return "TYPE_UINT16"
    elif np_dtype == np.float16:
        return "TYPE_FP16"
    elif np_dtype == np.float32:
        return "TYPE_FP32"
    elif np_dtype == np.float64:
        return "TYPE_FP64"
    elif np_dtype == np_dtype_string:
        return "TYPE_STRING"
    return None


def np_to_tf_dtype(np_dtype):
    if np_dtype == bool:
        return tf.bool
    elif np_dtype == np.int8:
        return tf.int8
    elif np_dtype == np.int16:
        return tf.int16
    elif np_dtype == np.int32:
        return tf.int32
    elif np_dtype == np.int64:
        return tf.int64
    elif np_dtype == np.uint8:
        return tf.uint8
    elif np_dtype == np.uint16:
        return tf.uint16
    elif np_dtype == np.float16:
        return tf.float16
    elif np_dtype == np.float32:
        return tf.float32
    elif np_dtype == np.float64:
        return tf.float64
    elif np_dtype == np_dtype_string:
        return tf.string
    return None


def np_to_trt_dtype(np_dtype):
    if np_dtype == bool:
        return trt.bool
    elif np_dtype == np.int8:
        return trt.int8
    elif np_dtype == np.int32:
        return trt.int32
    elif np_dtype == np.float16:
        return trt.float16
    elif np_dtype == np.float32:
        return trt.float32
    return None


def np_to_onnx_dtype(np_dtype):
    if np_dtype == bool:
        return onnx.TensorProto.BOOL
    elif np_dtype == np.int8:
        return onnx.TensorProto.INT8
    elif np_dtype == np.int16:
        return onnx.TensorProto.INT16
    elif np_dtype == np.int32:
        return onnx.TensorProto.INT32
    elif np_dtype == np.int64:
        return onnx.TensorProto.INT64
    elif np_dtype == np.uint8:
        return onnx.TensorProto.UINT8
    elif np_dtype == np.uint16:
        return onnx.TensorProto.UINT16
    elif np_dtype == np.float16:
        return onnx.TensorProto.FLOAT16
    elif np_dtype == np.float32:
        return onnx.TensorProto.FLOAT
    elif np_dtype == np.float64:
        return onnx.TensorProto.DOUBLE
    elif np_dtype == np_dtype_string:
        return onnx.TensorProto.STRING


def np_to_torch_dtype(np_dtype):
    if np_dtype == bool:
        return torch.bool
    elif np_dtype == np.int8:
        return torch.int8
    elif np_dtype == np.int16:
        return torch.int16
    elif np_dtype == np.int32:
        return torch.int
    elif np_dtype == np.int64:
        return torch.long
    elif np_dtype == np.uint8:
        return torch.uint8
    elif np_dtype == np.uint16:
        return None  # Not supported in Torch
    elif np_dtype == np.float16:
        return None
    elif np_dtype == np.float32:
        return torch.float
    elif np_dtype == np.float64:
        return torch.double
    elif np_dtype == np_dtype_string:
        return None  # Not supported in Torch
    return None


def create_tf_modelfile(create_savedmodel, models_dir, model_version, max_batch,
                        dtype, shape):

    if not tu.validate_for_tf_model(dtype, dtype, dtype, shape, shape, shape):
        return

    tf_input_dtype = np_to_tf_dtype(dtype)
    tf_dtype = tf_input_dtype

    # If the input is a string then use int32 for operation and just
    # cast to/from string for input and output.
    if tf_input_dtype == tf.string:
        tf_dtype = tf.int32

    # Create the model. If non-batching then don't include the batch
    # dimension.
    tf.reset_default_graph()
    if create_savedmodel and (max_batch == 0):
        input0 = tf.placeholder(tf_input_dtype, [
            1,
        ], "INPUT")
        if tf_input_dtype == tf.string:
            input0 = tf.strings.to_number(tf.strings.join(["0", input0]),
                                          tf_dtype)
        start0 = tf.placeholder(tf_dtype, [
            1,
        ], "START")
        end0 = tf.placeholder(tf_dtype, [
            1,
        ], "END")
        ready0 = tf.placeholder(tf_dtype, [
            1,
        ], "READY")
        corrid0 = tf.placeholder(tf.uint64, [
            1,
        ], "CORRID")
        corrid_cast0 = tf.cast(corrid0, tf_dtype)
        acc = tf.get_variable("ACC", [
            1,
        ], dtype=tf_dtype)
        tmp0 = tf.where(tf.equal(start0, 1), input0, tf.add(acc, input0))
        tmp1 = tf.where(tf.equal(end0, 1), tf.add(tmp0, corrid_cast0), tmp0)
        newacc = tf.where(tf.equal(ready0, 1), tmp1, acc)
        assign = tf.assign(acc, newacc)
        if tf_input_dtype == tf.string:
            output0 = tf.dtypes.as_string(assign, name="OUTPUT")
        else:
            output0 = tf.identity(assign, name="OUTPUT")
    else:
        # For batching we can't use a tf.variable to hold the
        # accumulated values since that forces the size of the output
        # to the size of the variable (which must be a max-batch-size
        # vector since require one accumulator each), instead of the
        # output shape being [None, 1]. So instead we just return 0 if
        # not-ready and 'INPUT'+'START'+('END'*'CORRID')
        # otherwise... the tests know to expect this.
        input0 = tf.placeholder(tf_input_dtype, [
            None,
        ] + tu.shape_to_tf_shape(shape), "INPUT")
        if tf_input_dtype == tf.string:
            input0 = tf.strings.to_number(tf.strings.join(["0", input0]),
                                          tf_dtype)
        start0 = tf.placeholder(tf_dtype, [None, 1], "START")
        end0 = tf.placeholder(tf_dtype, [None, 1], "END")
        ready0 = tf.placeholder(tf_dtype, [None, 1], "READY")
        corrid0 = tf.placeholder(tf.uint64, [None, 1], "CORRID")
        corrid_cast0 = tf.cast(corrid0, tf_dtype)
        tmp = tf.where(
            tf.equal(ready0, 1),
            tf.add(tf.add(start0, input0), tf.multiply(end0, corrid_cast0)),
            tf.zeros(tf.shape(input0), dtype=tf_dtype))
        if tf_input_dtype == tf.string:
            output0 = tf.dtypes.as_string(tmp, name="OUTPUT")
        else:
            output0 = tf.identity(tmp, name="OUTPUT")

    # Use a different model name for the non-batching variant
    if create_savedmodel:
        model_name = tu.get_dyna_sequence_model_name(
            "savedmodel_nobatch" if max_batch == 0 else "savedmodel", dtype)
    else:
        model_name = tu.get_dyna_sequence_model_name(
            "graphdef_nobatch" if max_batch == 0 else "graphdef", dtype)

    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    if create_savedmodel:
        with tf.Session() as sess:
            sess.run(tf.initializers.global_variables())
            input0_tensor = tf.get_default_graph().get_tensor_by_name("INPUT:0")
            start0_tensor = tf.get_default_graph().get_tensor_by_name("START:0")
            end0_tensor = tf.get_default_graph().get_tensor_by_name("END:0")
            ready0_tensor = tf.get_default_graph().get_tensor_by_name("READY:0")
            corrid0_tensor = tf.get_default_graph().get_tensor_by_name(
                "CORRID:0")
            output0_tensor = tf.get_default_graph().get_tensor_by_name(
                "OUTPUT:0")
            tf.saved_model.simple_save(sess,
                                       model_version_dir + "/model.savedmodel",
                                       inputs={
                                           "INPUT": input0_tensor,
                                           "START": start0_tensor,
                                           "END": end0_tensor,
                                           "READY": ready0_tensor,
                                           "CORRID": corrid0_tensor
                                       },
                                       outputs={"OUTPUT": output0_tensor})
    else:
        with tf.Session() as sess:
            sess.run(tf.initializers.global_variables())
            graph_io.write_graph(sess.graph.as_graph_def(),
                                 model_version_dir,
                                 "model.graphdef",
                                 as_text=False)


def create_tf_modelconfig(create_savedmodel, models_dir, model_version,
                          max_batch, dtype, shape):

    if not tu.validate_for_tf_model(dtype, dtype, dtype, shape, shape, shape):
        return

    # Use a different model name for the non-batching variant
    if create_savedmodel:
        model_name = tu.get_dyna_sequence_model_name(
            "savedmodel_nobatch" if max_batch == 0 else "savedmodel", dtype)
    else:
        model_name = tu.get_dyna_sequence_model_name(
            "graphdef_nobatch" if max_batch == 0 else "graphdef", dtype)

    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "{}"
max_batch_size: {}
sequence_batching {{
  max_sequence_idle_microseconds: 5000000
  {}
  control_input [
    {{
      name: "START"
      control [
        {{
          kind: CONTROL_SEQUENCE_START
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "END"
      control [
        {{
          kind: CONTROL_SEQUENCE_END
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "READY"
      control [
        {{
          kind: CONTROL_SEQUENCE_READY
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "CORRID"
      control [
        {{
          kind: CONTROL_SEQUENCE_CORRID
          data_type: TYPE_UINT64
        }}
      ]
    }}
  ]
}}
input [
  {{
    name: "INPUT"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT"
    data_type: {}
    dims: [ 1 ]
  }}
]
instance_group [
  {{
    kind: KIND_CPU
  }}
]
'''.format(
        model_name,
        "tensorflow_savedmodel" if create_savedmodel else "tensorflow_graphdef",
        max_batch,
        "oldest { max_candidate_sequences: 6\npreferred_batch_size: [ 4 ]\nmax_queue_delay_microseconds: 0\n}"
        if max_batch > 0 else "", "fp32" if dtype == np.float32 else "int32",
        "fp32" if dtype == np.float32 else "int32",
        "fp32" if dtype == np.float32 else "int32", np_to_model_dtype(dtype),
        tu.shape_to_dims_str(shape), np_to_model_dtype(dtype))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_plan_shape_tensor_modelfile(models_dir, model_version, max_batch,
                                       dtype, shape):
    # Note that resize layer does not support int tensors.
    # The model takes three inputs (INPUT, DUMMY_INPUT and SHAPE_INPUT)
    # and four control inputs(START, END, READY, CORR_ID).
    # In absence of proper accumulator,
    # OUTPUT : 0 if not-ready and 'DUMMY_INPUT'+'START'+('END'*'CORRID')
    #          otherwise
    # RESIZED_OUTPUT : Obtained after resizing 'INPUT' to shape specified
    #          in 'SHAPE_INPUT'
    # SHAPE_OUTPUT : The shape values of resized output

    trt_dtype = np_to_trt_dtype(dtype)
    trt_memory_format = trt.TensorFormat.LINEAR

    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

    unit_shape = ([1] * len(shape))
    dummy_shape = ([-1] * len(shape))
    if max_batch != 0:
        in0 = network.add_input("INPUT", trt.int32, [-1] + dummy_shape)
        dummy_in0 = network.add_input("DUMMY_INPUT", trt_dtype,
                                      [-1] + dummy_shape)
        shape_in0 = network.add_input("SHAPE_INPUT", trt.int32,
                                      [1 + len(shape)])
        start0 = network.add_input("START", trt.int32, [-1] + unit_shape)
        end0 = network.add_input("END", trt.int32, [-1] + unit_shape)
        ready0 = network.add_input("READY", trt.int32, [-1] + unit_shape)
        corrid0 = network.add_input("CORRID", trt.int32, [-1] + unit_shape)
    else:
        in0 = network.add_input("INPUT", trt.int32, dummy_shape)
        dummy_in0 = network.add_input("DUMMY_INPUT", trt_dtype, dummy_shape)
        shape_in0 = network.add_input("SHAPE_INPUT", trt.int32, [len(shape)])
        start0 = network.add_input("START", trt.int32, unit_shape)
        end0 = network.add_input("END", trt.int32, unit_shape)
        ready0 = network.add_input("READY", trt.int32, unit_shape)
        corrid0 = network.add_input("CORRID", trt.int32, unit_shape)

    add0 = network.add_elementwise(in0, start0, trt.ElementWiseOperation.SUM)
    mul0 = network.add_elementwise(end0, corrid0, trt.ElementWiseOperation.PROD)
    sum0 = network.add_elementwise(add0.get_output(0), mul0.get_output(0),
                                   trt.ElementWiseOperation.SUM)
    out0 = network.add_elementwise(sum0.get_output(0), ready0,
                                   trt.ElementWiseOperation.PROD).get_output(0)

    resize_layer = network.add_resize(dummy_in0)
    resize_layer.set_input(1, shape_in0)
    shape_out0 = network.add_shape(resize_layer.get_output(0))
    resized_out0 = resize_layer.get_output(0)

    shape_out0.get_output(0).name = "SHAPE_OUTPUT"
    shape_out0.get_output(0).dtype = trt.int32
    network.mark_output_for_shapes(shape_out0.get_output(0))

    out0.name = "OUTPUT"
    out0.dtype = trt.int32
    network.mark_output(out0)

    resized_out0.name = "RESIZED_OUTPUT"
    resized_out0.dtype = trt_dtype
    network.mark_output(resized_out0)

    shape_in0.allowed_formats = 1 << int(trt_memory_format)
    dummy_in0.allowed_formats = 1 << int(trt_memory_format)
    start0.allowed_formats = 1 << int(trt_memory_format)
    ready0.allowed_formats = 1 << int(trt_memory_format)
    out0.allowed_formats = 1 << int(trt_memory_format)
    shape_out0.get_output(0).allowed_formats = 1 << int(trt_memory_format)
    resized_out0.allowed_formats = 1 << int(trt_memory_format)

    if (trt_dtype == trt.int8):
        dummy_in0.dynamic_range = (-128.0, 127.0)
        resized_out0.dynamic_range = (-128.0, 127.0)
        start0.dynamic_range = (-128.0, 127.0)
        end0.dynamic_range = (-128.0, 127.0)
        ready0.dynamic_range = (-128.0, 127.0)

    flags = 1 << int(trt.BuilderFlag.STRICT_TYPES)
    if (trt_dtype == trt.int8):
        flags |= 1 << int(trt.BuilderFlag.INT8)
    elif (trt_dtype == trt.float16):
        flags |= 1 << int(trt.BuilderFlag.FP16)

    min_prefix = []
    opt_prefix = []
    max_prefix = []

    if max_batch != 0:
        min_prefix = [1]
        opt_prefix = [max(1, max_batch)]
        max_prefix = [max(1, max_batch)]

    min_shape = min_prefix + [1] * len(shape)
    opt_shape = opt_prefix + [8] * len(shape)
    max_shape = max_prefix + [32] * len(shape)

    profile = builder.create_optimization_profile()
    profile.set_shape("INPUT", min_shape, opt_shape, max_shape)
    profile.set_shape_input("SHAPE_INPUT", min_shape, opt_shape, max_shape)
    profile.set_shape("DUMMY_INPUT", min_shape, opt_shape, max_shape)
    profile.set_shape("START", min_prefix + unit_shape, opt_prefix + unit_shape,
                      max_prefix + unit_shape)
    profile.set_shape("END", min_prefix + unit_shape, opt_prefix + unit_shape,
                      max_prefix + unit_shape)
    profile.set_shape("READY", min_prefix + unit_shape, opt_prefix + unit_shape,
                      max_prefix + unit_shape)
    profile.set_shape("CORRID", min_prefix + unit_shape,
                      opt_prefix + unit_shape, max_prefix + unit_shape)

    config = builder.create_builder_config()
    config.flags = flags
    config.add_optimization_profile(profile)
    config.max_workspace_size = 1 << 20
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    model_name = tu.get_dyna_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_fixed_modelfile(models_dir, model_version, max_batch, dtype,
                                shape):
    trt_dtype = np_to_trt_dtype(dtype)
    # Create the model. For now don't implement a proper accumulator
    # just return 0 if not-ready and 'INPUT'+'START'+('END'*'CORRID')
    # otherwise...  the tests know to expect this.
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network()
    in0 = network.add_input("INPUT", trt_dtype, shape)
    start0 = network.add_input("START", trt_dtype, [1 for i in shape])
    end0 = network.add_input("END", trt_dtype, [1 for i in shape])
    ready0 = network.add_input("READY", trt_dtype, [1 for i in shape])
    corrid0 = network.add_input("CORRID", trt.int32, [1 for i in shape])
    add0 = network.add_elementwise(in0, start0, trt.ElementWiseOperation.SUM)
    mul0 = network.add_elementwise(end0, corrid0, trt.ElementWiseOperation.PROD)
    sum0 = network.add_elementwise(add0.get_output(0), mul0.get_output(0),
                                   trt.ElementWiseOperation.SUM)
    out0 = network.add_elementwise(sum0.get_output(0), ready0,
                                   trt.ElementWiseOperation.PROD)

    out0.get_output(0).name = "OUTPUT"
    network.mark_output(out0.get_output(0))

    config = builder.create_builder_config()
    config.max_workspace_size = 1 << 20
    builder.max_batch_size = max(1, max_batch)
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine
    del network

    model_name = tu.get_dyna_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_fixed_rf_modelfile(models_dir, model_version, max_batch, dtype,
                                   shape):
    trt_dtype = np_to_trt_dtype(dtype)
    trt_memory_format = trt.TensorFormat.LINEAR
    # Create the model. For now don't implement a proper accumulator
    # just return 0 if not-ready and 'INPUT'+'START'*('END'*'CORRID')
    # otherwise...  the tests know to expect this.
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network()
    in0 = network.add_input("INPUT", trt_dtype, shape)
    start0 = network.add_input("START", trt_dtype, [1 for i in shape])
    end0 = network.add_input("END", trt_dtype, [1 for i in shape])
    ready0 = network.add_input("READY", trt_dtype, [1 for i in shape])
    corrid0 = network.add_input("CORRID", trt.int32, [1 for i in shape])
    add0 = network.add_elementwise(in0, start0, trt.ElementWiseOperation.SUM)
    mul0 = network.add_elementwise(end0, corrid0, trt.ElementWiseOperation.PROD)
    sum0 = network.add_elementwise(add0.get_output(0), mul0.get_output(0),
                                   trt.ElementWiseOperation.SUM)
    out0 = network.add_elementwise(sum0.get_output(0), ready0,
                                   trt.ElementWiseOperation.PROD)

    out0.get_output(0).name = "OUTPUT"
    network.mark_output(out0.get_output(0))

    out0.get_output(0).dtype = trt_dtype

    in0.allowed_formats = 1 << int(trt_memory_format)
    start0.allowed_formats = 1 << int(trt_memory_format)
    end0.allowed_formats = 1 << int(trt_memory_format)
    ready0.allowed_formats = 1 << int(trt_memory_format)
    corrid0.allowed_formats = 1 << int(trt_memory_format)
    out0.get_output(0).allowed_formats = 1 << int(trt_memory_format)

    if (trt_dtype == trt.int8):
        in0.dynamic_range = (-128.0, 127.0)
        out0.dynamic_range = (-128.0, 127.0)
        start0.dynamic_range = (-128.0, 127.0)
        end0.dynamic_range = (-128.0, 127.0)
        ready0.dynamic_range = (-128.0, 127.0)
        corrid0.dynamic_range = (-128.0, 127.0)

    flags = 1 << int(trt.BuilderFlag.STRICT_TYPES)
    if (trt_dtype == trt.int8):
        flags |= 1 << int(trt.BuilderFlag.INT8)
    elif (trt_dtype == trt.float16):
        flags |= 1 << int(trt.BuilderFlag.FP16)

    config = builder.create_builder_config()
    config.flags = flags
    config.max_workspace_size = 1 << 20
    builder.max_batch_size = max(1, max_batch)
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    model_name = tu.get_dyna_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_dynamic_modelfile(models_dir, model_version, max_batch, dtype,
                                  shape):
    trt_dtype = np_to_trt_dtype(dtype)
    # Create the model. For now don't implement a proper accumulator
    # just return 0 if not-ready and 'INPUT'+'START'*('END'*'CORRID')
    # otherwise...  the tests know to expect this.
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

    unit_shape = ([1] * len(shape))
    if max_batch != 0:
        in0 = network.add_input("INPUT", trt_dtype, [-1] + shape)
        start0 = network.add_input("START", trt_dtype, [-1] + unit_shape)
        end0 = network.add_input("END", trt_dtype, [-1] + unit_shape)
        ready0 = network.add_input("READY", trt_dtype, [-1] + unit_shape)
        corrid0 = network.add_input("CORRID", trt.int32, [-1] + unit_shape)
    else:
        in0 = network.add_input("INPUT", trt_dtype, shape)
        start0 = network.add_input("START", trt_dtype, unit_shape)
        end0 = network.add_input("END", trt_dtype, unit_shape)
        ready0 = network.add_input("READY", trt_dtype, unit_shape)
        corrid0 = network.add_input("CORRID", trt.int32, unit_shape)

    add0 = network.add_elementwise(in0, start0, trt.ElementWiseOperation.SUM)
    mul0 = network.add_elementwise(end0, corrid0, trt.ElementWiseOperation.PROD)
    sum0 = network.add_elementwise(add0.get_output(0), mul0.get_output(0),
                                   trt.ElementWiseOperation.SUM)
    out0 = network.add_elementwise(sum0.get_output(0), ready0,
                                   trt.ElementWiseOperation.PROD)

    out0.get_output(0).name = "OUTPUT"
    network.mark_output(out0.get_output(0))

    min_shape = []
    opt_shape = []
    max_shape = []
    if max_batch != 0:
        min_shape = min_shape + [1]
        opt_shape = opt_shape + [max(1, max_batch)]
        max_shape = max_shape + [max(1, max_batch)]
    for i in shape:
        if i == -1:
            min_shape = min_shape + [1]
            opt_shape = opt_shape + [8]
            max_shape = max_shape + [32]
        else:
            min_shape = min_shape + [i]
            opt_shape = opt_shape + [i]
            max_shape = max_shape + [i]

    profile = builder.create_optimization_profile()
    profile.set_shape("INPUT", min_shape, opt_shape, max_shape)
    if max_batch != 0:
        profile.set_shape("START", [1] + unit_shape, [max_batch] + unit_shape,
                          [max_batch] + unit_shape)
        profile.set_shape("END", [1] + unit_shape, [max_batch] + unit_shape,
                          [max_batch] + unit_shape)
        profile.set_shape("READY", [1] + unit_shape, [max_batch] + unit_shape,
                          [max_batch] + unit_shape)
        profile.set_shape("CORRID", [1] + unit_shape, [max_batch] + unit_shape,
                          [max_batch] + unit_shape)
    else:
        profile.set_shape("START", unit_shape, unit_shape, unit_shape)
        profile.set_shape("END", unit_shape, unit_shape, unit_shape)
        profile.set_shape("READY", unit_shape, unit_shape, unit_shape)
        profile.set_shape("CORRID", unit_shape, unit_shape, unit_shape)
    config = builder.create_builder_config()
    config.add_optimization_profile(profile)

    config.max_workspace_size = 1 << 20
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    model_name = tu.get_dyna_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_dynamic_rf_modelfile(models_dir, model_version, max_batch,
                                     dtype, shape):
    trt_dtype = np_to_trt_dtype(dtype)
    trt_memory_format = trt.TensorFormat.LINEAR

    # Create the model. For now don't implement a proper accumulator
    # just return 0 if not-ready and 'INPUT'+'START'*('END'*'CORRID')
    # otherwise...  the tests know to expect this.
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

    unit_shape = ([1] * len(shape))
    if max_batch != 0:
        in0 = network.add_input("INPUT", trt_dtype, [-1] + shape)
        start0 = network.add_input("START", trt_dtype, [-1] + unit_shape)
        end0 = network.add_input("END", trt_dtype, [-1] + unit_shape)
        ready0 = network.add_input("READY", trt_dtype, [-1] + unit_shape)
        corrid0 = network.add_input("CORRID", trt.int32, [-1] + unit_shape)
    else:
        in0 = network.add_input("INPUT", trt_dtype, shape)
        start0 = network.add_input("START", trt_dtype, unit_shape)
        end0 = network.add_input("END", trt_dtype, unit_shape)
        ready0 = network.add_input("READY", trt_dtype, unit_shape)
        corrid0 = network.add_input("CORRID", trt.int32, unit_shape)

    add0 = network.add_elementwise(in0, start0, trt.ElementWiseOperation.SUM)
    mul0 = network.add_elementwise(end0, corrid0, trt.ElementWiseOperation.PROD)
    sum0 = network.add_elementwise(add0.get_output(0), mul0.get_output(0),
                                   trt.ElementWiseOperation.SUM)
    out0 = network.add_elementwise(sum0.get_output(0), ready0,
                                   trt.ElementWiseOperation.PROD)

    out0.get_output(0).name = "OUTPUT"
    network.mark_output(out0.get_output(0))

    out0.get_output(0).dtype = trt_dtype

    in0.allowed_formats = 1 << int(trt_memory_format)
    start0.allowed_formats = 1 << int(trt_memory_format)
    ready0.allowed_formats = 1 << int(trt_memory_format)
    out0.get_output(0).allowed_formats = 1 << int(trt_memory_format)

    if (trt_dtype == trt.int8):
        in0.dynamic_range = (-128.0, 127.0)
        out0.dynamic_range = (-128.0, 127.0)
        start0.dynamic_range = (-128.0, 127.0)
        end0.dynamic_range = (-128.0, 127.0)
        ready0.dynamic_range = (-128.0, 127.0)
        corrid0.dynamic_range = (-128.0, 127.0)

    flags = 1 << int(trt.BuilderFlag.STRICT_TYPES)
    if (trt_dtype == trt.int8):
        flags |= 1 << int(trt.BuilderFlag.INT8)
    elif (trt_dtype == trt.float16):
        flags |= 1 << int(trt.BuilderFlag.FP16)

    min_shape = []
    opt_shape = []
    max_shape = []
    if max_batch != 0:
        min_shape = min_shape + [1]
        opt_shape = opt_shape + [max(1, max_batch)]
        max_shape = max_shape + [max(1, max_batch)]
    for i in shape:
        if i == -1:
            min_shape = min_shape + [1]
            opt_shape = opt_shape + [8]
            max_shape = max_shape + [32]
        else:
            min_shape = min_shape + [i]
            opt_shape = opt_shape + [i]
            max_shape = max_shape + [i]

    profile = builder.create_optimization_profile()
    profile.set_shape("INPUT", min_shape, opt_shape, max_shape)
    if max_batch != 0:
        profile.set_shape("START", [1] + unit_shape, [max_batch] + unit_shape,
                          [max_batch] + unit_shape)
        profile.set_shape("END", [1] + unit_shape, [max_batch] + unit_shape,
                          [max_batch] + unit_shape)
        profile.set_shape("READY", [1] + unit_shape, [max_batch] + unit_shape,
                          [max_batch] + unit_shape)
        profile.set_shape("CORRID", [1] + unit_shape, [max_batch] + unit_shape,
                          [max_batch] + unit_shape)
    else:
        profile.set_shape("START", unit_shape, unit_shape, unit_shape)
        profile.set_shape("END", unit_shape, unit_shape, unit_shape)
        profile.set_shape("READY", unit_shape, unit_shape, unit_shape)
        profile.set_shape("CORRID", unit_shape, unit_shape, unit_shape)

    config = builder.create_builder_config()
    config.flags = flags
    config.add_optimization_profile(profile)
    config.max_workspace_size = 1 << 20
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    model_name = tu.get_dyna_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_modelfile(models_dir, model_version, max_batch, dtype, shape):

    if not tu.validate_for_trt_model(dtype, dtype, dtype, shape, shape, shape):
        return

    if dtype != np.float32:
        if (not tu.shape_is_fixed(shape)):
            create_plan_dynamic_rf_modelfile(models_dir, model_version,
                                             max_batch, dtype, shape)
        else:
            create_plan_fixed_rf_modelfile(models_dir, model_version, max_batch,
                                           dtype, shape)
    else:
        if (not tu.shape_is_fixed(shape)):
            create_plan_dynamic_modelfile(models_dir, model_version, max_batch,
                                          dtype, shape)
        else:
            create_plan_fixed_modelfile(models_dir, model_version, max_batch,
                                        dtype, shape)


def create_plan_modelconfig(models_dir, model_version, max_batch, dtype, shape):

    if not tu.validate_for_trt_model(dtype, dtype, dtype, shape, shape, shape):
        return

    model_name = tu.get_dyna_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    config_dir = models_dir + "/" + model_name

    if FLAGS.tensorrt_shape_io:
        shape_tensor_dim = len(shape)
        config = '''
name: "{}"
platform: "tensorrt_plan"
max_batch_size: {}
sequence_batching {{
  max_sequence_idle_microseconds: 5000000
  {}
  control_input [
    {{
      name: "START"
      control [
        {{
          kind: CONTROL_SEQUENCE_START
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "END"
      control [
        {{
          kind: CONTROL_SEQUENCE_END
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "READY"
      control [
        {{
          kind: CONTROL_SEQUENCE_READY
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "CORRID"
      control [
        {{
          kind: CONTROL_SEQUENCE_CORRID
          data_type: TYPE_INT32
        }}
      ]
    }}
  ]
}}
input [
  {{
    name: "INPUT"
    data_type: TYPE_INT32
    dims: [ {} ]
  }}
]
input [
  {{
    name: "DUMMY_INPUT"
    data_type: {}
    dims: [ {} ]
  }}
]
input [
  {{
    name: "SHAPE_INPUT"
    data_type: TYPE_INT32
    dims: [ {} ]
    is_shape_tensor: true
  }}
]
output [
  {{
    name: "OUTPUT"
    data_type: TYPE_INT32
    dims: [ {} ]
  }}
]
output [
  {{
    name: "RESIZED_OUTPUT"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "SHAPE_OUTPUT"
    data_type: TYPE_INT32
    dims: [ {} ]
    is_shape_tensor: true
  }}
]
instance_group [
  {{
    kind: KIND_GPU
  }}
]
'''.format(
            model_name, max_batch,
            "oldest { max_candidate_sequences: 6\npreferred_batch_size: [ 4 ]\nmax_queue_delay_microseconds: 0\n}"
            if max_batch > 0 else "", "int32", "int32", "int32",
            tu.shape_to_dims_str(shape), np_to_model_dtype(dtype),
            tu.shape_to_dims_str(shape), shape_tensor_dim,
            tu.shape_to_dims_str(shape), np_to_model_dtype(dtype),
            tu.shape_to_dims_str(shape), shape_tensor_dim)

    else:
        config = '''
name: "{}"
platform: "tensorrt_plan"
max_batch_size: {}
sequence_batching {{
  max_sequence_idle_microseconds: 5000000
  {}
  control_input [
    {{
      name: "START"
      control [
        {{
          kind: CONTROL_SEQUENCE_START
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "END"
      control [
        {{
          kind: CONTROL_SEQUENCE_END
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "READY"
      control [
        {{
          kind: CONTROL_SEQUENCE_READY
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "CORRID"
      control [
        {{
          kind: CONTROL_SEQUENCE_CORRID
          data_type: TYPE_INT32
        }}
      ]
    }}
  ]
}}
input [
  {{
    name: "INPUT"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT"
    data_type: {}
    dims: [ {} ]
  }}
]
instance_group [
  {{
    kind: KIND_GPU
  }}
]
'''.format(
            model_name, max_batch,
            "oldest { max_candidate_sequences: 6\npreferred_batch_size: [ 4 ]\nmax_queue_delay_microseconds: 0\n}"
            if max_batch > 0 else "", "int32" if dtype == np.int32 else "fp32",
            "int32" if dtype == np.int32 else "fp32",
            "int32" if dtype == np.int32 else "fp32", np_to_model_dtype(dtype),
            tu.shape_to_dims_str(shape), np_to_model_dtype(dtype),
            tu.shape_to_dims_str(shape))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_onnx_modelfile(models_dir, model_version, max_batch, dtype, shape):

    if not tu.validate_for_onnx_model(dtype, dtype, dtype, shape, shape, shape):
        return

    model_name = tu.get_dyna_sequence_model_name(
        "onnx_nobatch" if max_batch == 0 else "onnx", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    # Create the model. For now don't implement a proper accumulator
    # just return 0 if not-ready and 'INPUT'+'START'*('END'*'CORRID')
    # otherwise...  the tests know to expect this.
    onnx_dtype = np_to_onnx_dtype(dtype)
    onnx_input_shape, idx = tu.shape_to_onnx_shape(shape, 0)
    onnx_output_shape, idx = tu.shape_to_onnx_shape(shape, idx)

    # If the input is a string then use int32 for operation and just
    # cast to/from string for input and output.
    onnx_control_dtype = onnx_dtype
    if onnx_dtype == onnx.TensorProto.STRING:
        onnx_control_dtype = onnx.TensorProto.INT32

    batch_dim = [] if max_batch == 0 else [None]

    onnx_input = onnx.helper.make_tensor_value_info(
        "INPUT", onnx_dtype, batch_dim + onnx_input_shape)
    onnx_start = onnx.helper.make_tensor_value_info("START", onnx_control_dtype,
                                                    batch_dim + [1])
    onnx_end = onnx.helper.make_tensor_value_info("END", onnx_control_dtype,
                                                  batch_dim + [1])
    onnx_ready = onnx.helper.make_tensor_value_info("READY", onnx_control_dtype,
                                                    batch_dim + [1])
    onnx_corrid = onnx.helper.make_tensor_value_info("CORRID",
                                                     onnx.TensorProto.UINT64,
                                                     batch_dim + [1])
    onnx_output = onnx.helper.make_tensor_value_info(
        "OUTPUT", onnx_dtype, batch_dim + onnx_output_shape)

    internal_input = onnx.helper.make_node("Identity", ["INPUT"], ["_INPUT"])

    # cast int8, int16 input to higer precision int as Onnx Add/Sub operator doesn't support those type
    # Also casting String data type to int32
    if ((onnx_dtype == onnx.TensorProto.INT8) or
        (onnx_dtype == onnx.TensorProto.INT16) or
        (onnx_dtype == onnx.TensorProto.STRING)):
        internal_input = onnx.helper.make_node("Cast", ["INPUT"], ["_INPUT"],
                                               to=onnx.TensorProto.INT32)

    onnx_corrid_cast0 = onnx.helper.make_node("Cast", ["CORRID"],
                                              ["onnx_corrid_cast0"],
                                              to=onnx_control_dtype)
    add0 = onnx.helper.make_node("Add", ["_INPUT", "START"], ["add0"])
    mul0 = onnx.helper.make_node("Mul", ["END", "onnx_corrid_cast0"], ["mul0"])
    sum0 = onnx.helper.make_node("Add", ["add0", "mul0"], ["sum0"])
    res0 = onnx.helper.make_node("Mul", ["READY", "sum0"], ["CAST"])
    cast = onnx.helper.make_node("Cast", ["CAST"], ["OUTPUT"], to=onnx_dtype)

    # Avoid cast from float16 to float16
    # (bug in Onnx Runtime, cast from float16 to float16 will become cast from float16 to float32)
    if onnx_dtype == onnx.TensorProto.FLOAT16:
        cast = onnx.helper.make_node("Identity", ["CAST"], ["OUTPUT"])

    onnx_nodes = [
        internal_input, onnx_corrid_cast0, add0, mul0, sum0, res0, cast
    ]
    onnx_inputs = [onnx_input, onnx_start, onnx_end, onnx_ready, onnx_corrid]
    onnx_outputs = [onnx_output]

    graph_proto = onnx.helper.make_graph(onnx_nodes, model_name, onnx_inputs,
                                         onnx_outputs)
    if FLAGS.onnx_opset > 0:
        model_opset = onnx.helper.make_operatorsetid("", FLAGS.onnx_opset)
        model_def = onnx.helper.make_model(graph_proto,
                                           producer_name="triton",
                                           opset_imports=[model_opset])
    else:
        model_def = onnx.helper.make_model(graph_proto, producer_name="triton")

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    onnx.save(model_def, model_version_dir + "/model.onnx")


def create_onnx_modelconfig(models_dir, model_version, max_batch, dtype, shape):

    if not tu.validate_for_onnx_model(dtype, dtype, dtype, shape, shape, shape):
        return

    model_name = tu.get_dyna_sequence_model_name(
        "onnx_nobatch" if max_batch == 0 else "onnx", dtype)
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "onnxruntime_onnx"
max_batch_size: {}
sequence_batching {{
  max_sequence_idle_microseconds: 5000000
  {}
  control_input [
    {{
      name: "START"
      control [
        {{
          kind: CONTROL_SEQUENCE_START
          {type}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "END"
      control [
        {{
          kind: CONTROL_SEQUENCE_END
          {type}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "READY"
      control [
        {{
          kind: CONTROL_SEQUENCE_READY
          {type}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "CORRID"
      control [
        {{
          kind: CONTROL_SEQUENCE_CORRID
          data_type: TYPE_UINT64
        }}
      ]
    }}
  ]
}}
input [
  {{
    name: "INPUT"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT"
    data_type: {}
    dims: [ {} ]
  }}
]
instance_group [
  {{
    kind: KIND_CPU
  }}
]
'''.format(
        model_name,
        max_batch,
        "oldest { max_candidate_sequences: 6\npreferred_batch_size: [ 4 ]\nmax_queue_delay_microseconds: 0\n}"
        if max_batch > 0 else "",
        np_to_model_dtype(dtype),
        tu.shape_to_dims_str(shape),
        np_to_model_dtype(dtype),
        tu.shape_to_dims_str(shape),
        type="fp32" if dtype == np.float32 else "int32")

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_libtorch_modelfile(models_dir, model_version, max_batch, dtype,
                              shape):

    if not tu.validate_for_libtorch_model(dtype, dtype, dtype, shape, shape,
                                          shape):
        return

    torch_dtype = np_to_torch_dtype(dtype)

    model_name = tu.get_dyna_sequence_model_name(
        "libtorch_nobatch" if max_batch == 0 else "libtorch", dtype)
    # handle for -1 (when variable) since can't create tensor with shape of [-1]
    shape = [abs(ips) for ips in shape]

    class SequenceNet(nn.Module):

        def __init__(self):
            super(SequenceNet, self).__init__()

        def forward(self, input0, start0, end0, ready0, corrid0):
            tmp = input0 + start0 + (end0 * corrid0)
            return tmp * ready0

    sequenceModel = SequenceNet()
    example_input = torch.zeros(shape, dtype=torch_dtype)
    example_corrid_input = torch.zeros(shape, dtype=torch.long)
    traced = torch.jit.trace(sequenceModel,
                             (example_input, example_input, example_input,
                              example_input, example_corrid_input))

    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    traced.save(model_version_dir + "/model.pt")


def create_libtorch_modelconfig(models_dir, model_version, max_batch, dtype,
                                shape):

    if not tu.validate_for_libtorch_model(dtype, dtype, dtype, shape, shape,
                                          shape):
        return

    model_name = tu.get_dyna_sequence_model_name(
        "libtorch_nobatch" if max_batch == 0 else "libtorch", dtype)
    config_dir = models_dir + "/" + model_name
    #  FIX FOR LibTorch
    config = '''
name: "{}"
platform: "pytorch_libtorch"
max_batch_size: {}
sequence_batching {{
  max_sequence_idle_microseconds: 5000000
  {}
  control_input [
    {{
      name: "START__1"
      control [
        {{
          kind: CONTROL_SEQUENCE_START
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "END__2"
      control [
        {{
          kind: CONTROL_SEQUENCE_END
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "READY__3"
      control [
        {{
          kind: CONTROL_SEQUENCE_READY
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "CORRID__4"
      control [
        {{
          kind: CONTROL_SEQUENCE_CORRID
          data_type: TYPE_INT32
        }}
      ]
    }}
  ]
}}
input [
  {{
    name: "INPUT__0"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT__0"
    data_type: {}
    dims: [ 1 ]
  }}
]
instance_group [
  {{
    kind: KIND_CPU
  }}
]
'''.format(
        model_name, max_batch,
        "oldest { max_candidate_sequences: 6\npreferred_batch_size: [ 4 ]\nmax_queue_delay_microseconds: 0\n}"
        if max_batch > 0 else "", "int32" if dtype == np.int32 else "fp32",
        "int32" if dtype == np.int32 else "fp32",
        "int32" if dtype == np.int32 else "fp32", np_to_model_dtype(dtype),
        tu.shape_to_dims_str(shape), np_to_model_dtype(dtype))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_openvino_modelfile(models_dir, model_version, max_batch, dtype,
                              shape):

    batch_dim = [] if max_batch == 0 else [
        max_batch,
    ]
    if not tu.validate_for_openvino_model(dtype, dtype, dtype,
                                          batch_dim + shape, batch_dim + shape,
                                          batch_dim + shape):
        return

    model_name = tu.get_dyna_sequence_model_name(
        "openvino_nobatch" if max_batch == 0 else "openvino", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    in0 = ng.parameter(shape=batch_dim + shape, dtype=dtype, name="INPUT")
    start = ng.parameter(shape=batch_dim + shape, dtype=dtype, name="START")
    end = ng.parameter(shape=batch_dim + shape, dtype=dtype, name="END")
    ready = ng.parameter(shape=batch_dim + shape, dtype=dtype, name="READY")
    corrid = ng.parameter(shape=batch_dim + shape, dtype=dtype, name="CORRID")

    tmp1 = ng.add(in0, start)
    tmp2 = ng.multiply(end, corrid)
    tmp = ng.add(tmp1, tmp2)
    op0 = ng.multiply(tmp, ready, name="OUTPUT")

    function = ng.impl.Function([op0], [in0, start, end, ready, corrid],
                                model_name)
    ie_network = IENetwork(ng.impl.Function.to_capsule(function))

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    ie_network.serialize(model_version_dir + "/model.xml",
                         model_version_dir + "/model.bin")


def create_openvino_modelconfig(models_dir, model_version, max_batch, dtype,
                                shape):

    batch_dim = [] if max_batch == 0 else [
        max_batch,
    ]
    if not tu.validate_for_openvino_model(dtype, dtype, dtype,
                                          batch_dim + shape, batch_dim + shape,
                                          batch_dim + shape):
        return

    model_name = tu.get_dyna_sequence_model_name(
        "openvino_nobatch" if max_batch == 0 else "openvino", dtype)
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
backend: "openvino"
max_batch_size: {}
sequence_batching {{
  max_sequence_idle_microseconds: 5000000
  {}
  control_input [
    {{
      name: "START"
      control [
        {{
          kind: CONTROL_SEQUENCE_START
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "END"
      control [
        {{
          kind: CONTROL_SEQUENCE_END
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "READY"
      control [
        {{
          kind: CONTROL_SEQUENCE_READY
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "CORRID"
      control [
        {{
          kind: CONTROL_SEQUENCE_CORRID
          data_type: TYPE_INT32
        }}
      ]
    }}
  ]
}}
input [
  {{
    name: "INPUT"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT"
    data_type: {}
    dims: [ 1 ]
  }}
]
'''.format(
        model_name, max_batch,
        "oldest { max_candidate_sequences: 6\npreferred_batch_size: [ 4 ]\nmax_queue_delay_microseconds: 0\n}"
        if max_batch > 0 else "", "int32" if dtype == np.int32 else "fp32",
        "int32" if dtype == np.int32 else "fp32",
        "int32" if dtype == np.int32 else "fp32", np_to_model_dtype(dtype),
        tu.shape_to_dims_str(shape), np_to_model_dtype(dtype))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_shape_tensor_models(models_dir, dtype, shape, no_batch=True):
    model_version = 1

    create_plan_modelconfig(models_dir, model_version, 8, dtype, shape)
    create_plan_shape_tensor_modelfile(models_dir, model_version, 8, dtype,
                                       shape)
    if no_batch:
        create_plan_modelconfig(models_dir, model_version, 0, dtype, shape)
        create_plan_shape_tensor_modelfile(models_dir, model_version, 0, dtype,
                                           shape)


def create_models(models_dir, dtype, shape, no_batch=True):
    model_version = 1

    if FLAGS.graphdef:
        create_tf_modelconfig(False, models_dir, model_version, 8, dtype, shape)
        create_tf_modelfile(False, models_dir, model_version, 8, dtype, shape)
        if no_batch:
            create_tf_modelconfig(False, models_dir, model_version, 0, dtype,
                                  shape)
            create_tf_modelfile(False, models_dir, model_version, 0, dtype,
                                shape)

    if FLAGS.savedmodel:
        create_tf_modelconfig(True, models_dir, model_version, 8, dtype, shape)
        create_tf_modelfile(True, models_dir, model_version, 8, dtype, shape)
        if no_batch:
            create_tf_modelconfig(True, models_dir, model_version, 0, dtype,
                                  shape)
            create_tf_modelfile(True, models_dir, model_version, 0, dtype,
                                shape)

    if FLAGS.tensorrt:
        suffix = []
        if dtype == np.int8:
            suffix = [1, 1]

        create_plan_modelconfig(models_dir, model_version, 8, dtype,
                                shape + suffix)
        create_plan_modelfile(models_dir, model_version, 8, dtype,
                              shape + suffix)
        if no_batch:
            create_plan_modelconfig(models_dir, model_version, 0, dtype,
                                    shape + suffix)
            create_plan_modelfile(models_dir, model_version, 0, dtype,
                                  shape + suffix)

    if FLAGS.onnx:
        create_onnx_modelconfig(models_dir, model_version, 8, dtype, shape)
        create_onnx_modelfile(models_dir, model_version, 8, dtype, shape)
        if no_batch:
            create_onnx_modelconfig(models_dir, model_version, 0, dtype, shape)
            create_onnx_modelfile(models_dir, model_version, 0, dtype, shape)

    if FLAGS.libtorch:
        create_libtorch_modelconfig(models_dir, model_version, 8, dtype, shape)
        create_libtorch_modelfile(models_dir, model_version, 8, dtype, shape)
        if no_batch:
            create_libtorch_modelconfig(models_dir, model_version, 0, dtype,
                                        shape)
            create_libtorch_modelfile(models_dir, model_version, 0, dtype,
                                      shape)

    if FLAGS.openvino:
        create_openvino_modelconfig(models_dir, model_version, 8, dtype, shape)
        create_openvino_modelfile(models_dir, model_version, 8, dtype, shape)
        if no_batch:
            create_openvino_modelconfig(models_dir, model_version, 0, dtype,
                                        shape)
            create_openvino_modelfile(models_dir, model_version, 0, dtype,
                                      shape)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--models_dir',
                        type=str,
                        required=True,
                        help='Top-level model directory')
    parser.add_argument('--graphdef',
                        required=False,
                        action='store_true',
                        help='Generate GraphDef models')
    parser.add_argument('--savedmodel',
                        required=False,
                        action='store_true',
                        help='Generate SavedModel models')
    parser.add_argument('--tensorrt',
                        required=False,
                        action='store_true',
                        help='Generate TensorRT PLAN models')
    parser.add_argument(
        '--tensorrt-shape-io',
        required=False,
        action='store_true',
        help='Generate TensorRT PLAN models w/ shape tensor i/o')
    parser.add_argument('--onnx',
                        required=False,
                        action='store_true',
                        help='Generate Onnx models')
    parser.add_argument(
        '--onnx_opset',
        type=int,
        required=False,
        default=0,
        help='Opset used for Onnx models. Default is to use ONNXRT default')
    parser.add_argument('--libtorch',
                        required=False,
                        action='store_true',
                        help='Generate Pytorch LibTorch models')
    parser.add_argument('--openvino',
                        required=False,
                        action='store_true',
                        help='Generate OpenVino models')
    parser.add_argument('--variable',
                        required=False,
                        action='store_true',
                        help='Used variable-shape tensors for input/output')
    FLAGS, unparsed = parser.parse_known_args()

    if FLAGS.graphdef or FLAGS.savedmodel:
        import tensorflow as tf
        from tensorflow.python.framework import graph_io
    if FLAGS.tensorrt or FLAGS.tensorrt_shape_io:
        import tensorrt as trt
    if FLAGS.onnx:
        import onnx
    if FLAGS.libtorch:
        import torch
        from torch import nn
    if FLAGS.openvino:
        from openvino.inference_engine import IENetwork
        import ngraph as ng

    import test_util as tu

    if FLAGS.tensorrt_shape_io:
        create_shape_tensor_models(FLAGS.models_dir, np.float32, [
            -1,
        ])
    else:
        # Tests with models that accept fixed-shape input/output tensors
        if not FLAGS.variable:
            create_models(FLAGS.models_dir, np.int32, [
                1,
            ])

        # Tests with models that accept variable-shape input/output tensors
        if FLAGS.variable:
            create_models(FLAGS.models_dir, np.int32, [
                -1,
            ], False)


# qa/common/gen_qa_identity_models.py

# Copyright 2019-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
from builtins import range
import os
import numpy as np
import gen_ensemble_model_utils as emu

FLAGS = None
np_dtype_string = np.dtype(object)
from typing import List, Tuple


def np_to_model_dtype(np_dtype):
    if np_dtype == bool:
        return "TYPE_BOOL"
    elif np_dtype == np.int8:
        return "TYPE_INT8"
    elif np_dtype == np.int16:
        return "TYPE_INT16"
    elif np_dtype == np.int32:
        return "TYPE_INT32"
    elif np_dtype == np.int64:
        return "TYPE_INT64"
    elif np_dtype == np.uint8:
        return "TYPE_UINT8"
    elif np_dtype == np.uint16:
        return "TYPE_UINT16"
    elif np_dtype == np.float16:
        return "TYPE_FP16"
    elif np_dtype == np.float32:
        return "TYPE_FP32"
    elif np_dtype == np.float64:
        return "TYPE_FP64"
    elif np_dtype == np_dtype_string:
        return "TYPE_STRING"
    return None


def np_to_tf_dtype(np_dtype):
    if np_dtype == bool:
        return tf.bool
    elif np_dtype == np.int8:
        return tf.int8
    elif np_dtype == np.int16:
        return tf.int16
    elif np_dtype == np.int32:
        return tf.int32
    elif np_dtype == np.int64:
        return tf.int64
    elif np_dtype == np.uint8:
        return tf.uint8
    elif np_dtype == np.uint16:
        return tf.uint16
    elif np_dtype == np.float16:
        return tf.float16
    elif np_dtype == np.float32:
        return tf.float32
    elif np_dtype == np.float64:
        return tf.float64
    elif np_dtype == np_dtype_string:
        return tf.string
    return None


def np_to_trt_dtype(np_dtype):
    if np_dtype == bool:
        return trt.bool
    elif np_dtype == np.int8:
        return trt.int8
    elif np_dtype == np.int32:
        return trt.int32
    elif np_dtype == np.float16:
        return trt.float16
    elif np_dtype == np.float32:
        return trt.float32
    return None


def np_to_onnx_dtype(np_dtype):
    if np_dtype == bool:
        return onnx.TensorProto.BOOL
    elif np_dtype == np.int8:
        return onnx.TensorProto.INT8
    elif np_dtype == np.int16:
        return onnx.TensorProto.INT16
    elif np_dtype == np.int32:
        return onnx.TensorProto.INT32
    elif np_dtype == np.int64:
        return onnx.TensorProto.INT64
    elif np_dtype == np.uint8:
        return onnx.TensorProto.UINT8
    elif np_dtype == np.uint16:
        return onnx.TensorProto.UINT16
    elif np_dtype == np.float16:
        return onnx.TensorProto.FLOAT16
    elif np_dtype == np.float32:
        return onnx.TensorProto.FLOAT
    elif np_dtype == np.float64:
        return onnx.TensorProto.DOUBLE
    elif np_dtype == np_dtype_string:
        return onnx.TensorProto.STRING
    return None


def np_to_torch_dtype(np_dtype):
    if np_dtype == bool:
        return torch.bool
    elif np_dtype == np.int8:
        return torch.int8
    elif np_dtype == np.int16:
        return torch.int16
    elif np_dtype == np.int32:
        return torch.int
    elif np_dtype == np.int64:
        return torch.long
    elif np_dtype == np.uint8:
        return torch.uint8
    elif np_dtype == np.uint16:
        return None  # Not supported in Torch
    elif np_dtype == np.float16:
        return None
    elif np_dtype == np.float32:
        return torch.float
    elif np_dtype == np.float64:
        return torch.double
    elif np_dtype == np_dtype_string:
        return List[str]


def create_tf_modelfile(create_savedmodel, models_dir, model_version, io_cnt,
                        max_batch, dtype, shape):

    if not tu.validate_for_tf_model(dtype, dtype, dtype, shape, shape, shape):
        return

    tf_dtype = np_to_tf_dtype(dtype)

    # Create the model that copies inputs to corresponding outputs.
    tf.reset_default_graph()
    for io_num in range(io_cnt):
        input_name = "INPUT{}".format(io_num)
        output_name = "OUTPUT{}".format(io_num)
        if max_batch == 0:
            tin = tf.placeholder(tf_dtype, tu.shape_to_tf_shape(shape),
                                 input_name)
        else:
            tin = tf.placeholder(tf_dtype, [
                None,
            ] + tu.shape_to_tf_shape(shape), input_name)
        toutput = tf.identity(tin, name=output_name)

    # Use model name based on io_cnt and non-batching variant
    if create_savedmodel:
        model_name = tu.get_zero_model_name(
            "savedmodel_nobatch" if max_batch == 0 else "savedmodel", io_cnt,
            dtype)
    else:
        model_name = tu.get_zero_model_name(
            "graphdef_nobatch" if max_batch == 0 else "graphdef", io_cnt, dtype)

    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    if create_savedmodel:
        with tf.Session() as sess:
            input_dict = {}
            output_dict = {}
            for io_num in range(io_cnt):
                input_name = "INPUT{}".format(io_num)
                output_name = "OUTPUT{}".format(io_num)
                input_tensor = tf.get_default_graph().get_tensor_by_name(
                    input_name + ":0")
                output_tensor = tf.get_default_graph().get_tensor_by_name(
                    output_name + ":0")
                input_dict[input_name] = input_tensor
                output_dict[output_name] = output_tensor
            tf.saved_model.simple_save(sess,
                                       model_version_dir + "/model.savedmodel",
                                       inputs=input_dict,
                                       outputs=output_dict)
    else:
        with tf.Session() as sess:
            graph_io.write_graph(sess.graph.as_graph_def(),
                                 model_version_dir,
                                 "model.graphdef",
                                 as_text=False)


def create_tf_modelconfig(create_savedmodel, models_dir, model_version, io_cnt,
                          max_batch, dtype, shape):

    if not tu.validate_for_tf_model(dtype, dtype, dtype, shape, shape, shape):
        return

    shape_str = tu.shape_to_dims_str(shape)

    # Use a different model name for the non-batching variant
    if create_savedmodel:
        model_name = tu.get_zero_model_name(
            "savedmodel_nobatch" if max_batch == 0 else "savedmodel", io_cnt,
            dtype)
    else:
        model_name = tu.get_zero_model_name(
            "graphdef_nobatch" if max_batch == 0 else "graphdef", io_cnt, dtype)

    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "{}"
max_batch_size: {}
'''.format(
        model_name,
        "tensorflow_savedmodel" if create_savedmodel else "tensorflow_graphdef",
        max_batch)

    for io_num in range(io_cnt):
        config += '''
input [
  {{
    name: "INPUT{}"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT{}"
    data_type: {}
    dims: [ {} ]
  }}
]
'''.format(io_num, np_to_model_dtype(dtype), shape_str, io_num,
           np_to_model_dtype(dtype), shape_str)

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_ensemble_modelfile(create_savedmodel, models_dir, model_version,
                              io_cnt, max_batch, dtype, shape):
    if not tu.validate_for_ensemble_model("zero", dtype, dtype, dtype, shape,
                                          shape, shape):
        return

    emu.create_identity_ensemble_modelfile("zero", models_dir, model_version,
                                           max_batch, dtype, [shape] * io_cnt,
                                           [shape] * io_cnt)


def create_ensemble_modelconfig(create_savedmodel, models_dir, model_version,
                                io_cnt, max_batch, dtype, shape):
    if not tu.validate_for_ensemble_model("zero", dtype, dtype, dtype, shape,
                                          shape, shape):
        return

    emu.create_identity_ensemble_modelconfig("zero", models_dir, model_version,
                                             max_batch, dtype, [shape] * io_cnt,
                                             [shape] * io_cnt, [shape] * io_cnt,
                                             [shape] * io_cnt)


def create_onnx_modelfile(create_savedmodel, models_dir, model_version, io_cnt,
                          max_batch, dtype, shape):

    if not tu.validate_for_onnx_model(dtype, dtype, dtype, shape, shape, shape):
        return

    onnx_dtype = np_to_onnx_dtype(dtype)

    # Create the model
    model_name = tu.get_zero_model_name(
        "onnx_nobatch" if max_batch == 0 else "onnx", io_cnt, dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    batch_dim = [] if max_batch == 0 else [None]

    onnx_nodes = []
    onnx_inputs = []
    onnx_outputs = []
    idx = 0
    for io_num in range(io_cnt):
        # Repeat so that the variable dimension name is different
        in_shape, idx = tu.shape_to_onnx_shape(shape, idx)
        out_shape, idx = tu.shape_to_onnx_shape(shape, idx)
        in_name = "INPUT{}".format(io_num)
        out_name = "OUTPUT{}".format(io_num)

        onnx_inputs.append(
            onnx.helper.make_tensor_value_info(in_name, onnx_dtype,
                                               batch_dim + in_shape))
        onnx_outputs.append(
            onnx.helper.make_tensor_value_info(out_name, onnx_dtype,
                                               batch_dim + out_shape))
        onnx_nodes.append(
            onnx.helper.make_node("Identity", [in_name], [out_name]))

    graph_proto = onnx.helper.make_graph(onnx_nodes, model_name, onnx_inputs,
                                         onnx_outputs)
    if FLAGS.onnx_opset > 0:
        model_opset = onnx.helper.make_operatorsetid("", FLAGS.onnx_opset)
        model_def = onnx.helper.make_model(graph_proto,
                                           producer_name="triton",
                                           opset_imports=[model_opset])
    else:
        model_def = onnx.helper.make_model(graph_proto, producer_name="triton")

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    onnx.save(model_def, model_version_dir + "/model.onnx")


def create_onnx_modelconfig(create_savedmodel, models_dir, model_version,
                            io_cnt, max_batch, dtype, shape):

    if not tu.validate_for_onnx_model(dtype, dtype, dtype, shape, shape, shape):
        return

    # Use a different model name for the non-batching variant
    model_name = tu.get_zero_model_name(
        "onnx_nobatch" if max_batch == 0 else "onnx", io_cnt, dtype)
    config_dir = models_dir + "/" + model_name

    config = emu.create_general_modelconfig(model_name,
                                            "onnxruntime_onnx",
                                            max_batch,
                                            emu.repeat(dtype, io_cnt),
                                            emu.repeat(shape, io_cnt),
                                            emu.repeat(shape, io_cnt),
                                            emu.repeat(dtype, io_cnt),
                                            emu.repeat(shape, io_cnt),
                                            emu.repeat(shape, io_cnt),
                                            emu.repeat(None, io_cnt),
                                            force_tensor_number_suffix=True)

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_libtorch_modelfile(create_savedmodel, models_dir, model_version,
                              io_cnt, max_batch, dtype, shape):

    if not tu.validate_for_libtorch_model(dtype, dtype, dtype, shape, shape,
                                          shape, max_batch):
        return

    model_name = tu.get_zero_model_name(
        "libtorch_nobatch" if max_batch == 0 else "libtorch", io_cnt, dtype)

    # Create the model
    if io_cnt == 1:
        if (dtype == np_dtype_string):

            class IdentityNet(nn.Module):

                def __init__(self):
                    super(IdentityNet, self).__init__()

                def forward(self, input0: List[str]) -> List[str]:
                    return input0
        else:

            class IdentityNet(nn.Module):

                def __init__(self):
                    super(IdentityNet, self).__init__()

                def forward(self, input0):
                    return input0
    elif io_cnt == 2:
        if (dtype == np_dtype_string):

            class IdentityNet(nn.Module):

                def __init__(self):
                    super(IdentityNet, self).__init__()

                def forward(self, input0: List[str],
                            input1: List[str]) -> Tuple[List[str], List[str]]:
                    return input0, input1
        else:

            class IdentityNet(nn.Module):

                def __init__(self):
                    super(IdentityNet, self).__init__()

                def forward(self, input0, input1):
                    return input0, input1
    elif io_cnt == 3:
        if (dtype == np_dtype_string):

            class IdentityNet(nn.Module):

                def __init__(self):
                    super(IdentityNet, self).__init__()

                def forward(
                    self, input0: List[str], input1: List[str],
                    input2: List[str]
                ) -> Tuple[List[str], List[str], List[str]]:
                    return input0, input1, input2
        else:

            class IdentityNet(nn.Module):

                def __init__(self):
                    super(IdentityNet, self).__init__()

                def forward(self, input0, input1, input2):
                    return input0, input1, input2
    elif io_cnt == 4:
        if (dtype == np_dtype_string):

            class IdentityNet(nn.Module):

                def __init__(self):
                    super(IdentityNet, self).__init__()

                def forward(
                    self, input0: List[str], input1: List[str],
                    input2: List[str], input3: List[str]
                ) -> Tuple[List[str], List[str], List[str], List[str]]:
                    return input0, input1, input2, input3
        else:

            class IdentityNet(nn.Module):

                def __init__(self):
                    super(IdentityNet, self).__init__()

                def forward(self, input0, input1, input2, input3):
                    return input0, input1, input2, input3

    identityModel = IdentityNet()
    traced = torch.jit.script(identityModel)

    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    traced.save(model_version_dir + "/model.pt")


def create_libtorch_modelconfig(create_savedmodel, models_dir, model_version,
                                io_cnt, max_batch, dtype, shape):

    if not tu.validate_for_libtorch_model(dtype, dtype, dtype, shape, shape,
                                          shape, max_batch):
        return

    # Unpack version policy
    version_policy_str = "{ latest { num_versions: 1 }}"

    # Use a different model name for the non-batching variant
    model_name = tu.get_zero_model_name(
        "libtorch_nobatch" if max_batch == 0 else "libtorch", io_cnt, dtype)
    shape_str = tu.shape_to_dims_str(shape)

    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "pytorch_libtorch"
max_batch_size: {}
version_policy: {}
'''.format(model_name, max_batch, version_policy_str)

    for io_num in range(io_cnt):
        config += '''
input [
  {{
    name: "INPUT__{}"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT__{}"
    data_type: {}
    dims: [ {} ]
  }}
]
'''.format(io_num, np_to_model_dtype(dtype), shape_str, io_num,
           np_to_model_dtype(dtype), shape_str)

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_openvino_modelfile(models_dir, model_version, io_cnt, max_batch,
                              dtype, shape):

    batch_dim = [] if max_batch == 0 else [
        max_batch,
    ]
    if not tu.validate_for_openvino_model(dtype, dtype, dtype,
                                          batch_dim + shape, batch_dim + shape,
                                          batch_dim + shape):
        return

    # Create the model
    model_name = tu.get_zero_model_name(
        "openvino_nobatch" if max_batch == 0 else "openvino", io_cnt, dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    openvino_inputs = []
    openvino_outputs = []
    for io_num in range(io_cnt):
        in_name = "INPUT{}".format(io_num)
        out_name = "OUTPUT{}".format(io_num)
        openvino_inputs.append(
            ng.parameter(shape=batch_dim + shape, dtype=dtype, name=in_name))
        openvino_outputs.append(
            ng.result(openvino_inputs[io_num], name=out_name))

    function = ng.impl.Function(openvino_outputs, openvino_inputs, model_name)
    ie_network = IENetwork(ng.impl.Function.to_capsule(function))

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    ie_network.serialize(model_version_dir + "/model.xml",
                         model_version_dir + "/model.bin")


def create_openvino_modelconfig(models_dir, model_version, io_cnt, max_batch,
                                dtype, shape):

    batch_dim = [] if max_batch == 0 else [
        max_batch,
    ]
    if not tu.validate_for_openvino_model(dtype, dtype, dtype,
                                          batch_dim + shape, batch_dim + shape,
                                          batch_dim + shape):
        return

    # Unpack version policy
    version_policy_str = "{ latest { num_versions: 1 }}"

    # Use a different model name for the non-batching variant
    model_name = tu.get_zero_model_name(
        "openvino_nobatch" if max_batch == 0 else "openvino", io_cnt, dtype)
    shape_str = tu.shape_to_dims_str(shape)

    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
backend: "openvino"
max_batch_size: {}
version_policy: {}
'''.format(model_name, max_batch, version_policy_str)

    for io_num in range(io_cnt):
        config += '''
input [
  {{
    name: "INPUT__{}"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT__{}"
    data_type: {}
    dims: [ {} ]
  }}
]
'''.format(io_num, np_to_model_dtype(dtype), shape_str, io_num,
           np_to_model_dtype(dtype), shape_str)

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_plan_modelfile(create_savedmodel, models_dir, model_version, io_cnt,
                          max_batch, dtype, shape, profile_max_size):

    if not tu.validate_for_trt_model(dtype, dtype, dtype, shape, shape, shape):
        return

    # generate models with different configuration to ensure test coverage
    if dtype != np.float32:
        create_plan_dynamic_rf_modelfile(models_dir, model_version, io_cnt,
                                         max_batch, dtype, shape,
                                         profile_max_size)
    else:
        create_plan_dynamic_modelfile(models_dir, model_version, io_cnt,
                                      max_batch, dtype, shape, profile_max_size)


def create_plan_dynamic_rf_modelfile(models_dir, model_version, io_cnt,
                                     max_batch, dtype, shape, profile_max_size):
    # Create the model
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    if max_batch == 0:
        shape_with_batchsize = [i for i in shape]
    else:
        shape_with_batchsize = [-1] + [i for i in shape]

    trt_dtype = np_to_trt_dtype(dtype)
    trt_memory_format = trt.TensorFormat.LINEAR
    for io_num in range(io_cnt):
        in_node = network.add_input("INPUT{}".format(io_num), trt_dtype,
                                    shape_with_batchsize)
        in_node.allowed_formats = 1 << int(trt_memory_format)

        out_node = network.add_identity(in_node)

        out_node.get_output(0).name = "OUTPUT{}".format(io_num)
        out_node.get_output(0).dtype = trt_dtype
        network.mark_output(out_node.get_output(0))
        out_node.get_output(0).allowed_formats = 1 << int(trt_memory_format)

        if (trt_dtype == trt.int8):
            in_node.dynamic_range = (-128.0, 127.0)
            out_node.get_output(0).dynamic_range = (-128.0, 127.0)

    min_shape = []
    opt_shape = []
    max_shape = []
    if max_batch != 0:
        min_shape = min_shape + [1]
        opt_shape = opt_shape + [max(1, max_batch)]
        max_shape = max_shape + [max(1, max_batch)]
    for i in shape:
        if i == -1:
            # Generating a very generous optimization profile
            min_shape = min_shape + [1]
            opt_shape = opt_shape + [8]
            max_shape = max_shape + [profile_max_size]
        else:
            min_shape = min_shape + [i]
            opt_shape = opt_shape + [i]
            max_shape = max_shape + [i]

    profile = builder.create_optimization_profile()
    for io_num in range(io_cnt):
        profile.set_shape("INPUT{}".format(io_num), min_shape, opt_shape,
                          max_shape)

    flags = 1 << int(trt.BuilderFlag.STRICT_TYPES)
    datatype_set = set([trt_dtype])
    for dt in datatype_set:
        if (dt == trt.int8):
            flags |= 1 << int(trt.BuilderFlag.INT8)
        elif (dt == trt.float16):
            flags |= 1 << int(trt.BuilderFlag.FP16)
    config = builder.create_builder_config()
    config.flags = flags
    config.max_workspace_size = 1 << 20
    config.add_optimization_profile(profile)
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    model_name = tu.get_zero_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", io_cnt, dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_shape_tensor_modelfile(models_dir, model_version, io_cnt,
                                       max_batch, dtype, shape,
                                       profile_max_size):
    # Note that resize layer does not support int tensors.
    # The model takes two inputs (INPUT and DUMMY_INPUT)
    # and produce two outputs.
    # OUTPUT : The shape of resized output 'DUMMY_OUTPUT'.
    # DUMMY_OUTPUT : Obtained after resizing 'DUMMY_INPUT'
    # to shape specified in 'INPUT'.
    # Note that values of OUTPUT tensor must be identical
    # to INPUT values

    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    if max_batch == 0:
        shape_with_batchsize = len(shape)
        dummy_shape = [-1] * shape_with_batchsize
    else:
        shape_with_batchsize = len(shape) + 1
        dummy_shape = [-1] * shape_with_batchsize

    trt_dtype = np_to_trt_dtype(dtype)
    trt_memory_format = trt.TensorFormat.LINEAR
    for io_num in range(io_cnt):
        in_node = network.add_input("INPUT{}".format(io_num), trt.int32,
                                    [shape_with_batchsize])
        in_node.allowed_formats = 1 << int(trt_memory_format)
        dummy_in_node = network.add_input("DUMMY_INPUT{}".format(io_num),
                                          trt_dtype, dummy_shape)
        dummy_in_node.allowed_formats = 1 << int(trt_memory_format)
        resize_layer = network.add_resize(dummy_in_node)
        resize_layer.set_input(1, in_node)
        out_node = network.add_shape(resize_layer.get_output(0))

        dummy_out_node = resize_layer.get_output(0)
        out_node.get_output(0).name = "OUTPUT{}".format(io_num)

        dummy_out_node.name = "DUMMY_OUTPUT{}".format(io_num)

        dummy_out_node.dtype = trt_dtype
        network.mark_output(dummy_out_node)
        dummy_out_node.allowed_formats = 1 << int(trt_memory_format)

        out_node.get_output(0).dtype = trt.int32
        network.mark_output_for_shapes(out_node.get_output(0))
        out_node.get_output(0).allowed_formats = 1 << int(trt_memory_format)

        if (trt_dtype == trt.int8):
            in_node.dynamic_range = (-128.0, 127.0)
            out_node.get_output(0).dynamic_range = (-128.0, 127.0)

    config = builder.create_builder_config()

    min_prefix = []
    opt_prefix = []
    max_prefix = []

    if max_batch != 0:
        min_prefix = [1]
        opt_prefix = [max(1, max_batch)]
        max_prefix = [max(1, max_batch)]

    min_shape = min_prefix + [1] * len(shape)
    opt_shape = opt_prefix + [8] * len(shape)
    max_shape = max_prefix + [profile_max_size] * len(shape)

    profile = builder.create_optimization_profile()
    for io_num in range(io_cnt):
        profile.set_shape_input("INPUT{}".format(io_num), min_shape, opt_shape,
                                max_shape)
        profile.set_shape("DUMMY_INPUT{}".format(io_num), min_shape, opt_shape,
                          max_shape)

    config.add_optimization_profile(profile)

    flags = 1 << int(trt.BuilderFlag.STRICT_TYPES)
    datatype_set = set([trt_dtype])
    for dt in datatype_set:
        if (dt == trt.int8):
            flags |= 1 << int(trt.BuilderFlag.INT8)
        elif (dt == trt.float16):
            flags |= 1 << int(trt.BuilderFlag.FP16)
    config.flags = flags

    config.max_workspace_size = 1 << 20
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    model_name = tu.get_zero_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", io_cnt, dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_dynamic_modelfile(models_dir, model_version, io_cnt, max_batch,
                                  dtype, shape, profile_max_size):
    # Create the model
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    if max_batch == 0:
        shape_with_batchsize = [i for i in shape]
    else:
        shape_with_batchsize = [-1] + [i for i in shape]

    trt_dtype = np_to_trt_dtype(dtype)
    for io_num in range(io_cnt):
        in_node = network.add_input("INPUT{}".format(io_num), trt_dtype,
                                    shape_with_batchsize)
        out_node = network.add_identity(in_node)
        out_node.get_output(0).name = "OUTPUT{}".format(io_num)
        network.mark_output(out_node.get_output(0))

    min_shape = []
    opt_shape = []
    max_shape = []
    if max_batch != 0:
        min_shape = min_shape + [1]
        opt_shape = opt_shape + [max(1, max_batch)]
        max_shape = max_shape + [max(1, max_batch)]
    for i in shape:
        if i == -1:
            # Generating a very generous optimization profile
            min_shape = min_shape + [1]
            opt_shape = opt_shape + [8]
            max_shape = max_shape + [profile_max_size]
        else:
            min_shape = min_shape + [i]
            opt_shape = opt_shape + [i]
            max_shape = max_shape + [i]

    profile = builder.create_optimization_profile()
    for io_num in range(io_cnt):
        profile.set_shape("INPUT{}".format(io_num), min_shape, opt_shape,
                          max_shape)
    config = builder.create_builder_config()
    config.add_optimization_profile(profile)
    config.max_workspace_size = 1 << 20
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    model_name = tu.get_zero_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", io_cnt, dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_modelconfig(create_savedmodel, models_dir, model_version,
                            io_cnt, max_batch, dtype, shape):
    if not tu.validate_for_trt_model(dtype, dtype, dtype, shape, shape, shape):
        return

    shape_str = tu.shape_to_dims_str(shape)

    model_name = tu.get_zero_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", io_cnt, dtype)
    config_dir = models_dir + "/" + model_name

    if FLAGS.tensorrt_shape_io:
        shape_tensor_dim = len(shape)
        config = '''
name: "{}"
platform: "tensorrt_plan"
max_batch_size: {}
'''.format(model_name, max_batch)

        for io_num in range(io_cnt):
            config += '''
input [
  {{
    name: "DUMMY_INPUT{}"
    data_type: {}
    dims: [ {} ]
  }},
  {{
    name: "INPUT{}"
    data_type: TYPE_INT32
    dims: [ {} ]
    is_shape_tensor: true
  }}
]
output [
  {{
    name: "DUMMY_OUTPUT{}"
    data_type: {}
    dims: [ {} ]
  }},
  {{
    name: "OUTPUT{}"
    data_type: TYPE_INT32
    dims: [ {} ]
    is_shape_tensor: true
  }}
]
'''.format(io_num, np_to_model_dtype(dtype), shape_str, io_num,
            shape_tensor_dim, io_num, np_to_model_dtype(dtype), shape_str,
            io_num, shape_tensor_dim)

    else:
        config = '''
name: "{}"
platform: "tensorrt_plan"
max_batch_size: {}
'''.format(model_name, max_batch)

        for io_num in range(io_cnt):
            config += '''
input [
  {{
    name: "INPUT{}"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT{}"
    data_type: {}
    dims: [ {} ]
  }}
]
'''.format(io_num, np_to_model_dtype(dtype), shape_str, io_num,
            np_to_model_dtype(dtype), shape_str)

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_shape_tensor_models(models_dir,
                               dtype,
                               shape,
                               io_cnt=1,
                               no_batch=True):
    model_version = 1

    create_plan_modelconfig(True, models_dir, model_version, io_cnt, 8, dtype,
                            shape)
    create_plan_shape_tensor_modelfile(models_dir, model_version, io_cnt, 8,
                                       dtype, shape, 32)
    if no_batch:
        create_plan_modelconfig(True, models_dir, model_version, io_cnt, 0,
                                dtype, shape)
        create_plan_shape_tensor_modelfile(models_dir, model_version, io_cnt, 0,
                                           dtype, shape, 32)


def create_models(models_dir, dtype, shape, io_cnt=1, no_batch=True):
    model_version = 1

    if FLAGS.graphdef:
        create_tf_modelconfig(False, models_dir, model_version, io_cnt, 8,
                              dtype, shape)
        create_tf_modelfile(False, models_dir, model_version, io_cnt, 8, dtype,
                            shape)
        if no_batch:
            create_tf_modelconfig(False, models_dir, model_version, io_cnt, 0,
                                  dtype, shape)
            create_tf_modelfile(False, models_dir, model_version, io_cnt, 0,
                                dtype, shape)

    if FLAGS.savedmodel:
        create_tf_modelconfig(True, models_dir, model_version, io_cnt, 8, dtype,
                              shape)
        create_tf_modelfile(True, models_dir, model_version, io_cnt, 8, dtype,
                            shape)
        if no_batch:
            create_tf_modelconfig(True, models_dir, model_version, io_cnt, 0,
                                  dtype, shape)
            create_tf_modelfile(True, models_dir, model_version, io_cnt, 0,
                                dtype, shape)

    if FLAGS.onnx:
        create_onnx_modelconfig(True, models_dir, model_version, io_cnt, 8,
                                dtype, shape)
        create_onnx_modelfile(True, models_dir, model_version, io_cnt, 8, dtype,
                              shape)
        if no_batch:
            create_onnx_modelconfig(True, models_dir, model_version, io_cnt, 0,
                                    dtype, shape)
            create_onnx_modelfile(True, models_dir, model_version, io_cnt, 0,
                                  dtype, shape)

    if FLAGS.openvino:
        create_openvino_modelconfig(models_dir, model_version, io_cnt, 8, dtype,
                                    shape)
        create_openvino_modelfile(models_dir, model_version, io_cnt, 8, dtype,
                                  shape)
        if no_batch:
            create_openvino_modelconfig(models_dir, model_version, io_cnt, 0,
                                        dtype, shape)
            create_openvino_modelfile(models_dir, model_version, io_cnt, 0,
                                      dtype, shape)

    if FLAGS.libtorch:
        create_libtorch_modelconfig(True, models_dir, model_version, io_cnt, 8,
                                    dtype, shape)
        create_libtorch_modelfile(True, models_dir, model_version, io_cnt, 8,
                                  dtype, shape)
        if no_batch:
            create_libtorch_modelconfig(True, models_dir, model_version, io_cnt,
                                        0, dtype, shape)
            create_libtorch_modelfile(True, models_dir, model_version, io_cnt,
                                      0, dtype, shape)

    if FLAGS.tensorrt:
        create_plan_modelconfig(True, models_dir, model_version, io_cnt, 8,
                                dtype, shape)
        create_plan_modelfile(True, models_dir, model_version, io_cnt, 8, dtype,
                              shape, 32)
        if no_batch:
            create_plan_modelconfig(True, models_dir, model_version, io_cnt, 0,
                                    dtype, shape)
            create_plan_modelfile(True, models_dir, model_version, io_cnt, 0,
                                  dtype, shape, 32)

    if FLAGS.tensorrt_big:
        create_plan_modelconfig(True, models_dir, model_version, io_cnt, 8,
                                dtype, shape)
        create_plan_modelfile(True, models_dir, model_version, io_cnt, 8, dtype,
                              shape, 16 * 1024 * 1024)
        if no_batch:
            create_plan_modelconfig(True, models_dir, model_version, io_cnt, 0,
                                    dtype, shape)
            create_plan_modelfile(True, models_dir, model_version, io_cnt, 0,
                                  dtype, shape, 16 * 1024 * 1024)

    if FLAGS.ensemble:
        emu.create_nop_modelconfig(models_dir, shape, dtype)
        create_ensemble_modelconfig(True, models_dir, model_version, io_cnt, 8,
                                    dtype, shape)
        create_ensemble_modelfile(True, models_dir, model_version, io_cnt, 8,
                                  dtype, shape)
        if no_batch:
            create_ensemble_modelconfig(True, models_dir, model_version, io_cnt,
                                        0, dtype, shape)
            create_ensemble_modelfile(True, models_dir, model_version, io_cnt,
                                      0, dtype, shape)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--models_dir',
                        type=str,
                        required=True,
                        help='Top-level model directory')
    parser.add_argument('--graphdef',
                        required=False,
                        action='store_true',
                        help='Generate GraphDef models')
    parser.add_argument('--savedmodel',
                        required=False,
                        action='store_true',
                        help='Generate SavedModel models')
    parser.add_argument('--onnx',
                        required=False,
                        action='store_true',
                        help='Generate Onnx Runtime Onnx models')
    parser.add_argument(
        '--onnx_opset',
        type=int,
        required=False,
        default=0,
        help='Opset used for Onnx models. Default is to use ONNXRT default')
    parser.add_argument('--libtorch',
                        required=False,
                        action='store_true',
                        help='Generate Pytorch LibTorch models')
    parser.add_argument('--openvino',
                        required=False,
                        action='store_true',
                        help='Generate OpenVino models')
    parser.add_argument('--tensorrt',
                        required=False,
                        action='store_true',
                        help='Generate TensorRT PLAN models')
    parser.add_argument(
        '--tensorrt-big',
        required=False,
        action='store_true',
        help='Generate TensorRT PLAN models w/ opt profile with large max')
    parser.add_argument(
        '--tensorrt-shape-io',
        required=False,
        action='store_true',
        help='Generate TensorRT PLAN models w/ shape tensor i/o')
    parser.add_argument('--ensemble',
                        required=False,
                        action='store_true',
                        help='Generate ensemble models')
    FLAGS, unparsed = parser.parse_known_args()

    if FLAGS.graphdef or FLAGS.savedmodel:
        import tensorflow as tf
        from tensorflow.python.framework import graph_io
    if FLAGS.onnx:
        import onnx
    if FLAGS.libtorch:
        import torch
        from torch import nn
    if FLAGS.tensorrt or FLAGS.tensorrt_big or FLAGS.tensorrt_shape_io:
        import tensorrt as trt
    if FLAGS.openvino:
        from openvino.inference_engine import IENetwork
        import ngraph as ng

    import test_util as tu

    # Create models with variable-sized input and output. For big
    # TensorRT models only create the one needed for performance
    # testing
    if FLAGS.tensorrt_big:
        create_models(FLAGS.models_dir, np.float32, [-1], io_cnt=1)
    elif FLAGS.tensorrt_shape_io:
        create_shape_tensor_models(FLAGS.models_dir,
                                   np.float32, [-1, -1],
                                   io_cnt=1)
    else:
        create_models(FLAGS.models_dir, bool, [-1], io_cnt=1)
        create_models(FLAGS.models_dir, np.float32, [-1], io_cnt=1)
        create_models(FLAGS.models_dir, np.float32, [-1], io_cnt=3)
        create_models(FLAGS.models_dir, np.float16, [-1, -1], io_cnt=1)
        create_models(FLAGS.models_dir, np.float16, [-1, -1], io_cnt=3)
        create_models(FLAGS.models_dir, np_dtype_string, [-1], io_cnt=1)
        create_models(FLAGS.models_dir, np_dtype_string, [-1, -1], io_cnt=3)


# qa/common/gen_qa_implicit_models.py

# Copyright 2021-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
import os
import numpy as np
import gen_ensemble_model_utils as emu

FLAGS = None
np_dtype_string = np.dtype(object)


def np_to_model_dtype(np_dtype):
    if np_dtype == bool:
        return "TYPE_BOOL"
    elif np_dtype == np.int8:
        return "TYPE_INT8"
    elif np_dtype == np.int16:
        return "TYPE_INT16"
    elif np_dtype == np.int32:
        return "TYPE_INT32"
    elif np_dtype == np.int64:
        return "TYPE_INT64"
    elif np_dtype == np.uint8:
        return "TYPE_UINT8"
    elif np_dtype == np.uint16:
        return "TYPE_UINT16"
    elif np_dtype == np.float16:
        return "TYPE_FP16"
    elif np_dtype == np.float32:
        return "TYPE_FP32"
    elif np_dtype == np.float64:
        return "TYPE_FP64"
    elif np_dtype == np_dtype_string:
        return "TYPE_STRING"
    return None


def np_to_onnx_dtype(np_dtype):
    if np_dtype == bool:
        return onnx.TensorProto.BOOL
    elif np_dtype == np.int8:
        return onnx.TensorProto.INT8
    elif np_dtype == np.int16:
        return onnx.TensorProto.INT16
    elif np_dtype == np.int32:
        return onnx.TensorProto.INT32
    elif np_dtype == np.int64:
        return onnx.TensorProto.INT64
    elif np_dtype == np.uint8:
        return onnx.TensorProto.UINT8
    elif np_dtype == np.uint16:
        return onnx.TensorProto.UINT16
    elif np_dtype == np.float16:
        return onnx.TensorProto.FLOAT16
    elif np_dtype == np.float32:
        return onnx.TensorProto.FLOAT
    elif np_dtype == np.float64:
        return onnx.TensorProto.DOUBLE
    elif np_dtype == np_dtype_string:
        return onnx.TensorProto.STRING


def np_to_trt_dtype(np_dtype):
    if np_dtype == bool:
        return trt.bool
    elif np_dtype == np.int8:
        return trt.int8
    elif np_dtype == np.int32:
        return trt.int32
    elif np_dtype == np.float16:
        return trt.float16
    elif np_dtype == np.float32:
        return trt.float32
    return None


def create_onnx_modelfile_wo_initial_state(models_dir, model_version, max_batch,
                                           dtype, shape):

    if not tu.validate_for_onnx_model(dtype, dtype, dtype, shape, shape, shape):
        return

    model_name = tu.get_sequence_model_name(
        "onnx_nobatch" if max_batch == 0 else "onnx", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    onnx_dtype = np_to_onnx_dtype(dtype)
    onnx_control_dtype = onnx_dtype
    onnx_input_shape, idx = tu.shape_to_onnx_shape(shape, 0)
    onnx_output_shape, idx = tu.shape_to_onnx_shape(shape, idx)

    # If the input is a string then use int32 for operation and just
    # cast to/from string for input and output.
    if onnx_dtype == onnx.TensorProto.STRING:
        onnx_control_dtype = onnx.TensorProto.INT32

    # If input dtype is bool, then use bool type for control and
    # int32 type for input/output
    if onnx_dtype == onnx.TensorProto.BOOL:
        onnx_dtype = onnx.TensorProto.INT32

    batch_dim = [] if max_batch == 0 else [None]

    onnx_input = onnx.helper.make_tensor_value_info(
        "INPUT", onnx_dtype, batch_dim + onnx_input_shape)
    onnx_input_state = onnx.helper.make_tensor_value_info(
        "INPUT_STATE", onnx_dtype, batch_dim + onnx_input_shape)
    onnx_start = onnx.helper.make_tensor_value_info("START", onnx_control_dtype,
                                                    batch_dim + [1])
    onnx_ready = onnx.helper.make_tensor_value_info("READY", onnx_control_dtype,
                                                    batch_dim + [1])
    onnx_output = onnx.helper.make_tensor_value_info(
        "OUTPUT", onnx_dtype, batch_dim + onnx_output_shape)
    onnx_output_state = onnx.helper.make_tensor_value_info(
        "OUTPUT_STATE", onnx_dtype, batch_dim + onnx_output_shape)

    internal_input = onnx.helper.make_node("Identity", ["INPUT"], ["_INPUT"])
    internal_input_state = onnx.helper.make_node("Identity", ["INPUT_STATE"],
                                                 ["_INPUT_STATE"])
    # cast int8, int16 input to higer precision int as Onnx Add/Sub operator doesn't support those type
    if ((onnx_dtype == onnx.TensorProto.INT8) or
        (onnx_dtype == onnx.TensorProto.INT16)):
        internal_input = onnx.helper.make_node("Cast", ["INPUT"], ["_INPUT"],
                                               to=onnx.TensorProto.INT32)
        internal_input_state = onnx.helper.make_node("Cast", ["INPUT_STATE"],
                                                     ["_INPUT_STATE"],
                                                     to=onnx.TensorProto.INT32)

    # Convert boolean value to int32 value
    if onnx_control_dtype == onnx.TensorProto.BOOL:
        if onnx_dtype != onnx.TensorProto.STRING:
            internal_input1 = onnx.helper.make_node("Cast", ["START"],
                                                    ["_START"],
                                                    to=onnx.TensorProto.INT32)
            internal_input2 = onnx.helper.make_node("Cast", ["READY"],
                                                    ["_READY"],
                                                    to=onnx.TensorProto.INT32)
            not_start_cast = onnx.helper.make_node("Not", ["START"],
                                                   ["_NOT_START_CAST"])
            not_start = onnx.helper.make_node("Cast", ["_NOT_START_CAST"],
                                              ["_NOT_START"],
                                              to=onnx.TensorProto.INT32)
            not_ready_cast = onnx.helper.make_node("Not", ["START"],
                                                   ["_NOT_READY_CAST"])
            not_ready = onnx.helper.make_node("Cast", ["_NOT_READY_CAST"],
                                              ["_NOT_READY"],
                                              to=onnx.TensorProto.INT32)
            input_state_cond = onnx.helper.make_node(
                "And", ["READY", "_NOT_START_CAST"], ["input_state_cond"])
            input_state_cond_cast = onnx.helper.make_node(
                "Cast", ["input_state_cond"], ["input_state_cond_cast"],
                to=onnx.TensorProto.INT32)
            mul_state = onnx.helper.make_node(
                "Mul", ["_INPUT_STATE", "input_state_cond_cast"], ["mul_state"])
            add = onnx.helper.make_node("Add", ["_INPUT", "mul_state"],
                                        ["CAST"])

    else:

        if onnx_dtype != onnx.TensorProto.STRING:
            start_cast = onnx.helper.make_node("Cast", ["START"],
                                               ["_START_CAST"],
                                               to=onnx.TensorProto.BOOL)
            not_start_cast = onnx.helper.make_node("Not", ["_START_CAST"],
                                                   ["_NOT_START_CAST"])
            not_start = onnx.helper.make_node("Cast", ["_NOT_START_CAST"],
                                              ["_NOT_START"],
                                              to=onnx.TensorProto.INT32)

            ready_cast = onnx.helper.make_node("Cast", ["READY"],
                                               ["_READY_CAST"],
                                               to=onnx.TensorProto.BOOL)
            not_ready_cast = onnx.helper.make_node("Not", ["_READY_CAST"],
                                                   ["_NOT_READY_CAST"])
            not_ready = onnx.helper.make_node("Cast", ["_NOT_READY_CAST"],
                                              ["_NOT_READY"],
                                              to=onnx.TensorProto.INT32)
            # Take advantage of knowledge that the READY false value is 0 and true is 1
            input_state_cond = onnx.helper.make_node(
                "And", ["_NOT_START_CAST", "_READY_CAST"], ["input_state_cond"])
            input_state_cond_cast = onnx.helper.make_node(
                "Cast", ["input_state_cond"], ["input_state_cond_cast"],
                to=onnx.TensorProto.INT32)
            mul_state = onnx.helper.make_node(
                "Mul", ["_INPUT_STATE", "input_state_cond_cast"], ["mul_state"])
            add = onnx.helper.make_node("Add", ["_INPUT", "mul_state"],
                                        ["CAST"])

    if onnx_dtype == onnx.TensorProto.STRING:
        cast = onnx.helper.make_node("Identity", ["_INPUT"], ["OUTPUT"])
        cast_output_state = onnx.helper.make_node("Identity", ["_INPUT"],
                                                  ["OUTPUT_STATE"])
    elif onnx_dtype == onnx.TensorProto.FLOAT16:
        # Avoid cast from float16 to float16
        # (bug in Onnx Runtime, cast from float16 to float16 will become cast from float16 to float32)
        cast = onnx.helper.make_node("Identity", ["CAST"], ["OUTPUT"])
        cast_output_state = onnx.helper.make_node("Identity", ["CAST"],
                                                  ["OUTPUT_STATE"])
    else:
        cast = onnx.helper.make_node("Cast", ["CAST"], ["OUTPUT"],
                                     to=onnx_dtype)
        cast_output_state = onnx.helper.make_node("Cast", ["CAST"],
                                                  ["OUTPUT_STATE"],
                                                  to=onnx_dtype)

    if onnx_control_dtype == onnx.TensorProto.BOOL:
        if onnx_dtype != onnx.TensorProto.STRING:
            onnx_nodes = [
                internal_input, internal_input_state, internal_input1,
                internal_input2, not_start_cast, not_start, not_ready_cast,
                not_ready, input_state_cond, input_state_cond_cast, mul_state,
                add, cast, cast_output_state
            ]
        else:
            onnx_nodes = [
                internal_input, internal_input_state, cast, cast_output_state
            ]
    else:
        if onnx_dtype != onnx.TensorProto.STRING:
            onnx_nodes = [
                internal_input, internal_input_state, start_cast,
                not_start_cast, not_start, ready_cast, not_ready_cast,
                not_ready, input_state_cond, input_state_cond_cast, mul_state,
                add, cast, cast_output_state
            ]
        else:
            onnx_nodes = [
                internal_input, internal_input_state, cast, cast_output_state
            ]

    onnx_inputs = [onnx_input_state, onnx_input, onnx_start, onnx_ready]
    onnx_outputs = [onnx_output, onnx_output_state]
    graph_proto = onnx.helper.make_graph(onnx_nodes, model_name, onnx_inputs,
                                         onnx_outputs)

    if FLAGS.onnx_opset > 0:
        model_opset = onnx.helper.make_operatorsetid("", FLAGS.onnx_opset)
        model_def = onnx.helper.make_model(graph_proto,
                                           producer_name="triton",
                                           opset_imports=[model_opset])
    else:
        model_def = onnx.helper.make_model(graph_proto, producer_name="triton")

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    onnx.save(model_def, model_version_dir + "/model.onnx")


def create_onnx_modelfile_with_initial_state(models_dir, model_version,
                                             max_batch, dtype, shape):
    if not tu.validate_for_onnx_model(dtype, dtype, dtype, shape, shape, shape):
        return

    model_name = tu.get_sequence_model_name(
        "onnx_nobatch" if max_batch == 0 else "onnx", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    onnx_dtype = np_to_onnx_dtype(dtype)
    onnx_control_dtype = onnx_dtype
    onnx_input_shape, idx = tu.shape_to_onnx_shape(shape, 0)
    onnx_output_shape, idx = tu.shape_to_onnx_shape(shape, idx)

    # If the input is a string then use int32 for operation and just
    # cast to/from string for input and output.
    if onnx_dtype == onnx.TensorProto.STRING:
        onnx_control_dtype = onnx.TensorProto.INT32

    # If input dtype is bool, then use bool type for control and
    # int32 type for input/output
    if onnx_dtype == onnx.TensorProto.BOOL:
        onnx_dtype = onnx.TensorProto.INT32

    batch_dim = [] if max_batch == 0 else [None]

    onnx_input = onnx.helper.make_tensor_value_info(
        "INPUT", onnx_dtype, batch_dim + onnx_input_shape)
    onnx_input_state = onnx.helper.make_tensor_value_info(
        "INPUT_STATE", onnx_dtype, batch_dim + onnx_input_shape)
    onnx_start = onnx.helper.make_tensor_value_info("START", onnx_control_dtype,
                                                    batch_dim + [1])
    onnx_ready = onnx.helper.make_tensor_value_info("READY", onnx_control_dtype,
                                                    batch_dim + [1])
    onnx_output = onnx.helper.make_tensor_value_info(
        "OUTPUT", onnx_dtype, batch_dim + onnx_output_shape)
    onnx_output_state = onnx.helper.make_tensor_value_info(
        "OUTPUT_STATE", onnx_dtype, batch_dim + onnx_output_shape)

    internal_input = onnx.helper.make_node("Identity", ["INPUT"], ["_INPUT"])
    internal_input_state = onnx.helper.make_node("Identity", ["INPUT_STATE"],
                                                 ["_INPUT_STATE"])
    # cast int8, int16 input to higer precision int as Onnx Add/Sub operator doesn't support those type
    if ((onnx_dtype == onnx.TensorProto.INT8) or
        (onnx_dtype == onnx.TensorProto.INT16)):
        internal_input = onnx.helper.make_node("Cast", ["INPUT"], ["_INPUT"],
                                               to=onnx.TensorProto.INT32)
        internal_input_state = onnx.helper.make_node("Cast", ["INPUT_STATE"],
                                                     ["_INPUT_STATE"],
                                                     to=onnx.TensorProto.INT32)

    if onnx_dtype == onnx.TensorProto.STRING:
        identity = onnx.helper.make_node("Identity", ["_INPUT"], ["OUTPUT"])
        identity_output_state = onnx.helper.make_node("Identity", ["_INPUT"],
                                                      ["OUTPUT_STATE"])
    else:
        add = onnx.helper.make_node("Add", ["_INPUT", "_INPUT_STATE"], ["CAST"])
        cast = onnx.helper.make_node("Cast", ["CAST"], ["OUTPUT"],
                                     to=onnx_dtype)
        cast_output_state = onnx.helper.make_node("Cast", ["CAST"],
                                                  ["OUTPUT_STATE"],
                                                  to=onnx_dtype)

    # Avoid cast from float16 to float16
    # (bug in Onnx Runtime, cast from float16 to float16 will become cast from float16 to float32)
    if onnx_dtype == onnx.TensorProto.FLOAT16:
        cast = onnx.helper.make_node("Identity", ["CAST"], ["OUTPUT"])
        cast_output_state = onnx.helper.make_node("Identity", ["CAST"],
                                                  ["OUTPUT_STATE"])

    if onnx_dtype != onnx.TensorProto.STRING:
        onnx_nodes = [
            internal_input, internal_input_state, add, cast, cast_output_state
        ]
    else:
        onnx_nodes = [
            internal_input, internal_input_state, identity,
            identity_output_state
        ]

    onnx_inputs = [onnx_input_state, onnx_input, onnx_start, onnx_ready]
    onnx_outputs = [onnx_output, onnx_output_state]
    graph_proto = onnx.helper.make_graph(onnx_nodes, model_name, onnx_inputs,
                                         onnx_outputs)

    if FLAGS.onnx_opset > 0:
        model_opset = onnx.helper.make_operatorsetid("", FLAGS.onnx_opset)
        model_def = onnx.helper.make_model(graph_proto,
                                           producer_name="triton",
                                           opset_imports=[model_opset])
    else:
        model_def = onnx.helper.make_model(graph_proto, producer_name="triton")

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    onnx.save(model_def, model_version_dir + "/model.onnx")


def create_onnx_modelfile(models_dir, model_version, max_batch, dtype, shape,
                          initial_state):

    if initial_state is None:
        create_onnx_modelfile_wo_initial_state(models_dir, model_version,
                                               max_batch, dtype, shape)
    else:
        # This model assumes that the initial state contains correct data
        create_onnx_modelfile_with_initial_state(models_dir, model_version,
                                                 max_batch, dtype, shape)


def create_onnx_modelconfig(models_dir, model_version, max_batch, dtype, shape,
                            initial_state):

    if not tu.validate_for_onnx_model(dtype, dtype, dtype, shape, shape, shape):
        return

    model_name = tu.get_sequence_model_name(
        "onnx_nobatch" if max_batch == 0 else "onnx", dtype)
    config_dir = models_dir + "/" + model_name

    if dtype == np.float32:
        control_type = "fp32"
    elif dtype == bool:
        control_type = "bool"
        dtype = np.int32
    else:
        control_type = "int32"

    instance_group_string = '''
instance_group [
  {
    kind: KIND_GPU
  }
]
'''

    # [TODO] move create_general_modelconfig() out of emu as it is general
    # enough for all backends to use
    config = emu.create_general_modelconfig(
        model_name,
        "onnxruntime_onnx",
        max_batch, [dtype], [shape], [None], [dtype], [shape], [None], [None],
        force_tensor_number_suffix=False,
        instance_group_str=instance_group_string)

    # Prepare the shapes for initial state initialization
    shape_without_variable_dims = []
    for dim in shape:
        if dim == -1:
            shape_without_variable_dims.append(1)
        else:
            shape_without_variable_dims.append(dim)

    if initial_state is None:
        config += '''
    sequence_batching {{
      max_sequence_idle_microseconds: 5000000
      control_input [
        {{
          name: "START"
          control [
            {{
              kind: CONTROL_SEQUENCE_START
              {type}_false_true: [ 0, 1 ]
            }}
          ]
        }},
        {{
          name: "READY"
          control [
            {{
              kind: CONTROL_SEQUENCE_READY
              {type}_false_true: [ 0, 1 ]
            }}
          ]
        }}
      ]
      state [
        {{
          input_name: "INPUT_STATE"
          output_name: "OUTPUT_STATE"
          data_type: {dtype}
          dims: {dims}
        }} 
      ]
    }}
    '''.format(type=control_type,
               dims=tu.shape_to_dims_str(shape),
               dtype=emu.dtype_str(dtype))
    elif initial_state == 'zero':
        config += f'''
    sequence_batching {{
      max_sequence_idle_microseconds: 5000000
      control_input [
        {{
          name: "START"
          control [
            {{
              kind: CONTROL_SEQUENCE_START
              {control_type}_false_true: [ 0, 1 ]
            }}
          ]
        }},
        {{
          name: "READY"
          control [
            {{
              kind: CONTROL_SEQUENCE_READY
              {control_type}_false_true: [ 0, 1 ]
            }}
          ]
        }}
      ]
      state [
        {{
          input_name: "INPUT_STATE"
          output_name: "OUTPUT_STATE"
          data_type: {emu.dtype_str(dtype)}
          dims: {tu.shape_to_dims_str(shape)}
          initial_state: {{
              name: "state init"
              data_type: {emu.dtype_str(dtype)}
              dims: {tu.shape_to_dims_str(shape_without_variable_dims)}
              zero_data: true
          }}
        }} 
      ]
    }}
    '''
    elif initial_state == 'file':
        config += '''
    sequence_batching {{
      max_sequence_idle_microseconds: 5000000
      control_input [
        {{
          name: "START"
          control [
            {{
              kind: CONTROL_SEQUENCE_START
              {type}_false_true: [ 0, 1 ]
            }}
          ]
        }},
        {{
          name: "READY"
          control [
            {{
              kind: CONTROL_SEQUENCE_READY
              {type}_false_true: [ 0, 1 ]
            }}
          ]
        }}
      ]
      state [
        {{
          input_name: "INPUT_STATE"
          output_name: "OUTPUT_STATE"
          data_type: {dtype}
          dims: {dims}
          initial_state: {{
              name: "state init"
              data_type: {dtype}
              dims: {shape_without_variable_dims}
              data_file: input_state_data
          }}
        }} 
      ]
    }}
    '''.format(type=control_type,
               dims=tu.shape_to_dims_str(shape),
               dtype=emu.dtype_str(dtype),
               shape_without_variable_dims=tu.shape_to_dims_str(
                   shape_without_variable_dims))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_plan_fixed_modelfile(models_dir, model_version, max_batch, dtype,
                                shape):
    trt_dtype = np_to_trt_dtype(dtype)
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network()
    in0 = network.add_input("INPUT", trt_dtype, shape)
    in_state0 = network.add_input("INPUT_STATE", trt_dtype, shape)
    start0 = network.add_input("START", trt_dtype, [1 for i in shape])
    network.add_input("READY", trt_dtype, [1 for i in shape])
    constant_1_data = trt.Weights(np.ones([1 for i in shape], dtype=dtype))
    constant_1 = network.add_constant([1 for i in shape], constant_1_data)
    not_start = network.add_elementwise(constant_1.get_output(0), start0,
                                        trt.ElementWiseOperation.SUB)
    not_start.set_output_type(0, trt_dtype)
    internal_state = network.add_elementwise(in_state0, not_start.get_output(0),
                                             trt.ElementWiseOperation.PROD)
    out0 = network.add_elementwise(internal_state.get_output(0), in0,
                                   trt.ElementWiseOperation.SUM)
    out0_state = network.add_elementwise(internal_state.get_output(0), in0,
                                         trt.ElementWiseOperation.SUM)

    out0.get_output(0).name = "OUTPUT"
    network.mark_output(out0.get_output(0))

    out0_state.get_output(0).name = "OUTPUT_STATE"
    network.mark_output(out0_state.get_output(0))

    config = builder.create_builder_config()
    config.max_workspace_size = 1 << 20
    builder.max_batch_size = max(1, max_batch)
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine
    del network

    model_name = tu.get_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_fixed_rf_modelfile(models_dir, model_version, max_batch, dtype,
                                   shape):
    trt_dtype = np_to_trt_dtype(dtype)
    trt_memory_format = trt.TensorFormat.LINEAR

    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network()
    in0 = network.add_input("INPUT", trt_dtype, shape)
    in_state0 = network.add_input("INPUT_STATE", trt_dtype, shape)
    start0 = network.add_input("START", trt_dtype, [1 for i in shape])
    ready0 = network.add_input("READY", trt_dtype, [1 for i in shape])
    constant_1_data = trt.Weights(np.ones([1 for i in shape], dtype=dtype))
    constant_1 = network.add_constant([1 for i in shape], constant_1_data)
    not_start = network.add_elementwise(constant_1.get_output(0), start0,
                                        trt.ElementWiseOperation.SUB)
    not_start.set_output_type(0, trt_dtype)

    internal_state = network.add_elementwise(in_state0, not_start.get_output(0),
                                             trt.ElementWiseOperation.PROD)
    out0 = network.add_elementwise(internal_state.get_output(0), in0,
                                   trt.ElementWiseOperation.SUM)
    out0_state = network.add_elementwise(internal_state.get_output(0), in0,
                                         trt.ElementWiseOperation.SUM)

    out0.get_output(0).name = "OUTPUT"
    network.mark_output(out0.get_output(0))
    out0.get_output(0).dtype = trt_dtype

    out0_state.get_output(0).name = "OUTPUT_STATE"
    network.mark_output(out0_state.get_output(0))
    out0_state.get_output(0).dtype = trt_dtype

    in0.allowed_formats = 1 << int(trt_memory_format)
    in_state0.allowed_formats = 1 << int(trt_memory_format)
    start0.allowed_formats = 1 << int(trt_memory_format)
    ready0.allowed_formats = 1 << int(trt_memory_format)
    out0.get_output(0).allowed_formats = 1 << int(trt_memory_format)
    out0_state.get_output(0).allowed_formats = 1 << int(trt_memory_format)

    if (trt_dtype == trt.int8):
        in0.dynamic_range = (-128.0, 127.0)
        in_state0.dynamic_range = (-128.0, 127.0)
        out0.dynamic_range = (-128.0, 127.0)
        out0_state.dynamic_range = (-128.0, 127.0)
        start0.dynamic_range = (-128.0, 127.0)
        ready0.dynamic_range = (-128.0, 127.0)

    flags = 1 << int(trt.BuilderFlag.STRICT_TYPES)

    if (trt_dtype == trt.int8):
        flags |= 1 << int(trt.BuilderFlag.INT8)
    elif (trt_dtype == trt.float16):
        flags |= 1 << int(trt.BuilderFlag.FP16)

    config = builder.create_builder_config()
    config.flags = flags
    config.max_workspace_size = 1 << 20
    builder.max_batch_size = max(1, max_batch)
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    model_name = tu.get_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_dynamic_modelfile(models_dir, model_version, max_batch, dtype,
                                  shape):
    trt_dtype = np_to_trt_dtype(dtype)
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)

    # EXPLICIT_BATCH must be used when the dimension is variable
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

    unit_shape = ([1] * len(shape))
    if max_batch != 0:
        in0 = network.add_input("INPUT", trt_dtype, [-1] + shape)
        start0 = network.add_input("START", trt_dtype, [-1] + unit_shape)
        ready0 = network.add_input("READY", trt_dtype, [-1] + unit_shape)
        in_state0 = network.add_input("INPUT_STATE", trt_dtype, [-1] + shape)
    else:
        in0 = network.add_input("INPUT", trt_dtype, shape)
        start0 = network.add_input("START", trt_dtype, unit_shape)
        ready0 = network.add_input("READY", trt_dtype, unit_shape)
        in_state0 = network.add_input("INPUT_STATE", trt_dtype, shape)

    # Append the dimension by 1 so that broadcasting works properly
    constant_1_data = trt.Weights(np.ones(unit_shape + [1], dtype=dtype))
    constant_1 = network.add_constant(unit_shape + [1], constant_1_data)
    not_start = network.add_elementwise(constant_1.get_output(0), start0,
                                        trt.ElementWiseOperation.SUB)
    not_start.set_output_type(0, trt_dtype)
    internal_state = network.add_elementwise(in_state0, not_start.get_output(0),
                                             trt.ElementWiseOperation.PROD)
    out0 = network.add_elementwise(internal_state.get_output(0), in0,
                                   trt.ElementWiseOperation.SUM)
    out0_state = network.add_elementwise(internal_state.get_output(0), in0,
                                         trt.ElementWiseOperation.SUM)

    out0.get_output(0).name = "OUTPUT"
    network.mark_output(out0.get_output(0))

    out0_state.get_output(0).name = "OUTPUT_STATE"
    network.mark_output(out0_state.get_output(0))

    min_shape = []
    opt_shape = []
    max_shape = []
    if max_batch != 0:
        min_shape = min_shape + [1]
        opt_shape = opt_shape + [max(1, max_batch)]
        max_shape = max_shape + [max(1, max_batch)]
    for i in shape:
        if i == -1:
            min_shape = min_shape + [1]
            opt_shape = opt_shape + [8]
            max_shape = max_shape + [32]
        else:
            min_shape = min_shape + [i]
            opt_shape = opt_shape + [i]
            max_shape = max_shape + [i]

    profile = builder.create_optimization_profile()
    profile.set_shape("INPUT", min_shape, opt_shape, max_shape)
    profile.set_shape("INPUT_STATE", min_shape, opt_shape, max_shape)
    if max_batch != 0:
        profile.set_shape("START", [1] + unit_shape, [max_batch] + unit_shape,
                          [max_batch] + unit_shape)
        profile.set_shape("READY", [1] + unit_shape, [max_batch] + unit_shape,
                          [max_batch] + unit_shape)
    else:
        profile.set_shape("START", unit_shape, unit_shape, unit_shape)
        profile.set_shape("READY", unit_shape, unit_shape, unit_shape)
    config = builder.create_builder_config()
    config.add_optimization_profile(profile)

    config.max_workspace_size = 1 << 20
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    model_name = tu.get_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_dynamic_rf_modelfile(models_dir, model_version, max_batch,
                                     dtype, shape):
    trt_dtype = np_to_trt_dtype(dtype)
    trt_memory_format = trt.TensorFormat.LINEAR

    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)

    # EXPLICIT_BATCH must be used when the dimension is variable
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

    unit_shape = ([1] * len(shape))
    if max_batch != 0:
        in0 = network.add_input("INPUT", trt_dtype, [-1] + shape)
        start0 = network.add_input("START", trt_dtype, [-1] + unit_shape)
        ready0 = network.add_input("READY", trt_dtype, [-1] + unit_shape)
        in_state0 = network.add_input("INPUT_STATE", trt_dtype, [-1] + shape)
    else:
        in0 = network.add_input("INPUT", trt_dtype, shape)
        start0 = network.add_input("START", trt_dtype, unit_shape)
        ready0 = network.add_input("READY", trt_dtype, unit_shape)
        in_state0 = network.add_input("INPUT_STATE", trt_dtype, [-1] + shape)

    # Append the dimension by 1 so that broadcasting works properly
    constant_1_data = trt.Weights(np.ones(unit_shape + [1], dtype=dtype))
    constant_1 = network.add_constant(unit_shape + [1], constant_1_data)
    not_start = network.add_elementwise(constant_1.get_output(0), start0,
                                        trt.ElementWiseOperation.SUB)
    not_start.set_output_type(0, trt_dtype)
    internal_state = network.add_elementwise(in_state0, not_start.get_output(0),
                                             trt.ElementWiseOperation.PROD)
    out0 = network.add_elementwise(internal_state.get_output(0), in0,
                                   trt.ElementWiseOperation.SUM)
    out0_state = network.add_elementwise(internal_state.get_output(0), in0,
                                         trt.ElementWiseOperation.SUM)
    out0.get_output(0).name = "OUTPUT"
    network.mark_output(out0.get_output(0))

    out0_state.get_output(0).name = "OUTPUT_STATE"
    network.mark_output(out0_state.get_output(0))

    out0.get_output(0).dtype = trt_dtype
    out0_state.get_output(0).dtype = trt_dtype

    in0.allowed_formats = 1 << int(trt_memory_format)
    start0.allowed_formats = 1 << int(trt_memory_format)
    ready0.allowed_formats = 1 << int(trt_memory_format)
    out0.get_output(0).allowed_formats = 1 << int(trt_memory_format)

    if (trt_dtype == trt.int8):
        in0.dynamic_range = (-128.0, 127.0)
        in_state0.dynamic_range = (-128.0, 127.0)
        out0.dynamic_range = (-128.0, 127.0)
        start0.dynamic_range = (-128.0, 127.0)
        ready0.dynamic_range = (-128.0, 127.0)

    flags = 1 << int(trt.BuilderFlag.STRICT_TYPES)

    if (trt_dtype == trt.int8):
        flags |= 1 << int(trt.BuilderFlag.INT8)
    elif (trt_dtype == trt.float16):
        flags |= 1 << int(trt.BuilderFlag.FP16)

    min_shape = []
    opt_shape = []
    max_shape = []
    if max_batch != 0:
        min_shape = min_shape + [1]
        opt_shape = opt_shape + [max(1, max_batch)]
        max_shape = max_shape + [max(1, max_batch)]
    for i in shape:
        if i == -1:
            min_shape = min_shape + [1]
            opt_shape = opt_shape + [8]
            max_shape = max_shape + [32]
        else:
            min_shape = min_shape + [i]
            opt_shape = opt_shape + [i]
            max_shape = max_shape + [i]

    profile = builder.create_optimization_profile()
    profile.set_shape("INPUT", min_shape, opt_shape, max_shape)
    profile.set_shape("INPUT_STATE", min_shape, opt_shape, max_shape)
    if max_batch != 0:
        profile.set_shape("START", [1] + unit_shape, [max_batch] + unit_shape,
                          [max_batch] + unit_shape)
        profile.set_shape("READY", [1] + unit_shape, [max_batch] + unit_shape,
                          [max_batch] + unit_shape)
    else:
        profile.set_shape("START", unit_shape, unit_shape, unit_shape)
        profile.set_shape("READY", unit_shape, unit_shape, unit_shape)

    config = builder.create_builder_config()
    config.flags = flags
    config.add_optimization_profile(profile)

    config.max_workspace_size = 1 << 20
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    model_name = tu.get_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_modelfile(models_dir, model_version, max_batch, dtype, shape):

    if not tu.validate_for_trt_model(dtype, dtype, dtype, shape, shape, shape):
        return

    if dtype != np.float32:
        if (not tu.shape_is_fixed(shape)):
            create_plan_dynamic_rf_modelfile(models_dir, model_version,
                                             max_batch, dtype, shape)
        else:
            create_plan_fixed_rf_modelfile(models_dir, model_version, max_batch,
                                           dtype, shape)
    else:
        if (not tu.shape_is_fixed(shape)):
            create_plan_dynamic_modelfile(models_dir, model_version, max_batch,
                                          dtype, shape)
        else:
            create_plan_fixed_modelfile(models_dir, model_version, max_batch,
                                        dtype, shape)


def create_plan_modelconfig(models_dir, model_version, max_batch, dtype, shape):

    if not tu.validate_for_trt_model(dtype, dtype, dtype, shape, shape, shape):
        return

    model_name = tu.get_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "tensorrt_plan"
max_batch_size: {}
sequence_batching {{
  max_sequence_idle_microseconds: 5000000
  control_input [
    {{
      name: "START"
      control [
        {{
          kind: CONTROL_SEQUENCE_START
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "READY"
      control [
        {{
          kind: CONTROL_SEQUENCE_READY
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }}
  ]
  state [
    {{
      input_name: "INPUT_STATE"
      output_name: "OUTPUT_STATE"
      data_type: {dtype}
      dims: {shape}
    }} 
  ]
}}
input [
  {{
    name: "INPUT"
    data_type: {dtype}
    dims: [ {shape} ]
  }}
]
output [
  {{
    name: "OUTPUT"
    data_type: {dtype}
    dims: [ {shape} ]
  }}
]
instance_group [
  {{
    kind: KIND_GPU
  }}
]
'''.format(model_name,
           max_batch,
           "int32" if dtype == np.int32 else "fp32",
           "int32" if dtype == np.int32 else "fp32",
           dtype=np_to_model_dtype(dtype),
           shape=tu.shape_to_dims_str(shape))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_models(models_dir, dtype, shape, initial_state, no_batch=True):
    model_version = 1

    if FLAGS.onnx:
        create_onnx_modelconfig(models_dir, model_version, 8, dtype, shape,
                                initial_state)
        create_onnx_modelfile(models_dir, model_version, 8, dtype, shape,
                              initial_state)
        if no_batch:
            create_onnx_modelconfig(models_dir, model_version, 0, dtype, shape,
                                    initial_state)
            create_onnx_modelfile(models_dir, model_version, 0, dtype, shape,
                                  initial_state)

    if FLAGS.tensorrt:
        if dtype == bool:
            return
        suffix = []
        if dtype == np.int8:
            suffix = [1, 1]

        create_plan_modelconfig(models_dir, model_version, 8, dtype,
                                shape + suffix)
        create_plan_modelfile(models_dir, model_version, 8, dtype,
                              shape + suffix)
        if no_batch:
            create_plan_modelconfig(models_dir, model_version, 0, dtype,
                                    shape + suffix)
            create_plan_modelfile(models_dir, model_version, 0, dtype,
                                  shape + suffix)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--models_dir',
                        type=str,
                        required=True,
                        help='Top-level model directory')
    parser.add_argument('--graphdef',
                        required=False,
                        action='store_true',
                        help='Generate GraphDef models')
    parser.add_argument('--savedmodel',
                        required=False,
                        action='store_true',
                        help='Generate SavedModel models')
    parser.add_argument('--tensorrt',
                        required=False,
                        action='store_true',
                        help='Generate TensorRT PLAN models')
    parser.add_argument('--initial-state',
                        required=False,
                        choices=['zero', 'file'],
                        help='Generate models that rely on initial state.')
    parser.add_argument(
        '--tensorrt-shape-io',
        required=False,
        action='store_true',
        help='Generate TensorRT PLAN models w/ shape tensor i/o')
    parser.add_argument('--onnx',
                        required=False,
                        action='store_true',
                        help='Generate Onnx models')
    parser.add_argument(
        '--onnx_opset',
        type=int,
        required=False,
        default=0,
        help='Opset used for Onnx models. Default is to use ONNXRT default')
    parser.add_argument('--libtorch',
                        required=False,
                        action='store_true',
                        help='Generate Pytorch LibTorch models')
    parser.add_argument('--openvino',
                        required=False,
                        action='store_true',
                        help='Generate OpenVino models')
    parser.add_argument('--variable',
                        required=False,
                        action='store_true',
                        help='Used variable-shape tensors for input/output')
    parser.add_argument('--ensemble',
                        required=False,
                        action='store_true',
                        help='Generate ensemble models against the models' +
                        ' in all platforms. Note that the models generated' +
                        ' are not completed.')
    FLAGS, unparsed = parser.parse_known_args()

    if FLAGS.onnx:
        import onnx

    if FLAGS.tensorrt:
        import tensorrt as trt

    import test_util as tu

    # Tests with models that accept fixed-shape input/output tensors
    if not FLAGS.variable:
        create_models(FLAGS.models_dir, np.float32, [
            1,
        ], FLAGS.initial_state)
        create_models(FLAGS.models_dir, np.int32, [
            1,
        ], FLAGS.initial_state)
        create_models(FLAGS.models_dir, np_dtype_string, [
            1,
        ], FLAGS.initial_state)
        create_models(FLAGS.models_dir, bool, [
            1,
        ], FLAGS.initial_state)

    # Tests with models that accept variable-shape input/output tensors
    if FLAGS.variable:
        create_models(FLAGS.models_dir, np.int32, [
            -1,
        ], FLAGS.initial_state, False)
        create_models(FLAGS.models_dir, np.float32, [
            -1,
        ], FLAGS.initial_state, False)
        create_models(FLAGS.models_dir, np_dtype_string, [
            -1,
        ], FLAGS.initial_state, False)
        create_models(FLAGS.models_dir, bool, [
            -1,
        ], FLAGS.initial_state, False)


# qa/common/gen_qa_models.py

# Copyright 2018-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
from builtins import range
import os
import numpy as np
import gen_ensemble_model_utils as emu

FLAGS = None
np_dtype_string = np.dtype(object)
from typing import List, Tuple


def np_to_model_dtype(np_dtype):
    if np_dtype == bool:
        return "TYPE_BOOL"
    elif np_dtype == np.int8:
        return "TYPE_INT8"
    elif np_dtype == np.int16:
        return "TYPE_INT16"
    elif np_dtype == np.int32:
        return "TYPE_INT32"
    elif np_dtype == np.int64:
        return "TYPE_INT64"
    elif np_dtype == np.uint8:
        return "TYPE_UINT8"
    elif np_dtype == np.uint16:
        return "TYPE_UINT16"
    elif np_dtype == np.float16:
        return "TYPE_FP16"
    elif np_dtype == np.float32:
        return "TYPE_FP32"
    elif np_dtype == np.float64:
        return "TYPE_FP64"
    elif np_dtype == np_dtype_string:
        return "TYPE_STRING"
    return None


def np_to_tf_dtype(np_dtype):
    if np_dtype == bool:
        return tf.bool
    elif np_dtype == np.int8:
        return tf.int8
    elif np_dtype == np.int16:
        return tf.int16
    elif np_dtype == np.int32:
        return tf.int32
    elif np_dtype == np.int64:
        return tf.int64
    elif np_dtype == np.uint8:
        return tf.uint8
    elif np_dtype == np.uint16:
        return tf.uint16
    elif np_dtype == np.float16:
        return tf.float16
    elif np_dtype == np.float32:
        return tf.float32
    elif np_dtype == np.float64:
        return tf.float64
    elif np_dtype == np_dtype_string:
        return tf.string
    return None


def np_to_trt_dtype(np_dtype):
    if np_dtype == bool:
        return trt.bool
    elif np_dtype == np.int8:
        return trt.int8
    elif np_dtype == np.int32:
        return trt.int32
    elif np_dtype == np.uint8:
        return trt.uint8
    elif np_dtype == np.float16:
        return trt.float16
    elif np_dtype == np.float32:
        return trt.float32
    return None


def np_to_onnx_dtype(np_dtype):
    if np_dtype == bool:
        return onnx.TensorProto.BOOL
    elif np_dtype == np.int8:
        return onnx.TensorProto.INT8
    elif np_dtype == np.int16:
        return onnx.TensorProto.INT16
    elif np_dtype == np.int32:
        return onnx.TensorProto.INT32
    elif np_dtype == np.int64:
        return onnx.TensorProto.INT64
    elif np_dtype == np.uint8:
        return onnx.TensorProto.UINT8
    elif np_dtype == np.uint16:
        return onnx.TensorProto.UINT16
    elif np_dtype == np.float16:
        return onnx.TensorProto.FLOAT16
    elif np_dtype == np.float32:
        return onnx.TensorProto.FLOAT
    elif np_dtype == np.float64:
        return onnx.TensorProto.DOUBLE
    elif np_dtype == np_dtype_string:
        return onnx.TensorProto.STRING
    return None


def np_to_torch_dtype(np_dtype):
    if np_dtype == bool:
        return torch.bool
    elif np_dtype == np.int8:
        return torch.int8
    elif np_dtype == np.int16:
        return torch.int16
    elif np_dtype == np.int32:
        return torch.int
    elif np_dtype == np.int64:
        return torch.long
    elif np_dtype == np.uint8:
        return torch.uint8
    elif np_dtype == np.uint16:
        return None  # Not supported in Torch
    elif np_dtype == np.float16:
        return None
    elif np_dtype == np.float32:
        return torch.float
    elif np_dtype == np.float64:
        return torch.double
    elif np_dtype == np_dtype_string:
        return List[str]


def create_graphdef_modelfile(models_dir,
                              max_batch,
                              model_version,
                              input_shape,
                              output0_shape,
                              output1_shape,
                              input_dtype,
                              output0_dtype,
                              output1_dtype,
                              swap=False):

    if not tu.validate_for_tf_model(input_dtype, output0_dtype, output1_dtype,
                                    input_shape, output0_shape, output1_shape):
        return

    tf_input_dtype = np_to_tf_dtype(input_dtype)
    tf_output0_dtype = np_to_tf_dtype(output0_dtype)
    tf_output1_dtype = np_to_tf_dtype(output1_dtype)

    # Create the model. If non-batching then don't include the batch
    # dimension.
    tf.reset_default_graph()
    if max_batch == 0:
        in0 = tf.placeholder(tf_input_dtype, tu.shape_to_tf_shape(input_shape),
                             "INPUT0")
        in1 = tf.placeholder(tf_input_dtype, tu.shape_to_tf_shape(input_shape),
                             "INPUT1")
    else:
        in0 = tf.placeholder(tf_input_dtype, [
            None,
        ] + tu.shape_to_tf_shape(input_shape), "INPUT0")
        in1 = tf.placeholder(tf_input_dtype, [
            None,
        ] + tu.shape_to_tf_shape(input_shape), "INPUT1")

    # If the input is a string, then convert each string to the
    # equivalent int32 value.
    if tf_input_dtype == tf.string:
        in0 = tf.strings.to_number(in0, tf.int32)
        in1 = tf.strings.to_number(in1, tf.int32)

    add = tf.add(in0, in1, "ADD")
    sub = tf.subtract(in0, in1, "SUB")

    # Cast or convert result to the output dtype.
    if tf_output0_dtype == tf.string:
        cast0 = tf.dtypes.as_string(add if not swap else sub, name="TOSTR0")
    else:
        cast0 = tf.cast(add if not swap else sub, tf_output0_dtype, "CAST0")

    if tf_output1_dtype == tf.string:
        cast1 = tf.dtypes.as_string(sub if not swap else add, name="TOSTR1")
    else:
        cast1 = tf.cast(sub if not swap else add, tf_output1_dtype, "CAST1")

    out0 = tf.identity(cast0, "OUTPUT0")
    out1 = tf.identity(cast1, "OUTPUT1")

    # Use a different model name for the non-batching variant
    model_name = tu.get_model_name(
        "graphdef_nobatch" if max_batch == 0 else "graphdef", input_dtype,
        output0_dtype, output1_dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with tf.Session() as sess:
        graph_io.write_graph(sess.graph.as_graph_def(),
                             model_version_dir,
                             "model.graphdef",
                             as_text=False)


def create_graphdef_modelconfig(models_dir, max_batch, model_version,
                                input_shape, output0_shape, output1_shape,
                                input_dtype, output0_dtype, output1_dtype,
                                output0_label_cnt, version_policy):

    if not tu.validate_for_tf_model(input_dtype, output0_dtype, output1_dtype,
                                    input_shape, output0_shape, output1_shape):
        return

    # Unpack version policy
    version_policy_str = "{ latest { num_versions: 1 }}"
    if version_policy is not None:
        type, val = version_policy
        if type == 'latest':
            version_policy_str = "{{ latest {{ num_versions: {} }}}}".format(
                val)
        elif type == 'specific':
            version_policy_str = "{{ specific {{ versions: {} }}}}".format(val)
        else:
            version_policy_str = "{ all { }}"

    # Use a different model name for the non-batching variant
    model_name = tu.get_model_name(
        "graphdef_nobatch" if max_batch == 0 else "graphdef", input_dtype,
        output0_dtype, output1_dtype)
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "tensorflow_graphdef"
max_batch_size: {}
version_policy: {}
input [
  {{
    name: "INPUT0"
    data_type: {}
    dims: [ {} ]
  }},
  {{
    name: "INPUT1"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT0"
    data_type: {}
    dims: [ {} ]
    label_filename: "output0_labels.txt"
  }},
  {{
    name: "OUTPUT1"
    data_type: {}
    dims: [ {} ]
  }}
]
'''.format(model_name, max_batch, version_policy_str,
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(output0_dtype),
           tu.shape_to_dims_str(output0_shape),
           np_to_model_dtype(output1_dtype),
           tu.shape_to_dims_str(output1_shape))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)

    with open(config_dir + "/output0_labels.txt", "w") as lfile:
        for l in range(output0_label_cnt):
            lfile.write("label" + str(l) + "\n")


def create_savedmodel_modelfile(models_dir,
                                max_batch,
                                model_version,
                                input_shape,
                                output0_shape,
                                output1_shape,
                                input_dtype,
                                output0_dtype,
                                output1_dtype,
                                swap=False):

    if not tu.validate_for_tf_model(input_dtype, output0_dtype, output1_dtype,
                                    input_shape, output0_shape, output1_shape):
        return

    tf_input_dtype = np_to_tf_dtype(input_dtype)
    tf_output0_dtype = np_to_tf_dtype(output0_dtype)
    tf_output1_dtype = np_to_tf_dtype(output1_dtype)

    # Create the model. If non-batching then don't include the batch
    # dimension.
    tf.reset_default_graph()
    if max_batch == 0:
        in0 = tf.placeholder(tf_input_dtype, tu.shape_to_tf_shape(input_shape),
                             "TENSOR_INPUT0")
        in1 = tf.placeholder(tf_input_dtype, tu.shape_to_tf_shape(input_shape),
                             "TENSOR_INPUT1")
    else:
        in0 = tf.placeholder(tf_input_dtype, [
            None,
        ] + tu.shape_to_tf_shape(input_shape), "TENSOR_INPUT0")
        in1 = tf.placeholder(tf_input_dtype, [
            None,
        ] + tu.shape_to_tf_shape(input_shape), "TENSOR_INPUT1")

    # If the input is a string, then convert each string to the
    # equivalent float value.
    if tf_input_dtype == tf.string:
        in0 = tf.strings.to_number(in0, tf.int32)
        in1 = tf.strings.to_number(in1, tf.int32)

    add = tf.add(in0, in1, "ADD")
    sub = tf.subtract(in0, in1, "SUB")

    # Cast or convert result to the output dtype.
    if tf_output0_dtype == tf.string:
        cast0 = tf.dtypes.as_string(add if not swap else sub, name="TOSTR0")
    else:
        cast0 = tf.cast(add if not swap else sub, tf_output0_dtype, "CAST0")

    if tf_output1_dtype == tf.string:
        cast1 = tf.dtypes.as_string(sub if not swap else add, name="TOSTR1")
    else:
        cast1 = tf.cast(sub if not swap else add, tf_output1_dtype, "CAST1")

    out0 = tf.identity(cast0, "TENSOR_OUTPUT0")
    out1 = tf.identity(cast1, "TENSOR_OUTPUT1")

    # Use a different model name for the non-batching variant
    model_name = tu.get_model_name(
        "savedmodel_nobatch" if max_batch == 0 else "savedmodel", input_dtype,
        output0_dtype, output1_dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with tf.Session() as sess:
        input0_tensor = tf.get_default_graph().get_tensor_by_name(
            "TENSOR_INPUT0:0")
        input1_tensor = tf.get_default_graph().get_tensor_by_name(
            "TENSOR_INPUT1:0")
        output0_tensor = tf.get_default_graph().get_tensor_by_name(
            "TENSOR_OUTPUT0:0")
        output1_tensor = tf.get_default_graph().get_tensor_by_name(
            "TENSOR_OUTPUT1:0")
        tf.saved_model.simple_save(sess,
                                   model_version_dir + "/model.savedmodel",
                                   inputs={
                                       "INPUT0": input0_tensor,
                                       "INPUT1": input1_tensor
                                   },
                                   outputs={
                                       "OUTPUT0": output0_tensor,
                                       "OUTPUT1": output1_tensor
                                   })


def create_savedmodel_modelconfig(models_dir, max_batch, model_version,
                                  input_shape, output0_shape, output1_shape,
                                  input_dtype, output0_dtype, output1_dtype,
                                  output0_label_cnt, version_policy):

    if not tu.validate_for_tf_model(input_dtype, output0_dtype, output1_dtype,
                                    input_shape, output0_shape, output1_shape):
        return

    # Unpack version policy
    version_policy_str = "{ latest { num_versions: 1 }}"
    if version_policy is not None:
        type, val = version_policy
        if type == 'latest':
            version_policy_str = "{{ latest {{ num_versions: {} }}}}".format(
                val)
        elif type == 'specific':
            version_policy_str = "{{ specific {{ versions: {} }}}}".format(val)
        else:
            version_policy_str = "{ all { }}"

    # Use a different model name for the non-batching variant
    model_name = tu.get_model_name(
        "savedmodel_nobatch" if max_batch == 0 else "savedmodel", input_dtype,
        output0_dtype, output1_dtype)
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "tensorflow_savedmodel"
max_batch_size: {}
version_policy: {}
input [
  {{
    name: "INPUT0"
    data_type: {}
    dims: [ {} ]
  }},
  {{
    name: "INPUT1"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT0"
    data_type: {}
    dims: [ {} ]
    label_filename: "output0_labels.txt"
  }},
  {{
    name: "OUTPUT1"
    data_type: {}
    dims: [ {} ]
  }}
]
'''.format(model_name, max_batch, version_policy_str,
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(output0_dtype),
           tu.shape_to_dims_str(output0_shape),
           np_to_model_dtype(output1_dtype),
           tu.shape_to_dims_str(output1_shape))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)

    with open(config_dir + "/output0_labels.txt", "w") as lfile:
        for l in range(output0_label_cnt):
            lfile.write("label" + str(l) + "\n")


def create_plan_dynamic_rf_modelfile(models_dir, max_batch, model_version,
                                     input_shape, output0_shape, output1_shape,
                                     input_dtype, output0_dtype, output1_dtype,
                                     swap, min_dim, max_dim):
    trt_input_dtype = np_to_trt_dtype(input_dtype)
    trt_output0_dtype = np_to_trt_dtype(output0_dtype)
    trt_output1_dtype = np_to_trt_dtype(output1_dtype)
    trt_memory_format = trt.TensorFormat.LINEAR

    # Create the model
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    if max_batch == 0:
        input_with_batchsize = [i for i in input_shape]
    else:
        input_with_batchsize = [-1] + [i for i in input_shape]

    in0 = network.add_input("INPUT0", trt_input_dtype, input_with_batchsize)
    in1 = network.add_input("INPUT1", trt_input_dtype, input_with_batchsize)

    # TRT uint8 cannot be used to represent quantized floating-point value yet
    # uint8 must be converted to float16 or float32 before any operation
    # FIXME: Remove support check when jetson supports TRT 8.5 (DLIS-4256)
    if tu.support_trt_uint8():
        if trt_input_dtype == trt.uint8:
            in0_cast = network.add_identity(in0)
            in0_cast.set_output_type(0, trt.float32)
            in0 = in0_cast.get_output(0)
            in1_cast = network.add_identity(in1)
            in1_cast.set_output_type(0, trt.float32)
            in1 = in1_cast.get_output(0)

    add = network.add_elementwise(in0, in1, trt.ElementWiseOperation.SUM)
    sub = network.add_elementwise(in0, in1, trt.ElementWiseOperation.SUB)
    out0 = add if not swap else sub
    out1 = sub if not swap else add

    # uint8 conversion after operations
    # FIXME: Remove support check when jetson supports TRT 8.5 (DLIS-4256)
    if tu.support_trt_uint8():
        if trt_output0_dtype == trt.uint8:
            out0 = network.add_identity(out0.get_output(0))
            out0.set_output_type(0, trt.uint8)
        if trt_output1_dtype == trt.uint8:
            out1 = network.add_identity(out1.get_output(0))
            out1.set_output_type(0, trt.uint8)

    out0.get_output(0).name = "OUTPUT0"
    out1.get_output(0).name = "OUTPUT1"
    network.mark_output(out0.get_output(0))
    network.mark_output(out1.get_output(0))

    out0.get_output(0).dtype = trt_output0_dtype
    out1.get_output(0).dtype = trt_output1_dtype

    in0.allowed_formats = 1 << int(trt_memory_format)
    in1.allowed_formats = 1 << int(trt_memory_format)
    out0.get_output(0).allowed_formats = 1 << int(trt_memory_format)
    out1.get_output(0).allowed_formats = 1 << int(trt_memory_format)

    if (trt_input_dtype == trt.int8):
        in0.dynamic_range = (-128.0, 127.0)
        in1.dynamic_range = (-128.0, 127.0)
    if (trt_output0_dtype == trt.int8):
        out0.get_output(0).dynamic_range = (-128.0, 127.0)
    if (trt_output1_dtype == trt.int8):
        out1.get_output(0).dynamic_range = (-128.0, 127.0)

    min_shape = []
    opt_shape = []
    max_shape = []
    if max_batch != 0:
        min_shape = min_shape + [1]
        opt_shape = opt_shape + [max(1, max_batch)]
        max_shape = max_shape + [max(1, max_batch)]
    for i in input_shape:
        if i == -1:
            min_shape = min_shape + [min_dim]
            opt_shape = opt_shape + [int((max_dim + min_dim) / 2)]
            max_shape = max_shape + [max_dim]
        else:
            min_shape = min_shape + [i]
            opt_shape = opt_shape + [i]
            max_shape = max_shape + [i]

    profile = builder.create_optimization_profile()
    profile.set_shape("INPUT0", min_shape, opt_shape, max_shape)
    profile.set_shape("INPUT1", min_shape, opt_shape, max_shape)
    flags = 1 << int(trt.BuilderFlag.STRICT_TYPES)
    datatype_set = set([trt_input_dtype, trt_output0_dtype, trt_output1_dtype])
    for dt in datatype_set:
        if (dt == trt.int8):
            flags |= 1 << int(trt.BuilderFlag.INT8)
        elif (dt == trt.float16):
            flags |= 1 << int(trt.BuilderFlag.FP16)
    config = builder.create_builder_config()
    config.flags = flags
    config.add_optimization_profile(profile)
    config.max_workspace_size = 1 << 20
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    # Use a different model name for different kinds of models
    model_name = tu.get_model_name("plan_nobatch" if max_batch == 0 else "plan",
                                   input_dtype, output0_dtype, output1_dtype)
    if min_dim != 1 or max_dim != 32:
        model_name = "{}-{}-{}".format(model_name, min_dim, max_dim)

    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_dynamic_modelfile(models_dir, max_batch, model_version,
                                  input_shape, output0_shape, output1_shape,
                                  input_dtype, output0_dtype, output1_dtype,
                                  swap, min_dim, max_dim):
    trt_input_dtype = np_to_trt_dtype(input_dtype)
    trt_output0_dtype = np_to_trt_dtype(output0_dtype)
    trt_output1_dtype = np_to_trt_dtype(output1_dtype)

    # Create the model
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    if max_batch == 0:
        input_with_batchsize = [i for i in input_shape]
    else:
        input_with_batchsize = [-1] + [i for i in input_shape]

    in0 = network.add_input("INPUT0", trt_input_dtype, input_with_batchsize)
    in1 = network.add_input("INPUT1", trt_input_dtype, input_with_batchsize)
    add = network.add_elementwise(in0, in1, trt.ElementWiseOperation.SUM)
    sub = network.add_elementwise(in0, in1, trt.ElementWiseOperation.SUB)

    out0 = add if not swap else sub
    out1 = sub if not swap else add

    out0.get_output(0).name = "OUTPUT0"
    out1.get_output(0).name = "OUTPUT1"
    network.mark_output(out0.get_output(0))
    network.mark_output(out1.get_output(0))

    min_shape = []
    opt_shape = []
    max_shape = []
    for i in input_shape:
        if i == -1:
            min_shape = min_shape + [min_dim]
            opt_shape = opt_shape + [int((max_dim + min_dim) / 2)]
            max_shape = max_shape + [max_dim]
        else:
            min_shape = min_shape + [i]
            opt_shape = opt_shape + [i]
            max_shape = max_shape + [i]

    config = builder.create_builder_config()
    # create multiple profiles with same shape for testing
    # with decreasing batch sizes
    profile = []
    for i in range(4):
        profile.append(builder.create_optimization_profile())
        if max_batch == 0:
            profile[i].set_shape("INPUT0", min_shape, opt_shape, max_shape)
            profile[i].set_shape("INPUT1", min_shape, opt_shape, max_shape)
        else:
            bs = [max_batch - i if max_batch > i else 1]
            opt_bs = [1 + i if 1 + i < max_batch - 1 else max_batch - 1]
            # Hardcoded 'max_shape[0] += 1' in default profile for
            # L0_trt_dynamic_shape, to differentiate whether default profile
            # is used if no profile is specified
            max_shape_override = max_shape
            if i == 0 and (min_dim == 1 and max_dim == 32):
                max_shape_override[0] += 1

            profile[i].set_shape("INPUT0", [1] + min_shape, opt_bs + opt_shape,
                                 bs + max_shape_override)
            profile[i].set_shape("INPUT1", [1] + min_shape, opt_bs + opt_shape,
                                 bs + max_shape_override)
        config.add_optimization_profile(profile[i])
    # some profiles with non-one min shape for first dim to test autofiller
    for i in range(2):
        profile.append(builder.create_optimization_profile())
        if max_batch == 0:
            profile[i + 4].set_shape("INPUT0", min_shape, opt_shape, max_shape)
            profile[i + 4].set_shape("INPUT1", min_shape, opt_shape, max_shape)
        else:
            profile[i + 4].set_shape("INPUT0", [5 + i] + min_shape,
                                     [6] + opt_shape, [max_batch] + max_shape)
            profile[i + 4].set_shape("INPUT1", [5 + i] + min_shape,
                                     [6] + opt_shape, [max_batch] + max_shape)
        config.add_optimization_profile(profile[i + 4])
    # Will repeat another profile with same min and max shapes as the first profile to test non-zero profile
    # for infer_variable test.
    profile.append(builder.create_optimization_profile())
    if max_batch == 0:
        profile[6].set_shape("INPUT0", min_shape, opt_shape, max_shape)
        profile[6].set_shape("INPUT1", min_shape, opt_shape, max_shape)
    else:
        profile[6].set_shape("INPUT0", [1] + min_shape, [1] + opt_shape,
                             [max_batch] + max_shape)
        profile[6].set_shape("INPUT1", [1] + min_shape, [1] + opt_shape,
                             [max_batch] + max_shape)
    config.add_optimization_profile(profile[6])

    # Will add some profiles with static shapes to test the cases where min_shape=opt_shape=max_shape
    for i in range(3):
        profile.append(builder.create_optimization_profile())
        if max_batch == 0:
            static_shape = max_shape
            profile[7 + i].set_shape("INPUT0", static_shape, static_shape,
                                     static_shape)
            profile[7 + i].set_shape("INPUT1", static_shape, static_shape,
                                     static_shape)
        else:
            # Skipping alternate batch sizes for testing unsupported batches in L0_trt_dynamic_shape.
            full_static_shape = [1 + (2 * i)] + max_shape
            profile[7 + i].set_shape("INPUT0", full_static_shape,
                                     full_static_shape, full_static_shape)
            profile[7 + i].set_shape("INPUT1", full_static_shape,
                                     full_static_shape, full_static_shape)
        config.add_optimization_profile(profile[7 + i])

    # Add profiles where each profile supports a specific batch size
    if max_batch != 0:
        for i in range(max_batch):
            profile.append(builder.create_optimization_profile())
            profile[10 + i].set_shape("INPUT0", [1 + i] + min_shape,
                                      [1 + i] + opt_shape, [1 + i] + max_shape)
            profile[10 + i].set_shape("INPUT1", [1 + i] + min_shape,
                                      [1 + i] + opt_shape, [1 + i] + max_shape)
            config.add_optimization_profile(profile[10 + i])

    config.max_workspace_size = 1 << 20
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    # Use a different model name for different kinds of models
    model_name = tu.get_model_name("plan_nobatch" if max_batch == 0 else "plan",
                                   input_dtype, output0_dtype, output1_dtype)
    if min_dim != 1 or max_dim != 32:
        model_name = "{}-{}-{}".format(model_name, min_dim, max_dim)

    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_fixed_rf_modelfile(models_dir, max_batch, model_version,
                                   input_shape, output0_shape, output1_shape,
                                   input_dtype, output0_dtype, output1_dtype,
                                   swap):
    trt_input_dtype = np_to_trt_dtype(input_dtype)
    trt_output0_dtype = np_to_trt_dtype(output0_dtype)
    trt_output1_dtype = np_to_trt_dtype(output1_dtype)
    trt_memory_format = trt.TensorFormat.LINEAR

    # Create the model
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network()
    in0 = network.add_input("INPUT0", trt_input_dtype, input_shape)
    in1 = network.add_input("INPUT1", trt_input_dtype, input_shape)
    add = network.add_elementwise(in0, in1, trt.ElementWiseOperation.SUM)
    sub = network.add_elementwise(in0, in1, trt.ElementWiseOperation.SUB)

    out0 = add if not swap else sub
    out1 = sub if not swap else add

    out0.get_output(0).name = "OUTPUT0"
    out1.get_output(0).name = "OUTPUT1"
    network.mark_output(out0.get_output(0))
    network.mark_output(out1.get_output(0))

    out0.get_output(0).dtype = trt_output0_dtype
    out1.get_output(0).dtype = trt_output1_dtype

    in0.allowed_formats = 1 << int(trt_memory_format)
    in1.allowed_formats = 1 << int(trt_memory_format)
    out0.get_output(0).allowed_formats = 1 << int(trt_memory_format)
    out1.get_output(0).allowed_formats = 1 << int(trt_memory_format)

    if (trt_input_dtype == trt.int8):
        in0.dynamic_range = (-128.0, 127.0)
        in1.dynamic_range = (-128.0, 127.0)
    if (trt_output0_dtype == trt.int8):
        out0.get_output(0).dynamic_range = (-128.0, 127.0)
    if (trt_output1_dtype == trt.int8):
        out1.get_output(0).dynamic_range = (-128.0, 127.0)

    flags = 1 << int(trt.BuilderFlag.STRICT_TYPES)
    datatype_set = set([trt_input_dtype, trt_output0_dtype, trt_output1_dtype])
    for dt in datatype_set:
        if (dt == trt.int8):
            flags |= 1 << int(trt.BuilderFlag.INT8)
        elif (dt == trt.float16):
            flags |= 1 << int(trt.BuilderFlag.FP16)
    config = builder.create_builder_config()
    config.flags = flags
    config.max_workspace_size = 1 << 20
    builder.max_batch_size = max(1, max_batch)
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    model_name = tu.get_model_name("plan_nobatch" if max_batch == 0 else "plan",
                                   input_dtype, output0_dtype, output1_dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_fixed_modelfile(models_dir, max_batch, model_version,
                                input_shape, output0_shape, output1_shape,
                                input_dtype, output0_dtype, output1_dtype,
                                swap):
    trt_input_dtype = np_to_trt_dtype(input_dtype)
    trt_output0_dtype = np_to_trt_dtype(output0_dtype)
    trt_output1_dtype = np_to_trt_dtype(output1_dtype)

    # Create the model
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network()
    in0 = network.add_input("INPUT0", trt_input_dtype, input_shape)
    in1 = network.add_input("INPUT1", trt_input_dtype, input_shape)
    add = network.add_elementwise(in0, in1, trt.ElementWiseOperation.SUM)
    sub = network.add_elementwise(in0, in1, trt.ElementWiseOperation.SUB)

    out0 = add if not swap else sub
    out1 = sub if not swap else add

    out0.get_output(0).name = "OUTPUT0"
    out1.get_output(0).name = "OUTPUT1"
    network.mark_output(out0.get_output(0))
    network.mark_output(out1.get_output(0))

    config = builder.create_builder_config()
    config.max_workspace_size = 1 << 20
    builder.max_batch_size = max(1, max_batch)
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine
    del network

    model_name = tu.get_model_name("plan_nobatch" if max_batch == 0 else "plan",
                                   input_dtype, output0_dtype, output1_dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_modelfile(models_dir,
                          max_batch,
                          model_version,
                          input_shape,
                          output0_shape,
                          output1_shape,
                          input_dtype,
                          output0_dtype,
                          output1_dtype,
                          swap=False,
                          min_dim=1,
                          max_dim=32):

    if not tu.validate_for_trt_model(input_dtype, output0_dtype, output1_dtype,
                                     input_shape, output0_shape, output1_shape):
        return

    if input_dtype == np.uint8 or output0_dtype == np.uint8 or output1_dtype == np.uint8:
        # TRT uint8 cannot be used to represent quantized floating-point value yet
        # EXPLICIT_BATCH network and conversion are required to create models
        create_plan_dynamic_rf_modelfile(models_dir, max_batch, model_version,
                                         input_shape, output0_shape,
                                         output1_shape, input_dtype,
                                         output0_dtype, output1_dtype, swap,
                                         min_dim, max_dim)

    elif input_dtype != np.float32 or output0_dtype != np.float32 or output1_dtype != np.float32:
        if (not tu.shape_is_fixed(input_shape) or
                not tu.shape_is_fixed(output0_shape) or
                not tu.shape_is_fixed(output1_shape)):
            create_plan_dynamic_rf_modelfile(models_dir, max_batch,
                                             model_version, input_shape,
                                             output0_shape, output1_shape,
                                             input_dtype, output0_dtype,
                                             output1_dtype, swap, min_dim,
                                             max_dim)
        else:
            create_plan_fixed_rf_modelfile(models_dir, max_batch, model_version,
                                           input_shape, output0_shape,
                                           output1_shape, input_dtype,
                                           output0_dtype, output1_dtype, swap)

    else:
        if (not tu.shape_is_fixed(input_shape) or
                not tu.shape_is_fixed(output0_shape) or
                not tu.shape_is_fixed(output1_shape)):
            create_plan_dynamic_modelfile(models_dir, max_batch, model_version,
                                          input_shape, output0_shape,
                                          output1_shape, input_dtype,
                                          output0_dtype, output1_dtype, swap,
                                          min_dim, max_dim)
        else:
            create_plan_fixed_modelfile(models_dir, max_batch, model_version,
                                        input_shape, output0_shape,
                                        output1_shape, input_dtype,
                                        output0_dtype, output1_dtype, swap)


def create_plan_modelconfig(models_dir,
                            max_batch,
                            model_version,
                            input_shape,
                            output0_shape,
                            output1_shape,
                            input_dtype,
                            output0_dtype,
                            output1_dtype,
                            output0_label_cnt,
                            version_policy,
                            min_dim=1,
                            max_dim=32):

    if not tu.validate_for_trt_model(input_dtype, output0_dtype, output1_dtype,
                                     input_shape, output0_shape, output1_shape):
        return

    # Unpack version policy
    version_policy_str = "{ latest { num_versions: 1 }}"
    if version_policy is not None:
        type, val = version_policy
        if type == 'latest':
            version_policy_str = "{{ latest {{ num_versions: {} }}}}".format(
                val)
        elif type == 'specific':
            version_policy_str = "{{ specific {{ versions: {} }}}}".format(val)
        else:
            version_policy_str = "{ all { }}"

    # Use a different model name for different kinds of models
    model_name = tu.get_model_name("plan_nobatch" if max_batch == 0 else "plan",
                                   input_dtype, output0_dtype, output1_dtype)
    if min_dim != 1 or max_dim != 32:
        model_name = "{}-{}-{}".format(model_name, min_dim, max_dim)

    config_dir = models_dir + "/" + model_name
    if -1 in input_shape:
        # Selects the sixth profile for FP32 datatype
        # Note the min and max shapes of first and sixth
        # profile are identical.
        profile_index = 6 if input_dtype == np.float32 else 0
        config = '''
name: "{}"
platform: "tensorrt_plan"
max_batch_size: {}
version_policy: {}
input [
  {{
    name: "INPUT0"
    data_type: {}
    dims: [ {} ]
  }},
  {{
    name: "INPUT1"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT0"
    data_type: {}
    dims: [ {} ]
    label_filename: "output0_labels.txt"
   }},
  {{
    name: "OUTPUT1"
    data_type: {}
    dims: [ {} ]
  }}
]
instance_group [
  {{
      profile:"{}"
  }}
]
'''.format(model_name, max_batch, version_policy_str,
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(output0_dtype),
           tu.shape_to_dims_str(output0_shape),
           np_to_model_dtype(output1_dtype),
           tu.shape_to_dims_str(output1_shape), profile_index)
    else:
        config = '''
name: "{}"
platform: "tensorrt_plan"
max_batch_size: {}
version_policy: {}
input [
  {{
    name: "INPUT0"
    data_type: {}
    dims: [ {} ]
  }},
  {{
    name: "INPUT1"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT0"
    data_type: {}
    dims: [ {} ]
    label_filename: "output0_labels.txt"
   }},
  {{
    name: "OUTPUT1"
    data_type: {}
    dims: [ {} ]
  }}
]
'''.format(model_name, max_batch, version_policy_str,
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(output0_dtype),
           tu.shape_to_dims_str(output0_shape),
           np_to_model_dtype(output1_dtype),
           tu.shape_to_dims_str(output1_shape))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)

    with open(config_dir + "/output0_labels.txt", "w") as lfile:
        for l in range(output0_label_cnt):
            lfile.write("label" + str(l) + "\n")


def create_onnx_modelfile(models_dir,
                          max_batch,
                          model_version,
                          input_shape,
                          output0_shape,
                          output1_shape,
                          input_dtype,
                          output0_dtype,
                          output1_dtype,
                          swap=False):

    if not tu.validate_for_onnx_model(input_dtype, output0_dtype, output1_dtype,
                                      input_shape, output0_shape,
                                      output1_shape):
        return

    onnx_input_dtype = np_to_onnx_dtype(input_dtype)
    onnx_output0_dtype = np_to_onnx_dtype(output0_dtype)
    onnx_output1_dtype = np_to_onnx_dtype(output1_dtype)

    onnx_input_shape, idx = tu.shape_to_onnx_shape(input_shape, 0)
    onnx_output0_shape, idx = tu.shape_to_onnx_shape(input_shape, idx)
    onnx_output1_shape, idx = tu.shape_to_onnx_shape(input_shape, idx)

    # Create the model
    model_name = tu.get_model_name("onnx_nobatch" if max_batch == 0 else "onnx",
                                   input_dtype, output0_dtype, output1_dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    batch_dim = [] if max_batch == 0 else [None]

    in0 = onnx.helper.make_tensor_value_info("INPUT0", onnx_input_dtype,
                                             batch_dim + onnx_input_shape)
    in1 = onnx.helper.make_tensor_value_info("INPUT1", onnx_input_dtype,
                                             batch_dim + onnx_input_shape)

    out0 = onnx.helper.make_tensor_value_info("OUTPUT0", onnx_output0_dtype,
                                              batch_dim + onnx_output0_shape)
    out1 = onnx.helper.make_tensor_value_info("OUTPUT1", onnx_output1_dtype,
                                              batch_dim + onnx_output1_shape)

    internal_in0 = onnx.helper.make_node("Identity", ["INPUT0"], ["_INPUT0"])
    internal_in1 = onnx.helper.make_node("Identity", ["INPUT1"], ["_INPUT1"])

    # cast int8, int16 input to higer precision int as Onnx Add/Sub operator doesn't support those type
    # Also casting String data type to int32
    if ((onnx_input_dtype == onnx.TensorProto.INT8) or
        (onnx_input_dtype == onnx.TensorProto.INT16) or
        (onnx_input_dtype == onnx.TensorProto.STRING)):
        internal_in0 = onnx.helper.make_node("Cast", ["INPUT0"], ["_INPUT0"],
                                             to=onnx.TensorProto.INT32)
        internal_in1 = onnx.helper.make_node("Cast", ["INPUT1"], ["_INPUT1"],
                                             to=onnx.TensorProto.INT32)

    add = onnx.helper.make_node("Add", ["_INPUT0", "_INPUT1"],
                                ["CAST0" if not swap else "CAST1"])
    sub = onnx.helper.make_node("Sub", ["_INPUT0", "_INPUT1"],
                                ["CAST1" if not swap else "CAST0"])
    cast0 = onnx.helper.make_node("Cast", ["CAST0"], ["OUTPUT0"],
                                  to=onnx_output0_dtype)
    cast1 = onnx.helper.make_node("Cast", ["CAST1"], ["OUTPUT1"],
                                  to=onnx_output1_dtype)

    # Avoid cast from float16 to float16
    # (bug in Onnx Runtime, cast from float16 to float16 will become cast from float16 to float32)
    if onnx_input_dtype == onnx.TensorProto.FLOAT16:
        if onnx_output0_dtype == onnx_input_dtype:
            cast0 = onnx.helper.make_node("Identity", ["CAST0"], ["OUTPUT0"])
        if onnx_output1_dtype == onnx_input_dtype:
            cast1 = onnx.helper.make_node("Identity", ["CAST1"], ["OUTPUT1"])

    onnx_nodes = [internal_in0, internal_in1, add, sub, cast0, cast1]
    onnx_inputs = [in0, in1]
    onnx_outputs = [out0, out1]

    graph_proto = onnx.helper.make_graph(onnx_nodes, model_name, onnx_inputs,
                                         onnx_outputs)
    if FLAGS.onnx_opset > 0:
        model_opset = onnx.helper.make_operatorsetid("", FLAGS.onnx_opset)
        model_def = onnx.helper.make_model(graph_proto,
                                           producer_name="triton",
                                           opset_imports=[model_opset])
    else:
        model_def = onnx.helper.make_model(graph_proto, producer_name="triton")

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    onnx.save(model_def, model_version_dir + "/model.onnx")


def create_onnx_modelconfig(models_dir, max_batch, model_version, input_shape,
                            output0_shape, output1_shape, input_dtype,
                            output0_dtype, output1_dtype, output0_label_cnt,
                            version_policy):

    if not tu.validate_for_onnx_model(input_dtype, output0_dtype, output1_dtype,
                                      input_shape, output0_shape,
                                      output1_shape):
        return

    # Use a different model name for the non-batching variant
    model_name = tu.get_model_name("onnx_nobatch" if max_batch == 0 else "onnx",
                                   input_dtype, output0_dtype, output1_dtype)
    config_dir = models_dir + "/" + model_name

    # [TODO] move create_general_modelconfig() out of emu as it is general
    # enough for all backends to use
    config = emu.create_general_modelconfig(model_name,
                                            "onnxruntime_onnx",
                                            max_batch,
                                            emu.repeat(input_dtype, 2),
                                            emu.repeat(input_shape, 2),
                                            emu.repeat(None, 2),
                                            [output0_dtype, output1_dtype],
                                            [output0_shape, output1_shape],
                                            emu.repeat(None, 2),
                                            ["output0_labels.txt", None],
                                            version_policy=version_policy,
                                            force_tensor_number_suffix=True)

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)

    with open(config_dir + "/output0_labels.txt", "w") as lfile:
        for l in range(output0_label_cnt):
            lfile.write("label" + str(l) + "\n")


def create_libtorch_modelfile(models_dir,
                              max_batch,
                              model_version,
                              input_shape,
                              output0_shape,
                              output1_shape,
                              input_dtype,
                              output0_dtype,
                              output1_dtype,
                              swap=False):

    if not tu.validate_for_libtorch_model(
            input_dtype, output0_dtype, output1_dtype, input_shape,
            output0_shape, output1_shape, max_batch):
        return

    torch_output0_dtype = np_to_torch_dtype(output0_dtype)
    torch_output1_dtype = np_to_torch_dtype(output1_dtype)

    model_name = tu.get_model_name(
        "libtorch_nobatch" if max_batch == 0 else "libtorch", input_dtype,
        output0_dtype, output1_dtype)
    # handle for -1 (when variable) since can't create tensor with shape of [-1]
    input_shape = [abs(ips) for ips in input_shape]

    # Create the model
    if (input_dtype
            == np_dtype_string) and (output0_dtype != np_dtype_string) and (
                output1_dtype != np_dtype_string):

        class AddSubNet(nn.Module):

            def __init__(self, *args):
                self.output0_dtype = args[0][0]
                self.output1_dtype = args[0][1]
                self.swap = args[0][2]
                super(AddSubNet, self).__init__()

            def forward(self, INPUT0: List[str], INPUT1: List[str]):
                input0_int = torch.tensor([int(i) for i in INPUT0])
                input1_int = torch.tensor([int(i) for i in INPUT1])
                op0 = input0_int + input1_int if not self.swap else input0_int - input1_int
                op1 = input0_int - input1_int if not self.swap else input0_int + input1_int
                return op0.to(self.output0_dtype), op1.to(self.output1_dtype)
    elif (input_dtype
          == np_dtype_string) and (output0_dtype
                                   == np_dtype_string) and (output1_dtype
                                                            == np_dtype_string):

        class AddSubNet(nn.Module):

            def __init__(self, *args):
                self.output0_dtype = args[0][0]
                self.output1_dtype = args[0][1]
                self.swap = args[0][2]
                super(AddSubNet, self).__init__()

            def forward(self, INPUT0: List[str],
                        INPUT1: List[str]) -> Tuple[List[str], List[str]]:
                input0_int = torch.tensor([int(i) for i in INPUT0])
                input1_int = torch.tensor([int(i) for i in INPUT1])
                op0 = [
                    str(i.item())
                    for i in (input0_int +
                              input1_int if not self.swap else input0_int -
                              input1_int)
                ]
                op1 = [
                    str(i.item())
                    for i in (input0_int -
                              input1_int if not self.swap else input0_int +
                              input1_int)
                ]
                return op0, op1
    elif (input_dtype
          == np_dtype_string) and (output0_dtype == np_dtype_string) and (
              output1_dtype != np_dtype_string):

        class AddSubNet(nn.Module):

            def __init__(self, *args):
                self.output0_dtype = args[0][0]
                self.output1_dtype = args[0][1]
                self.swap = args[0][2]
                super(AddSubNet, self).__init__()

            def forward(self, INPUT0: List[str],
                        INPUT1: List[str]) -> Tuple[List[str], torch.Tensor]:
                input0_int = torch.tensor([int(i) for i in INPUT0])
                input1_int = torch.tensor([int(i) for i in INPUT1])
                op0 = [
                    str(i.item())
                    for i in (input0_int +
                              input1_int if not self.swap else input0_int -
                              input1_int)
                ]
                op1 = (input0_int -
                       input1_int if not self.swap else input0_int +
                       input1_int).to(self.output1_dtype)
                return op0, op1
    elif (input_dtype == np_dtype_string) and (
            output0_dtype != np_dtype_string) and (output1_dtype
                                                   == np_dtype_string):

        class AddSubNet(nn.Module):

            def __init__(self, *args):
                self.output0_dtype = args[0][0]
                self.output1_dtype = args[0][1]
                self.swap = args[0][2]
                super(AddSubNet, self).__init__()

            def forward(self, INPUT0: List[str],
                        INPUT1: List[str]) -> Tuple[torch.Tensor, List[str]]:
                input0_int = torch.tensor([int(i) for i in INPUT0])
                input1_int = torch.tensor([int(i) for i in INPUT1])
                op0 = (input0_int +
                       input1_int if not self.swap else input0_int -
                       input1_int).to(self.output0_dtype)
                op1 = [
                    str(i.item())
                    for i in (input0_int -
                              input1_int if not self.swap else input0_int +
                              input1_int)
                ]
                return op0, op1
    elif (input_dtype != np_dtype_string) and (
            output0_dtype == np_dtype_string) and (output1_dtype
                                                   == np_dtype_string):

        class AddSubNet(nn.Module):

            def __init__(self, *args):
                self.output0_dtype = args[0][0]
                self.output1_dtype = args[0][1]
                self.swap = args[0][2]
                super(AddSubNet, self).__init__()

            def forward(self, INPUT0, INPUT1) -> Tuple[List[str], List[str]]:
                op0 = [
                    str(i.item())
                    for i in (INPUT0 + INPUT1 if not self.swap else INPUT0 -
                              INPUT1)
                ]
                op1 = [
                    str(i.item())
                    for i in (INPUT0 - INPUT1 if not self.swap else INPUT0 +
                              INPUT1)
                ]
                return op0, op1
    elif (input_dtype != np_dtype_string) and (
            output0_dtype != np_dtype_string) and (output1_dtype
                                                   == np_dtype_string):

        class AddSubNet(nn.Module):

            def __init__(self, *args):
                self.output0_dtype = args[0][0]
                self.output1_dtype = args[0][1]
                self.swap = args[0][2]
                super(AddSubNet, self).__init__()

            def forward(self, INPUT0, INPUT1) -> Tuple[torch.Tensor, List[str]]:
                op0 = (INPUT0 + INPUT1 if not self.swap else INPUT0 -
                       INPUT1).to(self.output0_dtype)
                op1 = [
                    str(i.item())
                    for i in (INPUT0 - INPUT1 if not self.swap else INPUT0 +
                              INPUT1)
                ]
                return op0, op1
    elif (input_dtype != np_dtype_string) and (
            output0_dtype
            == np_dtype_string) and (output1_dtype != np_dtype_string):

        class AddSubNet(nn.Module):

            def __init__(self, *args):
                self.output0_dtype = args[0][0]
                self.output1_dtype = args[0][1]
                self.swap = args[0][2]
                super(AddSubNet, self).__init__()

            def forward(self, INPUT0, INPUT1) -> Tuple[List[str], torch.Tensor]:
                op0 = [
                    str(i.item())
                    for i in (INPUT0 + INPUT1 if not self.swap else INPUT0 -
                              INPUT1)
                ]
                op1 = (INPUT0 - INPUT1 if not self.swap else INPUT0 +
                       INPUT1).to(self.output1_dtype)
                return op0, op1
    else:

        class AddSubNet(nn.Module):

            def __init__(self, *args):
                self.output0_dtype = args[0][0]
                self.output1_dtype = args[0][1]
                self.swap = args[0][2]
                super(AddSubNet, self).__init__()

            def forward(self, INPUT0, INPUT1):
                op0 = (INPUT0 + INPUT1 if not self.swap else INPUT0 -
                       INPUT1).to(self.output0_dtype)
                op1 = (INPUT0 - INPUT1 if not self.swap else INPUT0 +
                       INPUT1).to(self.output1_dtype)
                return op0, op1

    addSubModel = AddSubNet((torch_output0_dtype, torch_output1_dtype, swap))
    traced = torch.jit.script(addSubModel)

    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    traced.save(model_version_dir + "/model.pt")


def create_libtorch_modelconfig(models_dir, max_batch, model_version,
                                input_shape, output0_shape, output1_shape,
                                input_dtype, output0_dtype, output1_dtype,
                                output0_label_cnt, version_policy):

    if not tu.validate_for_libtorch_model(
            input_dtype, output0_dtype, output1_dtype, input_shape,
            output0_shape, output1_shape, max_batch):
        return

    # Unpack version policy
    version_policy_str = "{ latest { num_versions: 1 }}"
    if version_policy is not None:
        type, val = version_policy
        if type == 'latest':
            version_policy_str = "{{ latest {{ num_versions: {} }}}}".format(
                val)
        elif type == 'specific':
            version_policy_str = "{{ specific {{ versions: {} }}}}".format(val)
        else:
            version_policy_str = "{ all { }}"

    # Use a different model name for the non-batching variant
    model_name = tu.get_model_name(
        "libtorch_nobatch" if max_batch == 0 else "libtorch", input_dtype,
        output0_dtype, output1_dtype)
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "pytorch_libtorch"
max_batch_size: {}
version_policy: {}
input [
  {{
    name: "INPUT0"
    data_type: {}
    dims: [ {} ]
  }},
  {{
    name: "INPUT1"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT__0"
    data_type: {}
    dims: [ {} ]
    label_filename: "output0_labels.txt"
  }},
  {{
    name: "OUTPUT__1"
    data_type: {}
    dims: [ {} ]
  }}
]
'''.format(model_name, max_batch, version_policy_str,
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(output0_dtype),
           tu.shape_to_dims_str(output0_shape),
           np_to_model_dtype(output1_dtype),
           tu.shape_to_dims_str(output1_shape))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)

    with open(config_dir + "/output0_labels.txt", "w") as lfile:
        for l in range(output0_label_cnt):
            lfile.write("label" + str(l) + "\n")


def create_openvino_modelfile(models_dir,
                              max_batch,
                              model_version,
                              input_shape,
                              output0_shape,
                              output1_shape,
                              input_dtype,
                              output0_dtype,
                              output1_dtype,
                              swap=False):

    batch_dim = () if max_batch == 0 else (max_batch,)
    if not tu.validate_for_openvino_model(
            input_dtype, output0_dtype, output1_dtype, batch_dim + input_shape,
            batch_dim + output0_shape, batch_dim + output1_shape):
        return

    # Create the model
    model_name = tu.get_model_name(
        "openvino_nobatch" if max_batch == 0 else "openvino", input_dtype,
        output0_dtype, output1_dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    in0 = ng.parameter(shape=batch_dim + input_shape,
                       dtype=input_dtype,
                       name="INPUT0")
    in1 = ng.parameter(shape=batch_dim + input_shape,
                       dtype=input_dtype,
                       name="INPUT1")

    r0 = ng.add(in0, in1) if not swap else ng.subtract(in0, in1)
    r1 = ng.subtract(in0, in1) if not swap else ng.add(in0, in1)

    result0 = ng.reshape(r0, batch_dim + output0_shape, special_zero=False)
    result1 = ng.reshape(r1, batch_dim + output1_shape, special_zero=False)

    op0 = ng.convert(result0, destination_type=output0_dtype, name="OUTPUT0")
    op1 = ng.convert(result1, destination_type=output1_dtype, name="OUTPUT1")

    function = ng.impl.Function([op0, op1], [in0, in1], model_name)
    ie_network = IENetwork(ng.impl.Function.to_capsule(function))

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    ie_network.serialize(model_version_dir + "/model.xml",
                         model_version_dir + "/model.bin")


def create_openvino_modelconfig(models_dir, max_batch, model_version,
                                input_shape, output0_shape, output1_shape,
                                input_dtype, output0_dtype, output1_dtype,
                                output0_label_cnt, version_policy):

    batch_dim = () if max_batch == 0 else (max_batch,)
    if not tu.validate_for_openvino_model(
            input_dtype, output0_dtype, output1_dtype, batch_dim + input_shape,
            batch_dim + output0_shape, batch_dim + output1_shape):
        return

    # Unpack version policy
    version_policy_str = "{ latest { num_versions: 1 }}"
    if version_policy is not None:
        type, val = version_policy
        if type == 'latest':
            version_policy_str = "{{ latest {{ num_versions: {} }}}}".format(
                val)
        elif type == 'specific':
            version_policy_str = "{{ specific {{ versions: {} }}}}".format(val)
        else:
            version_policy_str = "{ all { }}"

    # Use a different model name for the non-batching variant
    model_name = tu.get_model_name(
        "openvino_nobatch" if max_batch == 0 else "openvino", input_dtype,
        output0_dtype, output1_dtype)
    config_dir = models_dir + "/" + model_name

    # platform is empty and backend is 'openvino' for openvino model
    config = '''
name: "{}"
backend: "openvino"
max_batch_size: {}
version_policy: {}
input [
  {{
    name: "INPUT0"
    data_type: {}
    dims: [ {} ]
  }},
  {{
    name: "INPUT1"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT0"
    data_type: {}
    dims: [ {} ]
    label_filename: "output0_labels.txt"
   }},
  {{
    name: "OUTPUT1"
    data_type: {}
    dims: [ {} ]
  }}
]
'''.format(model_name, max_batch, version_policy_str,
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(output0_dtype),
           tu.shape_to_dims_str(output0_shape),
           np_to_model_dtype(output1_dtype),
           tu.shape_to_dims_str(output1_shape))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)

    with open(config_dir + "/output0_labels.txt", "w") as lfile:
        for l in range(output0_label_cnt):
            lfile.write("label" + str(l) + "\n")


def create_models(models_dir,
                  input_dtype,
                  output0_dtype,
                  output1_dtype,
                  input_shape,
                  output0_shape,
                  output1_shape,
                  output0_label_cnt,
                  version_policy=None):
    model_version = 1

    # Create two models, one that supports batching with a max-batch
    # of 8, and one that does not with a max-batch of 0
    if FLAGS.graphdef:
        # max-batch 8
        create_graphdef_modelconfig(models_dir, 8, model_version, input_shape,
                                    output0_shape, output1_shape, input_dtype,
                                    output0_dtype, output1_dtype,
                                    output0_label_cnt, version_policy)
        create_graphdef_modelfile(models_dir, 8, model_version, input_shape,
                                  output0_shape, output1_shape, input_dtype,
                                  output0_dtype, output1_dtype)
        # max-batch 0
        create_graphdef_modelconfig(models_dir, 0, model_version, input_shape,
                                    output0_shape, output1_shape, input_dtype,
                                    output0_dtype, output1_dtype,
                                    output0_label_cnt, version_policy)
        create_graphdef_modelfile(models_dir, 0, model_version, input_shape,
                                  output0_shape, output1_shape, input_dtype,
                                  output0_dtype, output1_dtype)

    if FLAGS.savedmodel:
        # max-batch 8
        create_savedmodel_modelconfig(models_dir, 8, model_version, input_shape,
                                      output0_shape, output1_shape, input_dtype,
                                      output0_dtype, output1_dtype,
                                      output0_label_cnt, version_policy)
        create_savedmodel_modelfile(models_dir, 8, model_version, input_shape,
                                    output0_shape, output1_shape, input_dtype,
                                    output0_dtype, output1_dtype)
        # max-batch 0
        create_savedmodel_modelconfig(models_dir, 0, model_version, input_shape,
                                      output0_shape, output1_shape, input_dtype,
                                      output0_dtype, output1_dtype,
                                      output0_label_cnt, version_policy)
        create_savedmodel_modelfile(models_dir, 0, model_version, input_shape,
                                    output0_shape, output1_shape, input_dtype,
                                    output0_dtype, output1_dtype)

    if FLAGS.tensorrt:
        # max-batch 8
        suffix = ()
        if input_dtype == np.int8 or output0_dtype == np.int8 or output1_dtype == np.int8:
            suffix = (1, 1)
        create_plan_modelconfig(models_dir, 8, model_version,
                                input_shape + suffix, output0_shape + suffix,
                                output1_shape + suffix, input_dtype,
                                output0_dtype, output1_dtype, output0_label_cnt,
                                version_policy)
        create_plan_modelfile(models_dir, 8, model_version,
                              input_shape + suffix, output0_shape + suffix,
                              output1_shape + suffix, input_dtype,
                              output0_dtype, output1_dtype)
        # max-batch 0
        create_plan_modelconfig(models_dir, 0, model_version,
                                input_shape + suffix, output0_shape + suffix,
                                output1_shape + suffix, input_dtype,
                                output0_dtype, output1_dtype, output0_label_cnt,
                                version_policy)
        create_plan_modelfile(models_dir, 0, model_version,
                              input_shape + suffix, output0_shape + suffix,
                              output1_shape + suffix, input_dtype,
                              output0_dtype, output1_dtype)

        if -1 in input_shape:
            # models for testing optimization profiles
            create_plan_modelconfig(models_dir,
                                    8,
                                    model_version,
                                    input_shape + suffix,
                                    output0_shape + suffix,
                                    output1_shape + suffix,
                                    input_dtype,
                                    output0_dtype,
                                    output1_dtype,
                                    output0_label_cnt,
                                    version_policy,
                                    min_dim=4,
                                    max_dim=32)
            create_plan_modelfile(models_dir,
                                  8,
                                  model_version,
                                  input_shape + suffix,
                                  output0_shape + suffix,
                                  output1_shape + suffix,
                                  input_dtype,
                                  output0_dtype,
                                  output1_dtype,
                                  min_dim=4,
                                  max_dim=32)

    if FLAGS.onnx:
        # max-batch 8
        create_onnx_modelconfig(models_dir, 8, model_version, input_shape,
                                output0_shape, output1_shape, input_dtype,
                                output0_dtype, output1_dtype, output0_label_cnt,
                                version_policy)
        create_onnx_modelfile(models_dir, 8, model_version, input_shape,
                              output0_shape, output1_shape, input_dtype,
                              output0_dtype, output1_dtype)
        # max-batch 0
        create_onnx_modelconfig(models_dir, 0, model_version, input_shape,
                                output0_shape, output1_shape, input_dtype,
                                output0_dtype, output1_dtype, output0_label_cnt,
                                version_policy)
        create_onnx_modelfile(models_dir, 0, model_version, input_shape,
                              output0_shape, output1_shape, input_dtype,
                              output0_dtype, output1_dtype)

    if FLAGS.libtorch:
        # max-batch 8
        create_libtorch_modelconfig(models_dir, 8, model_version, input_shape,
                                    output0_shape, output1_shape, input_dtype,
                                    output0_dtype, output1_dtype,
                                    output0_label_cnt, version_policy)
        create_libtorch_modelfile(models_dir, 8, model_version, input_shape,
                                  output0_shape, output1_shape, input_dtype,
                                  output0_dtype, output1_dtype)
        # max-batch 0
        create_libtorch_modelconfig(models_dir, 0, model_version, input_shape,
                                    output0_shape, output1_shape, input_dtype,
                                    output0_dtype, output1_dtype,
                                    output0_label_cnt, version_policy)
        create_libtorch_modelfile(models_dir, 0, model_version, input_shape,
                                  output0_shape, output1_shape, input_dtype,
                                  output0_dtype, output1_dtype)

    if FLAGS.openvino:
        # max-batch 8
        create_openvino_modelconfig(models_dir, 8, model_version, input_shape,
                                    output0_shape, output1_shape, input_dtype,
                                    output0_dtype, output1_dtype,
                                    output0_label_cnt, version_policy)
        create_openvino_modelfile(models_dir, 8, model_version, input_shape,
                                  output0_shape, output1_shape, input_dtype,
                                  output0_dtype, output1_dtype)
        # max-batch 0
        create_openvino_modelconfig(models_dir, 0, model_version, input_shape,
                                    output0_shape, output1_shape, input_dtype,
                                    output0_dtype, output1_dtype,
                                    output0_label_cnt, version_policy)
        create_openvino_modelfile(models_dir, 0, model_version, input_shape,
                                  output0_shape, output1_shape, input_dtype,
                                  output0_dtype, output1_dtype)

    if FLAGS.ensemble:
        for pair in emu.platform_types_and_validation():
            if not pair[1](input_dtype, output0_dtype, output1_dtype,
                           input_shape, output0_shape, output1_shape):
                continue

            config_input_shape = input_shape
            config_output0_shape = output0_shape
            config_output1_shape = output1_shape
            if pair[0] == "plan":
                if len(input_shape) == 1 and input_dtype == np.int8:
                    config_input_shape = (input_shape[0], 1, 1)
                if len(output0_shape) == 1 and output0_dtype == np.int8:
                    config_output0_shape = (output0_shape[0], 1, 1)
                if len(output1_shape) == 1 and output1_dtype == np.int8:
                    config_output1_shape = (output1_shape[0], 1, 1)

            # max-batch 0
            emu.create_ensemble_modelconfig(pair[0], models_dir, 0,
                                            model_version, config_input_shape,
                                            config_output0_shape,
                                            config_output1_shape, input_dtype,
                                            output0_dtype, output1_dtype,
                                            output0_label_cnt, version_policy)
            emu.create_ensemble_modelfile(pair[0], models_dir, 0, model_version,
                                          config_input_shape,
                                          config_output0_shape,
                                          config_output1_shape, input_dtype,
                                          output0_dtype, output1_dtype)

            # max-batch 8 (Skip for PyTorch models with String I/O)
            if (pair[0] == "libtorch") and not pair[1](
                    input_dtype, output0_dtype, output1_dtype, input_shape,
                    output0_shape, output1_shape, 8):
                continue

            emu.create_ensemble_modelconfig(pair[0], models_dir, 8,
                                            model_version, config_input_shape,
                                            config_output0_shape,
                                            config_output1_shape, input_dtype,
                                            output0_dtype, output1_dtype,
                                            output0_label_cnt, version_policy)
            emu.create_ensemble_modelfile(pair[0], models_dir, 8, model_version,
                                          config_input_shape,
                                          config_output0_shape,
                                          config_output1_shape, input_dtype,
                                          output0_dtype, output1_dtype)


def create_fixed_models(models_dir,
                        input_dtype,
                        output0_dtype,
                        output1_dtype,
                        version_policy=None):
    input_size = 16
    create_models(models_dir, input_dtype, output0_dtype, output1_dtype,
                  (input_size,), (input_size,), (input_size,), input_size,
                  version_policy)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--models_dir',
                        type=str,
                        required=True,
                        help='Top-level model directory')
    parser.add_argument('--graphdef',
                        required=False,
                        action='store_true',
                        help='Generate GraphDef models')
    parser.add_argument('--savedmodel',
                        required=False,
                        action='store_true',
                        help='Generate SavedModel models')
    parser.add_argument('--tensorrt',
                        required=False,
                        action='store_true',
                        help='Generate TensorRT PLAN models')
    parser.add_argument('--onnx',
                        required=False,
                        action='store_true',
                        help='Generate Onnx Runtime Onnx models')
    parser.add_argument(
        '--onnx_opset',
        type=int,
        required=False,
        default=0,
        help='Opset used for Onnx models. Default is to use ONNXRT default')
    parser.add_argument('--libtorch',
                        required=False,
                        action='store_true',
                        help='Generate Pytorch LibTorch models')
    parser.add_argument('--openvino',
                        required=False,
                        action='store_true',
                        help='Generate Openvino models')
    parser.add_argument('--variable',
                        required=False,
                        action='store_true',
                        help='Used variable-shape tensors for input/output')
    parser.add_argument('--ensemble',
                        required=False,
                        action='store_true',
                        help='Generate ensemble models against the models' +
                        ' in all platforms. Note that the models generated' +
                        ' are not completed.')
    FLAGS, unparsed = parser.parse_known_args()

    if FLAGS.graphdef or FLAGS.savedmodel:
        import tensorflow as tf
        from tensorflow.python.framework import graph_io
    if FLAGS.tensorrt:
        import tensorrt as trt
    if FLAGS.onnx:
        import onnx
    if FLAGS.libtorch:
        import torch
        from torch import nn
    if FLAGS.openvino:
        from openvino.inference_engine import IENetwork
        import ngraph as ng

    import test_util as tu

    # Tests with models that accept fixed-shape input/output tensors
    if not FLAGS.variable:
        create_fixed_models(FLAGS.models_dir, np.uint8, np.uint8, np.uint8,
                            ('latest', 3))
        create_fixed_models(FLAGS.models_dir, np.int8, np.int8, np.int8,
                            ('latest', 1))
        create_fixed_models(FLAGS.models_dir, np.int16, np.int16, np.int16,
                            ('latest', 2))
        create_fixed_models(FLAGS.models_dir, np.int32, np.int32, np.int32,
                            ('all', None))
        create_fixed_models(FLAGS.models_dir, np.int64, np.int64, np.int64)
        create_fixed_models(FLAGS.models_dir, np.float16, np.float16,
                            np.float16, ('specific', [
                                1,
                            ]))
        create_fixed_models(FLAGS.models_dir, np.float32, np.float32,
                            np.float32, ('specific', [1, 3]))
        create_fixed_models(FLAGS.models_dir, np.float16, np.float32,
                            np.float32)
        create_fixed_models(FLAGS.models_dir, np.int32, np.int8, np.int8)
        create_fixed_models(FLAGS.models_dir, np.int8, np.int32, np.int32)
        create_fixed_models(FLAGS.models_dir, np.int32, np.int8, np.int16)
        create_fixed_models(FLAGS.models_dir, np.float32, np.uint8, np.uint8)
        create_fixed_models(FLAGS.models_dir, np.uint8, np.float32, np.float32)
        create_fixed_models(FLAGS.models_dir, np.float32, np.uint8, np.float16)
        create_fixed_models(FLAGS.models_dir, np.int32, np.float32, np.float32)
        create_fixed_models(FLAGS.models_dir, np.float32, np.int32, np.int32)
        create_fixed_models(FLAGS.models_dir, np.int32, np.float16, np.int16)

        create_fixed_models(FLAGS.models_dir, np_dtype_string, np.int32,
                            np.int32)
        create_fixed_models(FLAGS.models_dir, np_dtype_string, np_dtype_string,
                            np_dtype_string)
        create_fixed_models(FLAGS.models_dir, np_dtype_string, np.int32,
                            np_dtype_string)
        create_fixed_models(FLAGS.models_dir, np_dtype_string, np_dtype_string,
                            np.int32)
        create_fixed_models(FLAGS.models_dir, np.int32, np_dtype_string,
                            np_dtype_string)
        create_fixed_models(FLAGS.models_dir, np.int32, np.int32,
                            np_dtype_string)
        create_fixed_models(FLAGS.models_dir, np.int32, np_dtype_string,
                            np.int32)

        # Make multiple versions of some models for version testing
        # (they use different version policies when created above)
        if FLAGS.graphdef:
            for vt in [np.float16, np.float32, np.int8, np.int16, np.int32]:
                create_graphdef_modelfile(FLAGS.models_dir,
                                          8,
                                          2, (16,), (16,), (16,),
                                          vt,
                                          vt,
                                          vt,
                                          swap=True)
                create_graphdef_modelfile(FLAGS.models_dir,
                                          8,
                                          3, (16,), (16,), (16,),
                                          vt,
                                          vt,
                                          vt,
                                          swap=True)
                create_graphdef_modelfile(FLAGS.models_dir,
                                          0,
                                          2, (16,), (16,), (16,),
                                          vt,
                                          vt,
                                          vt,
                                          swap=True)
                create_graphdef_modelfile(FLAGS.models_dir,
                                          0,
                                          3, (16,), (16,), (16,),
                                          vt,
                                          vt,
                                          vt,
                                          swap=True)

        if FLAGS.savedmodel:
            for vt in [np.float16, np.float32, np.int8, np.int16, np.int32]:
                create_savedmodel_modelfile(FLAGS.models_dir,
                                            8,
                                            2, (16,), (16,), (16,),
                                            vt,
                                            vt,
                                            vt,
                                            swap=True)
                create_savedmodel_modelfile(FLAGS.models_dir,
                                            8,
                                            3, (16,), (16,), (16,),
                                            vt,
                                            vt,
                                            vt,
                                            swap=True)
                create_savedmodel_modelfile(FLAGS.models_dir,
                                            0,
                                            2, (16,), (16,), (16,),
                                            vt,
                                            vt,
                                            vt,
                                            swap=True)
                create_savedmodel_modelfile(FLAGS.models_dir,
                                            0,
                                            3, (16,), (16,), (16,),
                                            vt,
                                            vt,
                                            vt,
                                            swap=True)

        if FLAGS.tensorrt:
            for vt in [np.float32, np.float16, np.int32, np.uint8]:
                create_plan_modelfile(FLAGS.models_dir,
                                      8,
                                      2, (16,), (16,), (16,),
                                      vt,
                                      vt,
                                      vt,
                                      swap=True)
                create_plan_modelfile(FLAGS.models_dir,
                                      8,
                                      3, (16,), (16,), (16,),
                                      vt,
                                      vt,
                                      vt,
                                      swap=True)
                create_plan_modelfile(FLAGS.models_dir,
                                      0,
                                      2, (16,), (16,), (16,),
                                      vt,
                                      vt,
                                      vt,
                                      swap=True)
                create_plan_modelfile(FLAGS.models_dir,
                                      0,
                                      3, (16,), (16,), (16,),
                                      vt,
                                      vt,
                                      vt,
                                      swap=True)

            vt = np.int8
            #handle INT8 separately as it doesn't allow 1d tensors
            create_plan_modelfile(FLAGS.models_dir,
                                  8,
                                  2, (16, 1, 1), (16, 1, 1), (16, 1, 1),
                                  vt,
                                  vt,
                                  vt,
                                  swap=True)
            create_plan_modelfile(FLAGS.models_dir,
                                  8,
                                  3, (16, 1, 1), (16, 1, 1), (16, 1, 1),
                                  vt,
                                  vt,
                                  vt,
                                  swap=True)
            create_plan_modelfile(FLAGS.models_dir,
                                  0,
                                  2, (16, 1, 1), (16, 1, 1), (16, 1, 1),
                                  vt,
                                  vt,
                                  vt,
                                  swap=True)
            create_plan_modelfile(FLAGS.models_dir,
                                  0,
                                  3, (16, 1, 1), (16, 1, 1), (16, 1, 1),
                                  vt,
                                  vt,
                                  vt,
                                  swap=True)

        if FLAGS.onnx:
            for vt in [np.float16, np.float32, np.int8, np.int16, np.int32]:
                create_onnx_modelfile(FLAGS.models_dir,
                                      8,
                                      2, (16,), (16,), (16,),
                                      vt,
                                      vt,
                                      vt,
                                      swap=True)
                create_onnx_modelfile(FLAGS.models_dir,
                                      8,
                                      3, (16,), (16,), (16,),
                                      vt,
                                      vt,
                                      vt,
                                      swap=True)
                create_onnx_modelfile(FLAGS.models_dir,
                                      0,
                                      2, (16,), (16,), (16,),
                                      vt,
                                      vt,
                                      vt,
                                      swap=True)
                create_onnx_modelfile(FLAGS.models_dir,
                                      0,
                                      3, (16,), (16,), (16,),
                                      vt,
                                      vt,
                                      vt,
                                      swap=True)
        if FLAGS.libtorch:
            for vt in [np.float32, np.int32, np.int16, np.int8]:
                create_libtorch_modelfile(FLAGS.models_dir,
                                          8,
                                          2, (16,), (16,), (16,),
                                          vt,
                                          vt,
                                          vt,
                                          swap=True)
                create_libtorch_modelfile(FLAGS.models_dir,
                                          8,
                                          3, (16,), (16,), (16,),
                                          vt,
                                          vt,
                                          vt,
                                          swap=True)
                create_libtorch_modelfile(FLAGS.models_dir,
                                          0,
                                          2, (16,), (16,), (16,),
                                          vt,
                                          vt,
                                          vt,
                                          swap=True)
                create_libtorch_modelfile(FLAGS.models_dir,
                                          0,
                                          3, (16,), (16,), (16,),
                                          vt,
                                          vt,
                                          vt,
                                          swap=True)
        if FLAGS.openvino:
            for vt in [np.float16, np.float32, np.int8, np.int16, np.int32]:
                create_openvino_modelfile(FLAGS.models_dir,
                                          8,
                                          2, (16,), (16,), (16,),
                                          vt,
                                          vt,
                                          vt,
                                          swap=True)
                create_openvino_modelfile(FLAGS.models_dir,
                                          8,
                                          3, (16,), (16,), (16,),
                                          vt,
                                          vt,
                                          vt,
                                          swap=True)
                create_openvino_modelfile(FLAGS.models_dir,
                                          0,
                                          2, (16,), (16,), (16,),
                                          vt,
                                          vt,
                                          vt,
                                          swap=True)
                create_openvino_modelfile(FLAGS.models_dir,
                                          0,
                                          3, (16,), (16,), (16,),
                                          vt,
                                          vt,
                                          vt,
                                          swap=True)

        if FLAGS.ensemble:
            for pair in emu.platform_types_and_validation():
                for vt in [np.float16, np.float32, np.int8, np.int16, np.int32]:
                    shape = (16, 1, 1) if (pair[0] == "plan" and
                                           vt == np.int8) else (16,)
                    if not pair[1](vt, vt, vt, shape, shape, shape):
                        continue
                    emu.create_ensemble_modelfile(pair[0],
                                                  FLAGS.models_dir,
                                                  8,
                                                  2,
                                                  shape,
                                                  shape,
                                                  shape,
                                                  vt,
                                                  vt,
                                                  vt,
                                                  swap=True)
                    emu.create_ensemble_modelfile(pair[0],
                                                  FLAGS.models_dir,
                                                  8,
                                                  3,
                                                  shape,
                                                  shape,
                                                  shape,
                                                  vt,
                                                  vt,
                                                  vt,
                                                  swap=True)
                    emu.create_ensemble_modelfile(pair[0],
                                                  FLAGS.models_dir,
                                                  0,
                                                  2,
                                                  shape,
                                                  shape,
                                                  shape,
                                                  vt,
                                                  vt,
                                                  vt,
                                                  swap=True)
                    emu.create_ensemble_modelfile(pair[0],
                                                  FLAGS.models_dir,
                                                  0,
                                                  3,
                                                  shape,
                                                  shape,
                                                  shape,
                                                  vt,
                                                  vt,
                                                  vt,
                                                  swap=True)

    # Tests with models that accept variable-shape input/output tensors
    if FLAGS.variable:
        create_models(FLAGS.models_dir, np.float32, np.float32, np.float32,
                      (-1,), (-1,), (-1,), 16)
        create_models(FLAGS.models_dir, np.float32, np.int32, np.int32,
                      (-1, -1), (-1, -1), (-1, -1), 16)
        create_models(FLAGS.models_dir, np.float32, np.int64, np.int64, (8, -1),
                      (8, -1), (8, -1), 32)
        create_models(FLAGS.models_dir, np.float32, np.int32, np.int64,
                      (-1, 8, -1), (-1, 8, -1), (-1, 8, -1), 32)
        create_models(FLAGS.models_dir, np.float32, np.float32, np.int32, (-1,),
                      (-1,), (-1,), 16)
        create_models(FLAGS.models_dir, np.int32, np.int32, np.int32, (-1, -1),
                      (-1, -1), (-1, -1), 16)
        create_models(FLAGS.models_dir, np.int32, np.int32, np.float32,
                      (-1, 8, -1), (-1, 8, -1), (-1, 8, -1), 32)

        create_models(FLAGS.models_dir, np_dtype_string, np_dtype_string,
                      np_dtype_string, (-1,), (-1,), (-1,), 16)
        create_models(FLAGS.models_dir, np_dtype_string, np.int32, np.int32,
                      (-1, -1), (-1, -1), (-1, -1), 16)
        create_models(FLAGS.models_dir, np_dtype_string, np_dtype_string,
                      np.int32, (8, -1), (8, -1), (8, -1), 32)
        create_models(FLAGS.models_dir, np_dtype_string, np.int32,
                      np_dtype_string, (-1, 8, -1), (-1, 8, -1), (-1, 8, -1),
                      32)

    if FLAGS.ensemble:
        # Create utility models used in ensemble
        # nop (only creates model config, should add model file before use)
        model_dtypes = ["TYPE_BOOL", "TYPE_STRING"]
        for s in [8, 16, 32, 64]:
            for t in ["INT", "UINT", "FP"]:
                if t == "FP" and s == 8:
                    continue
                model_dtypes.append("TYPE_{}{}".format(t, s))

        for model_dtype in model_dtypes:
            # Use variable size to handle all shape. Note: piping variable size output
            # to fixed size model is not safe but doable
            for model_shape in [(-1,), (-1, -1), (-1, -1, -1)]:
                emu.create_nop_modelconfig(FLAGS.models_dir, model_shape,
                                           model_dtype)


# qa/common/gen_qa_noshape_models.py

# Copyright (c) 2018-2019, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
from builtins import range
import os
import numpy as np

FLAGS = None
np_dtype_string = np.dtype(object)


def np_to_model_dtype(np_dtype):
    if np_dtype == bool:
        return "TYPE_BOOL"
    elif np_dtype == np.int8:
        return "TYPE_INT8"
    elif np_dtype == np.int16:
        return "TYPE_INT16"
    elif np_dtype == np.int32:
        return "TYPE_INT32"
    elif np_dtype == np.int64:
        return "TYPE_INT64"
    elif np_dtype == np.uint8:
        return "TYPE_UINT8"
    elif np_dtype == np.uint16:
        return "TYPE_UINT16"
    elif np_dtype == np.float16:
        return "TYPE_FP16"
    elif np_dtype == np.float32:
        return "TYPE_FP32"
    elif np_dtype == np.float64:
        return "TYPE_FP64"
    elif np_dtype == np_dtype_string:
        return "TYPE_STRING"
    return None


def np_to_tf_dtype(np_dtype):
    if np_dtype == bool:
        return tf.bool
    elif np_dtype == np.int8:
        return tf.int8
    elif np_dtype == np.int16:
        return tf.int16
    elif np_dtype == np.int32:
        return tf.int32
    elif np_dtype == np.int64:
        return tf.int64
    elif np_dtype == np.uint8:
        return tf.uint8
    elif np_dtype == np.uint16:
        return tf.uint16
    elif np_dtype == np.float16:
        return tf.float16
    elif np_dtype == np.float32:
        return tf.float32
    elif np_dtype == np.float64:
        return tf.float64
    elif np_dtype == np_dtype_string:
        return tf.string
    return None


def create_savedmodel_modelfile(models_dir,
                                max_batch,
                                model_version,
                                input_shape,
                                output0_shape,
                                output1_shape,
                                input_dtype,
                                output0_dtype,
                                output1_dtype,
                                swap=False):

    if not tu.validate_for_tf_model(input_dtype, output0_dtype, output1_dtype,
                                    input_shape, output0_shape, output1_shape):
        return

    tf_input_dtype = np_to_tf_dtype(input_dtype)
    tf_output0_dtype = np_to_tf_dtype(output0_dtype)
    tf_output1_dtype = np_to_tf_dtype(output1_dtype)

    # Create the model. If non-batching then don't include the batch
    # dimension.
    tf.reset_default_graph()
    if max_batch == 0:
        in0 = tf.placeholder(tf_input_dtype, tu.shape_to_tf_shape([]),
                             "TENSOR_INPUT0")
        in1 = tf.placeholder(tf_input_dtype, tu.shape_to_tf_shape(input_shape),
                             "TENSOR_INPUT1")
    else:
        in0 = tf.placeholder(tf_input_dtype, tu.shape_to_tf_shape([]),
                             "TENSOR_INPUT0")
        in1 = tf.placeholder(tf_input_dtype, [
            None,
        ] + tu.shape_to_tf_shape(input_shape), "TENSOR_INPUT1")

    # If the input is a string, then convert each string to the
    # equivalent float value.
    if tf_input_dtype == tf.string:
        in0 = tf.strings.to_number(in0, tf.int32)
        in1 = tf.strings.to_number(in1, tf.int32)

    add = tf.add(in0, in1, "ADD")
    sub = tf.subtract(in0, in1, "SUB")

    # Cast or convert result to the output dtype.
    if tf_output0_dtype == tf.string:
        cast0 = tf.dtypes.as_string(add if not swap else sub, name="TOSTR0")
    else:
        cast0 = tf.cast(add if not swap else sub, tf_output0_dtype, "CAST0")

    if tf_output1_dtype == tf.string:
        cast1 = tf.dtypes.as_string(sub if not swap else add, name="TOSTR1")
    else:
        cast1 = tf.cast(sub if not swap else add, tf_output1_dtype, "CAST1")

    out0 = tf.identity(cast0, "TENSOR_OUTPUT0")
    out1 = tf.identity(cast1, "TENSOR_OUTPUT1")

    # Use a different model name for the non-batching variant
    model_name = tu.get_model_name(
        "savedmodel_nobatch" if max_batch == 0 else "savedmodel", input_dtype,
        output0_dtype, output1_dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with tf.Session() as sess:
        input0_tensor = tf.get_default_graph().get_tensor_by_name(
            "TENSOR_INPUT0:0")
        input1_tensor = tf.get_default_graph().get_tensor_by_name(
            "TENSOR_INPUT1:0")
        output0_tensor = tf.get_default_graph().get_tensor_by_name(
            "TENSOR_OUTPUT0:0")
        output1_tensor = tf.get_default_graph().get_tensor_by_name(
            "TENSOR_OUTPUT1:0")
        tf.saved_model.simple_save(sess,
                                   model_version_dir + "/model.savedmodel",
                                   inputs={
                                       "INPUT0": input0_tensor,
                                       "INPUT1": input1_tensor
                                   },
                                   outputs={
                                       "OUTPUT0": output0_tensor,
                                       "OUTPUT1": output1_tensor
                                   })


def create_savedmodel_modelconfig(models_dir, max_batch, model_version,
                                  input_shape, output0_shape, output1_shape,
                                  input_dtype, output0_dtype, output1_dtype,
                                  output0_label_cnt, version_policy):

    if not tu.validate_for_tf_model(input_dtype, output0_dtype, output1_dtype,
                                    input_shape, output0_shape, output1_shape):
        return

    # Unpack version policy
    version_policy_str = "{ latest { num_versions: 1 }}"
    if version_policy is not None:
        type, val = version_policy
        if type == 'latest':
            version_policy_str = "{{ latest {{ num_versions: {} }}}}".format(
                val)
        elif type == 'specific':
            version_policy_str = "{{ specific {{ versions: {} }}}}".format(val)
        else:
            version_policy_str = "{ all { }}"

    # Use a different model name for the non-batching variant
    model_name = tu.get_model_name(
        "savedmodel_nobatch" if max_batch == 0 else "savedmodel", input_dtype,
        output0_dtype, output1_dtype)
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "tensorflow_savedmodel"
max_batch_size: {}
version_policy: {}
input [
  {{
    name: "INPUT0"
    data_type: {}
    dims: [ {} ]
  }},
  {{
    name: "INPUT1"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT0"
    data_type: {}
    dims: [ {} ]
    label_filename: "output0_labels.txt"
  }},
  {{
    name: "OUTPUT1"
    data_type: {}
    dims: [ {} ]
  }}
]
'''.format(model_name, max_batch, version_policy_str,
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(output0_dtype),
           tu.shape_to_dims_str(output0_shape),
           np_to_model_dtype(output1_dtype),
           tu.shape_to_dims_str(output1_shape))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)

    with open(config_dir + "/output0_labels.txt", "w") as lfile:
        for l in range(output0_label_cnt):
            lfile.write("label" + str(l) + "\n")


def create_models(models_dir,
                  input_dtype,
                  output0_dtype,
                  output1_dtype,
                  input_shape,
                  output0_shape,
                  output1_shape,
                  output0_label_cnt,
                  version_policy=None):
    model_version = 1

    # Create two models, one that supports batching with a max-batch
    # of 8, and one that does not with a max-batch of 0

    if FLAGS.savedmodel:
        # max-batch 8
        create_savedmodel_modelconfig(models_dir, 8, model_version, input_shape,
                                      output0_shape, output1_shape, input_dtype,
                                      output0_dtype, output1_dtype,
                                      output0_label_cnt, version_policy)
        create_savedmodel_modelfile(models_dir, 8, model_version, input_shape,
                                    output0_shape, output1_shape, input_dtype,
                                    output0_dtype, output1_dtype)
        # max-batch 0
        create_savedmodel_modelconfig(models_dir, 0, model_version, input_shape,
                                      output0_shape, output1_shape, input_dtype,
                                      output0_dtype, output1_dtype,
                                      output0_label_cnt, version_policy)
        create_savedmodel_modelfile(models_dir, 0, model_version, input_shape,
                                    output0_shape, output1_shape, input_dtype,
                                    output0_dtype, output1_dtype)


def create_fixed_models(models_dir,
                        input_dtype,
                        output0_dtype,
                        output1_dtype,
                        version_policy=None):
    input_size = 16

    create_models(models_dir, input_dtype, output0_dtype, output1_dtype,
                  (input_size,), (input_size,), (input_size,), input_size,
                  version_policy)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--models_dir',
                        type=str,
                        required=True,
                        help='Top-level model directory')
    parser.add_argument('--graphdef',
                        required=False,
                        action='store_true',
                        help='Generate GraphDef models')
    parser.add_argument('--savedmodel',
                        required=False,
                        action='store_true',
                        help='Generate SavedModel models')
    parser.add_argument('--tensorrt',
                        required=False,
                        action='store_true',
                        help='Generate TensorRT PLAN models')
    parser.add_argument('--onnx',
                        required=False,
                        action='store_true',
                        help='Generate Onnx Runtime Onnx models')
    parser.add_argument('--libtorch',
                        required=False,
                        action='store_true',
                        help='Generate Pytorch LibTorch models')
    parser.add_argument('--variable',
                        required=False,
                        action='store_true',
                        help='Used variable-shape tensors for input/output')
    parser.add_argument('--ensemble',
                        required=False,
                        action='store_true',
                        help='Generate ensemble models against the models' +
                        ' in all platforms. Note that the models generated' +
                        ' are not completed.')
    FLAGS, unparsed = parser.parse_known_args()

    if FLAGS.savedmodel:
        import tensorflow as tf

    import test_util as tu

    # Tests with models that accept fixed-shape input/output tensors
    if not FLAGS.variable:
        create_fixed_models(FLAGS.models_dir, np.int8, np.int8, np.int8,
                            ('latest', 1))
        create_fixed_models(FLAGS.models_dir, np.int16, np.int16, np.int16,
                            ('latest', 2))
        create_fixed_models(FLAGS.models_dir, np.int32, np.int32, np.int32,
                            ('all', None))
        create_fixed_models(FLAGS.models_dir, np.int64, np.int64, np.int64)
        create_fixed_models(FLAGS.models_dir, np.float16, np.float16,
                            np.float16, ('specific', [
                                1,
                            ]))
        create_fixed_models(FLAGS.models_dir, np.float32, np.float32,
                            np.float32, ('specific', [1, 3]))
        create_fixed_models(FLAGS.models_dir, np.float16, np.float32,
                            np.float32)
        create_fixed_models(FLAGS.models_dir, np.int32, np.int8, np.int8)
        create_fixed_models(FLAGS.models_dir, np.int8, np.int32, np.int32)
        create_fixed_models(FLAGS.models_dir, np.int32, np.int8, np.int16)
        create_fixed_models(FLAGS.models_dir, np.int32, np.float32, np.float32)
        create_fixed_models(FLAGS.models_dir, np.float32, np.int32, np.int32)
        create_fixed_models(FLAGS.models_dir, np.int32, np.float16, np.int16)

        if FLAGS.savedmodel:
            for vt in [np.float16, np.float32, np.int8, np.int16, np.int32]:
                create_savedmodel_modelfile(FLAGS.models_dir,
                                            8,
                                            2, (16,), (16,), (16,),
                                            vt,
                                            vt,
                                            vt,
                                            swap=True)
                create_savedmodel_modelfile(FLAGS.models_dir,
                                            8,
                                            3, (16,), (16,), (16,),
                                            vt,
                                            vt,
                                            vt,
                                            swap=True)
                create_savedmodel_modelfile(FLAGS.models_dir,
                                            0,
                                            2, (16,), (16,), (16,),
                                            vt,
                                            vt,
                                            vt,
                                            swap=True)
                create_savedmodel_modelfile(FLAGS.models_dir,
                                            0,
                                            3, (16,), (16,), (16,),
                                            vt,
                                            vt,
                                            vt,
                                            swap=True)


# qa/common/gen_qa_ragged_models.py

# Copyright 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
import os

import numpy as np


def np_to_model_dtype(np_dtype):
    if np_dtype == bool:
        return "TYPE_BOOL"
    elif np_dtype == np.int8:
        return "TYPE_INT8"
    elif np_dtype == np.int16:
        return "TYPE_INT16"
    elif np_dtype == np.int32:
        return "TYPE_INT32"
    elif np_dtype == np.int64:
        return "TYPE_INT64"
    elif np_dtype == np.uint8:
        return "TYPE_UINT8"
    elif np_dtype == np.uint16:
        return "TYPE_UINT16"
    elif np_dtype == np.float16:
        return "TYPE_FP16"
    elif np_dtype == np.float32:
        return "TYPE_FP32"
    elif np_dtype == np.float64:
        return "TYPE_FP64"
    elif np_dtype == np_dtype_string:
        return "TYPE_STRING"
    return None


def np_to_trt_dtype(np_dtype):
    if np_dtype == bool:
        return trt.bool
    elif np_dtype == np.int8:
        return trt.int8
    elif np_dtype == np.int32:
        return trt.int32
    elif np_dtype == np.float16:
        return trt.float16
    elif np_dtype == np.float32:
        return trt.float32
    return None


def np_to_tf_dtype(np_dtype):
    if np_dtype == bool:
        return tf.bool
    elif np_dtype == np.int8:
        return tf.int8
    elif np_dtype == np.int16:
        return tf.int16
    elif np_dtype == np.int32:
        return tf.int32
    elif np_dtype == np.int64:
        return tf.int64
    elif np_dtype == np.uint8:
        return tf.uint8
    elif np_dtype == np.uint16:
        return tf.uint16
    elif np_dtype == np.float16:
        return tf.float16
    elif np_dtype == np.float32:
        return tf.float32
    elif np_dtype == np.float64:
        return tf.float64
    elif np_dtype == np_dtype_string:
        return tf.string
    return None


def np_to_onnx_dtype(np_dtype):
    if np_dtype == bool:
        return onnx.TensorProto.BOOL
    elif np_dtype == np.int8:
        return onnx.TensorProto.INT8
    elif np_dtype == np.int16:
        return onnx.TensorProto.INT16
    elif np_dtype == np.int32:
        return onnx.TensorProto.INT32
    elif np_dtype == np.int64:
        return onnx.TensorProto.INT64
    elif np_dtype == np.uint8:
        return onnx.TensorProto.UINT8
    elif np_dtype == np.uint16:
        return onnx.TensorProto.UINT16
    elif np_dtype == np.float16:
        return onnx.TensorProto.FLOAT16
    elif np_dtype == np.float32:
        return onnx.TensorProto.FLOAT
    elif np_dtype == np.float64:
        return onnx.TensorProto.DOUBLE
    elif np_dtype == np_dtype_string:
        return onnx.TensorProto.STRING
    return None


def create_savedmodel_modelfile(models_dir, model_version, dtype):
    # Create special identity model for batch input testing.
    # Because the ragged input and batch input are one dimensional vector
    # when passing to the model, the model must generate output with batch
    # dimension so that Triton can scatter it to different responses along
    # the batch dimension.
    # 'BATCH_AND_SIZE_INPUT' is also used as a hint to generate output with
    # batch dimension, 'BATCH_AND_SIZE_INPUT' must have shape [batch_size].
    # Each output corresponds to the input with the same name, so if there
    # are two requests, one has "RAGGED_INPUT" [2, 4] and the other has [1],
    # since the input is ragged, the model sees the input as [2, 4, 1], and
    # "BATCH_AND_SIZE_INPUT" will have shape [2]. Then the model output will
    # be [[2, 4, 1], [2, 4, 1]] and Triton will send responses that each has
    # value [[2, 4, 1]].
    # For "BATCH_INPUT", the input tensor must only have one variable dimension
    # to be broadcasted along the batch dimension properly, thus the currently
    # allowed batch input types are:
    # - BATCH_ACCUMULATED_ELEMENT_COUNT
    # - BATCH_ACCUMULATED_ELEMENT_COUNT_WITH_ZERO
    # - BATCH_MAX_ELEMENT_COUNT_AS_SHAPE
    # - BATCH_ITEM_SHAPE_FLATTEN

    tf_dtype = np_to_tf_dtype(dtype)

    tf.reset_default_graph()
    in_node = tf.placeholder(tf_dtype, tu.shape_to_tf_shape([-1]),
                             "TENSOR_RAGGED_INPUT")
    bs_node = tf.placeholder(tf_dtype, tu.shape_to_tf_shape([-1]),
                             "TENSOR_BATCH_AND_SIZE_INPUT")
    batch_node = tf.placeholder(tf_dtype, tu.shape_to_tf_shape([-1]),
                                "TENSOR_BATCH_INPUT")

    in_mat = tf.reshape(in_node, [1, -1])
    bs_mat = tf.reshape(bs_node, [1, -1])
    batch_mat = tf.reshape(batch_node, [1, -1])

    output_expander = tf.reshape(tf.divide(bs_node, bs_node), [-1, 1])

    out_node = tf.matmul(output_expander, in_mat, name="TENSOR_RAGGED_OUTPUT")
    bs_out_node = tf.matmul(output_expander,
                            bs_mat,
                            name="TENSOR_BATCH_AND_SIZE_OUTPUT")
    batch_out_node = tf.matmul(output_expander,
                               batch_mat,
                               name="TENSOR_BATCH_OUTPUT")

    model_name = "savedmodel_batch_input"
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with tf.Session() as sess:
        in_tensor = tf.get_default_graph().get_tensor_by_name(
            "TENSOR_RAGGED_INPUT:0")
        bs_tensor = tf.get_default_graph().get_tensor_by_name(
            "TENSOR_BATCH_AND_SIZE_INPUT:0")
        batch_tensor = tf.get_default_graph().get_tensor_by_name(
            "TENSOR_BATCH_INPUT:0")
        out_tensor = tf.get_default_graph().get_tensor_by_name(
            "TENSOR_RAGGED_OUTPUT:0")
        bs_out_tensor = tf.get_default_graph().get_tensor_by_name(
            "TENSOR_BATCH_AND_SIZE_OUTPUT:0")
        batch_out_tensor = tf.get_default_graph().get_tensor_by_name(
            "TENSOR_BATCH_OUTPUT:0")
        tf.saved_model.simple_save(sess,
                                   model_version_dir + "/model.savedmodel",
                                   inputs={
                                       "RAGGED_INPUT": in_tensor,
                                       "BATCH_AND_SIZE_INPUT": bs_tensor,
                                       "BATCH_INPUT": batch_tensor,
                                   },
                                   outputs={
                                       "RAGGED_OUTPUT": out_tensor,
                                       "BATCH_AND_SIZE_OUTPUT": bs_out_tensor,
                                       "BATCH_OUTPUT": batch_out_tensor,
                                   })


def create_plan_modelfile(models_dir, model_version, dtype):
    # Create special identity model for batch input testing.
    # Because the ragged input and batch input are one dimensional vector
    # when passing to the model, the model must generate output with batch
    # dimension so that Triton can scatter it to different responses along
    # the batch dimension.
    # 'BATCH_AND_SIZE_INPUT' is also used as a hint to generate output with
    # batch dimension, 'BATCH_AND_SIZE_INPUT' must have shape [batch_size].
    # Each output corresponds to the input with the same name, so if there
    # are two requests, one has "RAGGED_INPUT" [2, 4] and the other has [1],
    # since the input is ragged, the model sees the input as [2, 4, 1], and
    # "BATCH_AND_SIZE_INPUT" will have shape [2]. Then the model output will
    # be [[2, 4, 1], [2, 4, 1]] and Triton will send responses that each has
    # value [[2, 4, 1]].
    # For "BATCH_INPUT", the input tensor must only have one variable dimension
    # to be broadcasted along the batch dimension properly, thus the currently
    # allowed batch input types are:
    # - BATCH_ACCUMULATED_ELEMENT_COUNT
    # - BATCH_ACCUMULATED_ELEMENT_COUNT_WITH_ZERO
    # - BATCH_MAX_ELEMENT_COUNT_AS_SHAPE
    # - BATCH_ITEM_SHAPE_FLATTEN

    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    trt_dtype = np_to_trt_dtype(dtype)

    in_node = network.add_input("RAGGED_INPUT", trt_dtype, [-1])
    bs_node = network.add_input("BATCH_AND_SIZE_INPUT", trt_dtype, [-1])
    batch_node = network.add_input("BATCH_INPUT", trt_dtype, [-1])

    reshape_dims = trt.Dims([-1, 1])
    in_mat = network.add_shuffle(in_node)
    in_mat.reshape_dims = reshape_dims
    bs_mat = network.add_shuffle(bs_node)
    bs_mat.reshape_dims = reshape_dims
    batch_mat = network.add_shuffle(batch_node)
    batch_mat.reshape_dims = reshape_dims

    batch_entry = network.add_elementwise(bs_mat.get_output(0),
                                          bs_mat.get_output(0),
                                          trt.ElementWiseOperation.DIV)
    out_node = network.add_matrix_multiply(batch_entry.get_output(0),
                                           trt.MatrixOperation.NONE,
                                           in_mat.get_output(0),
                                           trt.MatrixOperation.TRANSPOSE)
    bs_out_node = network.add_matrix_multiply(batch_entry.get_output(0),
                                              trt.MatrixOperation.NONE,
                                              bs_mat.get_output(0),
                                              trt.MatrixOperation.TRANSPOSE)
    batch_out_node = network.add_matrix_multiply(batch_entry.get_output(0),
                                                 trt.MatrixOperation.NONE,
                                                 batch_mat.get_output(0),
                                                 trt.MatrixOperation.TRANSPOSE)
    out_node.get_output(0).name = "RAGGED_OUTPUT"
    bs_out_node.get_output(0).name = "BATCH_AND_SIZE_OUTPUT"
    batch_out_node.get_output(0).name = "BATCH_OUTPUT"
    network.mark_output(out_node.get_output(0))
    network.mark_output(bs_out_node.get_output(0))
    network.mark_output(batch_out_node.get_output(0))

    # Hard coded optimization profile
    min_shape = [1]
    opt_shape = [8]
    max_shape = [32]

    profile = builder.create_optimization_profile()
    for input_name in ["RAGGED_INPUT", "BATCH_AND_SIZE_INPUT", "BATCH_INPUT"]:
        profile.set_shape("{}".format(input_name), min_shape, opt_shape,
                          max_shape)
    config = builder.create_builder_config()
    config.add_optimization_profile(profile)
    config.max_workspace_size = 1 << 20
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    model_name = "plan_batch_input"
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_onnx_modelfile(models_dir, model_version, dtype):
    # Create special identity model for batch input testing.
    # Because the ragged input and batch input are one dimensional vector
    # when passing to the model, the model must generate output with batch
    # dimension so that Triton can scatter it to different responses along
    # the batch dimension.
    # 'BATCH_AND_SIZE_INPUT' is also used as a hint to generate output with
    # batch dimension, 'BATCH_AND_SIZE_INPUT' must have shape [batch_size].
    # Each output corresponds to the input with the same name, so if there
    # are two requests, one has "RAGGED_INPUT" [2, 4] and the other has [1],
    # since the input is ragged, the model sees the input as [2, 4, 1], and
    # "BATCH_AND_SIZE_INPUT" will have shape [2]. Then the model output will
    # be [[2, 4, 1], [2, 4, 1]] and Triton will send responses that each has
    # value [[2, 4, 1]].
    # For "BATCH_INPUT", the input tensor must only have one variable dimension
    # to be broadcasted along the batch dimension properly, thus the currently
    # allowed batch input types are:
    # - BATCH_ACCUMULATED_ELEMENT_COUNT
    # - BATCH_ACCUMULATED_ELEMENT_COUNT_WITH_ZERO
    # - BATCH_MAX_ELEMENT_COUNT_AS_SHAPE
    # - BATCH_ITEM_SHAPE_FLATTEN

    onnx_dtype = np_to_onnx_dtype(dtype)

    # Create the model
    model_name = "onnx_batch_input"
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    in0_shape, idx = tu.shape_to_onnx_shape([-1], 0)
    bs_shape, idx = tu.shape_to_onnx_shape([-1], 0)
    batch_shape, idx = tu.shape_to_onnx_shape([-1], 0)

    in0 = onnx.helper.make_tensor_value_info("RAGGED_INPUT", onnx_dtype,
                                             in0_shape)
    bs_in = onnx.helper.make_tensor_value_info("BATCH_AND_SIZE_INPUT",
                                               onnx_dtype, bs_shape)
    batch_in = onnx.helper.make_tensor_value_info("BATCH_INPUT", onnx_dtype,
                                                  batch_shape)

    out_shape, idx = tu.shape_to_onnx_shape([-1, -1], idx)
    bs_out_shape, idx = tu.shape_to_onnx_shape([-1, -1], idx)
    batch_out_shape, idx = tu.shape_to_onnx_shape([-1, -1], idx)

    out = onnx.helper.make_tensor_value_info("RAGGED_OUTPUT", onnx_dtype,
                                             out_shape)
    bs_out = onnx.helper.make_tensor_value_info("BATCH_AND_SIZE_OUTPUT",
                                                onnx_dtype, bs_out_shape)
    batch_out = onnx.helper.make_tensor_value_info("BATCH_OUTPUT", onnx_dtype,
                                                   batch_out_shape)

    const_node_shape = onnx.helper.make_node(
        'Constant', [], ["shape"],
        value=onnx.helper.make_tensor("const_shape", onnx.TensorProto.INT64,
                                      [2], [1, -1]))

    const_node_expander_shape = onnx.helper.make_node(
        'Constant', [], ["expander_shape"],
        value=onnx.helper.make_tensor("const_expander_shape",
                                      onnx.TensorProto.INT64, [2], [-1, 1]))

    in0_mat_node = onnx.helper.make_node("Reshape", ["RAGGED_INPUT", "shape"],
                                         ["in_mat"])
    bs_mat_node = onnx.helper.make_node("Reshape",
                                        ["BATCH_AND_SIZE_INPUT", "shape"],
                                        ["bs_mat"])
    batch_mat_node = onnx.helper.make_node("Reshape", ["BATCH_INPUT", "shape"],
                                           ["batch_mat"])

    internal_node_div = onnx.helper.make_node(
        "Div", ["BATCH_AND_SIZE_INPUT", "BATCH_AND_SIZE_INPUT"],
        ["output_expander_int"])
    internal_node_reshape = onnx.helper.make_node(
        "Reshape", ["output_expander_int", "expander_shape"],
        ["output_expander"])

    out_node = onnx.helper.make_node("MatMul", ["output_expander", "in_mat"],
                                     ["RAGGED_OUTPUT"])
    bs_out_node = onnx.helper.make_node("MatMul", ["output_expander", "bs_mat"],
                                        ["BATCH_AND_SIZE_OUTPUT"])
    batch_out_node = onnx.helper.make_node("MatMul",
                                           ["output_expander", "batch_mat"],
                                           ["BATCH_OUTPUT"])

    onnx_nodes = [
        const_node_shape, const_node_expander_shape, in0_mat_node, bs_mat_node,
        batch_mat_node, internal_node_div, internal_node_reshape, out_node,
        bs_out_node, batch_out_node
    ]
    onnx_inputs = [in0, bs_in, batch_in]
    onnx_outputs = [out, bs_out, batch_out]

    graph_proto = onnx.helper.make_graph(onnx_nodes, model_name, onnx_inputs,
                                         onnx_outputs)
    if FLAGS.onnx_opset > 0:
        model_opset = onnx.helper.make_operatorsetid("", FLAGS.onnx_opset)
        model_def = onnx.helper.make_model(graph_proto,
                                           producer_name="triton",
                                           opset_imports=[model_opset])
    else:
        model_def = onnx.helper.make_model(graph_proto, producer_name="triton")

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    onnx.save(model_def, model_version_dir + "/model.onnx")


def create_modelconfig(models_dir, max_batch, model_version, dtype, backend,
                       platform):
    version_policy_str = "{ latest { num_versions: 1 }}"

    backend_spec = '''
backend: "{}"
'''.format(backend)
    if backend == "tensorflow":
        backend_spec += '''
platform: "{}_{}"
'''.format(backend, platform)

    model_name = "{}_batch_input".format(platform)
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
{}
max_batch_size: {}
version_policy: {}
input [
  {{
    name: "RAGGED_INPUT"
    data_type: {data_type}
    dims: [ -1 ]
    allow_ragged_batch: true
  }}
]
output [
  {{
    name: "RAGGED_OUTPUT"
    data_type: {data_type}
    dims: [ -1 ]
   }}
]
output [
  {{
    name: "BATCH_AND_SIZE_OUTPUT"
    data_type: {data_type}
    dims: [ -1 ]
   }}
]
output [
  {{
    name: "BATCH_OUTPUT"
    data_type: {data_type}
    dims: [ -1 ]
   }}
]
batch_input [
  {{
    kind: BATCH_ELEMENT_COUNT
    target_name: "BATCH_AND_SIZE_INPUT"
    data_type: {data_type}
    source_input: "RAGGED_INPUT"
  }},
  {{
    kind: BATCH_ACCUMULATED_ELEMENT_COUNT_WITH_ZERO
    target_name: "BATCH_INPUT"
    data_type: {data_type}
    source_input: "RAGGED_INPUT"
  }}
]
dynamic_batching {{
  max_queue_delay_microseconds: 1000000
}}
'''.format(model_name,
           backend_spec,
           max_batch,
           version_policy_str,
           data_type=np_to_model_dtype(dtype))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_savedmodel_itemshape_modelfile(models_dir, model_version, dtype):
    # Create special identity model for batch input 'BATCH_ITEM_SHAPE' testing,
    # such model has one ragged input and one batch input, and one output to
    # return the batch input directly. Because 'BATCH_ITEM_SHAPE' should be
    # generated to have matching batch dimension, the output can be produced
    # via identity op and expect Triton will scatter the output properly.

    tf_dtype = np_to_tf_dtype(dtype)

    tf.reset_default_graph()
    in_node = tf.placeholder(tf_dtype, tu.shape_to_tf_shape([-1]),
                             "TENSOR_RAGGED_INPUT")
    # Shape is predefined
    batch_node = tf.placeholder(tf_dtype, tu.shape_to_tf_shape([-1, 2]),
                                "TENSOR_BATCH_INPUT")
    batch_output_node = tf.identity(batch_node, name="TENSOR_BATCH_OUTPUT")

    model_name = "savedmodel_batch_item"
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with tf.Session() as sess:
        in_tensor = tf.get_default_graph().get_tensor_by_name(
            "TENSOR_RAGGED_INPUT:0")
        batch_tensor = tf.get_default_graph().get_tensor_by_name(
            "TENSOR_BATCH_INPUT:0")
        batch_out_tensor = tf.get_default_graph().get_tensor_by_name(
            "TENSOR_BATCH_OUTPUT:0")
        tf.saved_model.simple_save(sess,
                                   model_version_dir + "/model.savedmodel",
                                   inputs={
                                       "RAGGED_INPUT": in_tensor,
                                       "BATCH_INPUT": batch_tensor,
                                   },
                                   outputs={
                                       "BATCH_OUTPUT": batch_out_tensor,
                                   })


def create_plan_itemshape_modelfile(models_dir, model_version, dtype):
    # Create special identity model for batch input 'BATCH_ITEM_SHAPE' testing,
    # such model has one ragged input and one batch input, and one output to
    # return the batch input directly. Because 'BATCH_ITEM_SHAPE' should be
    # generated to have matching batch dimension, the output can be produced
    # via identity op and expect Triton will scatter the output properly.

    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    trt_dtype = np_to_trt_dtype(dtype)

    in_node = network.add_input("RAGGED_INPUT", trt_dtype, [-1])
    batch_node = network.add_input("BATCH_INPUT", trt_dtype, [-1, 2])

    batch_out_node = network.add_identity(batch_node)
    batch_out_node.get_output(0).name = "BATCH_OUTPUT"
    network.mark_output(batch_out_node.get_output(0))

    # Hard coded optimization profile
    min_shape = [1]
    opt_shape = [8]
    max_shape = [32]

    profile = builder.create_optimization_profile()
    profile.set_shape("RAGGED_INPUT", min_shape, opt_shape, max_shape)
    profile.set_shape("BATCH_INPUT", min_shape + [2], opt_shape + [2],
                      max_shape + [2])
    config = builder.create_builder_config()
    config.add_optimization_profile(profile)
    config.max_workspace_size = 1 << 20
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    model_name = "plan_batch_item"
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_onnx_itemshape_modelfile(models_dir, model_version, dtype):
    # Create special identity model for batch input 'BATCH_ITEM_SHAPE' testing,
    # such model has one ragged input and one batch input, and one output to
    # return the batch input directly. Because 'BATCH_ITEM_SHAPE' should be
    # generated to have matching batch dimension, the output can be produced
    # via identity op and expect Triton will scatter the output properly.

    onnx_dtype = np_to_onnx_dtype(dtype)

    # Create the model
    model_name = "onnx_batch_item"
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    in0_shape, idx = tu.shape_to_onnx_shape([-1], 0)
    batch_shape, idx = tu.shape_to_onnx_shape([-1, 2], 0)

    in0 = onnx.helper.make_tensor_value_info("RAGGED_INPUT", onnx_dtype,
                                             in0_shape)
    batch_in = onnx.helper.make_tensor_value_info("BATCH_INPUT", onnx_dtype,
                                                  batch_shape)

    batch_out_shape, idx = tu.shape_to_onnx_shape([-1, -1], idx)
    batch_out = onnx.helper.make_tensor_value_info("BATCH_OUTPUT", onnx_dtype,
                                                   batch_out_shape)

    onnx_nodes = [
        onnx.helper.make_node("Identity", ["BATCH_INPUT"], ["BATCH_OUTPUT"])
    ]
    onnx_inputs = [in0, batch_in]
    onnx_outputs = [batch_out]

    graph_proto = onnx.helper.make_graph(onnx_nodes, model_name, onnx_inputs,
                                         onnx_outputs)
    if FLAGS.onnx_opset > 0:
        model_opset = onnx.helper.make_operatorsetid("", FLAGS.onnx_opset)
        model_def = onnx.helper.make_model(graph_proto,
                                           producer_name="triton",
                                           opset_imports=[model_opset])
    else:
        model_def = onnx.helper.make_model(graph_proto, producer_name="triton")

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    onnx.save(model_def, model_version_dir + "/model.onnx")


def create_itemshape_modelconfig(models_dir, max_batch, model_version, dtype,
                                 backend, platform):
    version_policy_str = "{ latest { num_versions: 1 }}"

    backend_spec = '''
backend: "{}"
'''.format(backend)
    if backend == "tensorflow":
        backend_spec += '''
platform: "{}_{}"
'''.format(backend, platform)

    model_name = "{}_batch_item".format(platform)
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
{}
max_batch_size: {}
version_policy: {}
input [
  {{
    name: "RAGGED_INPUT"
    data_type: {data_type}
    dims: [ -1, -1 ]
    allow_ragged_batch: true
  }}
]
output [
  {{
    name: "BATCH_OUTPUT"
    data_type: {data_type}
    dims: [ 2 ]
   }}
]
batch_input [
  {{
    kind: BATCH_ITEM_SHAPE
    target_name: "BATCH_INPUT"
    data_type: {data_type}
    source_input: "RAGGED_INPUT"
  }}
]
dynamic_batching {{
  max_queue_delay_microseconds: 1000000
}}
'''.format(model_name,
           backend_spec,
           max_batch,
           version_policy_str,
           data_type=np_to_model_dtype(dtype))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_batch_input_models(models_dir):
    model_version = 1
    if FLAGS.tensorrt:
        create_modelconfig(models_dir, 4, model_version, np.float32, "tensorrt",
                           "plan")
        create_plan_modelfile(models_dir, model_version, np.float32)
        create_itemshape_modelconfig(models_dir, 4, model_version, np.float32,
                                     "tensorrt", "plan")
        create_plan_itemshape_modelfile(models_dir, model_version, np.float32)
    if FLAGS.savedmodel:
        create_modelconfig(models_dir, 4, model_version, np.float32,
                           "tensorflow", "savedmodel")
        create_savedmodel_modelfile(models_dir, model_version, np.float32)
        create_itemshape_modelconfig(models_dir, 4, model_version, np.float32,
                                     "tensorflow", "savedmodel")
        create_savedmodel_itemshape_modelfile(models_dir, model_version,
                                              np.float32)
    if FLAGS.onnx:
        create_modelconfig(models_dir, 4, model_version, np.float32,
                           "onnxruntime", "onnx")
        create_onnx_modelfile(models_dir, model_version, np.float32)
        create_itemshape_modelconfig(models_dir, 4, model_version, np.float32,
                                     "onnxruntime", "onnx")
        create_onnx_itemshape_modelfile(models_dir, model_version, np.float32)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--models_dir',
                        type=str,
                        required=True,
                        help='Top-level model directory')
    parser.add_argument('--tensorrt',
                        required=False,
                        action='store_true',
                        help='Generate TensorRT PLAN models')
    parser.add_argument('--savedmodel',
                        required=False,
                        action='store_true',
                        help='Generate SavedModel models')
    parser.add_argument('--graphdef',
                        required=False,
                        action='store_true',
                        help='Generate GraphDef models')
    parser.add_argument('--onnx',
                        required=False,
                        action='store_true',
                        help='Generate Onnx Runtime Onnx models')
    parser.add_argument(
        '--onnx_opset',
        type=int,
        required=False,
        default=0,
        help='Opset used for Onnx models. Default is to use ONNXRT default')

    FLAGS, unparsed = parser.parse_known_args()

    import test_util as tu
    if FLAGS.tensorrt:
        import tensorrt as trt
    if FLAGS.graphdef or FLAGS.savedmodel:
        import tensorflow as tf
    if FLAGS.onnx:
        import onnx

    create_batch_input_models(FLAGS.models_dir)


# qa/common/gen_qa_reshape_models.py

# Copyright 2019-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
from builtins import range
import os
import numpy as np
import gen_ensemble_model_utils as emu

FLAGS = None
np_dtype_string = np.dtype(object)
from typing import List


def np_to_model_dtype(np_dtype):
    if np_dtype == bool:
        return "TYPE_BOOL"
    elif np_dtype == np.int8:
        return "TYPE_INT8"
    elif np_dtype == np.int16:
        return "TYPE_INT16"
    elif np_dtype == np.int32:
        return "TYPE_INT32"
    elif np_dtype == np.int64:
        return "TYPE_INT64"
    elif np_dtype == np.uint8:
        return "TYPE_UINT8"
    elif np_dtype == np.uint16:
        return "TYPE_UINT16"
    elif np_dtype == np.float16:
        return "TYPE_FP16"
    elif np_dtype == np.float32:
        return "TYPE_FP32"
    elif np_dtype == np.float64:
        return "TYPE_FP64"
    elif np_dtype == np_dtype_string:
        return "TYPE_STRING"
    return None


def np_to_tf_dtype(np_dtype):
    if np_dtype == bool:
        return tf.bool
    elif np_dtype == np.int8:
        return tf.int8
    elif np_dtype == np.int16:
        return tf.int16
    elif np_dtype == np.int32:
        return tf.int32
    elif np_dtype == np.int64:
        return tf.int64
    elif np_dtype == np.uint8:
        return tf.uint8
    elif np_dtype == np.uint16:
        return tf.uint16
    elif np_dtype == np.float16:
        return tf.float16
    elif np_dtype == np.float32:
        return tf.float32
    elif np_dtype == np.float64:
        return tf.float64
    elif np_dtype == np_dtype_string:
        return tf.string
    return None


def np_to_trt_dtype(np_dtype):
    if np_dtype == bool:
        return trt.bool
    elif np_dtype == np.int8:
        return trt.int8
    elif np_dtype == np.int32:
        return trt.int32
    elif np_dtype == np.float16:
        return trt.float16
    elif np_dtype == np.float32:
        return trt.float32
    return None


def np_to_onnx_dtype(np_dtype):
    if np_dtype == bool:
        return onnx.TensorProto.BOOL
    elif np_dtype == np.int8:
        return onnx.TensorProto.INT8
    elif np_dtype == np.int16:
        return onnx.TensorProto.INT16
    elif np_dtype == np.int32:
        return onnx.TensorProto.INT32
    elif np_dtype == np.int64:
        return onnx.TensorProto.INT64
    elif np_dtype == np.uint8:
        return onnx.TensorProto.UINT8
    elif np_dtype == np.uint16:
        return onnx.TensorProto.UINT16
    elif np_dtype == np.float16:
        return onnx.TensorProto.FLOAT16
    elif np_dtype == np.float32:
        return onnx.TensorProto.FLOAT
    elif np_dtype == np.float64:
        return onnx.TensorProto.DOUBLE
    elif np_dtype == np_dtype_string:
        return onnx.TensorProto.STRING
    return None


def np_to_torch_dtype(np_dtype):
    if np_dtype == bool:
        return torch.bool
    elif np_dtype == np.int8:
        return torch.int8
    elif np_dtype == np.int16:
        return torch.int16
    elif np_dtype == np.int32:
        return torch.int
    elif np_dtype == np.int64:
        return torch.long
    elif np_dtype == np.uint8:
        return torch.uint8
    elif np_dtype == np.uint16:
        return None  # Not supported in Torch
    elif np_dtype == np.float16:
        return None
    elif np_dtype == np.float32:
        return torch.float
    elif np_dtype == np.float64:
        return torch.double
    elif np_dtype == np_dtype_string:
        return List[str]
    return None


def create_tf_modelfile(create_savedmodel, models_dir, model_version, max_batch,
                        dtype, input_shapes, output_shapes):

    assert len(input_shapes) == len(output_shapes)
    if not tu.validate_for_tf_model(dtype, dtype, dtype, input_shapes[0],
                                    input_shapes[0], input_shapes[0]):
        return

    tf_dtype = np_to_tf_dtype(dtype)
    io_cnt = len(input_shapes)

    # Create the model that copies inputs to corresponding outputs.
    tf.reset_default_graph()
    for io_num in range(io_cnt):
        input_name = "INPUT{}".format(io_num)
        output_name = "OUTPUT{}".format(io_num)
        if max_batch == 0:
            tin = tf.placeholder(tf_dtype,
                                 tu.shape_to_tf_shape(input_shapes[io_num]),
                                 input_name)
        else:
            tin = tf.placeholder(tf_dtype, [
                None,
            ] + tu.shape_to_tf_shape(input_shapes[io_num]), input_name)

        if input_shapes == output_shapes:
            toutput = tf.identity(tin, name=output_name)
        else:
            if max_batch == 0:
                toutput = tf.reshape(tin,
                                     output_shapes[io_num],
                                     name=output_name)
            else:
                toutput = tf.reshape(tin, [
                    -1,
                ] + output_shapes[io_num],
                                     name=output_name)

    # Use model name based on input/output count and non-batching variant
    if create_savedmodel:
        model_name = tu.get_zero_model_name(
            "savedmodel_nobatch" if max_batch == 0 else "savedmodel", io_cnt,
            dtype)
    else:
        model_name = tu.get_zero_model_name(
            "graphdef_nobatch" if max_batch == 0 else "graphdef", io_cnt, dtype)

    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    if create_savedmodel:
        with tf.Session() as sess:
            input_dict = {}
            output_dict = {}
            for io_num in range(io_cnt):
                input_name = "INPUT{}".format(io_num)
                output_name = "OUTPUT{}".format(io_num)
                input_tensor = tf.get_default_graph().get_tensor_by_name(
                    input_name + ":0")
                output_tensor = tf.get_default_graph().get_tensor_by_name(
                    output_name + ":0")
                input_dict[input_name] = input_tensor
                output_dict[output_name] = output_tensor
            tf.saved_model.simple_save(sess,
                                       model_version_dir + "/model.savedmodel",
                                       inputs=input_dict,
                                       outputs=output_dict)
    else:
        with tf.Session() as sess:
            graph_io.write_graph(sess.graph.as_graph_def(),
                                 model_version_dir,
                                 "model.graphdef",
                                 as_text=False)


def create_tf_modelconfig(create_savedmodel, models_dir, model_version,
                          max_batch, dtype, input_shapes, input_model_shapes,
                          output_shapes, output_model_shapes):

    assert len(input_shapes) == len(input_model_shapes)
    assert len(output_shapes) == len(output_model_shapes)
    assert len(input_shapes) == len(output_shapes)
    if not tu.validate_for_tf_model(dtype, dtype, dtype, input_shapes[0],
                                    input_shapes[0], input_shapes[0]):
        return

    io_cnt = len(input_shapes)

    # Use a different model name for the non-batching variant
    if create_savedmodel:
        model_name = tu.get_zero_model_name(
            "savedmodel_nobatch" if max_batch == 0 else "savedmodel", io_cnt,
            dtype)
    else:
        model_name = tu.get_zero_model_name(
            "graphdef_nobatch" if max_batch == 0 else "graphdef", io_cnt, dtype)

    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "{}"
max_batch_size: {}
'''.format(
        model_name,
        "tensorflow_savedmodel" if create_savedmodel else "tensorflow_graphdef",
        max_batch)

    for io_num in range(io_cnt):
        config += '''
input [
  {{
    name: "INPUT{}"
    data_type: {}
    dims: [ {} ]
    {}
  }}
]
output [
  {{
    name: "OUTPUT{}"
    data_type: {}
    dims: [ {} ]
    {}
  }}
]
'''.format(
            io_num, np_to_model_dtype(dtype),
            tu.shape_to_dims_str(input_shapes[io_num]),
            "reshape: {{ shape: [ {} ] }}".format(
                tu.shape_to_dims_str(input_model_shapes[io_num]))
            if input_shapes[io_num] != input_model_shapes[io_num] else "",
            io_num, np_to_model_dtype(dtype),
            tu.shape_to_dims_str(output_shapes[io_num]),
            "reshape: {{ shape: [ {} ] }}".format(
                tu.shape_to_dims_str(output_model_shapes[io_num]))
            if output_shapes[io_num] != output_model_shapes[io_num] else "")

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_plan_modelfile(models_dir, model_version, max_batch, dtype,
                          input_shapes, output_shapes):

    assert len(input_shapes) == len(output_shapes)
    if not tu.validate_for_trt_model(dtype, dtype, dtype, input_shapes[0],
                                     input_shapes[0], input_shapes[0]):
        return

    trt_dtype = np_to_trt_dtype(dtype)
    io_cnt = len(input_shapes)

    # Create the model that copies inputs to corresponding outputs.
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network()

    for io_num in range(io_cnt):
        input_name = "INPUT{}".format(io_num)
        output_name = "OUTPUT{}".format(io_num)
        in0 = network.add_input(input_name, trt_dtype, input_shapes[io_num])
        if input_shapes == output_shapes:
            out0 = network.add_identity(in0)
        else:
            out0 = network.add_shuffle(in0)
            out0.set_reshape_dimensions(output_shapes[io_num])

        out0.get_output(0).name = output_name
        network.mark_output(out0.get_output(0))

    config = builder.create_builder_config()
    config.max_workspace_size = 1 << 20
    builder.max_batch_size = max(1, max_batch)
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine
    del network

    model_name = tu.get_zero_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", io_cnt, dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_modelconfig(models_dir, model_version, max_batch, dtype,
                            input_shapes, input_model_shapes, output_shapes,
                            output_model_shapes):

    assert len(input_shapes) == len(input_model_shapes)
    assert len(output_shapes) == len(output_model_shapes)
    assert len(input_shapes) == len(output_shapes)
    if not tu.validate_for_trt_model(dtype, dtype, dtype, input_shapes[0],
                                     input_shapes[0], input_shapes[0]):
        return

    io_cnt = len(input_shapes)

    model_name = tu.get_zero_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", io_cnt, dtype)
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "tensorrt_plan"
max_batch_size: {}
'''.format(model_name, max_batch)

    for io_num in range(io_cnt):
        config += '''
input [
  {{
    name: "INPUT{}"
    data_type: {}
    dims: [ {} ]
    {}
  }}
]
output [
  {{
    name: "OUTPUT{}"
    data_type: {}
    dims: [ {} ]
    {}
  }}
]
'''.format(
            io_num, np_to_model_dtype(dtype),
            tu.shape_to_dims_str(input_shapes[io_num]),
            "reshape: {{ shape: [ {} ] }}".format(
                tu.shape_to_dims_str(input_model_shapes[io_num]))
            if input_shapes[io_num] != input_model_shapes[io_num] else "",
            io_num, np_to_model_dtype(dtype),
            tu.shape_to_dims_str(output_shapes[io_num]),
            "reshape: {{ shape: [ {} ] }}".format(
                tu.shape_to_dims_str(output_model_shapes[io_num]))
            if output_shapes[io_num] != output_model_shapes[io_num] else "")

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_libtorch_modelfile(models_dir, model_version, max_batch, dtype,
                              input_shapes, output_shapes):

    assert len(input_shapes) == len(output_shapes)
    if not tu.validate_for_libtorch_model(dtype,
                                          dtype,
                                          dtype,
                                          input_shapes[0],
                                          input_shapes[0],
                                          input_shapes[0],
                                          max_batch,
                                          reshape=True):
        return

    torch_dtype = np_to_torch_dtype(dtype)
    io_cnt = len(input_shapes)
    model_name = tu.get_zero_model_name(
        "libtorch_nobatch" if max_batch == 0 else "libtorch", io_cnt, dtype)

    # Create the model that reshapes inputs to corresponding outputs
    # Note that string I/O is supported only for 1-dimensional inputs/outputs.
    # Use identity model for string I/O models and add 'reshape' field with
    # empty shape so that batching is supported and the full shape becomes [-1].
    if io_cnt == 1:
        if (dtype == np_dtype_string):

            class IdentityNet(nn.Module):

                def __init__(self):
                    super(IdentityNet, self).__init__()

                def forward(self, input0: List[str]) -> List[str]:
                    return input0
        else:

            class ReshapeNet(nn.Module):

                def __init__(self, *args):
                    super(ReshapeNet, self).__init__()
                    self.shape = args[0][0]
                    self.max_batch = args[0][1]

                def forward(self, input0):
                    if self.max_batch == 0:
                        return input0.view(self.shape[0])
                    else:
                        return input0.view([
                            -1,
                        ] + self.shape[0])
    elif io_cnt == 2:
        if (dtype == np_dtype_string):

            class IdentityNet(nn.Module):

                def __init__(self):
                    super(IdentityNet, self).__init__()

                def forward(self, input0: List[str],
                            input1: List[str]) -> Tuple[List[str], List[str]]:
                    return input0, input1
        else:

            class ReshapeNet(nn.Module):

                def __init__(self, *args):
                    super(ReshapeNet, self).__init__()
                    self.shape = args[0][0]
                    self.max_batch = args[0][1]

                def forward(self, input0, input1):
                    if self.max_batch == 0:
                        return input0.view(self.shape[0]), input1.view(
                            self.shape[1])
                    else:
                        return input0.view([
                            -1,
                        ] + self.shape[0]), input1.view([
                            -1,
                        ] + self.shape[1])
    elif io_cnt == 3:
        if (dtype == np_dtype_string):

            class IdentityNet(nn.Module):

                def __init__(self):
                    super(IdentityNet, self).__init__()

                def forward(
                    self, input0: List[str], input1: List[str],
                    input2: List[str]
                ) -> Tuple[List[str], List[str], List[str]]:
                    return input0, input1, input2
        else:

            class ReshapeNet(nn.Module):

                def __init__(self, *args):
                    super(ReshapeNet, self).__init__()
                    self.shape = args[0][0]
                    self.max_batch = args[0][1]

                def forward(self, input0, input1, input2):
                    if self.max_batch == 0:
                        return input0.view(self.shape[0]), input1.view(
                            self.shape[1]), input2.view(self.shape[2])
                    else:
                        return input0.view([-1,]+self.shape[0]), input1.view([-1,]+self.shape[1]), \
                            input2.view([-1,]+self.shape[2])
    elif io_cnt == 4:
        if (dtype == np_dtype_string):

            class IdentityNet(nn.Module):

                def __init__(self):
                    super(IdentityNet, self).__init__()

                def forward(
                    self, input0: List[str], input1: List[str],
                    input2: List[str], input3: List[str]
                ) -> Tuple[List[str], List[str], List[str], List[str]]:
                    return input0, input1, input2, input3
        else:

            class ReshapeNet(nn.Module):

                def __init__(self, *args):
                    super(ReshapeNet, self).__init__()
                    self.shape = args[0][0]
                    self.max_batch = args[0][1]

                def forward(self, input0, input1, input2, input3):
                    if self.max_batch == 0:
                        return input0.view(self.shape[0]), input1.view(self.shape[1]), input2.view(self.shape[2]), \
                            input3.view(self.shape[3])
                    else:
                        return input0.view([-1,]+self.shape[0]), input1.view([-1,]+self.shape[1]), \
                            input2.view([-1,]+self.shape[2]), input3.view([-1,]+self.shape[3])

    if (dtype == np_dtype_string):
        identityModel = IdentityNet()
        traced = torch.jit.script(identityModel)
    else:
        reshapeModel = ReshapeNet([[op_shape for op_shape in output_shapes],
                                   max_batch])
        traced = torch.jit.script(reshapeModel)

    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    traced.save(model_version_dir + "/model.pt")


def create_libtorch_modelconfig(models_dir, model_version, max_batch, dtype,
                                input_shapes, input_model_shapes, output_shapes,
                                output_model_shapes):

    assert len(input_shapes) == len(input_model_shapes)
    assert len(output_shapes) == len(output_model_shapes)
    assert len(input_shapes) == len(output_shapes)
    if not tu.validate_for_libtorch_model(dtype,
                                          dtype,
                                          dtype,
                                          input_shapes[0],
                                          input_shapes[0],
                                          input_shapes[0],
                                          max_batch,
                                          reshape=True):
        return

    io_cnt = len(input_shapes)

    model_name = tu.get_zero_model_name(
        "libtorch_nobatch" if max_batch == 0 else "libtorch", io_cnt, dtype)
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "pytorch_libtorch"
max_batch_size: {}
'''.format(model_name, max_batch)

    for io_num in range(io_cnt):
        config += '''
input [
  {{
    name: "INPUT__{}"
    data_type: {}
    dims: [ {} ]
    {}
  }}
]
output [
  {{
    name: "OUTPUT__{}"
    data_type: {}
    dims: [ {} ]
    {}
  }}
]
'''.format(
            io_num, np_to_model_dtype(dtype),
            tu.shape_to_dims_str(input_shapes[io_num]),
            "reshape: {{ shape: [ {} ] }}".format(
                tu.shape_to_dims_str(input_model_shapes[io_num]))
            if input_shapes[io_num] != input_model_shapes[io_num] else "",
            io_num, np_to_model_dtype(dtype),
            tu.shape_to_dims_str(output_shapes[io_num]),
            "reshape: {{ shape: [ {} ] }}".format(
                tu.shape_to_dims_str(output_model_shapes[io_num]))
            if output_shapes[io_num] != output_model_shapes[io_num] else "")

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_ensemble_modelfile(models_dir, model_version, max_batch, dtype,
                              input_shapes, output_shapes):

    assert len(input_shapes) == len(output_shapes)
    if not tu.validate_for_ensemble_model("reshape", dtype, dtype, dtype,
                                          input_shapes[0], input_shapes[0],
                                          input_shapes[0]):
        return

    emu.create_identity_ensemble_modelfile("reshape", models_dir, model_version,
                                           max_batch, dtype, input_shapes,
                                           output_shapes)


def create_ensemble_modelconfig(models_dir, model_version, max_batch, dtype,
                                input_shapes, input_model_shapes, output_shapes,
                                output_model_shapes):

    assert len(input_shapes) == len(input_model_shapes)
    assert len(output_shapes) == len(output_model_shapes)
    assert len(input_shapes) == len(output_shapes)
    if not tu.validate_for_ensemble_model("reshape", dtype, dtype, dtype,
                                          input_shapes[0], input_shapes[0],
                                          input_shapes[0]):
        return

    # No reason to reshape ensemble inputs / outputs to empty as the inner models
    # have to have non-empty shapes for inputs / outputs.
    input_model_shapes_list = []
    output_model_shapes_list = []
    for idx in range(len(input_shapes)):
        if len(input_model_shapes[idx]) == 0:
            input_model_shapes_list.append(input_shapes[idx])
        else:
            input_model_shapes_list.append(input_model_shapes[idx])
        if len(output_model_shapes[idx]) == 0:
            output_model_shapes_list.append(output_shapes[idx])
        else:
            output_model_shapes_list.append(output_model_shapes[idx])

    emu.create_identity_ensemble_modelconfig("reshape", models_dir,
                                             model_version, max_batch, dtype,
                                             input_shapes,
                                             tuple(input_model_shapes_list),
                                             output_shapes,
                                             tuple(output_model_shapes_list))


def create_onnx_modelfile(models_dir, model_version, max_batch, dtype,
                          input_shapes, output_shapes):

    assert len(input_shapes) == len(output_shapes)
    if not tu.validate_for_onnx_model(dtype, dtype, dtype, input_shapes[0],
                                      input_shapes[0], input_shapes[0]):
        return

    onnx_dtype = np_to_onnx_dtype(dtype)
    io_cnt = len(input_shapes)

    # Create the model
    model_name = tu.get_zero_model_name(
        "onnx_nobatch" if max_batch == 0 else "onnx", io_cnt, dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    batch_dim = [] if max_batch == 0 else [None]

    onnx_nodes = []
    onnx_inputs = []
    onnx_outputs = []
    idx = 0
    for io_num in range(io_cnt):
        # Repeat so that the variable dimension name is different
        in_shape, idx = tu.shape_to_onnx_shape(input_shapes[io_num], idx)
        out_shape, idx = tu.shape_to_onnx_shape(output_shapes[io_num], idx)
        in_name = "INPUT{}".format(io_num)
        out_name = "OUTPUT{}".format(io_num)
        out_shape_name = out_name + "_shape"

        onnx_inputs.append(
            onnx.helper.make_tensor_value_info(in_name, onnx_dtype,
                                               batch_dim + in_shape))
        onnx_outputs.append(
            onnx.helper.make_tensor_value_info(out_name, onnx_dtype,
                                               batch_dim + out_shape))

        if input_shapes == output_shapes:
            onnx_nodes.append(
                onnx.helper.make_node("Identity", [in_name], [out_name]))
        else:
            onnx_nodes.append(
                onnx.helper.make_node("Shape", [out_name], [out_shape_name]))
            onnx_nodes.append(
                onnx.helper.make_node("Reshape", [in_name, out_shape_name],
                                      [out_name]))

    graph_proto = onnx.helper.make_graph(onnx_nodes, model_name, onnx_inputs,
                                         onnx_outputs)
    if FLAGS.onnx_opset > 0:
        model_opset = onnx.helper.make_operatorsetid("", FLAGS.onnx_opset)
        model_def = onnx.helper.make_model(graph_proto,
                                           producer_name="triton",
                                           opset_imports=[model_opset])
    else:
        model_def = onnx.helper.make_model(graph_proto, producer_name="triton")

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    onnx.save(model_def, model_version_dir + "/model.onnx")


def create_onnx_modelconfig(models_dir, model_version, max_batch, dtype,
                            input_shapes, input_model_shapes, output_shapes,
                            output_model_shapes):

    assert len(input_shapes) == len(input_model_shapes)
    assert len(output_shapes) == len(output_model_shapes)
    assert len(input_shapes) == len(output_shapes)
    if not tu.validate_for_onnx_model(dtype, dtype, dtype, input_shapes[0],
                                      input_shapes[0], input_shapes[0]):
        return

    io_cnt = len(input_shapes)

    # Use a different model name for the non-batching variant
    model_name = tu.get_zero_model_name(
        "onnx_nobatch" if max_batch == 0 else "onnx", io_cnt, dtype)
    config_dir = models_dir + "/" + model_name

    config = emu.create_general_modelconfig(model_name,
                                            "onnxruntime_onnx",
                                            max_batch,
                                            emu.repeat(dtype, io_cnt),
                                            input_shapes,
                                            input_model_shapes,
                                            emu.repeat(dtype, io_cnt),
                                            output_shapes,
                                            output_model_shapes,
                                            emu.repeat(None, io_cnt),
                                            force_tensor_number_suffix=True)

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_openvino_modelfile(models_dir, model_version, max_batch, dtype,
                              input_shapes, output_shapes):

    assert len(input_shapes) == len(output_shapes)
    batch_dim = [] if max_batch == 0 else [
        max_batch,
    ]
    if not tu.validate_for_openvino_model(
            dtype, dtype, dtype, batch_dim + input_shapes[0],
            batch_dim + input_shapes[0], batch_dim + input_shapes[0]):
        return

    io_cnt = len(input_shapes)

    # Create the model
    model_name = tu.get_zero_model_name(
        "openvino_nobatch" if max_batch == 0 else "openvino", io_cnt, dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    openvino_inputs = []
    openvino_outputs = []
    for io_num in range(io_cnt):
        in_name = "INPUT{}".format(io_num)
        out_name = "OUTPUT{}".format(io_num)
        openvino_inputs.append(
            ng.parameter(shape=batch_dim + input_shapes[io_num],
                         dtype=dtype,
                         name=in_name))

        openvino_outputs.append(
            ng.reshape(openvino_inputs[io_num],
                       batch_dim + output_shapes[io_num],
                       name=out_name,
                       special_zero=False))

    function = ng.impl.Function(openvino_outputs, openvino_inputs, model_name)
    ie_network = IENetwork(ng.impl.Function.to_capsule(function))

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    ie_network.serialize(model_version_dir + "/model.xml",
                         model_version_dir + "/model.bin")


def create_openvino_modelconfig(models_dir, model_version, max_batch, dtype,
                                input_shapes, input_model_shapes, output_shapes,
                                output_model_shapes):

    assert len(input_shapes) == len(input_model_shapes)
    assert len(output_shapes) == len(output_model_shapes)
    assert len(input_shapes) == len(output_shapes)
    batch_dim = [] if max_batch == 0 else [
        max_batch,
    ]
    if not tu.validate_for_openvino_model(
            dtype, dtype, dtype, batch_dim + input_shapes[0],
            batch_dim + input_shapes[0], batch_dim + input_shapes[0]):
        return

    io_cnt = len(input_shapes)

    # Use a different model name for the non-batching variant
    model_name = tu.get_zero_model_name(
        "openvino_nobatch" if max_batch == 0 else "openvino", io_cnt, dtype)
    config_dir = models_dir + "/" + model_name

    config = '''
name: "{}"
backend: "openvino"
max_batch_size: {}
'''.format(model_name, max_batch)

    for io_num in range(io_cnt):
        config += '''
input [
  {{
    name: "INPUT{}"
    data_type: {}
    dims: [ {} ]
    {}
  }}
]
output [
  {{
    name: "OUTPUT{}"
    data_type: {}
    dims: [ {} ]
    {}
  }}
]
'''.format(
            io_num, np_to_model_dtype(dtype),
            tu.shape_to_dims_str(input_shapes[io_num]),
            "reshape: {{ shape: [ {} ] }}".format(
                tu.shape_to_dims_str(input_model_shapes[io_num]))
            if input_shapes[io_num] != input_model_shapes[io_num] else "",
            io_num, np_to_model_dtype(dtype),
            tu.shape_to_dims_str(output_shapes[io_num]),
            "reshape: {{ shape: [ {} ] }}".format(
                tu.shape_to_dims_str(output_model_shapes[io_num]))
            if output_shapes[io_num] != output_model_shapes[io_num] else "")

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_models(models_dir,
                  dtype,
                  input_shapes,
                  input_model_shapes,
                  output_shapes=None,
                  output_model_shapes=None,
                  no_batch=True):
    model_version = 1
    if output_shapes is None:
        output_shapes = input_shapes
    if output_model_shapes is None:
        output_model_shapes = input_model_shapes

    if FLAGS.graphdef:
        create_tf_modelconfig(False, models_dir, model_version, 8, dtype,
                              input_shapes, input_model_shapes, output_shapes,
                              output_model_shapes)
        create_tf_modelfile(False, models_dir, model_version, 8, dtype,
                            input_model_shapes, output_model_shapes)
        if no_batch:
            create_tf_modelconfig(False, models_dir, model_version, 0, dtype,
                                  input_shapes, input_model_shapes,
                                  output_shapes, output_model_shapes)
            create_tf_modelfile(False, models_dir, model_version, 0, dtype,
                                input_model_shapes, output_model_shapes)

    if FLAGS.savedmodel:
        create_tf_modelconfig(True, models_dir, model_version, 8, dtype,
                              input_shapes, input_model_shapes, output_shapes,
                              output_model_shapes)
        create_tf_modelfile(True, models_dir, model_version, 8, dtype,
                            input_model_shapes, output_model_shapes)
        if no_batch:
            create_tf_modelconfig(True, models_dir, model_version, 0, dtype,
                                  input_shapes, input_model_shapes,
                                  output_shapes, output_model_shapes)
            create_tf_modelfile(True, models_dir, model_version, 0, dtype,
                                input_model_shapes, output_model_shapes)

    if FLAGS.onnx:
        create_onnx_modelconfig(models_dir, model_version, 8, dtype,
                                input_shapes, input_model_shapes, output_shapes,
                                output_model_shapes)
        create_onnx_modelfile(models_dir, model_version, 8, dtype,
                              input_model_shapes, output_model_shapes)
        if no_batch:
            create_onnx_modelconfig(models_dir, model_version, 0, dtype,
                                    input_shapes, input_model_shapes,
                                    output_shapes, output_model_shapes)
            create_onnx_modelfile(models_dir, model_version, 0, dtype,
                                  input_model_shapes, output_model_shapes)

    # Shouldn't create ensembles that reshape to zero-sized tensors. Reshaping
    # from / to zero dimension is not allow as ensemble inputs / outputs
    # are passed from / to other model AS IF direct inference from client.
    # But create it anyway, expecting that the ensemble models can be served but
    # they will always return error message.
    if FLAGS.ensemble:
        # Create fixed size nop for ensemble models
        for shape in input_model_shapes:
            emu.create_nop_modelconfig(models_dir, shape, np.float32)
            emu.create_nop_tunnel_modelconfig(models_dir, shape, np.float32)
            emu.create_nop_modelconfig(models_dir, [-1], np.float32)
        create_ensemble_modelconfig(models_dir, model_version, 8, dtype,
                                    input_shapes, input_model_shapes,
                                    output_shapes, output_model_shapes)
        create_ensemble_modelfile(models_dir, model_version, 8, dtype,
                                  input_model_shapes, output_model_shapes)
        if no_batch:
            create_ensemble_modelconfig(models_dir, model_version, 0, dtype,
                                        input_shapes, input_model_shapes,
                                        output_shapes, output_model_shapes)
            create_ensemble_modelfile(models_dir, model_version, 0, dtype,
                                      input_model_shapes, output_model_shapes)


def create_trt_models(models_dir,
                      dtype,
                      input_shapes,
                      input_model_shapes,
                      output_shapes=None,
                      output_model_shapes=None,
                      no_batch=True):
    model_version = 1
    if output_shapes is None:
        output_shapes = input_shapes
    if output_model_shapes is None:
        output_model_shapes = input_model_shapes

    if FLAGS.tensorrt:
        create_plan_modelconfig(models_dir, model_version, 8, dtype,
                                input_shapes, input_model_shapes, output_shapes,
                                output_model_shapes)
        create_plan_modelfile(models_dir, model_version, 8, dtype,
                              input_model_shapes, output_model_shapes)
        if no_batch:
            create_plan_modelconfig(models_dir, model_version, 0, dtype,
                                    input_shapes, input_model_shapes,
                                    output_shapes, output_model_shapes)
            create_plan_modelfile(models_dir, model_version, 0, dtype,
                                  input_model_shapes, output_model_shapes)


def create_libtorch_models(models_dir,
                           dtype,
                           input_shapes,
                           input_model_shapes,
                           output_shapes=None,
                           output_model_shapes=None,
                           no_batch=True):
    model_version = 1
    if output_shapes is None:
        output_shapes = input_shapes
    if output_model_shapes is None:
        output_model_shapes = input_model_shapes

    if FLAGS.libtorch:
        create_libtorch_modelconfig(models_dir, model_version, 8, dtype,
                                    input_shapes, input_model_shapes,
                                    output_shapes, output_model_shapes)
        create_libtorch_modelfile(models_dir, model_version, 8, dtype,
                                  input_model_shapes, output_model_shapes)
        # skip for libtorch string I/O
        if no_batch and (dtype != np_dtype_string):
            create_libtorch_modelconfig(models_dir, model_version, 0, dtype,
                                        input_shapes, input_model_shapes,
                                        output_shapes, output_model_shapes)
            create_libtorch_modelfile(models_dir, model_version, 0, dtype,
                                      input_model_shapes, output_model_shapes)


def create_openvino_models(models_dir,
                           dtype,
                           input_shapes,
                           input_model_shapes,
                           output_shapes=None,
                           output_model_shapes=None,
                           no_batch=True):
    model_version = 1
    if output_shapes is None:
        output_shapes = input_shapes
    if output_model_shapes is None:
        output_model_shapes = input_model_shapes

    if FLAGS.openvino:
        create_openvino_modelconfig(models_dir, model_version, 8, dtype,
                                    input_shapes, input_model_shapes,
                                    output_shapes, output_model_shapes)
        create_openvino_modelfile(models_dir, model_version, 8, dtype,
                                  input_model_shapes, output_model_shapes)
        if no_batch:
            create_openvino_modelconfig(models_dir, model_version, 0, dtype,
                                        input_shapes, input_model_shapes,
                                        output_shapes, output_model_shapes)
            create_openvino_modelfile(models_dir, model_version, 0, dtype,
                                      input_model_shapes, output_model_shapes)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--models_dir',
                        type=str,
                        required=True,
                        help='Top-level model directory')
    parser.add_argument('--graphdef',
                        required=False,
                        action='store_true',
                        help='Generate GraphDef models')
    parser.add_argument('--savedmodel',
                        required=False,
                        action='store_true',
                        help='Generate SavedModel models')
    parser.add_argument('--tensorrt',
                        required=False,
                        action='store_true',
                        help='Generate TensorRT PLAN models')
    parser.add_argument('--onnx',
                        required=False,
                        action='store_true',
                        help='Generate Onnx Runtime Onnx models')
    parser.add_argument(
        '--onnx_opset',
        type=int,
        required=False,
        default=0,
        help='Opset used for Onnx models. Default is to use ONNXRT default')
    parser.add_argument('--libtorch',
                        required=False,
                        action='store_true',
                        help='Generate Pytorch LibTorch models')
    parser.add_argument('--openvino',
                        required=False,
                        action='store_true',
                        help='Generate OpenVino models')
    parser.add_argument('--ensemble',
                        required=False,
                        action='store_true',
                        help='Generate ensemble models')
    parser.add_argument('--variable',
                        required=False,
                        action='store_true',
                        help='Used variable-shape tensors for input/output')
    FLAGS, unparsed = parser.parse_known_args()

    if FLAGS.graphdef or FLAGS.savedmodel:
        import tensorflow as tf
        from tensorflow.python.framework import graph_io
    if FLAGS.tensorrt:
        import tensorrt as trt
    if FLAGS.onnx:
        import onnx
    if FLAGS.libtorch:
        import torch
        from torch import nn
    if FLAGS.openvino:
        from openvino.inference_engine import IENetwork
        import ngraph as ng

    import test_util as tu

    # TensorRT, OpenVino and LibTorch must be handled separately since they
    # don't support zero-sized tensors.
    create_models(FLAGS.models_dir,
                  np_dtype_string, ([1],), ([],),
                  no_batch=False)
    create_models(FLAGS.models_dir, np.float32, ([1],), ([],), no_batch=False)
    create_models(FLAGS.models_dir,
                  np.float32, ([1], [8]), ([], [4, 1, 2]),
                  no_batch=False)
    create_models(FLAGS.models_dir, np.float32, ([4, 4], [2], [2, 2, 3]),
                  ([16], [1, 2], [3, 2, 2]))
    create_libtorch_models(FLAGS.models_dir,
                           np.float32, ([1],), ([1, 1, 1],),
                           no_batch=False)
    create_libtorch_models(FLAGS.models_dir,
                           np.float32, ([1], [8]), ([1, 1, 1], [4, 1, 2]),
                           no_batch=False)
    create_libtorch_models(FLAGS.models_dir, np.float32,
                           ([4, 4], [2], [2, 2, 3]), ([16], [1, 2], [3, 2, 2]))
    create_libtorch_models(FLAGS.models_dir,
                           np_dtype_string, ([1],), ([],),
                           no_batch=False)
    create_openvino_models(FLAGS.models_dir,
                           np.float32, ([1],), ([1, 1, 1],),
                           no_batch=False)
    create_openvino_models(FLAGS.models_dir,
                           np.float32, ([1], [8]), ([1, 1, 1], [4, 1, 2]),
                           no_batch=False)
    create_openvino_models(FLAGS.models_dir, np.float32,
                           ([4, 4], [2], [2, 2, 3]), ([16], [1, 2], [3, 2, 2]))
    create_trt_models(FLAGS.models_dir, np.float32, ([1], [8]),
                      ([1, 1, 1], [4, 1, 2]))

    # Models that reshape only the input, not the output.
    create_models(FLAGS.models_dir,
                  np.float32, ([4, 4], [2], [2, 2, 3], [1]),
                  ([16], [1, 2], [3, 2, 2], [1]),
                  output_shapes=([16], [1, 2], [3, 2, 2], [1]),
                  output_model_shapes=([16], [1, 2], [3, 2, 2], [1]))

    create_libtorch_models(FLAGS.models_dir,
                           np.float32, ([4, 4], [2], [2, 2, 3], [1]),
                           ([16], [1, 2], [3, 2, 2], [1]),
                           output_shapes=([16], [1, 2], [3, 2, 2], [1]),
                           output_model_shapes=([16], [1, 2], [3, 2, 2], [1]))

    create_openvino_models(FLAGS.models_dir,
                           np.float32, ([4, 4], [2], [2, 2, 3], [1]),
                           ([16], [1, 2], [3, 2, 2], [1]),
                           output_shapes=([16], [1, 2], [3, 2, 2], [1]),
                           output_model_shapes=([16], [1, 2], [3, 2, 2], [1]))

    create_trt_models(FLAGS.models_dir,
                      np.float32, ([4, 4], [2], [2, 2, 3], [1]),
                      ([2, 2, 4], [1, 2, 1], [3, 2, 2], [1, 1, 1]),
                      output_shapes=([2, 2, 4], [1, 2, 1], [3, 2, 2], [1, 1,
                                                                       1]),
                      output_model_shapes=([2, 2, 4], [1, 2, 1], [3, 2,
                                                                  2], [1, 1,
                                                                       1]))

    # Tests with models that accept variable-shape input/output tensors and reshape
    # TensorRT is ignored as it only allows fixed-shape tensors
    # PyTorch is ignored as "tensor.view()" is shape dependent (shape is fixed
    # based on input used for tracing), need to find equivalent operation that
    # is not shape dependent.
    if FLAGS.variable:
        create_models(FLAGS.models_dir, np.int32, ([2, 4, -1, 6],),
                      ([8, -1, 1, 6],))
        create_models(FLAGS.models_dir, np.int32, ([1, -1, 1], [-1], [2, 2, 3]),
                      ([-1], [1, -1, 1], [3, 2, 2]))
        create_models(FLAGS.models_dir,
                      np.int32, ([-1, 1], [2]), ([1, -1], [1, 2]),
                      output_shapes=([1, -1], [1, 2]),
                      output_model_shapes=([1, -1], [1, 2]))

    # TRT plan that reshapes neither input nor output. Needed for
    # L0_perflab_nomodel.
    create_trt_models(FLAGS.models_dir, np.float32, ([1],), ([1],))


# qa/common/gen_qa_sequence_models.py

# Copyright 2019-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
import os
import numpy as np
import gen_ensemble_model_utils as emu

FLAGS = None
np_dtype_string = np.dtype(object)


def np_to_model_dtype(np_dtype):
    if np_dtype == bool:
        return "TYPE_BOOL"
    elif np_dtype == np.int8:
        return "TYPE_INT8"
    elif np_dtype == np.int16:
        return "TYPE_INT16"
    elif np_dtype == np.int32:
        return "TYPE_INT32"
    elif np_dtype == np.int64:
        return "TYPE_INT64"
    elif np_dtype == np.uint8:
        return "TYPE_UINT8"
    elif np_dtype == np.uint16:
        return "TYPE_UINT16"
    elif np_dtype == np.float16:
        return "TYPE_FP16"
    elif np_dtype == np.float32:
        return "TYPE_FP32"
    elif np_dtype == np.float64:
        return "TYPE_FP64"
    elif np_dtype == np_dtype_string:
        return "TYPE_STRING"
    return None


def np_to_tf_dtype(np_dtype):
    if np_dtype == bool:
        return tf.bool
    elif np_dtype == np.int8:
        return tf.int8
    elif np_dtype == np.int16:
        return tf.int16
    elif np_dtype == np.int32:
        return tf.int32
    elif np_dtype == np.int64:
        return tf.int64
    elif np_dtype == np.uint8:
        return tf.uint8
    elif np_dtype == np.uint16:
        return tf.uint16
    elif np_dtype == np.float16:
        return tf.float16
    elif np_dtype == np.float32:
        return tf.float32
    elif np_dtype == np.float64:
        return tf.float64
    elif np_dtype == np_dtype_string:
        return tf.string
    return None


def np_to_trt_dtype(np_dtype):
    if np_dtype == bool:
        return trt.bool
    elif np_dtype == np.int8:
        return trt.int8
    elif np_dtype == np.int32:
        return trt.int32
    elif np_dtype == np.float16:
        return trt.float16
    elif np_dtype == np.float32:
        return trt.float32
    return None


def np_to_onnx_dtype(np_dtype):
    if np_dtype == bool:
        return onnx.TensorProto.BOOL
    elif np_dtype == np.int8:
        return onnx.TensorProto.INT8
    elif np_dtype == np.int16:
        return onnx.TensorProto.INT16
    elif np_dtype == np.int32:
        return onnx.TensorProto.INT32
    elif np_dtype == np.int64:
        return onnx.TensorProto.INT64
    elif np_dtype == np.uint8:
        return onnx.TensorProto.UINT8
    elif np_dtype == np.uint16:
        return onnx.TensorProto.UINT16
    elif np_dtype == np.float16:
        return onnx.TensorProto.FLOAT16
    elif np_dtype == np.float32:
        return onnx.TensorProto.FLOAT
    elif np_dtype == np.float64:
        return onnx.TensorProto.DOUBLE
    elif np_dtype == np_dtype_string:
        return onnx.TensorProto.STRING


def np_to_torch_dtype(np_dtype):
    if np_dtype == bool:
        return torch.bool
    elif np_dtype == np.int8:
        return torch.int8
    elif np_dtype == np.int16:
        return torch.int16
    elif np_dtype == np.int32:
        return torch.int
    elif np_dtype == np.int64:
        return torch.long
    elif np_dtype == np.uint8:
        return torch.uint8
    elif np_dtype == np.uint16:
        return None  # Not supported in Torch
    elif np_dtype == np.float16:
        return None
    elif np_dtype == np.float32:
        return torch.float
    elif np_dtype == np.float64:
        return torch.double
    elif np_dtype == np_dtype_string:
        return None  # Not supported in Torch
    return None


def create_tf_modelfile(create_savedmodel, models_dir, model_version, max_batch,
                        dtype, shape):

    if not tu.validate_for_tf_model(dtype, dtype, dtype, shape, shape, shape):
        return

    tf_input_dtype = np_to_tf_dtype(dtype)
    tf_dtype = tf_input_dtype
    tf_control_type = tf_input_dtype

    # If the input is a string then use int32 for operation and just
    # cast to/from string for input and output.
    if tf_input_dtype == tf.string:
        tf_dtype = tf.int32
        tf_control_type = tf.int32

    # If input dtype is bool, then use bool type for control and
    # int32 type for input/output
    if tf_input_dtype == tf.bool:
        tf_control_type = tf.bool
        tf_dtype = tf.int32
        tf_input_dtype = tf.int32

    # Create the model. If non-batching then don't include the batch
    # dimension.
    tf.reset_default_graph()
    if create_savedmodel and (max_batch == 0):
        input0 = tf.placeholder(tf_input_dtype, [
            1,
        ], "INPUT")
        if tf_input_dtype == tf.string:
            input0 = tf.strings.to_number(tf.strings.join(["0", input0]),
                                          tf_dtype)
        start0 = tf.placeholder(tf_control_type, [
            1,
        ], "START")
        ready0 = tf.placeholder(tf_control_type, [
            1,
        ], "READY")
        acc = tf.get_variable("ACC", [
            1,
        ], dtype=tf_dtype)

        # Convert boolean value to int32 value
        if tf_control_type == tf.bool:
            start0 = tf.cast(start0, tf.int32)
            ready0 = tf.cast(ready0, tf.int32)

        tmp = tf.where(tf.equal(start0, 1), input0, tf.add(acc, input0))
        newacc = tf.where(tf.equal(ready0, 1), tmp, acc)

        assign = tf.assign(acc, newacc)
        if tf_input_dtype == tf.string:
            output0 = tf.dtypes.as_string(assign, name="OUTPUT")
        else:
            output0 = tf.identity(assign, name="OUTPUT")
    else:
        # For batching we can't use a tf.variable to hold the
        # accumulated values since that forces the size of the output
        # to the size of the variable (which must be a max-batch-size
        # vector since require one accumulator each), instead of the
        # output shape being [None, 1]. So instead we just return 0 if
        # not-ready and 'INPUT'+'START' otherwise... the tests know to
        # expect this.
        input0 = tf.placeholder(tf_input_dtype, [
            None,
        ] + tu.shape_to_tf_shape(shape), "INPUT")
        if tf_input_dtype == tf.string:
            input0 = tf.strings.to_number(tf.strings.join(["0", input0]),
                                          tf_dtype)
        start0 = tf.placeholder(tf_control_type, [None, 1], "START")
        ready0 = tf.placeholder(tf_control_type, [None, 1], "READY")

        # Convert boolean value to int32 value
        if tf_control_type == tf.bool:
            start0 = tf.cast(start0, tf.int32)
            ready0 = tf.cast(ready0, tf.int32)

        tmp = tf.where(tf.equal(ready0, 1), tf.add(start0, input0),
                       tf.zeros(tf.shape(input0), dtype=tf_dtype))

        if tf_input_dtype == tf.string:
            output0 = tf.dtypes.as_string(tmp, name="OUTPUT")
        else:
            output0 = tf.identity(tmp, name="OUTPUT")

    # Use a different model name for the non-batching variant
    if create_savedmodel:
        model_name = tu.get_sequence_model_name(
            "savedmodel_nobatch" if max_batch == 0 else "savedmodel", dtype)
    else:
        model_name = tu.get_sequence_model_name(
            "graphdef_nobatch" if max_batch == 0 else "graphdef", dtype)

    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    if create_savedmodel:
        with tf.Session() as sess:
            sess.run(tf.initializers.global_variables())
            input0_tensor = tf.get_default_graph().get_tensor_by_name("INPUT:0")
            start0_tensor = tf.get_default_graph().get_tensor_by_name("START:0")
            ready0_tensor = tf.get_default_graph().get_tensor_by_name("READY:0")
            output0_tensor = tf.get_default_graph().get_tensor_by_name(
                "OUTPUT:0")
            tf.saved_model.simple_save(sess,
                                       model_version_dir + "/model.savedmodel",
                                       inputs={
                                           "INPUT": input0_tensor,
                                           "START": start0_tensor,
                                           "READY": ready0_tensor
                                       },
                                       outputs={"OUTPUT": output0_tensor})
    else:
        with tf.Session() as sess:
            sess.run(tf.initializers.global_variables())
            graph_io.write_graph(sess.graph.as_graph_def(),
                                 model_version_dir,
                                 "model.graphdef",
                                 as_text=False)


def create_tf_modelconfig(create_savedmodel, models_dir, model_version,
                          max_batch, dtype, shape):

    if not tu.validate_for_tf_model(dtype, dtype, dtype, shape, shape, shape):
        return

    # Use a different model name for the non-batching variant
    if create_savedmodel:
        model_name = tu.get_sequence_model_name(
            "savedmodel_nobatch" if max_batch == 0 else "savedmodel", dtype)
    else:
        model_name = tu.get_sequence_model_name(
            "graphdef_nobatch" if max_batch == 0 else "graphdef", dtype)

    if dtype == np.float32:
        control_type = "fp32"
    elif dtype == bool:
        control_type = "bool"
        dtype = np.int32
    else:
        control_type = "int32"

    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "{}"
max_batch_size: {}
sequence_batching {{
  max_sequence_idle_microseconds: 5000000
  control_input [
    {{
      name: "START"
      control [
        {{
          kind: CONTROL_SEQUENCE_START
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "READY"
      control [
        {{
          kind: CONTROL_SEQUENCE_READY
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }}
  ]
}}
input [
  {{
    name: "INPUT"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT"
    data_type: {}
    dims: [ 1 ]
  }}
]
instance_group [
  {{
    kind: KIND_GPU
  }}
]
'''.format(
        model_name,
        "tensorflow_savedmodel" if create_savedmodel else "tensorflow_graphdef",
        max_batch, control_type, control_type, np_to_model_dtype(dtype),
        tu.shape_to_dims_str(shape), np_to_model_dtype(dtype))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_plan_shape_tensor_modelfile(models_dir, model_version, max_batch,
                                       dtype, shape):
    # Note that resize layer does not support int tensors.
    # The model takes two inputs (INPUT and SHAPE_INPUT)
    # and two control inputs(START and READY).
    # In absence of proper accumulator,
    # OUTPUT : 0 if not-ready and 'INPUT'+'START' otherwise
    # RESIZED_OUTPUT : Obtained after resizing 'INPUT' to shape specified
    #          in 'SHAPE_INPUT'
    # SHAPE_OUTPUT : The shape values of resized output

    trt_dtype = np_to_trt_dtype(dtype)
    trt_memory_format = trt.TensorFormat.LINEAR

    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

    unit_shape = ([1] * len(shape))
    if max_batch != 0:
        shape_in0 = network.add_input("SHAPE_INPUT", trt.int32,
                                      [1 + len(shape)])
        in0 = network.add_input("INPUT", trt_dtype, [-1] + shape)
        start0 = network.add_input("START", trt_dtype, [-1] + unit_shape)
        ready0 = network.add_input("READY", trt_dtype, [-1] + unit_shape)
    else:
        shape_in0 = network.add_input("SHAPE_INPUT", trt.int32, [len(shape)])
        in0 = network.add_input("INPUT", trt_dtype, shape)
        start0 = network.add_input("START", trt_dtype, unit_shape)
        ready0 = network.add_input("READY", trt_dtype, unit_shape)

    add = network.add_elementwise(in0, start0, trt.ElementWiseOperation.SUM)
    out0 = network.add_elementwise(add.get_output(0), ready0,
                                   trt.ElementWiseOperation.PROD).get_output(0)

    resize_layer = network.add_resize(input=in0)
    resize_layer.set_input(1, shape_in0)
    resized_out0 = resize_layer.get_output(0)
    shape_out0 = network.add_shape(resized_out0)

    shape_out0.get_output(0).name = "SHAPE_OUTPUT"
    shape_out0.get_output(0).dtype = trt.int32
    network.mark_output_for_shapes(shape_out0.get_output(0))

    out0.name = "OUTPUT"
    out0.dtype = trt_dtype
    network.mark_output(out0)

    resized_out0.name = "RESIZED_OUTPUT"
    resized_out0.dtype = trt_dtype
    network.mark_output(resized_out0)

    in0.allowed_formats = 1 << int(trt_memory_format)
    shape_in0.allowed_formats = 1 << int(trt_memory_format)
    start0.allowed_formats = 1 << int(trt_memory_format)
    ready0.allowed_formats = 1 << int(trt_memory_format)
    out0.allowed_formats = 1 << int(trt_memory_format)
    shape_out0.get_output(0).allowed_formats = 1 << int(trt_memory_format)
    resized_out0.allowed_formats = 1 << int(trt_memory_format)

    if (trt_dtype == trt.int8):
        in0.dynamic_range = (-128.0, 127.0)
        out0.dynamic_range = (-128.0, 127.0)
        resized_out0.dynamic_range = (-128.0, 127.0)
        start0.dynamic_range = (-128.0, 127.0)
        ready0.dynamic_range = (-128.0, 127.0)

    flags = 1 << int(trt.BuilderFlag.STRICT_TYPES)

    if (trt_dtype == trt.int8):
        flags |= 1 << int(trt.BuilderFlag.INT8)
    elif (trt_dtype == trt.float16):
        flags |= 1 << int(trt.BuilderFlag.FP16)

    min_prefix = []
    opt_prefix = []
    max_prefix = []

    if max_batch != 0:
        min_prefix = [1]
        opt_prefix = [max(1, max_batch)]
        max_prefix = [max(1, max_batch)]

    min_shape = min_prefix + [1] * len(shape)
    opt_shape = opt_prefix + [8] * len(shape)
    max_shape = max_prefix + [32] * len(shape)

    profile = builder.create_optimization_profile()
    profile.set_shape_input("SHAPE_INPUT", min_shape, opt_shape, max_shape)
    profile.set_shape("INPUT", min_shape, opt_shape, max_shape)
    profile.set_shape("START", min_prefix + unit_shape, opt_prefix + unit_shape,
                      opt_prefix + unit_shape)
    profile.set_shape("READY", min_prefix + unit_shape, opt_prefix + unit_shape,
                      opt_prefix + unit_shape)

    config = builder.create_builder_config()
    config.flags = flags
    config.add_optimization_profile(profile)

    config.max_workspace_size = 1 << 20
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    model_name = tu.get_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_fixed_modelfile(models_dir, model_version, max_batch, dtype,
                                shape):
    trt_dtype = np_to_trt_dtype(dtype)
    # Create the model. For now don't implement a proper accumulator
    # just return 0 if not-ready and 'INPUT'+'START' otherwise...  the
    # tests know to expect this.
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network()
    in0 = network.add_input("INPUT", trt_dtype, shape)
    start0 = network.add_input("START", trt_dtype, [1 for i in shape])
    ready0 = network.add_input("READY", trt_dtype, [1 for i in shape])
    add = network.add_elementwise(in0, start0, trt.ElementWiseOperation.SUM)
    out0 = network.add_elementwise(add.get_output(0), ready0,
                                   trt.ElementWiseOperation.PROD)

    out0.get_output(0).name = "OUTPUT"
    network.mark_output(out0.get_output(0))

    config = builder.create_builder_config()
    config.max_workspace_size = 1 << 20
    builder.max_batch_size = max(1, max_batch)
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine
    del network

    model_name = tu.get_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_fixed_rf_modelfile(models_dir, model_version, max_batch, dtype,
                                   shape):
    trt_dtype = np_to_trt_dtype(dtype)
    trt_memory_format = trt.TensorFormat.LINEAR
    # Create the model. For now don't implement a proper accumulator
    # just return 0 if not-ready and 'INPUT'+'START' otherwise...  the
    # tests know to expect this.
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network()
    in0 = network.add_input("INPUT", trt_dtype, shape)
    start0 = network.add_input("START", trt_dtype, [1 for i in shape])
    ready0 = network.add_input("READY", trt_dtype, [1 for i in shape])
    add = network.add_elementwise(in0, start0, trt.ElementWiseOperation.SUM)
    out0 = network.add_elementwise(add.get_output(0), ready0,
                                   trt.ElementWiseOperation.PROD)

    out0.get_output(0).name = "OUTPUT"
    network.mark_output(out0.get_output(0))

    out0.get_output(0).dtype = trt_dtype

    in0.allowed_formats = 1 << int(trt_memory_format)
    start0.allowed_formats = 1 << int(trt_memory_format)
    ready0.allowed_formats = 1 << int(trt_memory_format)
    out0.get_output(0).allowed_formats = 1 << int(trt_memory_format)

    if (trt_dtype == trt.int8):
        in0.dynamic_range = (-128.0, 127.0)
        out0.dynamic_range = (-128.0, 127.0)
        start0.dynamic_range = (-128.0, 127.0)
        ready0.dynamic_range = (-128.0, 127.0)

    flags = 1 << int(trt.BuilderFlag.STRICT_TYPES)

    if (trt_dtype == trt.int8):
        flags |= 1 << int(trt.BuilderFlag.INT8)
    elif (trt_dtype == trt.float16):
        flags |= 1 << int(trt.BuilderFlag.FP16)

    config = builder.create_builder_config()
    config.flags = flags
    config.max_workspace_size = 1 << 20
    builder.max_batch_size = max(1, max_batch)
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    model_name = tu.get_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_dynamic_modelfile(models_dir, model_version, max_batch, dtype,
                                  shape):
    trt_dtype = np_to_trt_dtype(dtype)
    # Create the model. For now don't implement a proper accumulator
    # just return 0 if not-ready and 'INPUT'+'START' otherwise...  the
    # tests know to expect this.
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

    unit_shape = ([1] * len(shape))
    if max_batch != 0:
        in0 = network.add_input("INPUT", trt_dtype, [-1] + shape)
        start0 = network.add_input("START", trt_dtype, [-1] + unit_shape)
        ready0 = network.add_input("READY", trt_dtype, [-1] + unit_shape)
    else:
        in0 = network.add_input("INPUT", trt_dtype, shape)
        start0 = network.add_input("START", trt_dtype, unit_shape)
        ready0 = network.add_input("READY", trt_dtype, unit_shape)

    add = network.add_elementwise(in0, start0, trt.ElementWiseOperation.SUM)
    out0 = network.add_elementwise(add.get_output(0), ready0,
                                   trt.ElementWiseOperation.PROD)

    out0.get_output(0).name = "OUTPUT"
    network.mark_output(out0.get_output(0))

    min_shape = []
    opt_shape = []
    max_shape = []
    if max_batch != 0:
        min_shape = min_shape + [1]
        opt_shape = opt_shape + [max(1, max_batch)]
        max_shape = max_shape + [max(1, max_batch)]
    for i in shape:
        if i == -1:
            min_shape = min_shape + [1]
            opt_shape = opt_shape + [8]
            max_shape = max_shape + [32]
        else:
            min_shape = min_shape + [i]
            opt_shape = opt_shape + [i]
            max_shape = max_shape + [i]

    profile = builder.create_optimization_profile()
    profile.set_shape("INPUT", min_shape, opt_shape, max_shape)
    if max_batch != 0:
        profile.set_shape("START", [1] + unit_shape, [max_batch] + unit_shape,
                          [max_batch] + unit_shape)
        profile.set_shape("READY", [1] + unit_shape, [max_batch] + unit_shape,
                          [max_batch] + unit_shape)
    else:
        profile.set_shape("START", unit_shape, unit_shape, unit_shape)
        profile.set_shape("READY", unit_shape, unit_shape, unit_shape)
    config = builder.create_builder_config()
    config.add_optimization_profile(profile)

    config.max_workspace_size = 1 << 20
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    model_name = tu.get_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_dynamic_rf_modelfile(models_dir, model_version, max_batch,
                                     dtype, shape):
    trt_dtype = np_to_trt_dtype(dtype)
    trt_memory_format = trt.TensorFormat.LINEAR

    # Create the model. For now don't implement a proper accumulator
    # just return 0 if not-ready and 'INPUT'+'START' otherwise...  the
    # tests know to expect this.
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

    unit_shape = ([1] * len(shape))
    if max_batch != 0:
        in0 = network.add_input("INPUT", trt_dtype, [-1] + shape)
        start0 = network.add_input("START", trt_dtype, [-1] + unit_shape)
        ready0 = network.add_input("READY", trt_dtype, [-1] + unit_shape)
    else:
        in0 = network.add_input("INPUT", trt_dtype, shape)
        start0 = network.add_input("START", trt_dtype, unit_shape)
        ready0 = network.add_input("READY", trt_dtype, unit_shape)

    add = network.add_elementwise(in0, start0, trt.ElementWiseOperation.SUM)
    out0 = network.add_elementwise(add.get_output(0), ready0,
                                   trt.ElementWiseOperation.PROD)

    out0.get_output(0).name = "OUTPUT"
    network.mark_output(out0.get_output(0))

    out0.get_output(0).dtype = trt_dtype

    in0.allowed_formats = 1 << int(trt_memory_format)
    start0.allowed_formats = 1 << int(trt_memory_format)
    ready0.allowed_formats = 1 << int(trt_memory_format)
    out0.get_output(0).allowed_formats = 1 << int(trt_memory_format)

    if (trt_dtype == trt.int8):
        in0.dynamic_range = (-128.0, 127.0)
        out0.dynamic_range = (-128.0, 127.0)
        start0.dynamic_range = (-128.0, 127.0)
        ready0.dynamic_range = (-128.0, 127.0)

    flags = 1 << int(trt.BuilderFlag.STRICT_TYPES)

    if (trt_dtype == trt.int8):
        flags |= 1 << int(trt.BuilderFlag.INT8)
    elif (trt_dtype == trt.float16):
        flags |= 1 << int(trt.BuilderFlag.FP16)

    min_shape = []
    opt_shape = []
    max_shape = []
    if max_batch != 0:
        min_shape = min_shape + [1]
        opt_shape = opt_shape + [max(1, max_batch)]
        max_shape = max_shape + [max(1, max_batch)]
    for i in shape:
        if i == -1:
            min_shape = min_shape + [1]
            opt_shape = opt_shape + [8]
            max_shape = max_shape + [32]
        else:
            min_shape = min_shape + [i]
            opt_shape = opt_shape + [i]
            max_shape = max_shape + [i]

    profile = builder.create_optimization_profile()
    profile.set_shape("INPUT", min_shape, opt_shape, max_shape)
    if max_batch != 0:
        profile.set_shape("START", [1] + unit_shape, [max_batch] + unit_shape,
                          [max_batch] + unit_shape)
        profile.set_shape("READY", [1] + unit_shape, [max_batch] + unit_shape,
                          [max_batch] + unit_shape)
    else:
        profile.set_shape("START", unit_shape, unit_shape, unit_shape)
        profile.set_shape("READY", unit_shape, unit_shape, unit_shape)

    config = builder.create_builder_config()
    config.flags = flags
    config.add_optimization_profile(profile)

    config.max_workspace_size = 1 << 20
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    model_name = tu.get_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_modelfile(models_dir, model_version, max_batch, dtype, shape):

    if not tu.validate_for_trt_model(dtype, dtype, dtype, shape, shape, shape):
        return

    if dtype != np.float32:
        if (not tu.shape_is_fixed(shape)):
            create_plan_dynamic_rf_modelfile(models_dir, model_version,
                                             max_batch, dtype, shape)
        else:
            create_plan_fixed_rf_modelfile(models_dir, model_version, max_batch,
                                           dtype, shape)
    else:
        if (not tu.shape_is_fixed(shape)):
            create_plan_dynamic_modelfile(models_dir, model_version, max_batch,
                                          dtype, shape)
        else:
            create_plan_fixed_modelfile(models_dir, model_version, max_batch,
                                        dtype, shape)


def create_plan_modelconfig(models_dir, model_version, max_batch, dtype, shape):

    if not tu.validate_for_trt_model(dtype, dtype, dtype, shape, shape, shape):
        return

    model_name = tu.get_sequence_model_name(
        "plan_nobatch" if max_batch == 0 else "plan", dtype)
    config_dir = models_dir + "/" + model_name
    if FLAGS.tensorrt_shape_io:
        shape_tensor_dim = len(shape)
        config = '''
name: "{}"
platform: "tensorrt_plan"
max_batch_size: {}
sequence_batching {{
  max_sequence_idle_microseconds: 5000000
  control_input [
    {{
      name: "START"
      control [
        {{
          kind: CONTROL_SEQUENCE_START
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "READY"
      control [
        {{
          kind: CONTROL_SEQUENCE_READY
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }}
  ]
}}
input [
  {{
    name: "INPUT"
    data_type: {}
    dims: [ {} ]
  }}
]
input [
  {{
    name: "SHAPE_INPUT"
    data_type: TYPE_INT32
    dims: [ {} ]
    is_shape_tensor: true
  }}
]
output [
  {{
    name: "OUTPUT"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "RESIZED_OUTPUT"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "SHAPE_OUTPUT"
    data_type: TYPE_INT32
    dims: [ {} ]
    is_shape_tensor: true
  }}
]
instance_group [
  {{
    kind: KIND_GPU
  }}
]
'''.format(model_name, max_batch, "int32" if dtype == np.int32 else "fp32",
           "int32" if dtype == np.int32 else "fp32", np_to_model_dtype(dtype),
           tu.shape_to_dims_str(shape), shape_tensor_dim,
           np_to_model_dtype(dtype), tu.shape_to_dims_str(shape),
           np_to_model_dtype(dtype), tu.shape_to_dims_str(shape),
           shape_tensor_dim)

    else:
        config = '''
name: "{}"
platform: "tensorrt_plan"
max_batch_size: {}
sequence_batching {{
  max_sequence_idle_microseconds: 5000000
  control_input [
    {{
      name: "START"
      control [
        {{
          kind: CONTROL_SEQUENCE_START
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "READY"
      control [
        {{
          kind: CONTROL_SEQUENCE_READY
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }}
  ]
}}
input [
  {{
    name: "INPUT"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT"
    data_type: {}
    dims: [ {} ]
  }}
]
instance_group [
  {{
    kind: KIND_GPU
  }}
]
'''.format(model_name, max_batch, "int32" if dtype == np.int32 else "fp32",
           "int32" if dtype == np.int32 else "fp32", np_to_model_dtype(dtype),
           tu.shape_to_dims_str(shape), np_to_model_dtype(dtype),
           tu.shape_to_dims_str(shape))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_onnx_modelfile(models_dir, model_version, max_batch, dtype, shape):

    if not tu.validate_for_onnx_model(dtype, dtype, dtype, shape, shape, shape):
        return

    model_name = tu.get_sequence_model_name(
        "onnx_nobatch" if max_batch == 0 else "onnx", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    # Create the model. For now don't implement a proper accumulator
    # just return 0 if not-ready and 'INPUT'+'START' otherwise...  the
    # tests know to expect this.
    onnx_dtype = np_to_onnx_dtype(dtype)
    onnx_control_dtype = onnx_dtype
    onnx_input_shape, idx = tu.shape_to_onnx_shape(shape, 0)
    onnx_output_shape, idx = tu.shape_to_onnx_shape(shape, idx)

    # If the input is a string then use int32 for operation and just
    # cast to/from string for input and output.
    if onnx_dtype == onnx.TensorProto.STRING:
        onnx_control_dtype = onnx.TensorProto.INT32

    # If input dtype is bool, then use bool type for control and
    # int32 type for input/output
    if onnx_dtype == onnx.TensorProto.BOOL:
        onnx_dtype = onnx.TensorProto.INT32

    batch_dim = [] if max_batch == 0 else [None]

    onnx_input = onnx.helper.make_tensor_value_info(
        "INPUT", onnx_dtype, batch_dim + onnx_input_shape)
    onnx_start = onnx.helper.make_tensor_value_info("START", onnx_control_dtype,
                                                    batch_dim + [1])
    onnx_ready = onnx.helper.make_tensor_value_info("READY", onnx_control_dtype,
                                                    batch_dim + [1])
    onnx_output = onnx.helper.make_tensor_value_info(
        "OUTPUT", onnx_dtype, batch_dim + onnx_output_shape)

    internal_input = onnx.helper.make_node("Identity", ["INPUT"], ["_INPUT"])

    # cast int8, int16 input to higer precision int as Onnx Add/Sub operator doesn't support those type
    # Also casting String data type to int32
    if ((onnx_dtype == onnx.TensorProto.INT8) or
        (onnx_dtype == onnx.TensorProto.INT16) or
        (onnx_dtype == onnx.TensorProto.STRING)):
        internal_input = onnx.helper.make_node("Cast", ["INPUT"], ["_INPUT"],
                                               to=onnx.TensorProto.INT32)

    # Convert boolean value to int32 value
    if onnx_control_dtype == onnx.TensorProto.BOOL:
        internal_input1 = onnx.helper.make_node("Cast", ["START"], ["_START"],
                                                to=onnx.TensorProto.INT32)
        internal_input2 = onnx.helper.make_node("Cast", ["READY"], ["_READY"],
                                                to=onnx.TensorProto.INT32)
        add = onnx.helper.make_node("Add", ["_INPUT", "_START"], ["add"])
        # Take advantage of knowledge that the READY false value is 0 and true is 1
        mul = onnx.helper.make_node("Mul", ["_READY", "add"], ["CAST"])

    else:
        add = onnx.helper.make_node("Add", ["_INPUT", "START"], ["add"])
        # Take advantage of knowledge that the READY false value is 0 and true is 1
        mul = onnx.helper.make_node("Mul", ["READY", "add"], ["CAST"])

    cast = onnx.helper.make_node("Cast", ["CAST"], ["OUTPUT"], to=onnx_dtype)

    # Avoid cast from float16 to float16
    # (bug in Onnx Runtime, cast from float16 to float16 will become cast from float16 to float32)
    if onnx_dtype == onnx.TensorProto.FLOAT16:
        cast = onnx.helper.make_node("Identity", ["CAST"], ["OUTPUT"])

    if onnx_control_dtype == onnx.TensorProto.BOOL:
        onnx_nodes = [
            internal_input, internal_input1, internal_input2, add, mul, cast
        ]
    else:
        onnx_nodes = [internal_input, add, mul, cast]
    onnx_inputs = [onnx_input, onnx_start, onnx_ready]
    onnx_outputs = [onnx_output]

    graph_proto = onnx.helper.make_graph(onnx_nodes, model_name, onnx_inputs,
                                         onnx_outputs)
    if FLAGS.onnx_opset > 0:
        model_opset = onnx.helper.make_operatorsetid("", FLAGS.onnx_opset)
        model_def = onnx.helper.make_model(graph_proto,
                                           producer_name="triton",
                                           opset_imports=[model_opset])
    else:
        model_def = onnx.helper.make_model(graph_proto, producer_name="triton")

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    onnx.save(model_def, model_version_dir + "/model.onnx")


def create_onnx_modelconfig(models_dir, model_version, max_batch, dtype, shape):

    if not tu.validate_for_onnx_model(dtype, dtype, dtype, shape, shape, shape):
        return

    model_name = tu.get_sequence_model_name(
        "onnx_nobatch" if max_batch == 0 else "onnx", dtype)
    config_dir = models_dir + "/" + model_name

    if dtype == np.float32:
        control_type = "fp32"
    elif dtype == bool:
        control_type = "bool"
        dtype = np.int32
    else:
        control_type = "int32"

    instance_group_string = '''
instance_group [
  {
    kind: KIND_GPU
  }
]
'''

    # [TODO] move create_general_modelconfig() out of emu as it is general
    # enough for all backends to use
    config = emu.create_general_modelconfig(
        model_name,
        "onnxruntime_onnx",
        max_batch, [dtype], [shape], [None], [dtype], [shape], [None], [None],
        force_tensor_number_suffix=False,
        instance_group_str=instance_group_string)

    config += '''
sequence_batching {{
  max_sequence_idle_microseconds: 5000000
  control_input [
    {{
      name: "START"
      control [
        {{
          kind: CONTROL_SEQUENCE_START
          {type}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "READY"
      control [
        {{
          kind: CONTROL_SEQUENCE_READY
          {type}_false_true: [ 0, 1 ]
        }}
      ]
    }}
  ]
}}
'''.format(type=control_type)

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_libtorch_modelfile(models_dir, model_version, max_batch, dtype,
                              shape):

    if not tu.validate_for_libtorch_model(dtype, dtype, dtype, shape, shape,
                                          shape, max_batch):
        return

    torch_dtype = np_to_torch_dtype(dtype)
    torch_control_type = torch_dtype

    # If input dtype is bool, then use bool type for control and
    # int32 type for input/output
    if torch_dtype == torch.bool:
        torch_dtype = torch.int32

    model_name = tu.get_sequence_model_name(
        "libtorch_nobatch" if max_batch == 0 else "libtorch", dtype)
    # handle for -1 (when variable) since can't create tensor with shape of [-1]
    shape = [abs(ips) for ips in shape]

    class SequenceNet(nn.Module):

        def __init__(self):
            super(SequenceNet, self).__init__()

        def forward(self, input0, start0, ready0):
            tmp = input0 + start0
            return tmp * ready0

    sequenceModel = SequenceNet()
    example_input0 = torch.zeros(shape, dtype=torch_dtype)
    example_input1 = torch.zeros(shape, dtype=torch_control_type)
    example_input2 = torch.zeros(shape, dtype=torch_control_type)

    # Convert boolean value to int32 value
    if torch_control_type == torch.bool:
        example_input1 = example_input1.long()
        example_input2 = example_input2.long()

    traced = torch.jit.trace(sequenceModel,
                             (example_input0, example_input1, example_input2))

    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    traced.save(model_version_dir + "/model.pt")


def create_libtorch_modelconfig(models_dir, model_version, max_batch, dtype,
                                shape):

    if not tu.validate_for_libtorch_model(dtype, dtype, dtype, shape, shape,
                                          shape, max_batch):
        return

    model_name = tu.get_sequence_model_name(
        "libtorch_nobatch" if max_batch == 0 else "libtorch", dtype)
    config_dir = models_dir + "/" + model_name

    if dtype == np.float32:
        control_type = "fp32"
    elif dtype == bool:
        control_type = "bool"
        dtype = np.int32
    else:
        control_type = "int32"

    #  FIX FOR LibTorch
    config = '''
name: "{}"
platform: "pytorch_libtorch"
max_batch_size: {}
sequence_batching {{
  max_sequence_idle_microseconds: 5000000
  control_input [
    {{
      name: "START__1"
      control [
        {{
          kind: CONTROL_SEQUENCE_START
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "READY__2"
      control [
        {{
          kind: CONTROL_SEQUENCE_READY
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }}
  ]
}}
input [
  {{
    name: "INPUT__0"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT__0"
    data_type: {}
    dims: [ 1 ]
  }}
]
instance_group [
  {{
    kind: KIND_GPU
  }}
]
'''.format(model_name, max_batch, control_type, control_type,
           np_to_model_dtype(dtype), tu.shape_to_dims_str(shape),
           np_to_model_dtype(dtype))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_openvino_modelfile(models_dir, model_version, max_batch, dtype,
                              shape):

    batch_dim = [] if max_batch == 0 else [
        max_batch,
    ]
    if not tu.validate_for_openvino_model(dtype, dtype, dtype,
                                          batch_dim + shape, batch_dim + shape,
                                          batch_dim + shape):
        return

    model_name = tu.get_sequence_model_name(
        "openvino_nobatch" if max_batch == 0 else "openvino", dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    in0 = ng.parameter(shape=batch_dim + shape, dtype=dtype, name="INPUT")
    start = ng.parameter(shape=batch_dim + shape, dtype=dtype, name="START")
    ready = ng.parameter(shape=batch_dim + shape, dtype=dtype, name="READY")

    tmp = ng.add(in0, start)
    op0 = ng.multiply(tmp, ready, name="OUTPUT")

    function = ng.impl.Function([op0], [in0, start, ready], model_name)
    ie_network = IENetwork(ng.impl.Function.to_capsule(function))

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    ie_network.serialize(model_version_dir + "/model.xml",
                         model_version_dir + "/model.bin")


def create_openvino_modelconfig(models_dir, model_version, max_batch, dtype,
                                shape):

    batch_dim = [] if max_batch == 0 else [
        max_batch,
    ]
    if not tu.validate_for_openvino_model(dtype, dtype, dtype,
                                          batch_dim + shape, batch_dim + shape,
                                          batch_dim + shape):
        return

    model_name = tu.get_sequence_model_name(
        "openvino_nobatch" if max_batch == 0 else "openvino", dtype)
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
backend: "openvino"
max_batch_size: {}
sequence_batching {{
  max_sequence_idle_microseconds: 5000000
  control_input [
    {{
      name: "START"
      control [
        {{
          kind: CONTROL_SEQUENCE_START
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }},
    {{
      name: "READY"
      control [
        {{
          kind: CONTROL_SEQUENCE_READY
          {}_false_true: [ 0, 1 ]
        }}
      ]
    }}
  ]
}}
input [
  {{
    name: "INPUT"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT"
    data_type: {}
    dims: [ 1 ]
  }}
]
'''.format(model_name, max_batch, "int32" if dtype == np.int32 else "fp32",
           "int32" if dtype == np.int32 else "fp32", np_to_model_dtype(dtype),
           tu.shape_to_dims_str(shape), np_to_model_dtype(dtype))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_shape_tensor_models(models_dir, dtype, shape, no_batch=True):
    model_version = 1

    create_plan_modelconfig(models_dir, model_version, 8, dtype, shape)
    create_plan_shape_tensor_modelfile(models_dir, model_version, 8, dtype,
                                       shape)
    if no_batch:
        create_plan_modelconfig(models_dir, model_version, 0, dtype, shape)
        create_plan_shape_tensor_modelfile(models_dir, model_version, 0, dtype,
                                           shape)


def create_models(models_dir, dtype, shape, no_batch=True):
    model_version = 1

    if FLAGS.graphdef:
        create_tf_modelconfig(False, models_dir, model_version, 8, dtype, shape)
        create_tf_modelfile(False, models_dir, model_version, 8, dtype, shape)
        if no_batch:
            create_tf_modelconfig(False, models_dir, model_version, 0, dtype,
                                  shape)
            create_tf_modelfile(False, models_dir, model_version, 0, dtype,
                                shape)

    if FLAGS.savedmodel:
        create_tf_modelconfig(True, models_dir, model_version, 8, dtype, shape)
        create_tf_modelfile(True, models_dir, model_version, 8, dtype, shape)
        if no_batch:
            create_tf_modelconfig(True, models_dir, model_version, 0, dtype,
                                  shape)
            create_tf_modelfile(True, models_dir, model_version, 0, dtype,
                                shape)

    if FLAGS.tensorrt:
        if dtype == bool:
            return
        suffix = []
        if dtype == np.int8:
            suffix = [1, 1]

        create_plan_modelconfig(models_dir, model_version, 8, dtype,
                                shape + suffix)
        create_plan_modelfile(models_dir, model_version, 8, dtype,
                              shape + suffix)
        if no_batch:
            create_plan_modelconfig(models_dir, model_version, 0, dtype,
                                    shape + suffix)
            create_plan_modelfile(models_dir, model_version, 0, dtype,
                                  shape + suffix)

    if FLAGS.onnx:
        create_onnx_modelconfig(models_dir, model_version, 8, dtype, shape)
        create_onnx_modelfile(models_dir, model_version, 8, dtype, shape)
        if no_batch:
            create_onnx_modelconfig(models_dir, model_version, 0, dtype, shape)
            create_onnx_modelfile(models_dir, model_version, 0, dtype, shape)

    # Skip for PyTorch String I/O
    if FLAGS.libtorch and (dtype != np_dtype_string):
        create_libtorch_modelconfig(models_dir, model_version, 8, dtype, shape)
        create_libtorch_modelfile(models_dir, model_version, 8, dtype, shape)
        if no_batch:
            create_libtorch_modelconfig(models_dir, model_version, 0, dtype,
                                        shape)
            create_libtorch_modelfile(models_dir, model_version, 0, dtype,
                                      shape)

    if FLAGS.openvino:
        create_openvino_modelconfig(models_dir, model_version, 8, dtype, shape)
        create_openvino_modelfile(models_dir, model_version, 8, dtype, shape)
        if no_batch:
            create_openvino_modelconfig(models_dir, model_version, 0, dtype,
                                        shape)
            create_openvino_modelfile(models_dir, model_version, 0, dtype,
                                      shape)

    if FLAGS.ensemble:
        if dtype == bool:
            return
        for pair in emu.platform_types_and_validation():
            config_shape = shape
            if pair[0] == "plan" and dtype == np.int8:
                config_shape = shape + [1, 1]
            if not pair[1](dtype, dtype, dtype, config_shape, config_shape,
                           config_shape):
                continue

            emu.create_sequence_ensemble_modelconfig(pair[0], models_dir, 8,
                                                     model_version,
                                                     config_shape, dtype)
            emu.create_sequence_ensemble_modelfile(pair[0], models_dir, 8,
                                                   model_version, config_shape,
                                                   dtype)
            if no_batch:
                emu.create_sequence_ensemble_modelconfig(
                    pair[0], models_dir, 0, model_version, config_shape, dtype)
                emu.create_sequence_ensemble_modelfile(pair[0], models_dir, 0,
                                                       model_version,
                                                       config_shape, dtype)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--models_dir',
                        type=str,
                        required=True,
                        help='Top-level model directory')
    parser.add_argument('--graphdef',
                        required=False,
                        action='store_true',
                        help='Generate GraphDef models')
    parser.add_argument('--savedmodel',
                        required=False,
                        action='store_true',
                        help='Generate SavedModel models')
    parser.add_argument('--tensorrt',
                        required=False,
                        action='store_true',
                        help='Generate TensorRT PLAN models')
    parser.add_argument(
        '--tensorrt-shape-io',
        required=False,
        action='store_true',
        help='Generate TensorRT PLAN models w/ shape tensor i/o')
    parser.add_argument('--onnx',
                        required=False,
                        action='store_true',
                        help='Generate Onnx models')
    parser.add_argument(
        '--onnx_opset',
        type=int,
        required=False,
        default=0,
        help='Opset used for Onnx models. Default is to use ONNXRT default')
    parser.add_argument('--libtorch',
                        required=False,
                        action='store_true',
                        help='Generate Pytorch LibTorch models')
    parser.add_argument('--openvino',
                        required=False,
                        action='store_true',
                        help='Generate OpenVino models')
    parser.add_argument('--variable',
                        required=False,
                        action='store_true',
                        help='Used variable-shape tensors for input/output')
    parser.add_argument('--ensemble',
                        required=False,
                        action='store_true',
                        help='Generate ensemble models against the models' +
                        ' in all platforms. Note that the models generated' +
                        ' are not completed.')
    FLAGS, unparsed = parser.parse_known_args()

    if FLAGS.graphdef or FLAGS.savedmodel:
        import tensorflow as tf
        from tensorflow.python.framework import graph_io
    if FLAGS.tensorrt or FLAGS.tensorrt_shape_io:
        import tensorrt as trt
    if FLAGS.onnx:
        import onnx
    if FLAGS.libtorch:
        import torch
        from torch import nn
    if FLAGS.openvino:
        from openvino.inference_engine import IENetwork
        import ngraph as ng

    import test_util as tu

    if FLAGS.tensorrt_shape_io:
        create_shape_tensor_models(FLAGS.models_dir, np.float32, [
            -1,
        ])
    else:
        # Tests with models that accept fixed-shape input/output tensors
        if not FLAGS.variable:
            create_models(FLAGS.models_dir, np.float32, [
                1,
            ])
            create_models(FLAGS.models_dir, np.int32, [
                1,
            ])
            create_models(FLAGS.models_dir, np_dtype_string, [
                1,
            ])
            create_models(FLAGS.models_dir, bool, [
                1,
            ])

        # Tests with models that accept variable-shape input/output tensors
        if FLAGS.variable:
            create_models(FLAGS.models_dir, np.int32, [
                -1,
            ], False)
            create_models(FLAGS.models_dir, np.float32, [
                -1,
            ], False)
            create_models(FLAGS.models_dir, np_dtype_string, [
                -1,
            ], False)
            create_models(FLAGS.models_dir, bool, [
                -1,
            ], False)

        if FLAGS.ensemble:
            # Create nop models used in ensemble
            for model_dtype in ["TYPE_INT32", "TYPE_FP32"]:
                for model_shape in [(-1,)]:
                    emu.create_nop_modelconfig(FLAGS.models_dir, model_shape,
                                               model_dtype)


# qa/common/gen_qa_tf_parameters.py

# Copyright 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../common")

import os
import tensorflow as tf
from tensorflow.python.framework import graph_io


def create_graphdefmodel(models_dir, model_name, model_version=1):
    """A simple tensorflow model that accumulates the INPUT with internal
    model parameter named VARIABLE and produces OUTPUT.
    """

    tf.reset_default_graph()
    input0 = tf.placeholder(tf.int32, [
        1,
    ], "INPUT")
    variable = tf.get_variable("VARIABLE", [
        1,
    ],
                               initializer=tf.zeros_initializer(),
                               dtype=tf.int32)
    tf.add(variable, input0, name="OUTPUT")
    tf.global_variables_initializer()
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with tf.Session() as sess:
        graph_io.write_graph(sess.graph.as_graph_def(),
                             model_version_dir,
                             "model.graphdef",
                             as_text=False)


def create_graphdef_modelconfig(models_dir, model_name):
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "tensorflow_graphdef"
input [
  {{
    name: "INPUT"
    data_type: TYPE_INT32
    dims: [ 1 ]
  }}
]
output [
  {{
    name: "OUTPUT"
    data_type: TYPE_INT32
    dims: [ 1 ]
  }}
]
'''.format(model_name)

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--models_dir',
                        type=str,
                        required=True,
                        help='Top-level model directory')
    args = parser.parse_args()

    model_name = "graphdef_variable"
    create_graphdefmodel(args.models_dir, model_name)
    create_graphdef_modelconfig(args.models_dir, model_name=model_name)


# qa/common/gen_qa_torchtrt_models.py

# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
import os

import torch
import torchvision
import torch_tensorrt


def create_resnet50_torchtrt(models_dir, max_batch):
    model = torchvision.models.resnet50(pretrained=True)
    model.eval()
    example_input = torch.rand(1, 3, 224, 224, dtype=torch.float)

    resnet50_ts = torch.jit.trace(model, example_input)

    trt_ts_module = torch_tensorrt.compile(
        resnet50_ts,
        inputs=[
            torch_tensorrt.Input(min_shape=[1, 3, 224, 224],
                                 opt_shape=[1, 3, 224, 224],
                                 max_shape=[max_batch, 3, 224, 224],
                                 dtype=torch.float)
        ],
        enabled_precisions={torch.float},
    )

    model_name = "resnet50_libtorch"

    model_version = 1
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    torch.jit.save(trt_ts_module, model_version_dir + "/model.pt")


def create_resnet50_torchtrt_modelconfig(models_dir, max_batch):

    model_name = "resnet50_libtorch"
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
backend: "pytorch"
max_batch_size: {}
input [
  {{
    name: "INPUT__0"
    data_type: TYPE_FP32
    format: FORMAT_NCHW
    dims: [ 3, 224, 224 ]
  }}
]
output [
  {{
    name: "OUTPUT__0"
    data_type: TYPE_FP32
    dims: [ 1000 ]
    label_filename: "resnet50_labels.txt"
  }}
]
'''.format(model_name, max_batch)

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--models_dir',
                        type=str,
                        required=True,
                        help='Top-level model directory')
    FLAGS, unparsed = parser.parse_known_args()

    create_resnet50_torchtrt(FLAGS.models_dir, 128)
    create_resnet50_torchtrt_modelconfig(FLAGS.models_dir, 128)


# qa/common/gen_qa_trt_data_dependent_shape.py

# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
import os
import numpy as np
import tensorrt as trt
import test_util as tu

def np_to_model_dtype(np_dtype):
    if np_dtype == bool:
        return "TYPE_BOOL"
    elif np_dtype == np.int8:
        return "TYPE_INT8"
    elif np_dtype == np.int16:
        return "TYPE_INT16"
    elif np_dtype == np.int32:
        return "TYPE_INT32"
    elif np_dtype == np.int64:
        return "TYPE_INT64"
    elif np_dtype == np.uint8:
        return "TYPE_UINT8"
    elif np_dtype == np.uint16:
        return "TYPE_UINT16"
    elif np_dtype == np.float16:
        return "TYPE_FP16"
    elif np_dtype == np.float32:
        return "TYPE_FP32"
    elif np_dtype == np.float64:
        return "TYPE_FP64"
    elif np_dtype == np_dtype_string:
        return "TYPE_STRING"
    return None


def np_to_trt_dtype(np_dtype):
    if np_dtype == bool:
        return trt.bool
    elif np_dtype == np.int8:
        return trt.int8
    elif np_dtype == np.int32:
        return trt.int32
    elif np_dtype == np.float16:
        return trt.float16
    elif np_dtype == np.float32:
        return trt.float32
    return None

# The 'nonzero' model that we use for data dependent shape is naturally
# not support batching, because the layer output is not trivially separable
# based on the request batch size.
# input_shape is config shape
def create_data_dependent_modelfile(models_dir,
                                    model_name,
                                    input_shape,
                                    input_dtype=np.int32,
                                    min_dim=1,
                                    max_dim=32):
    trt_input_dtype = np_to_trt_dtype(input_dtype)

    # Create the model
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

    # input
    in0 = network.add_input("INPUT", trt_input_dtype, input_shape)

    # layers
    non_zero = network.add_non_zero(in0)

    # configure output
    out0 = non_zero.get_output(0)
    out0.name = "OUTPUT"
    network.mark_output(out0)

    # optimization profile
    min_shape = []
    opt_shape = []
    max_shape = []
    for i in input_shape:
        if i == -1:
            min_shape = min_shape + [min_dim]
            opt_shape = opt_shape + [int((max_dim + min_dim) / 2)]
            max_shape = max_shape + [max_dim]
        else:
            min_shape = min_shape + [i]
            opt_shape = opt_shape + [i]
            max_shape = max_shape + [i]

    profile = builder.create_optimization_profile()
    profile.set_shape("INPUT", min_shape, opt_shape, max_shape)
    config = builder.create_builder_config()
    config.add_optimization_profile(profile)
    config.max_workspace_size = 1 << 20
    
    # serialized model
    engine_bytes = builder.build_serialized_network(network, config)

    model_version_dir = models_dir + "/" + model_name + "/1"
    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)

def create_data_dependent_modelconfig(models_dir,
                                    model_name,
                                      input_shape,
                                      input_dtype=np.int32):
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "tensorrt_plan"
max_batch_size: 0
input [
  {{
    name: "INPUT"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT"
    data_type: {}
    dims: [ {} ]
   }}
]
'''.format(model_name,
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(np.int32),
           tu.shape_to_dims_str((len(input_shape), -1)))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--models_dir',
                        type=str,
                        required=True,
                        help='Top-level model directory')
    FLAGS, unparsed = parser.parse_known_args()

    # Fixed input shape
    create_data_dependent_modelfile(FLAGS.models_dir, "plan_nobatch_nonzero_fixed", (4, 4))
    create_data_dependent_modelconfig(FLAGS.models_dir, "plan_nobatch_nonzero_fixed", (4, 4))

    # Dynamic input shape
    create_data_dependent_modelfile(FLAGS.models_dir, "plan_nobatch_nonzero_dynamic", (-1, -1))
    create_data_dependent_modelconfig(FLAGS.models_dir, "plan_nobatch_nonzero_dynamic", (-1, -1))



# qa/common/gen_qa_trt_format_models.py

# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
import os
import numpy as np
import tensorrt as trt
import test_util as tu


def np_to_model_dtype(np_dtype):
    if np_dtype == bool:
        return "TYPE_BOOL"
    elif np_dtype == np.int8:
        return "TYPE_INT8"
    elif np_dtype == np.int16:
        return "TYPE_INT16"
    elif np_dtype == np.int32:
        return "TYPE_INT32"
    elif np_dtype == np.int64:
        return "TYPE_INT64"
    elif np_dtype == np.uint8:
        return "TYPE_UINT8"
    elif np_dtype == np.uint16:
        return "TYPE_UINT16"
    elif np_dtype == np.float16:
        return "TYPE_FP16"
    elif np_dtype == np.float32:
        return "TYPE_FP32"
    elif np_dtype == np.float64:
        return "TYPE_FP64"
    elif np_dtype == np_dtype_string:
        return "TYPE_STRING"
    return None


def np_to_trt_dtype(np_dtype):
    if np_dtype == bool:
        return trt.bool
    elif np_dtype == np.int8:
        return trt.int8
    elif np_dtype == np.int32:
        return trt.int32
    elif np_dtype == np.float16:
        return trt.float16
    elif np_dtype == np.float32:
        return trt.float32
    return None


def trt_format_to_string(trt_format):
    # FIXME uncomment the following formats once TRT used is up-to-date
    # if trt_format == trt.TensorFormat.CDHW32:
    #     return "CDHW32"
    # if trt_format == trt.TensorFormat.DHWC8:
    #     return "DHWC8"
    # if trt_format == trt.TensorFormat.HWC:
    #     return "HWC"
    if trt_format == trt.TensorFormat.CHW2:
        return "CHW2"
    if trt_format == trt.TensorFormat.CHW32:
        return "CHW32"
    if trt_format == trt.TensorFormat.LINEAR:
        return "LINEAR"
    if trt_format == trt.TensorFormat.CHW4:
        return "CHW4"
    if trt_format == trt.TensorFormat.HWC8:
        return "HWC8"
    if trt_format == trt.TensorFormat.CHW16:
        return "CHW16"
    return "INVALID"


def create_plan_dynamic_modelfile(models_dir,
                                  max_batch,
                                  model_version,
                                  input_shape,
                                  output0_shape,
                                  output1_shape,
                                  input_dtype,
                                  output0_dtype,
                                  output1_dtype,
                                  input_memory_format,
                                  output_memory_format,
                                  min_dim=1,
                                  max_dim=64):
    trt_input_dtype = np_to_trt_dtype(input_dtype)
    trt_output0_dtype = np_to_trt_dtype(output0_dtype)
    trt_output1_dtype = np_to_trt_dtype(output1_dtype)
    trt_input_memory_format = input_memory_format
    trt_output_memory_format = output_memory_format

    # Create the model
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    if max_batch == 0:
        input_with_batchsize = [i for i in input_shape]
    else:
        input_with_batchsize = [-1] + [i for i in input_shape]

    in0 = network.add_input("INPUT0", trt_input_dtype, input_with_batchsize)
    in1 = network.add_input("INPUT1", trt_input_dtype, input_with_batchsize)
    add = network.add_elementwise(in0, in1, trt.ElementWiseOperation.SUM)
    sub = network.add_elementwise(in0, in1, trt.ElementWiseOperation.SUB)

    out0 = network.add_identity(add.get_output(0))
    out1 = network.add_identity(sub.get_output(0))

    out0.get_output(0).name = "OUTPUT0"
    out1.get_output(0).name = "OUTPUT1"
    network.mark_output(out0.get_output(0))
    network.mark_output(out1.get_output(0))

    out0.get_output(0).dtype = trt_output0_dtype
    out1.get_output(0).dtype = trt_output1_dtype

    in0.allowed_formats = 1 << int(trt_input_memory_format)
    in1.allowed_formats = 1 << int(trt_input_memory_format)
    out0.get_output(0).allowed_formats = 1 << int(trt_output_memory_format)
    out1.get_output(0).allowed_formats = 1 << int(trt_output_memory_format)

    if (trt_input_dtype == trt.int8):
        in0.dynamic_range = (-128.0, 127.0)
        in1.dynamic_range = (-128.0, 127.0)
    if (trt_output0_dtype == trt.int8):
        out0.get_output(0).dynamic_range = (-128.0, 127.0)
    if (trt_output1_dtype == trt.int8):
        out1.get_output(0).dynamic_range = (-128.0, 127.0)

    min_shape = []
    opt_shape = []
    max_shape = []
    if max_batch != 0:
        min_shape = min_shape + [1]
        opt_shape = opt_shape + [max(1, max_batch)]
        max_shape = max_shape + [max(1, max_batch)]
    for i in input_shape:
        if i == -1:
            min_shape = min_shape + [min_dim]
            opt_shape = opt_shape + [int((max_dim + min_dim) / 2)]
            max_shape = max_shape + [max_dim]
        else:
            min_shape = min_shape + [i]
            opt_shape = opt_shape + [i]
            max_shape = max_shape + [i]

    profile = builder.create_optimization_profile()
    profile.set_shape("INPUT0", min_shape, opt_shape, max_shape)
    profile.set_shape("INPUT1", min_shape, opt_shape, max_shape)
    flags = 1 << int(trt.BuilderFlag.STRICT_TYPES)
    datatype_set = set([trt_input_dtype, trt_output0_dtype, trt_output1_dtype])
    for dt in datatype_set:
        if (dt == trt.int8):
            flags |= 1 << int(trt.BuilderFlag.INT8)
        elif (dt == trt.float16):
            flags |= 1 << int(trt.BuilderFlag.FP16)
    config = builder.create_builder_config()
    config.flags = flags
    config.add_optimization_profile(profile)
    config.max_workspace_size = 1 << 20
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    # Use a different model name for different kinds of models
    base_name = "plan_nobatch" if max_batch == 0 else "plan"
    base_name += "_" + trt_format_to_string(
        input_memory_format) + "_" + trt_format_to_string(output_memory_format)
    model_name = tu.get_model_name(base_name, input_dtype, output0_dtype,
                                   output1_dtype)

    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_fixed_modelfile(models_dir, max_batch, model_version,
                                input_shape, output0_shape, output1_shape,
                                input_dtype, output0_dtype, output1_dtype,
                                input_memory_format, output_memory_format):
    trt_input_dtype = np_to_trt_dtype(input_dtype)
    trt_output0_dtype = np_to_trt_dtype(output0_dtype)
    trt_output1_dtype = np_to_trt_dtype(output1_dtype)
    trt_input_memory_format = input_memory_format
    trt_output_memory_format = output_memory_format

    # Create the model
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network()
    in0 = network.add_input("INPUT0", trt_input_dtype, input_shape)
    in1 = network.add_input("INPUT1", trt_input_dtype, input_shape)
    add = network.add_elementwise(in0, in1, trt.ElementWiseOperation.SUM)
    sub = network.add_elementwise(in0, in1, trt.ElementWiseOperation.SUB)

    out0 = network.add_identity(add.get_output(0))
    out1 = network.add_identity(sub.get_output(0))

    out0.get_output(0).name = "OUTPUT0"
    out1.get_output(0).name = "OUTPUT1"
    network.mark_output(out0.get_output(0))
    network.mark_output(out1.get_output(0))

    out0.get_output(0).dtype = trt_output0_dtype
    out1.get_output(0).dtype = trt_output1_dtype

    in0.allowed_formats = 1 << int(trt_input_memory_format)
    in1.allowed_formats = 1 << int(trt_input_memory_format)
    out0.get_output(0).allowed_formats = 1 << int(trt_output_memory_format)
    out1.get_output(0).allowed_formats = 1 << int(trt_output_memory_format)

    if (trt_input_dtype == trt.int8):
        in0.dynamic_range = (-128.0, 127.0)
        in1.dynamic_range = (-128.0, 127.0)
    if (trt_output0_dtype == trt.int8):
        out0.get_output(0).dynamic_range = (-128.0, 127.0)
    if (trt_output1_dtype == trt.int8):
        out1.get_output(0).dynamic_range = (-128.0, 127.0)

    flags = 1 << int(trt.BuilderFlag.STRICT_TYPES)
    datatype_set = set([trt_input_dtype, trt_output0_dtype, trt_output1_dtype])
    for dt in datatype_set:
        if (dt == trt.int8):
            flags |= 1 << int(trt.BuilderFlag.INT8)
        elif (dt == trt.float16):
            flags |= 1 << int(trt.BuilderFlag.FP16)
    config = builder.create_builder_config()
    config.flags = flags
    config.max_workspace_size = 1 << 20
    builder.max_batch_size = max(1, max_batch)
    try:
        engine_bytes = builder.build_serialized_network(network, config)
    except AttributeError:
        engine = builder.build_engine(network, config)
        engine_bytes = engine.serialize()
        del engine

    base_name = "plan_nobatch" if max_batch == 0 else "plan"
    base_name += "_" + trt_format_to_string(
        input_memory_format) + "_" + trt_format_to_string(output_memory_format)
    model_name = tu.get_model_name(base_name, input_dtype, output0_dtype,
                                   output1_dtype)
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(model_version_dir + "/model.plan", "wb") as f:
        f.write(engine_bytes)


def create_plan_modelconfig(models_dir, max_batch, model_version, input_shape,
                            output0_shape, output1_shape, input_dtype,
                            output0_dtype, output1_dtype, input_memory_format,
                            output_memory_format, version_policy):

    if not tu.validate_for_trt_model(input_dtype, output0_dtype, output1_dtype,
                                     input_shape, output0_shape, output1_shape):
        return

    # Unpack version policy
    version_policy_str = "{ latest { num_versions: 1 }}"
    if version_policy is not None:
        type, val = version_policy
        if type == 'latest':
            version_policy_str = "{{ latest {{ num_versions: {} }}}}".format(
                val)
        elif type == 'specific':
            version_policy_str = "{{ specific {{ versions: {} }}}}".format(val)
        else:
            version_policy_str = "{ all { }}"

    # Use a different model name for different kinds of models
    base_name = "plan_nobatch" if max_batch == 0 else "plan"
    base_name += "_" + trt_format_to_string(
        input_memory_format) + "_" + trt_format_to_string(output_memory_format)
    model_name = tu.get_model_name(base_name, input_dtype, output0_dtype,
                                   output1_dtype)

    config_dir = models_dir + "/" + model_name
    if -1 in input_shape:
        profile_index = 0
        config = '''
name: "{}"
platform: "tensorrt_plan"
max_batch_size: {}
version_policy: {}
input [
  {{
    name: "INPUT0"
    data_type: {}
    dims: [ {} ]
  }},
  {{
    name: "INPUT1"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT0"
    data_type: {}
    dims: [ {} ]
   }},
  {{
    name: "OUTPUT1"
    data_type: {}
    dims: [ {} ]
  }}
]
instance_group [
  {{
      profile:"{}"
  }}
]
'''.format(model_name, max_batch, version_policy_str,
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(output0_dtype),
           tu.shape_to_dims_str(output0_shape),
           np_to_model_dtype(output1_dtype),
           tu.shape_to_dims_str(output1_shape), profile_index)
    else:
        config = '''
name: "{}"
platform: "tensorrt_plan"
max_batch_size: {}
version_policy: {}
input [
  {{
    name: "INPUT0"
    data_type: {}
    dims: [ {} ]
  }},
  {{
    name: "INPUT1"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT0"
    data_type: {}
    dims: [ {} ]
   }},
  {{
    name: "OUTPUT1"
    data_type: {}
    dims: [ {} ]
  }}
]
'''.format(model_name, max_batch, version_policy_str,
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(output0_dtype),
           tu.shape_to_dims_str(output0_shape),
           np_to_model_dtype(output1_dtype),
           tu.shape_to_dims_str(output1_shape))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_plan_model(models_dir, max_batch, model_version, input_shape,
                      output0_shape, output1_shape, input_dtype, output0_dtype,
                      output1_dtype, input_memory_format, output_memory_format):

    if not tu.validate_for_trt_model(input_dtype, output0_dtype, output1_dtype,
                                     input_shape, output0_shape, output1_shape):
        return

    create_plan_modelconfig(models_dir, max_batch, model_version, input_shape,
                            output0_shape, output1_shape, input_dtype,
                            output0_dtype, output1_dtype, input_memory_format,
                            output_memory_format, None)

    if (not tu.shape_is_fixed(input_shape) or
            not tu.shape_is_fixed(output0_shape) or
            not tu.shape_is_fixed(output1_shape)):
        create_plan_dynamic_modelfile(models_dir, max_batch, model_version,
                                      input_shape, output0_shape, output1_shape,
                                      input_dtype, output0_dtype, output1_dtype,
                                      input_memory_format, output_memory_format)
    else:
        create_plan_fixed_modelfile(models_dir, max_batch, model_version,
                                    input_shape, output0_shape, output1_shape,
                                    input_dtype, output0_dtype, output1_dtype,
                                    input_memory_format, output_memory_format)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--models_dir',
                        type=str,
                        required=True,
                        help='Top-level model directory')
    FLAGS, unparsed = parser.parse_known_args()

    # reformat-free input
    # Fixed shape
    create_plan_model(FLAGS.models_dir, 0, 1, (13, 2, 1), (13, 2, 1),
                      (13, 2, 1), np.float16, np.float16, np.float16,
                      trt.TensorFormat.CHW2, trt.TensorFormat.LINEAR)
    create_plan_model(FLAGS.models_dir, 8, 1, (13, 2, 1), (13, 2, 1),
                      (13, 2, 1), np.float16, np.float16, np.float16,
                      trt.TensorFormat.CHW2, trt.TensorFormat.LINEAR)

    # Dynamic shape
    create_plan_model(FLAGS.models_dir, 0, 1, (-1, 2, 1), (-1, 2, 1),
                      (-1, 2, 1), np.float32, np.float32, np.float32,
                      trt.TensorFormat.CHW32, trt.TensorFormat.LINEAR)
    create_plan_model(FLAGS.models_dir, 8, 1, (-1, 2, 1), (-1, 2, 1),
                      (-1, 2, 1), np.float32, np.float32, np.float32,
                      trt.TensorFormat.CHW32, trt.TensorFormat.LINEAR)

    # reformat-free output
    # reformat-free I/O


# qa/common/gen_qa_trt_plugin_models.py

# Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
import os

import numpy as np
import tensorrt as trt

TRT_LOGGER = trt.Logger()

trt.init_libnvinfer_plugins(TRT_LOGGER, '')
PLUGIN_CREATORS = trt.get_plugin_registry().plugin_creator_list


def np_to_model_dtype(np_dtype):
    if np_dtype == bool:
        return "TYPE_BOOL"
    elif np_dtype == np.int8:
        return "TYPE_INT8"
    elif np_dtype == np.int16:
        return "TYPE_INT16"
    elif np_dtype == np.int32:
        return "TYPE_INT32"
    elif np_dtype == np.int64:
        return "TYPE_INT64"
    elif np_dtype == np.uint8:
        return "TYPE_UINT8"
    elif np_dtype == np.uint16:
        return "TYPE_UINT16"
    elif np_dtype == np.float16:
        return "TYPE_FP16"
    elif np_dtype == np.float32:
        return "TYPE_FP32"
    elif np_dtype == np.float64:
        return "TYPE_FP64"
    elif np_dtype == np_dtype_string:
        return "TYPE_STRING"
    return None


def np_to_trt_dtype(np_dtype):
    if np_dtype == bool:
        return trt.bool
    elif np_dtype == np.int8:
        return trt.int8
    elif np_dtype == np.int32:
        return trt.int32
    elif np_dtype == np.float16:
        return trt.float16
    elif np_dtype == np.float32:
        return trt.float32
    return None


def get_trt_plugin(plugin_name):
    plugin = None
    field_collection = None
    for plugin_creator in PLUGIN_CREATORS:
        if (plugin_creator.name == "Normalize_TRT") and \
                (plugin_name == "Normalize_TRT"):
            nbWeights = trt.PluginField("nbWeights",
                                        np.array([1], dtype=np.int32),
                                        trt.PluginFieldType.INT32)
            eps = trt.PluginField("eps", np.array([0.00001], dtype=np.float32),
                                  trt.PluginFieldType.FLOAT32)
            weights = trt.PluginField('weights',
                                      np.array([1] * 16, dtype=np.float32),
                                      trt.PluginFieldType.FLOAT32)
            field_collection = trt.PluginFieldCollection(
                [weights, eps, nbWeights])
            break
        elif (plugin_creator.name
              == "CustomGeluPluginDynamic") and (plugin_name
                                                 == "CustomGeluPluginDynamic"):
            type_id = trt.PluginField("type_id", np.array([0], np.int32),
                                      trt.PluginFieldType.INT32)
            bias = trt.PluginField("bias", np.array([[[1]]], np.float32),
                                   trt.PluginFieldType.FLOAT32)
            field_collection = trt.PluginFieldCollection([type_id, bias])
            break
        elif (plugin_creator.name
              == "CustomClipPlugin") and (plugin_name == "CustomClipPlugin"):
            min_clip = trt.PluginField("clipMin", np.array([0.1],\
                dtype=np.float32), trt.PluginFieldType.FLOAT32)
            max_clip = trt.PluginField("clipMax", np.array([0.5],\
                dtype=np.float32), trt.PluginFieldType.FLOAT32)
            field_collection = trt.PluginFieldCollection([min_clip, max_clip])
            break

    if field_collection is None:
        raise RuntimeError("Plugin not found: " + plugin_name)
    plugin = plugin_creator.create_plugin(name=plugin_name,
                                          field_collection=field_collection)

    return plugin


def create_plan_modelfile(models_dir, max_batch, model_version, plugin_name,
                          input_shape, output0_shape, input_dtype,
                          output0_dtype):

    if not tu.validate_for_trt_model(input_dtype, output0_dtype, output0_dtype,
                                     input_shape, output0_shape, output0_shape):
        return

    trt_input_dtype = np_to_trt_dtype(input_dtype)

    model_name = tu.get_model_name("plan_nobatch" if max_batch == 0 else "plan",
                                   input_dtype, output0_dtype,
                                   output0_dtype) + '_' + plugin_name

    # using explicit batch is necessary for CustomGeluPluginDynamic
    if plugin_name == "CustomGeluPluginDynamic":
        explicit_batch = 1 << (int)(
            trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
        with trt.Builder(TRT_LOGGER) as builder, builder.create_network(
                explicit_batch) as network:
            input_layer = network.add_input(name="INPUT0",
                                            dtype=trt_input_dtype,
                                            shape=input_shape)
            plugin_layer = network.add_plugin_v2(
                inputs=[input_layer], plugin=get_trt_plugin(plugin_name))
            plugin_layer.get_output(0).name = "OUTPUT0"
            network.mark_output(plugin_layer.get_output(0))

            config = builder.create_builder_config()
            config.max_workspace_size = 1 << 20

            try:
                engine_bytes = builder.build_serialized_network(network, config)
            except AttributeError:
                engine = builder.build_engine(network, config)
                engine_bytes = engine.serialize()
                del engine

            model_version_dir = models_dir + "/" + model_name + "/" + str(
                model_version)

            try:
                os.makedirs(model_version_dir)
            except OSError as ex:
                pass  # ignore existing dir

            with open(model_version_dir + "/model.plan", "wb") as f:
                f.write(engine_bytes)
    else:
        with trt.Builder(
                TRT_LOGGER) as builder, builder.create_network() as network:
            input_layer = network.add_input(name="INPUT0",
                                            dtype=trt_input_dtype,
                                            shape=input_shape)
            plugin_layer = network.add_plugin_v2(
                inputs=[input_layer], plugin=get_trt_plugin(plugin_name))
            plugin_layer.get_output(0).name = "OUTPUT0"
            network.mark_output(plugin_layer.get_output(0))

            config = builder.create_builder_config()
            config.max_workspace_size = 1 << 20
            builder.max_batch_size = max(1, max_batch)

            try:
                engine_bytes = builder.build_serialized_network(network, config)
            except AttributeError:
                engine = builder.build_engine(network, config)
                engine_bytes = engine.serialize()
                del engine

            model_version_dir = models_dir + "/" + model_name + "/" + str(
                model_version)

            try:
                os.makedirs(model_version_dir)
            except OSError as ex:
                pass  # ignore existing dir

            with open(model_version_dir + "/model.plan", "wb") as f:
                f.write(engine_bytes)


def create_plan_modelconfig(models_dir, max_batch, model_version, plugin_name,
                            input_shape, output0_shape, input_dtype,
                            output0_dtype):

    if not tu.validate_for_trt_model(input_dtype, output0_dtype, output0_dtype,
                                     input_shape, output0_shape, output0_shape):
        return

    version_policy_str = "{ latest { num_versions: 1 }}"

    # Use a different model name for the non-batching variant
    model_name = tu.get_model_name("plan_nobatch" if max_batch == 0 else "plan",
                                   input_dtype, output0_dtype,
                                   output0_dtype) + '_' + plugin_name
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "tensorrt_plan"
max_batch_size: {}
version_policy: {}
input [
  {{
    name: "INPUT0"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT0"
    data_type: {}
    dims: [ {} ]
   }}
]
'''.format(model_name, max_batch, version_policy_str,
           np_to_model_dtype(input_dtype), tu.shape_to_dims_str(input_shape),
           np_to_model_dtype(output0_dtype),
           tu.shape_to_dims_str(output0_shape))

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


def create_plugin_models(models_dir):
    model_version = 1

    # custom CustomClipPlugin
    create_plan_modelconfig(models_dir, 8, model_version, "CustomClipPlugin",
                            (16,), (16,), np.float32, np.float32)
    create_plan_modelfile(models_dir, 8, model_version, "CustomClipPlugin",
                          (16,), (16,), np.float32, np.float32)

    # default CustomGeluPluginDynamic plugin
    create_plan_modelconfig(models_dir, 0, model_version,
                            "CustomGeluPluginDynamic", (16, 1, 1), (16, 1, 1),
                            np.float32, np.float32)
    create_plan_modelfile(models_dir, 0, model_version,
                          "CustomGeluPluginDynamic", (16, 1, 1), (16, 1, 1),
                          np.float32, np.float32)

    # default Normalize_TRT
    create_plan_modelconfig(models_dir, 8, model_version, "Normalize_TRT", (
        16,
        16,
        16,
    ), (
        16,
        16,
        16,
    ), np.float32, np.float32)
    create_plan_modelfile(models_dir, 8, model_version, "Normalize_TRT", (
        16,
        16,
        16,
    ), (
        16,
        16,
        16,
    ), np.float32, np.float32)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--models_dir',
                        type=str,
                        required=True,
                        help='Top-level model directory')
    FLAGS, unparsed = parser.parse_known_args()

    import test_util as tu

    create_plugin_models(FLAGS.models_dir)


# qa/common/gen_tag_sigdef.py

# Copyright 2019-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../common")

import os
from builtins import range

from tensorflow.python.framework import ops
from tensorflow.python.saved_model import builder
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants
import tensorflow.compat.v1 as tf
import gen_ensemble_model_utils as gu
"""Create SavedModels that contains multiple tags and multiple signature defs"""


def create_savedmodel(models_dir,
                      model_version=1,
                      dims=16,
                      model_name="sig_tag",
                      tag_name="testTag",
                      signature_def_name="testSigDef",
                      different_io=False):
    """
    Creates one SavedModel with four variants of the model based on provided tag and signature_def.
    The models multiplies the input tensor by a multiplier and the multiplier value is different for each variant.
    Naming convention and config:
    <model_name>_0: tag: "serve",    signature_def: "serving_default",    multiplier 1
    <model_name>_1: tag: "serve",    signature_def: <signature_def_name>, multiplier 2
    <model_name>_2: tag: <tag_name>, signature_def: "serving_default",    multiplier 3
    <model_name>_3: tag: <tag_name>, signature_def: <signature_def_name>, multiplier 4

    If different_io is true, there will be two variants of the model created.
    The variants will have different numbers of inputs and outputs.
    Alternate naming convention and config:
    <model_name>0: tag: "serve",    signature_def: "serving_default",    two inputs/outputs
    <model_name>1: tag: <tag_name>, signature_def: <signature_def_name>, one input/output
    """
    model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)

    try:
        os.makedirs(model_version_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with tf.Session() as sess:
        input_tensor = tf.placeholder(tf.float32, [dims], "TENSOR_INPUT")

        # tag:"serve", signature_def:"serving_default"
        multiplier_0 = tf.constant(1.0, name="multiplier_0")
        # tag:"serve", signature_def:signature_def_name
        multiplier_1 = tf.constant(2.0, name="multiplier_1")
        # tag:tag_name, signature_def:"serving_default"
        multiplier_2 = tf.constant(3.0, name="multiplier_2")
        # tag:tag_name, signature_def:signature_def_name
        multiplier_3 = tf.constant(4.0, name="multiplier_3")

        output_tensor_0 = tf.multiply(multiplier_0,
                                      input_tensor,
                                      name="TENSOR_OUTPUT")
        output_tensor_1 = tf.multiply(multiplier_1,
                                      input_tensor,
                                      name="TENSOR_OUTPUT")
        output_tensor_2 = tf.multiply(multiplier_2,
                                      input_tensor,
                                      name="TENSOR_OUTPUT")
        output_tensor_3 = tf.multiply(multiplier_3,
                                      input_tensor,
                                      name="TENSOR_OUTPUT")

        # build_tensor_info_op could be used if build_tensor_info is deprecated
        input_tensor_info = tf.saved_model.utils.build_tensor_info(input_tensor)
        output_tensor_info_0 = tf.saved_model.utils.build_tensor_info(
            output_tensor_0)
        output_tensor_info_1 = tf.saved_model.utils.build_tensor_info(
            output_tensor_1)
        output_tensor_info_2 = tf.saved_model.utils.build_tensor_info(
            output_tensor_2)
        output_tensor_info_3 = tf.saved_model.utils.build_tensor_info(
            output_tensor_3)

        # Using predict method name because simple save uses it
        # tag:"serve", signature_def:"serving_default"
        signature_0 = tf.saved_model.signature_def_utils.build_signature_def(
            inputs={"INPUT": input_tensor_info},
            outputs={"OUTPUT": output_tensor_info_0},
            method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME)
        # tag:"serve", signature_def:signature_def_name
        signature_1 = tf.saved_model.signature_def_utils.build_signature_def(
            inputs={"INPUT": input_tensor_info},
            outputs={"OUTPUT": output_tensor_info_1},
            method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME)
        # tag:tag_name, signature_def:"serving_default"
        signature_2 = tf.saved_model.signature_def_utils.build_signature_def(
            inputs={"INPUT": input_tensor_info},
            outputs={"OUTPUT": output_tensor_info_2},
            method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME)
        # tag:tag_name, signature_def:signature_def_name
        signature_3 = tf.saved_model.signature_def_utils.build_signature_def(
            inputs={"INPUT": input_tensor_info},
            outputs={"OUTPUT": output_tensor_info_3},
            method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME)
        # tag:tag_name, signature_def:signature_def_name, two inputs/outputs
        signature_4 = tf.saved_model.signature_def_utils.build_signature_def(
            inputs={
                "INPUT": input_tensor_info,
                "INPUT1": input_tensor_info
            },
            outputs={
                "OUTPUT": output_tensor_info_0,
                "OUTPUT1": output_tensor_info_1
            },
            method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME)

        b = builder.SavedModelBuilder(model_version_dir + "/model.savedmodel")

        if different_io:
            b.add_meta_graph_and_variables(
                sess,
                tags=[tag_name],
                signature_def_map={signature_def_name: signature_0},
                assets_collection=ops.get_collection(
                    ops.GraphKeys.ASSET_FILEPATHS),
                clear_devices=True)
            b.add_meta_graph(
                tags=[tag_constants.SERVING],
                signature_def_map={
                    signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                        signature_4
                },
                assets_collection=ops.get_collection(
                    ops.GraphKeys.ASSET_FILEPATHS),
                clear_devices=True)
        else:
            signature_def_map_0 = {
                signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                    signature_0,
                signature_def_name:
                    signature_1
            }
            signature_def_map_1 = {
                signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                    signature_2,
                signature_def_name:
                    signature_3
            }

            b.add_meta_graph_and_variables(
                sess,
                tags=[tag_constants.SERVING],
                signature_def_map=signature_def_map_0,
                assets_collection=ops.get_collection(
                    ops.GraphKeys.ASSET_FILEPATHS),
                clear_devices=True)
            b.add_meta_graph(tags=[tag_name],
                             signature_def_map=signature_def_map_1,
                             assets_collection=ops.get_collection(
                                 ops.GraphKeys.ASSET_FILEPATHS),
                             clear_devices=True)

        b.save()


def create_savedmodel_modelconfig(models_dir,
                                  model_version=1,
                                  dims=16,
                                  model_name="sig_tag",
                                  tag_name="testTag",
                                  signature_def_name="testSigDef"):
    config_dir = models_dir + "/" + model_name
    config = '''
name: "{}"
platform: "tensorflow_savedmodel"
input [
  {{
    name: "INPUT"
    data_type: {}
    dims: [ {} ]
  }}
]
output [
  {{
    name: "OUTPUT"
    data_type: {}
    dims: [ {} ]
  }}
]
parameters: {{
key: "TF_GRAPH_TAG"
value: {{
string_value: "{}"
}}
}}
parameters: {{
key: "TF_SIGNATURE_DEF"
value: {{
string_value: "{}"
}}
}}
'''.format(model_name, gu.np_to_model_dtype(tf.float32), str(dims),
           gu.np_to_model_dtype(tf.float32), str(dims), tag_name,
           signature_def_name)

    try:
        os.makedirs(config_dir)
    except OSError as ex:
        pass  # ignore existing dir

    with open(config_dir + "/config.pbtxt", "w") as cfile:
        cfile.write(config)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='getting model output dir')
    parser.add_argument('--dir',
                        help='directory to run model in',
                        required=True)
    args = parser.parse_args()
    base_dir = args.dir
    base_model_name = "sig_tag"
    base_tag = "serve"
    test_tag = "testTag"
    base_sig_def = "serving_default"
    test_sig_def = "testSigDef"

    for i in range(4):
        model_name = base_model_name + str(i)
        create_savedmodel(base_dir,
                          model_name=model_name,
                          tag_name=test_tag,
                          signature_def_name=test_sig_def)
    create_savedmodel(base_dir,
                      model_name=base_model_name + "_different_io",
                      tag_name=test_tag,
                      signature_def_name=test_sig_def,
                      different_io=True)
    create_savedmodel_modelconfig(base_dir,
                                  model_name="sig_tag0",
                                  tag_name=base_tag,
                                  signature_def_name=base_sig_def)
    create_savedmodel_modelconfig(base_dir,
                                  model_name="sig_tag1",
                                  tag_name=base_tag,
                                  signature_def_name=test_sig_def)
    create_savedmodel_modelconfig(base_dir,
                                  model_name="sig_tag2",
                                  tag_name=test_tag,
                                  signature_def_name=base_sig_def)
    create_savedmodel_modelconfig(base_dir,
                                  model_name="sig_tag3",
                                  tag_name=test_tag,
                                  signature_def_name=test_sig_def)
    create_savedmodel_modelconfig(base_dir,
                                  model_name="sig_tag_different_io",
                                  tag_name=test_tag,
                                  signature_def_name=test_sig_def)


# qa/common/infer_test.py

# Copyright 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../common")

from builtins import range
from future.utils import iteritems
import unittest
import numpy as np
import infer_util as iu
import test_util as tu
import os

np_dtype_string = np.dtype(object)

# Allow caller to setup specific set of backends to test
DEFAULT_BACKENDS="graphdef savedmodel plan onnx libtorch"
TEST_BACKENDS = os.environ.get("BACKENDS", DEFAULT_BACKENDS).split()

class InferTest(tu.TestResultCollector):

    def _full_exact(self, input_dtype, output0_dtype, output1_dtype,
                    output0_raw, output1_raw, swap):

        def _infer_exact_helper(tester,
                                pf,
                                tensor_shape,
                                batch_size,
                                input_dtype,
                                output0_dtype,
                                output1_dtype,
                                output0_raw=True,
                                output1_raw=True,
                                model_version=None,
                                swap=False,
                                outputs=("OUTPUT0", "OUTPUT1"),
                                use_http=True,
                                use_grpc=True,
                                skip_request_id_check=False,
                                use_streaming=True,
                                correlation_id=0):
            for bs in (1, batch_size):
                iu.infer_exact(tester,
                               pf, (bs,) + tensor_shape,
                               bs,
                               input_dtype,
                               output0_dtype,
                               output1_dtype,
                               output0_raw=output0_raw,
                               output1_raw=output1_raw,
                               model_version=model_version,
                               swap=swap,
                               outputs=outputs,
                               use_http=use_http,
                               use_grpc=use_grpc,
                               skip_request_id_check=skip_request_id_check,
                               use_streaming=use_streaming,
                               correlation_id=correlation_id)


        input_size = 16

        if tu.validate_for_tf_model(input_dtype, output0_dtype, output1_dtype,
                                    (input_size,), (input_size,),
                                    (input_size,)):
            for pf in ["graphdef", "savedmodel"]:
                if pf in TEST_BACKENDS:
                    _infer_exact_helper(self,
                                        pf, (input_size,),
                                        8,
                                        input_dtype,
                                        output0_dtype,
                                        output1_dtype,
                                        output0_raw=output0_raw,
                                        output1_raw=output1_raw,
                                        swap=swap)

        if tu.validate_for_trt_model(input_dtype, output0_dtype, output1_dtype,
                                     (input_size, 1, 1), (input_size, 1, 1),
                                     (input_size, 1, 1)):
            if "plan" in TEST_BACKENDS:
                if input_dtype == np.int8:
                    shape = (input_size, 1, 1)
                else:
                    shape = (input_size,)
                _infer_exact_helper(self,
                                    'plan', shape,
                                    8,
                                    input_dtype,
                                    output0_dtype,
                                    output1_dtype,
                                    output0_raw=output0_raw,
                                    output1_raw=output1_raw,
                                    swap=swap)

        if tu.validate_for_onnx_model(input_dtype, output0_dtype, output1_dtype,
                                      (input_size,), (input_size,),
                                      (input_size,)):
            if "onnx" in TEST_BACKENDS:
                _infer_exact_helper(self,
                                    'onnx', (input_size,),
                                    8,
                                    input_dtype,
                                    output0_dtype,
                                    output1_dtype,
                                    output0_raw=output0_raw,
                                    output1_raw=output1_raw,
                                    swap=swap)

        # Skip for batched string I/O
        if tu.validate_for_libtorch_model(input_dtype, output0_dtype,
                                          output1_dtype, (input_size,),
                                          (input_size,), (input_size,), 8):
            if "libtorch" in TEST_BACKENDS:
                _infer_exact_helper(self,
                                    'libtorch', (input_size,),
                                    8,
                                    input_dtype,
                                    output0_dtype,
                                    output1_dtype,
                                    output0_raw=output0_raw,
                                    output1_raw=output1_raw,
                                    swap=swap)

    def test_raw_fff(self):
        self._full_exact(np.float32,
                         np.float32,
                         np.float32,
                         output0_raw=True,
                         output1_raw=True,
                         swap=True)

    def test_raw_ooo(self):
        self._full_exact(np_dtype_string,
                         np_dtype_string,
                         np_dtype_string,
                         output0_raw=True,
                         output1_raw=True,
                         swap=False)

    def test_class_fff(self):
        self._full_exact(np.float32,
                         np.float32,
                         np.float32,
                         output0_raw=False,
                         output1_raw=False,
                         swap=True)


if __name__ == '__main__':
    unittest.main()


# qa/common/infer_util.py

# Copyright 2018-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import os
import sys
import numpy as np
from functools import partial
import test_util as tu
import tritonclient.grpc as grpcclient
import tritonclient.http as httpclient
import shm_util as su
from tritonclient.utils import *

if sys.version_info >= (3, 0):
    import queue
else:
    import Queue as queue

# unicode() doesn't exist on python3, for how we use it the
# corresponding function is bytes()
if sys.version_info.major == 3:
    unicode = bytes

_seen_request_ids = set()

# By default, find tritonserver on "localhost", but can be overridden
# with TRITONSERVER_IPADDR envvar
_tritonserver_ipaddr = os.environ.get('TRITONSERVER_IPADDR', 'localhost')


def _unique_request_id():
    if len(_seen_request_ids) == 0:
        return 1
    else:
        return max(_seen_request_ids) + 1


def _range_repr_dtype(dtype):
    if dtype == np.float64:
        return np.int32
    elif dtype == np.float32:
        return np.int16
    elif dtype == np.float16:
        return np.int8
    elif dtype == np.object_:  # TYPE_STRING
        return np.int32
    return dtype


def serialize_byte_tensor_list(tensor_values):
    tensor_list = []
    for tensor_value in tensor_values:
        tensor_list.append(serialize_byte_tensor(tensor_value))
    return tensor_list


class UserData:

    def __init__(self):
        self._completed_requests = queue.Queue()


# Callback function used for async_stream_infer()
def completion_callback(user_data, result, error):
    # passing error raise and handling out
    user_data._completed_requests.put((result, error))


# Perform inference using an "addsum" type verification backend.
def infer_exact(
        tester,
        pf,
        tensor_shape,
        batch_size,
        input_dtype,
        output0_dtype,
        output1_dtype,
        output0_raw=True,
        output1_raw=True,
        model_version=None,
        swap=False,
        outputs=("OUTPUT0", "OUTPUT1"),
        use_http=True,
        use_grpc=True,
        use_http_json_tensors=True,
        skip_request_id_check=False,
        use_streaming=True,
        correlation_id=0,
        shm_region_names=None,
        precreated_shm_regions=None,
        use_system_shared_memory=False,
        use_cuda_shared_memory=False,
        priority=0,
        # 60 sec is the default value for L0_infer_valgrind
        network_timeout=60.0):
    # Lazy shm imports...
    if use_system_shared_memory or use_cuda_shared_memory:
        import tritonclient.utils.shared_memory as shm
        import tritonclient.utils.cuda_shared_memory as cudashm

    tester.assertTrue(use_http or use_grpc or use_streaming)
    # configs [ url, protocol, async stream, binary data ]
    configs = []
    if use_http:
        configs.append((f"{_tritonserver_ipaddr}:8000", "http", False, True))
        if output0_raw == output1_raw:
            # Float16 not supported for Input and Output via JSON
            if use_http_json_tensors and (input_dtype != np.float16) and \
               (output0_dtype != np.float16) and (output1_dtype != np.float16):
                configs.append(
                    (f"{_tritonserver_ipaddr}:8000", "http", False, False))
    if use_grpc:
        configs.append((f"{_tritonserver_ipaddr}:8001", "grpc", False, False))
    if use_streaming:
        configs.append((f"{_tritonserver_ipaddr}:8001", "grpc", True, False))

    # outputs are sum and difference of inputs so set max input
    # values so that they will not overflow the output. This
    # allows us to do an exact match. For float types use 8, 16,
    # 32 int range for fp 16, 32, 64 respectively. When getting
    # class outputs the result value/probability is returned as a
    # float so must use fp32 range in that case.
    rinput_dtype = _range_repr_dtype(input_dtype)
    routput0_dtype = _range_repr_dtype(
        output0_dtype if output0_raw else np.float32)
    routput1_dtype = _range_repr_dtype(
        output1_dtype if output1_raw else np.float32)
    val_min = max(
        np.iinfo(rinput_dtype).min,
        np.iinfo(routput0_dtype).min,
        np.iinfo(routput1_dtype).min) / 2
    val_max = min(
        np.iinfo(rinput_dtype).max,
        np.iinfo(routput0_dtype).max,
        np.iinfo(routput1_dtype).max) / 2

    input0_array = np.random.randint(low=val_min,
                                     high=val_max,
                                     size=tensor_shape,
                                     dtype=rinput_dtype)
    input1_array = np.random.randint(low=val_min,
                                     high=val_max,
                                     size=tensor_shape,
                                     dtype=rinput_dtype)
    if input_dtype != np.object_:
        input0_array = input0_array.astype(input_dtype)
        input1_array = input1_array.astype(input_dtype)

    # for unsigned data type, the value being subtracted must be less than the
    # value it is subtracted from, to avoid overflow.
    if val_min == 0:
        # swap element if the element at input 0 < input 1
        tmp = np.where(input0_array < input1_array, input1_array, input0_array)
        input1_array = np.where(input0_array < input1_array, input0_array,
                                input1_array)
        input0_array = tmp

    if not swap:
        output0_array = input0_array + input1_array
        output1_array = input0_array - input1_array
    else:
        output0_array = input0_array - input1_array
        output1_array = input0_array + input1_array

    if output0_dtype == np.object_:
        output0_array = np.array([
            unicode(str(x), encoding='utf-8') for x in (output0_array.flatten())
        ],
                                 dtype=object).reshape(output0_array.shape)
    else:
        output0_array = output0_array.astype(output0_dtype)
    if output1_dtype == np.object_:
        output1_array = np.array([
            unicode(str(x), encoding='utf-8') for x in (output1_array.flatten())
        ],
                                 dtype=object).reshape(output1_array.shape)
    else:
        output1_array = output1_array.astype(output1_dtype)

    if input_dtype == np.object_:
        in0n = np.array(
            [str(x) for x in input0_array.reshape(input0_array.size)],
            dtype=object)
        input0_array = in0n.reshape(input0_array.shape)
        in1n = np.array(
            [str(x) for x in input1_array.reshape(input1_array.size)],
            dtype=object)
        input1_array = in1n.reshape(input1_array.shape)

    # prepend size of string to output string data
    if output0_dtype == np.object_:
        if batch_size == 1:
            output0_array_tmp = serialize_byte_tensor_list([output0_array])
        else:
            output0_array_tmp = serialize_byte_tensor_list(output0_array)
    else:
        output0_array_tmp = output0_array

    if output1_dtype == np.object_:
        if batch_size == 1:
            output1_array_tmp = serialize_byte_tensor_list([output1_array])
        else:
            output1_array_tmp = serialize_byte_tensor_list(output1_array)
    else:
        output1_array_tmp = output1_array

    if output0_dtype == np.object_:
        output0_byte_size = sum(
            [serialized_byte_size(o0) for o0 in output0_array_tmp])
    else:
        output0_byte_size = sum([o0.nbytes for o0 in output0_array_tmp])

    if output1_dtype == np.object_:
        output1_byte_size = sum(
            [serialized_byte_size(o1) for o1 in output1_array_tmp])
    else:
        output1_byte_size = sum([o1.nbytes for o1 in output1_array_tmp])

    if batch_size == 1:
        input0_list = [input0_array]
        input1_list = [input1_array]
    else:
        input0_list = [x for x in input0_array]
        input1_list = [x for x in input1_array]

    # Serialization of string tensors in the case of shared memory must be done manually
    if input_dtype == np.object_:
        input0_list_tmp = serialize_byte_tensor_list(input0_list)
        input1_list_tmp = serialize_byte_tensor_list(input1_list)
    else:
        input0_list_tmp = input0_list
        input1_list_tmp = input1_list

    if input_dtype == np.object_:
        input0_byte_size = sum(
            [serialized_byte_size(i0) for i0 in input0_list_tmp])
        input1_byte_size = sum(
            [serialized_byte_size(i1) for i1 in input1_list_tmp])
    else:
        input0_byte_size = sum([i0.nbytes for i0 in input0_list_tmp])
        input1_byte_size = sum([i1.nbytes for i1 in input1_list_tmp])

    if model_version is not None:
        model_version = str(model_version)
    else:
        model_version = ""

    # Run inference and check results for each config
    inferAndCheckResults(tester, configs, pf, batch_size, model_version,
                         input_dtype, output0_dtype, output1_dtype,
                         tensor_shape, input0_array, input1_array,
                         output0_array, output1_array, output0_raw, output1_raw,
                         outputs, precreated_shm_regions, input0_list_tmp,
                         input1_list_tmp, shm_region_names, input0_byte_size,
                         input1_byte_size, output0_byte_size, output1_byte_size,
                         use_system_shared_memory, use_cuda_shared_memory,
                         network_timeout, skip_request_id_check)


def inferAndCheckResults(tester, configs, pf, batch_size, model_version,
                         input_dtype, output0_dtype, output1_dtype,
                         tensor_shape, input0_array, input1_array,
                         output0_array, output1_array, output0_raw, output1_raw,
                         outputs, precreated_shm_regions, input0_list_tmp,
                         input1_list_tmp, shm_region_names, input0_byte_size,
                         input1_byte_size, output0_byte_size, output1_byte_size,
                         use_system_shared_memory, use_cuda_shared_memory,
                         network_timeout, skip_request_id_check):
    # Lazy shm imports...
    if use_system_shared_memory or use_cuda_shared_memory:
        import tritonclient.utils.shared_memory as shm
        import tritonclient.utils.cuda_shared_memory as cudashm
    num_classes = 3

    # Get model platform
    model_name = tu.get_model_name(pf, input_dtype, output0_dtype,
                                   output1_dtype)
    if configs[0][1] == "http":
        metadata_client = httpclient.InferenceServerClient(configs[0][0],
                                                           verbose=True)
        metadata = metadata_client.get_model_metadata(model_name)
        platform = metadata["platform"]
    else:
        metadata_client = grpcclient.InferenceServerClient(configs[0][0],
                                                           verbose=True)
        metadata = metadata_client.get_model_metadata(model_name)
        platform = metadata.platform

    INPUT0 = "INPUT0"
    INPUT1 = "INPUT1"

    if platform == "pytorch_libtorch":
        OUTPUT0 = "OUTPUT__0"
        OUTPUT1 = "OUTPUT__1"
    else:
        OUTPUT0 = "OUTPUT0"
        OUTPUT1 = "OUTPUT1"

    # Create system/cuda shared memory regions if needed
    shm_regions, shm_handles = su.create_set_shm_regions(
        input0_list_tmp, input1_list_tmp, output0_byte_size, output1_byte_size,
        outputs, shm_region_names, precreated_shm_regions,
        use_system_shared_memory, use_cuda_shared_memory)
    for config in configs:
        model_name = tu.get_model_name(pf, input_dtype, output0_dtype,
                                       output1_dtype)

        if config[1] == "http":
            triton_client = httpclient.InferenceServerClient(
                config[0], verbose=True, network_timeout=network_timeout)
        else:
            triton_client = grpcclient.InferenceServerClient(config[0],
                                                             verbose=True)

        inputs = []
        if config[1] == "http":
            inputs.append(
                httpclient.InferInput(INPUT0, tensor_shape,
                                      np_to_triton_dtype(input_dtype)))
            inputs.append(
                httpclient.InferInput(INPUT1, tensor_shape,
                                      np_to_triton_dtype(input_dtype)))
        else:
            inputs.append(
                grpcclient.InferInput(INPUT0, tensor_shape,
                                      np_to_triton_dtype(input_dtype)))
            inputs.append(
                grpcclient.InferInput(INPUT1, tensor_shape,
                                      np_to_triton_dtype(input_dtype)))

        if not (use_cuda_shared_memory or use_system_shared_memory):
            if config[1] == "http":
                inputs[0].set_data_from_numpy(input0_array,
                                              binary_data=config[3])
                inputs[1].set_data_from_numpy(input1_array,
                                              binary_data=config[3])
            else:
                inputs[0].set_data_from_numpy(input0_array)
                inputs[1].set_data_from_numpy(input1_array)
        else:
            # Register necessary shared memory regions/handles
            su.register_add_shm_regions(inputs, outputs, shm_regions,
                                        precreated_shm_regions, shm_handles,
                                        input0_byte_size, input1_byte_size,
                                        output0_byte_size, output1_byte_size,
                                        use_system_shared_memory,
                                        use_cuda_shared_memory, triton_client)

        if batch_size == 1:
            expected0_sort_idx = [
                np.flip(np.argsort(x.flatten()), 0)
                for x in output0_array.reshape((1,) + tensor_shape)
            ]
            expected1_sort_idx = [
                np.flip(np.argsort(x.flatten()), 0)
                for x in output1_array.reshape((1,) + tensor_shape)
            ]
        else:
            expected0_sort_idx = [
                np.flip(np.argsort(x.flatten()), 0)
                for x in output0_array.reshape(tensor_shape)
            ]
            expected1_sort_idx = [
                np.flip(np.argsort(x.flatten()), 0)
                for x in output1_array.reshape(tensor_shape)
            ]

        # Force binary_data = False for shared memory and class
        output_req = []
        i = 0
        if "OUTPUT0" in outputs:
            if len(shm_regions) != 0:
                if config[1] == "http":
                    output_req.append(
                        httpclient.InferRequestedOutput(OUTPUT0,
                                                        binary_data=config[3]))
                else:
                    output_req.append(grpcclient.InferRequestedOutput(OUTPUT0))

                output_req[-1].set_shared_memory(shm_regions[2] + '_data',
                                                 output0_byte_size)
            else:
                if output0_raw:
                    if config[1] == "http":
                        output_req.append(
                            httpclient.InferRequestedOutput(
                                OUTPUT0, binary_data=config[3]))
                    else:
                        output_req.append(
                            grpcclient.InferRequestedOutput(OUTPUT0))
                else:
                    if config[1] == "http":
                        output_req.append(
                            httpclient.InferRequestedOutput(
                                OUTPUT0,
                                binary_data=config[3],
                                class_count=num_classes))
                    else:
                        output_req.append(
                            grpcclient.InferRequestedOutput(
                                OUTPUT0, class_count=num_classes))
            i += 1
        if "OUTPUT1" in outputs:
            if len(shm_regions) != 0:
                if config[1] == "http":
                    output_req.append(
                        httpclient.InferRequestedOutput(OUTPUT1,
                                                        binary_data=config[3]))
                else:
                    output_req.append(grpcclient.InferRequestedOutput(OUTPUT1))

                output_req[-1].set_shared_memory(shm_regions[2 + i] + '_data',
                                                 output1_byte_size)
            else:
                if output1_raw:
                    if config[1] == "http":
                        output_req.append(
                            httpclient.InferRequestedOutput(
                                OUTPUT1, binary_data=config[3]))
                    else:
                        output_req.append(
                            grpcclient.InferRequestedOutput(OUTPUT1))
                else:
                    if config[1] == "http":
                        output_req.append(
                            httpclient.InferRequestedOutput(
                                OUTPUT1,
                                binary_data=config[3],
                                class_count=num_classes))
                    else:
                        output_req.append(
                            grpcclient.InferRequestedOutput(
                                OUTPUT1, class_count=num_classes))

        if config[2]:
            user_data = UserData()
            triton_client.start_stream(partial(completion_callback, user_data))
            try:
                results = triton_client.async_stream_infer(
                    model_name,
                    inputs,
                    model_version=model_version,
                    outputs=output_req,
                    request_id=str(_unique_request_id()))
            except Exception as e:
                triton_client.stop_stream()
                raise e
            triton_client.stop_stream()
            (results, error) = user_data._completed_requests.get()
            if error is not None:
                raise error
        else:
            results = triton_client.infer(model_name,
                                          inputs,
                                          model_version=model_version,
                                          outputs=output_req,
                                          request_id=str(_unique_request_id()))

        last_response = results.get_response()

        if not skip_request_id_check:
            global _seen_request_ids
            if config[1] == "http":
                request_id = int(last_response["id"])
            else:
                request_id = int(last_response.id)
            tester.assertFalse(request_id in _seen_request_ids,
                               "request_id: {}".format(request_id))
            _seen_request_ids.add(request_id)

        if config[1] == "http":
            response_model_name = last_response["model_name"]
            if model_version != "":
                response_model_version = last_response["model_version"]
            response_outputs = last_response["outputs"]
        else:
            response_model_name = last_response.model_name
            if model_version != "":
                response_model_version = last_response.model_version
            response_outputs = last_response.outputs

        tester.assertEqual(response_model_name, model_name)

        if model_version != "":
            tester.assertEqual(str(response_model_version), model_version)

        tester.assertEqual(len(response_outputs), len(outputs))

        for result in response_outputs:
            if config[1] == "http":
                result_name = result["name"]
            else:
                result_name = result.name

            if ((result_name == OUTPUT0 and output0_raw) or
                (result_name == OUTPUT1 and output1_raw)):
                if use_system_shared_memory or use_cuda_shared_memory:
                    if result_name == OUTPUT0:
                        shm_handle = shm_handles[2]
                    else:
                        shm_handle = shm_handles[3]

                    output = results.get_output(result_name)
                    if config[1] == "http":
                        output_datatype = output['datatype']
                        output_shape = output['shape']
                    else:
                        output_datatype = output.datatype
                        output_shape = output.shape
                    output_dtype = triton_to_np_dtype(output_datatype)
                if use_system_shared_memory:
                    output_data = shm.get_contents_as_numpy(
                        shm_handle, output_dtype, output_shape)
                elif use_cuda_shared_memory:
                    output_data = cudashm.get_contents_as_numpy(
                        shm_handle, output_dtype, output_shape)
                else:
                    output_data = results.as_numpy(result_name)
                    if (output_data.dtype == np.object_) and (not config[3]):
                        if config[1] == 'http':
                            output_data = np.array([
                                unicode(str(x), encoding='utf-8')
                                for x in (output_data.flatten())
                            ],
                                                   dtype=np.object_).reshape(
                                                       output_data.shape)
                        elif config[1] == 'grpc':
                            output_data = np.array(
                                [x for x in (output_data.flatten())],
                                dtype=np.object_).reshape(output_data.shape)

                if result_name == OUTPUT0:
                    tester.assertTrue(
                        np.array_equal(output_data, output0_array),
                        "{}, {} expected: {}, got {}".format(
                            model_name, OUTPUT0, output0_array, output_data))
                elif result_name == OUTPUT1:
                    tester.assertTrue(
                        np.array_equal(output_data, output1_array),
                        "{}, {} expected: {}, got {}".format(
                            model_name, OUTPUT1, output1_array, output_data))
                else:
                    tester.assertTrue(
                        False, "unexpected raw result {}".format(result_name))
            else:
                for b in range(batch_size):
                    # num_classes values must be returned and must
                    # match expected top values
                    if "nobatch" in pf:
                        class_list = results.as_numpy(result_name)
                    else:
                        class_list = results.as_numpy(result_name)[b]

                    tester.assertEqual(len(class_list), num_classes)
                    if batch_size == 1:
                        expected0_flatten = output0_array.flatten()
                        expected1_flatten = output1_array.flatten()
                    else:
                        expected0_flatten = output0_array[b].flatten()
                        expected1_flatten = output1_array[b].flatten()

                    for idx, class_label in enumerate(class_list):
                        # can't compare indices since could have different
                        # indices with the same value/prob, so check that
                        # the value of each index equals the expected value.
                        # Only compare labels when the indices are equal.
                        if type(class_label) == str:
                            ctuple = class_label.split(':')
                        else:
                            ctuple = "".join(
                                chr(x) for x in class_label).split(':')
                        cval = float(ctuple[0])
                        cidx = int(ctuple[1])
                        if result_name == OUTPUT0:
                            tester.assertEqual(cval, expected0_flatten[cidx])
                            tester.assertEqual(
                                cval,
                                expected0_flatten[expected0_sort_idx[b][idx]])
                            if cidx == expected0_sort_idx[b][idx]:
                                tester.assertEqual(
                                    ctuple[2], 'label{}'.format(
                                        expected0_sort_idx[b][idx]))
                        elif result_name == OUTPUT1:
                            tester.assertEqual(cval, expected1_flatten[cidx])
                            tester.assertEqual(
                                cval,
                                expected1_flatten[expected1_sort_idx[b][idx]])
                        else:
                            tester.assertTrue(
                                False, "unexpected class result {}".format(
                                    result_name))

    # Unregister system/cuda shared memory regions if they exist
    su.unregister_cleanup_shm_regions(shm_regions, shm_handles,
                                      precreated_shm_regions, outputs,
                                      use_system_shared_memory,
                                      use_cuda_shared_memory)

    return results


# resize the dummy tensor with the provided values in the shape tensor and finally
# return the shape of the resized tensor.
def infer_shape_tensor(tester,
                       pf,
                       tensor_dtype,
                       input_shape_values,
                       dummy_input_shapes,
                       use_http=True,
                       use_grpc=True,
                       use_streaming=True,
                       shm_suffix="",
                       use_system_shared_memory=False,
                       priority=0,
                       timeout_us=0,
                       batch_size=1):
    # Lazy shm imports...
    if use_system_shared_memory:
        import tritonclient.utils.shared_memory as shm

    tester.assertTrue(use_http or use_grpc or use_streaming)
    tester.assertTrue(pf == "plan" or pf == "plan_nobatch")
    tester.assertEqual(len(input_shape_values), len(dummy_input_shapes))

    configs = []
    if use_http:
        configs.append((f"{_tritonserver_ipaddr}:8000", "http", False))
    if use_grpc:
        configs.append((f"{_tritonserver_ipaddr}:8001", "grpc", False))
    if use_streaming:
        configs.append((f"{_tritonserver_ipaddr}:8001", "grpc", True))

    io_cnt = len(input_shape_values)

    # FIXME wrap up shm handle cleanup
    # item is (handle, byte_size)
    input_shm_handle_list = []
    output_shm_handle_list = []
    dummy_input_list = []
    input_list = []
    expected_dict = dict()
    # Prepare IO in advance
    for io_num in range(io_cnt):
        dummy_input_name = "DUMMY_INPUT{}".format(io_num)
        input_name = "INPUT{}".format(io_num)
        dummy_output_name = "DUMMY_OUTPUT{}".format(io_num)
        output_name = "OUTPUT{}".format(io_num)

        # Prepare the dummy tensor
        rtensor_dtype = _range_repr_dtype(tensor_dtype)
        if (rtensor_dtype != bool):
            dummy_in0 = np.random.randint(low=np.iinfo(rtensor_dtype).min,
                                          high=np.iinfo(rtensor_dtype).max,
                                          size=dummy_input_shapes[io_num],
                                          dtype=rtensor_dtype)
        else:
            dummy_in0 = np.random.choice(a=[False, True],
                                         size=dummy_input_shapes[io_num])
        if tensor_dtype != np.object_:
            dummy_in0 = dummy_in0.astype(tensor_dtype)
        else:
            dummy_in0 = np.array([str(x) for x in dummy_in0.flatten()],
                                 dtype=object).reshape(dummy_in0.shape)
        dummy_input_list.append(dummy_in0)

        # Prepare shape input tensor
        in0 = np.asarray(input_shape_values[io_num], dtype=np.int32)
        input_list.append(in0)

        # Prepare the expected value for the output. Skip dummy output as we
        # only care about its shape (== value of OUTPUT*)
        expected_dict[output_name] = np.ndarray.copy(in0)

        # Only need to create region once
        input_byte_size = in0.size * np.dtype(np.int32).itemsize
        output_byte_size = input_byte_size * batch_size
        if use_system_shared_memory:
            input_shm_handle_list.append(
                (shm.create_shared_memory_region(input_name + shm_suffix,
                                                 '/' + input_name + shm_suffix,
                                                 input_byte_size),
                 input_byte_size))
            output_shm_handle_list.append(
                (shm.create_shared_memory_region(output_name + shm_suffix,
                                                 '/' + output_name + shm_suffix,
                                                 output_byte_size),
                 output_byte_size))
            shm.set_shared_memory_region(input_shm_handle_list[-1][0], [
                in0,
            ])

    model_name = tu.get_zero_model_name(pf, io_cnt, tensor_dtype)
    # Run inference and check results for each config
    for config in configs:
        client_utils = grpcclient if config[1] == "grpc" else httpclient
        triton_client = client_utils.InferenceServerClient(config[0],
                                                           verbose=True)

        inputs = []
        outputs = []

        # Set IOs
        for io_num in range(io_cnt):
            dummy_input_name = "DUMMY_INPUT{}".format(io_num)
            input_name = "INPUT{}".format(io_num)
            dummy_output_name = "DUMMY_OUTPUT{}".format(io_num)
            output_name = "OUTPUT{}".format(io_num)

            inputs.append(
                client_utils.InferInput(dummy_input_name,
                                        dummy_input_shapes[io_num],
                                        np_to_triton_dtype(tensor_dtype)))
            inputs.append(
                client_utils.InferInput(input_name, input_list[io_num].shape,
                                        "INT32"))
            outputs.append(client_utils.InferRequestedOutput(dummy_output_name))
            outputs.append(client_utils.InferRequestedOutput(output_name))

            # -2: dummy; -1: input
            inputs[-2].set_data_from_numpy(dummy_input_list[io_num])
            if (not use_system_shared_memory):
                inputs[-1].set_data_from_numpy(input_list[io_num])
            else:
                input_byte_size = input_shm_handle_list[io_num][1]
                output_byte_size = output_shm_handle_list[io_num][1]
                triton_client.register_system_shared_memory(
                    input_name + shm_suffix, "/" + input_name + shm_suffix,
                    input_byte_size)
                triton_client.register_system_shared_memory(
                    output_name + shm_suffix, "/" + output_name + shm_suffix,
                    output_byte_size)
                inputs[-1].set_shared_memory(input_name + shm_suffix,
                                             input_byte_size)
                outputs[-1].set_shared_memory(output_name + shm_suffix,
                                              output_byte_size)

        if config[2]:
            user_data = UserData()
            triton_client.start_stream(partial(completion_callback, user_data))
            try:
                results = triton_client.async_stream_infer(model_name,
                                                           inputs,
                                                           outputs=outputs,
                                                           priority=priority,
                                                           timeout=timeout_us)
            except Exception as e:
                triton_client.stop_stream()
                raise e
            triton_client.stop_stream()
            (results, error) = user_data._completed_requests.get()
            if error is not None:
                raise error
        else:
            results = triton_client.infer(model_name,
                                          inputs,
                                          outputs=outputs,
                                          priority=priority,
                                          timeout=timeout_us)

        for io_num in range(io_cnt):
            output_name = "OUTPUT{}".format(io_num)
            dummy_output_name = "DUMMY_OUTPUT{}".format(io_num)
            expected = expected_dict[output_name]

            # get outputs as numpy array
            dummy_out = results.as_numpy(dummy_output_name)
            if (not use_system_shared_memory):
                out = results.as_numpy(output_name)
            else:
                output = results.get_output(output_name)
                if config[1] == "grpc":
                    output_shape = output.shape
                else:
                    output_shape = output["shape"]
                out = shm.get_contents_as_numpy(
                    output_shm_handle_list[io_num][0], np.int32, output_shape)

            # if out shape is 2D, it is batched
            if (len(out.shape) == 2):
                # The shape of the dummy output should be equal to the shape values
                # specified in the shape tensor
                tester.assertTrue(
                    np.array_equal(dummy_out.shape[1:], out[0]),
                    "{}, {} shape, expected: {}, got {}".format(
                        model_name, dummy_output_name, out[0],
                        dummy_out.shape[1:]))
                for b in range(1, out.shape[0]):
                    tester.assertTrue(
                        np.array_equal(out[b - 1], out[b]),
                        "expect shape tensor has consistent value, "
                        "expected: {}, got {}".format(out[b - 1], out[b]))
                out = out[0]
            else:
                tester.assertTrue(
                    np.array_equal(dummy_out.shape, out),
                    "{}, {} shape, expected: {}, got {}".format(
                        model_name, dummy_output_name, out, dummy_out.shape))
            tester.assertTrue(
                np.array_equal(out, expected),
                "{}, {}, expected: {}, got {}".format(model_name, output_name,
                                                      expected, out))

            # unregister shared memory region for next config
            if use_system_shared_memory:
                triton_client.unregister_system_shared_memory(input_name +
                                                              shm_suffix)
                triton_client.unregister_system_shared_memory(output_name +
                                                              shm_suffix)

    for handle in input_shm_handle_list:
        shm.destroy_shared_memory_region(handle[0])
    for handle in output_shm_handle_list:
        shm.destroy_shared_memory_region(handle[0])


# Perform inference using a "nop" model that expects some form or
# zero-sized input/output tensor.
# FIXME Support for empty tensors using non-empty shared memory regions.
# Currently shared memory support is broken for empty input/outputs tensors.
def infer_zero(tester,
               pf,
               batch_size,
               tensor_dtype,
               input_shapes,
               output_shapes,
               model_version=None,
               use_http=True,
               use_grpc=True,
               use_http_json_tensors=True,
               use_streaming=True,
               shm_region_name_prefix=None,
               use_system_shared_memory=False,
               use_cuda_shared_memory=False,
               priority=0,
               timeout_us=0,
               override_model_name=None,
               override_input_names=[],
               override_output_names=[]):
    # Lazy shm imports...
    if use_system_shared_memory or use_cuda_shared_memory:
        import tritonclient.utils.shared_memory as shm
        import tritonclient.utils.cuda_shared_memory as cudashm

    tester.assertTrue(use_http or use_grpc or use_streaming)
    configs = []
    if use_http:
        configs.append((f"{_tritonserver_ipaddr}:8000", "http", False, True))
        if use_http_json_tensors and (tensor_dtype != np.float16):
            configs.append(
                (f"{_tritonserver_ipaddr}:8000", "http", False, False))
    if use_grpc:
        configs.append((f"{_tritonserver_ipaddr}:8001", "grpc", False, False))
    if use_streaming:
        configs.append((f"{_tritonserver_ipaddr}:8001", "grpc", True, False))
    tester.assertEqual(len(input_shapes), len(output_shapes))
    io_cnt = len(input_shapes)

    if shm_region_name_prefix is None:
        shm_region_name_prefix = ["input", "output"]

    input_dict = {}
    expected_dict = {}
    shm_ip_handles = list()
    shm_op_handles = list()

    # Get model platform
    if override_model_name is None:
        model_name = tu.get_zero_model_name(pf, io_cnt, tensor_dtype)
    else:
        model_name = override_model_name
    if configs[0][1] == "http":
        metadata_client = httpclient.InferenceServerClient(configs[0][0],
                                                           verbose=True)
        metadata = metadata_client.get_model_metadata(model_name)
        platform = metadata["platform"]
    else:
        metadata_client = grpcclient.InferenceServerClient(configs[0][0],
                                                           verbose=True)
        metadata = metadata_client.get_model_metadata(model_name)
        platform = metadata.platform

    for io_num in range(io_cnt):
        if override_input_names:
            input_name = override_input_names[io_num]
        else:
            if platform == "pytorch_libtorch":
                input_name = "INPUT__{}".format(io_num)
            else:
                input_name = "INPUT{}".format(io_num)

        if override_output_names:
            output_name = override_output_names[io_num]
        else:
            if platform == "pytorch_libtorch":
                output_name = "OUTPUT__{}".format(io_num)
            else:
                output_name = "OUTPUT{}".format(io_num)

        input_shape = input_shapes[io_num]
        output_shape = output_shapes[io_num]

        rtensor_dtype = _range_repr_dtype(tensor_dtype)
        if (rtensor_dtype != bool):
            input_array = np.random.randint(low=np.iinfo(rtensor_dtype).min,
                                            high=np.iinfo(rtensor_dtype).max,
                                            size=input_shape,
                                            dtype=rtensor_dtype)
        else:
            input_array = np.random.choice(a=[False, True], size=input_shape)
        if tensor_dtype != np.object_:
            input_array = input_array.astype(tensor_dtype)
            expected_array = np.ndarray.copy(input_array)
        else:
            expected_array = np.array([
                unicode(str(x), encoding='utf-8')
                for x in input_array.flatten()
            ],
                                      dtype=object)
            input_array = np.array([str(x) for x in input_array.flatten()],
                                   dtype=object).reshape(input_array.shape)

        expected_array = expected_array.reshape(output_shape)
        expected_dict[output_name] = expected_array

        if tensor_dtype == np.object_:
            output_byte_size = serialized_byte_size(expected_array)
        else:
            output_byte_size = expected_array.nbytes

        if batch_size == 1:
            input_list = [input_array]
        else:
            input_list = [x for x in input_array]

        # Serialization of string tensors in the case of shared memory must be done manually
        if tensor_dtype == np.object_:
            input_list_tmp = serialize_byte_tensor_list(input_list)
        else:
            input_list_tmp = input_list

        if tensor_dtype == np.object_:
            input_byte_size = sum(
                [serialized_byte_size(ip) for ip in input_list_tmp])
        else:
            input_byte_size = sum([ip.nbytes for ip in input_list_tmp])

        # create and register shared memory region for inputs and outputs
        shm_io_handles = su.create_set_either_shm_region(
            [
                shm_region_name_prefix[0] + str(io_num),
                shm_region_name_prefix[1] + str(io_num)
            ], input_list_tmp, input_byte_size, output_byte_size,
            use_system_shared_memory, use_cuda_shared_memory)

        if len(shm_io_handles) != 0:
            shm_ip_handles.append(shm_io_handles[0])
            shm_op_handles.append(shm_io_handles[1])
        input_dict[input_name] = input_array

    if model_version is not None:
        model_version = str(model_version)
    else:
        model_version = ""

    # Run inference and check results for each config
    for config in configs:
        if config[1] == "http":
            triton_client = httpclient.InferenceServerClient(config[0],
                                                             verbose=True)
        else:
            triton_client = grpcclient.InferenceServerClient(config[0],
                                                             verbose=True)

        inputs = []
        output_req = []
        for io_num, (input_name, output_name) in enumerate(
                zip(input_dict.keys(), expected_dict.keys())):
            input_data = input_dict[input_name]
            output_data = expected_dict[output_name]
            if tensor_dtype == np.object_:
                input_byte_size = serialized_byte_size(
                    serialize_byte_tensor(input_data))
                output_byte_size = serialized_byte_size(
                    serialize_byte_tensor(output_data))
            else:
                input_byte_size = input_data.nbytes
                output_byte_size = output_data.nbytes
            if config[1] == "http":
                inputs.append(
                    httpclient.InferInput(input_name, input_data.shape,
                                          np_to_triton_dtype(tensor_dtype)))
                output_req.append(
                    httpclient.InferRequestedOutput(output_name,
                                                    binary_data=config[3]))
            else:
                inputs.append(
                    grpcclient.InferInput(input_name, input_data.shape,
                                          np_to_triton_dtype(tensor_dtype)))
                output_req.append(grpcclient.InferRequestedOutput(output_name))

            if not (use_cuda_shared_memory or use_system_shared_memory):
                if config[1] == "http":
                    inputs[-1].set_data_from_numpy(input_data,
                                                   binary_data=config[3])
                else:
                    inputs[-1].set_data_from_numpy(input_data)
            else:
                # Register necessary shared memory regions/handles
                su.register_add_either_shm_regions(
                    inputs, output_req, shm_region_name_prefix,
                    (shm_ip_handles, shm_op_handles), io_num, input_byte_size,
                    output_byte_size, use_system_shared_memory,
                    use_cuda_shared_memory, triton_client)

        if config[2]:
            user_data = UserData()
            triton_client.start_stream(partial(completion_callback, user_data))
            try:
                results = triton_client.async_stream_infer(
                    model_name,
                    inputs,
                    model_version=model_version,
                    outputs=output_req,
                    request_id=str(_unique_request_id()),
                    priority=priority,
                    timeout=timeout_us)
            except Exception as e:
                triton_client.stop_stream()
                raise e
            triton_client.stop_stream()
            (results, error) = user_data._completed_requests.get()
            if error is not None:
                raise error
        else:
            results = triton_client.infer(model_name,
                                          inputs,
                                          model_version=model_version,
                                          outputs=output_req,
                                          request_id=str(_unique_request_id()),
                                          priority=priority,
                                          timeout=timeout_us)

        last_response = results.get_response()

        if config[1] == "http":
            response_model_name = last_response["model_name"]
            if model_version != "":
                response_model_version = last_response["model_version"]
            response_outputs = last_response["outputs"]
        else:
            response_model_name = last_response.model_name
            if model_version != "":
                response_model_version = last_response.model_version
            response_outputs = last_response.outputs

        tester.assertEqual(response_model_name, model_name)

        if model_version != "":
            tester.assertEqual(response_model_version, model_version)

        tester.assertEqual(len(response_outputs), io_cnt)

        for result in response_outputs:
            if config[1] == "http":
                result_name = result["name"]
            else:
                result_name = result.name

            tester.assertIn(result_name, expected_dict)
            if use_system_shared_memory or use_cuda_shared_memory:
                if platform == "pytorch_libtorch":
                    io_num = int(result_name.split("OUTPUT__")[1])
                else:
                    io_num = int(result_name.split("OUTPUT")[1])
                shm_handle = shm_op_handles[io_num]

                output = results.get_output(result_name)
                if config[1] == "http":
                    output_datatype = output['datatype']
                    output_shape = output['shape']
                else:
                    output_datatype = output.datatype
                    output_shape = output.shape
                output_dtype = triton_to_np_dtype(output_datatype)
            if use_system_shared_memory:
                output_data = shm.get_contents_as_numpy(shm_handle,
                                                        output_dtype,
                                                        output_shape)
            elif use_cuda_shared_memory:
                output_data = cudashm.get_contents_as_numpy(
                    shm_handle, output_dtype, output_shape)
            else:
                output_data = results.as_numpy(result_name)

                if (output_data.dtype == np.object_) and (config[3] == False):
                    if config[1] == 'http':
                        output_data = np.array([
                            unicode(str(x), encoding='utf-8')
                            for x in (output_data.flatten())
                        ],
                                               dtype=np.object_).reshape(
                                                   output_data.shape)
                    elif config[1] == 'grpc':
                        output_data = np.array(
                            [x for x in (output_data.flatten())],
                            dtype=np.object_).reshape(output_data.shape)

            expected = expected_dict[result_name]
            tester.assertEqual(output_data.shape, expected.shape)
            tester.assertTrue(
                np.array_equal(output_data, expected),
                "{}, {}, expected: {}, got {}".format(model_name, result_name,
                                                      expected, output_data))

    if len(shm_ip_handles) != 0:
        for io_num in range(io_cnt):
            if use_cuda_shared_memory:
                triton_client.unregister_cuda_shared_memory(
                    shm_region_name_prefix[0] + str(io_num) + '_data')
                triton_client.unregister_cuda_shared_memory(
                    shm_region_name_prefix[0] + str(io_num) + '_data')
                cudashm.destroy_shared_memory_region(shm_ip_handles[io_num])
                cudashm.destroy_shared_memory_region(shm_op_handles[io_num])
            else:
                triton_client.unregister_system_shared_memory(
                    shm_region_name_prefix[1] + str(io_num) + '_data')
                triton_client.unregister_system_shared_memory(
                    shm_region_name_prefix[1] + str(io_num) + '_data')
                shm.destroy_shared_memory_region(shm_ip_handles[io_num])
                shm.destroy_shared_memory_region(shm_op_handles[io_num])

    return results


# qa/common/libtorch_infer_client.py

#!/usr/bin/env python
# Copyright 2021-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import os
import sys

sys.path.append("../common")

import unittest
import numpy as np
import test_util as tu

import tritonclient.http as httpclient
from tritonclient.utils import InferenceServerException

# By default, find tritonserver on "localhost", but can be overridden
# with TRITONSERVER_IPADDR envvar
_tritonserver_ipaddr = os.environ.get('TRITONSERVER_IPADDR', 'localhost')


class InferTest(tu.TestResultCollector):

    def test_infer(self):
        try:
            triton_client = httpclient.InferenceServerClient(
                url=f"{_tritonserver_ipaddr}:8000")
        except Exception as e:
            print("channel creation failed: " + str(e))
            sys.exit(1)

        model_name = "libtorch_int32_int32_int32"

        inputs = []
        outputs = []
        inputs.append(httpclient.InferInput('INPUT0', [1, 16], "INT32"))
        inputs.append(httpclient.InferInput('INPUT1', [1, 16], "INT32"))

        # Create the data for the two input tensors. Initialize the first
        # to unique integers and the second to all ones.
        input0_data = np.arange(start=0, stop=16, dtype=np.int32)
        input0_data = np.expand_dims(input0_data, axis=0)
        input1_data = np.full(shape=(1, 16), fill_value=-1, dtype=np.int32)

        # Initialize the data
        inputs[0].set_data_from_numpy(input0_data, binary_data=True)
        inputs[1].set_data_from_numpy(input1_data, binary_data=True)

        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT__0', binary_data=True))
        outputs.append(
            httpclient.InferRequestedOutput('OUTPUT__1', binary_data=True))

        results = triton_client.infer(model_name, inputs, outputs=outputs)

        output0_data = results.as_numpy('OUTPUT__0')
        output1_data = results.as_numpy('OUTPUT__1')

        # Validate the results by comparing with precomputed values.
        for i in range(16):
            print(
                str(input0_data[0][i]) + " - " + str(input1_data[0][i]) +
                " = " + str(output0_data[0][i]))
            print(
                str(input0_data[0][i]) + " + " + str(input1_data[0][i]) +
                " = " + str(output1_data[0][i]))
            if (input0_data[0][i] - input1_data[0][i]) != output0_data[0][i]:
                print("sync infer error: incorrect difference")
                sys.exit(1)
            if (input0_data[0][i] + input1_data[0][i]) != output1_data[0][i]:
                print("sync infer error: incorrect sum")
                sys.exit(1)


if __name__ == '__main__':
    unittest.main()


# qa/common/nightly_email_helper.py

#!/usr/bin/env python
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

from email import encoders
import smtplib
from email.mime.multipart import MIMEMultipart
from email.mime.base import MIMEBase
from email.mime.text import MIMEText
import glob
import os
import sys
import tarfile


def send(subject: str,
         content: str,
         attachments=None,
         files_to_tar=None,
         is_html=False):
    FROM = os.environ.get('TRITON_FROM', '')
    TO = os.environ.get('TRITON_TO_DL', '')
    if FROM == '' or TO == '':
        print('Must set TRITON_FROM and TRITON_TO_DL env variables')
        sys.exit(1)

    msg = MIMEMultipart('alternative')
    msg['Subject'] = subject
    msg['From'] = FROM
    msg['To'] = TO
    if is_html:
        mime_text = MIMEText(content, 'html')
    else:
        mime_text = MIMEText(content)
    msg.attach(mime_text)

    if attachments is None:
        attachments = []

    if files_to_tar is not None:
        with tarfile.open(subject + ".tgz", "w:gz") as csv_tar:
            for filename in glob.glob(files_to_tar):
                csv_tar.add(filename)
        attachments.append(subject + ".tgz")

    for fname in attachments:
        p = MIMEBase('application', 'octet-stream')
        with open(fname, "rb") as attachment:
            p.set_payload((attachment).read())
        encoders.encode_base64(p)
        p.add_header('Content-Disposition',
                     "attachment; filename= %s" % (fname))
        msg.attach(p)

    mailServer = smtplib.SMTP("mailgw.nvidia.com")
    mailServer.send_message(msg)
    mailServer.quit()


# qa/common/reporter.py

#!/usr/bin/python

# Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
import csv
import json
import os
import requests
import socket

FLAGS = None

ENVS = [
    "CUDA_DRIVER_VERSION", "CUDA_VERSION", "TRITON_SERVER_VERSION",
    "NVIDIA_TRITON_SERVER_VERSION", "TRT_VERSION", "CUDNN_VERSION",
    "CUBLAS_VERSION", "BENCHMARK_PIPELINE", "BENCHMARK_REPO_BRANCH",
    "BENCHMARK_REPO_COMMIT", "BENCHMARK_CLUSTER", "BENCHMARK_GPU_COUNT"
]


def annotate(datas):
    # Add all interesting envvar values
    for data in datas:
        for env in ENVS:
            if env in os.environ:
                val = os.environ[env]
                data['s_' + env.lower()] = val

        # Add this system's name. If running within slurm use
        # SLURM_JOB_NODELIST as the name (this assumes that the slurm
        # job was scheduled on a single node, otherwise
        # SLURM_JOB_NODELIST will list multiple nodes).
        if 'SLURM_JOB_NODELIST' in os.environ:
            data['s_benchmark_system'] = os.environ['SLURM_JOB_NODELIST']
        else:
            data['s_benchmark_system'] = socket.gethostname()


def annotate_csv(data, csv_file):
    csv_reader = csv.reader(csv_file, delimiter=',')
    linenum = 0
    header_row = None
    concurrency_row = None
    for row in csv_reader:
        if linenum == 0:
            header_row = row
        else:
            concurrency_row = row
            break
        linenum += 1

    if (header_row is not None) and (concurrency_row is not None):
        avg_latency_us = 0
        for header, result in zip(header_row, concurrency_row):
            if header == 'Inferences/Second':
                data['d_infer_per_sec'] = float(result)
            elif ((header == 'Client Send') or
                  (header == 'Network+Server Send/Recv') or
                  (header == 'Server Queue') or
                  (header == 'Server Compute Input') or
                  (header == 'Server Compute Output') or
                  (header == 'Server Compute Infer') or
                  (header == 'Client Recv')):
                avg_latency_us += float(result)
            elif header == 'p50 latency':
                data['d_latency_p50_ms'] = float(result) / 1000.0
            elif header == 'p90 latency':
                data['d_latency_p90_ms'] = float(result) / 1000.0
            elif header == 'p95 latency':
                data['d_latency_p95_ms'] = float(result) / 1000.0
            elif header == 'p99 latency':
                data['d_latency_p99_ms'] = float(result) / 1000.0

        data['d_latency_avg_ms'] = avg_latency_us / 1000.0


def post_to_url(url, data):
    headers = {'Content-Type': 'application/json', 'Accept-Charset': 'UTF-8'}
    r = requests.post(url, data=data, headers=headers)
    r.raise_for_status()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-v',
                        '--verbose',
                        action="store_true",
                        required=False,
                        default=False,
                        help='Enable verbose output')
    parser.add_argument('-o',
                        '--output',
                        type=str,
                        required=False,
                        help='Output filename')
    parser.add_argument('-u',
                        '--url',
                        type=str,
                        required=False,
                        help='Post results to a URL')
    parser.add_argument('--csv',
                        type=argparse.FileType('r'),
                        required=False,
                        help='perf_analyzer generated CSV')
    parser.add_argument('file', type=argparse.FileType('r'))
    FLAGS = parser.parse_args()

    data = json.loads(FLAGS.file.read())

    if FLAGS.verbose:
        print("*** Load json ***")
        print(json.dumps(data, sort_keys=True, indent=2))

    if FLAGS.csv is not None:
        if len(data) != 1:
            raise Exception(
                "--csv requires that json data have a single array entry")
        annotate_csv(data[0], FLAGS.csv)
        if FLAGS.verbose:
            print("*** Annotate CSV ***")
            print(json.dumps(data, sort_keys=True, indent=2))

    annotate(data)

    if FLAGS.verbose:
        print("*** Post Annotate ***")
        print(json.dumps(data, sort_keys=True, indent=2))

    if FLAGS.output is not None:
        with open(FLAGS.output, "w") as f:
            f.write(json.dumps(data))
            f.write("\n")

    if FLAGS.url is not None:
        post_to_url(FLAGS.url, json.dumps(data))


# qa/common/sequence_util.py

# Copyright 2019-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

from builtins import range
from builtins import str
import os
import sys
import time
import threading
import numpy as np
import infer_util as iu
import test_util as tu
from functools import partial

import tritonclient.grpc as grpcclient
import tritonclient.http as httpclient
from tritonclient.utils import *

if sys.version_info >= (3, 0):
    import queue
else:
    import Queue as queue

# By default, find tritonserver on "localhost", but can be overridden
# with TRITONSERVER_IPADDR envvar
_tritonserver_ipaddr = os.environ.get('TRITONSERVER_IPADDR', 'localhost')

_test_system_shared_memory = bool(
    int(os.environ.get('TEST_SYSTEM_SHARED_MEMORY', 0)))
_test_cuda_shared_memory = bool(
    int(os.environ.get('TEST_CUDA_SHARED_MEMORY', 0)))

if _test_system_shared_memory:
    import tritonclient.utils.shared_memory as shm
if _test_cuda_shared_memory:
    import tritonclient.utils.cuda_shared_memory as cudashm

_test_valgrind = bool(int(os.environ.get('TEST_VALGRIND', 0)))
_test_jetson = bool(int(os.environ.get('TEST_JETSON', 0)))

_max_sequence_idle_ms = 5000
_valgrind_delay_ms = bool(int(os.environ.get('TEST_DELAY_MS', 50)))

_deferred_exceptions_lock = threading.Lock()
_deferred_exceptions = None
_jetson_slowdown_factor = 3


class UserData:

    def __init__(self):
        self._completed_requests = queue.Queue()


# Callback function used for async_stream_infer()
def completion_callback(user_data, result, error):
    # passing error raise and handling out
    user_data._completed_requests.put((result, error))


class SequenceBatcherTestUtil(tu.TestResultCollector):

    def setUp(self):
        # The helper client for setup will be GRPC for simplicity.
        self.triton_client_ = grpcclient.InferenceServerClient(
            f"{_tritonserver_ipaddr}:8001")
        self.clear_deferred_exceptions()

    def clear_deferred_exceptions(self):
        global _deferred_exceptions
        with _deferred_exceptions_lock:
            _deferred_exceptions = []

    def add_deferred_exception(self, ex):
        global _deferred_exceptions
        with _deferred_exceptions_lock:
            _deferred_exceptions.append(ex)

    def check_deferred_exception(self):
        # Just raise one of the exceptions...
        with _deferred_exceptions_lock:
            if len(_deferred_exceptions) > 0:
                raise _deferred_exceptions[0]

    def check_failure(self):
        # Check securely whether a failure has been registered
        # This is generic because the failure behavior is undefined
        # for ragged batches.
        with _deferred_exceptions_lock:
            if len(_deferred_exceptions) == 0:
                raise Exception("Unexpected inference success")

    def precreate_register_regions(self,
                                   value_list,
                                   dtype,
                                   i,
                                   batch_size=1,
                                   tensor_shape=(1,)):
        if _test_system_shared_memory or _test_cuda_shared_memory:
            shm_region_handles = []
            for j, value in enumerate(value_list):
                # For string we can't know the size of the output
                # so we conservatively assume 64 bytes for each
                # element of the output
                if dtype == np.object_:
                    output_byte_size = 4  # size of empty string
                else:
                    output_byte_size = 0

                # create data
                input_list = list()
                for b in range(batch_size):
                    if dtype == np.object_:
                        in0 = np.full(tensor_shape, value, dtype=np.int32)
                        in0n = np.array([
                            str(x).encode('utf-8')
                            for x in in0.reshape(in0.size)
                        ],
                                        dtype=object)
                        in0 = in0n.reshape(tensor_shape)
                        output_byte_size += 64 * in0.size
                    else:
                        in0 = np.full(tensor_shape, value, dtype=dtype)
                        output_byte_size += np.dtype(dtype).itemsize * in0.size
                    input_list.append(in0)

                if dtype == np.object_:
                    input_list_tmp = iu.serialize_byte_tensor_list(input_list)
                    input_byte_size = sum(
                        [serialized_byte_size(i0) for i0 in input_list_tmp])
                else:
                    input_list_tmp = input_list
                    input_byte_size = sum([i0.nbytes for i0 in input_list_tmp])

                # create shared memory regions and copy data for input values
                ip_name = 'ip{}{}'.format(i, j)
                op_name = 'op{}{}_data'.format(i, j)
                if _test_system_shared_memory:
                    shm_ip_handle = shm.create_shared_memory_region(
                        ip_name, '/' + ip_name, input_byte_size)
                    shm_op_handle = shm.create_shared_memory_region(
                        op_name, '/' + op_name, output_byte_size)
                    shm.set_shared_memory_region(shm_ip_handle, input_list_tmp)
                    self.triton_client_.register_system_shared_memory(
                        ip_name, '/' + ip_name, input_byte_size)
                    self.triton_client_.register_system_shared_memory(
                        op_name, '/' + op_name, output_byte_size)
                elif _test_cuda_shared_memory:
                    shm_ip_handle = cudashm.create_shared_memory_region(
                        ip_name, input_byte_size, 0)
                    shm_op_handle = cudashm.create_shared_memory_region(
                        op_name, output_byte_size, 0)
                    cudashm.set_shared_memory_region(shm_ip_handle,
                                                     input_list_tmp)
                    self.triton_client_.register_cuda_shared_memory(
                        ip_name, cudashm.get_raw_handle(shm_ip_handle), 0,
                        input_byte_size)
                    self.triton_client_.register_cuda_shared_memory(
                        op_name, cudashm.get_raw_handle(shm_op_handle), 0,
                        output_byte_size)
                shm_region_handles.append(
                    (ip_name, input_byte_size, shm_ip_handle))
                shm_region_handles.append(
                    (op_name, output_byte_size, shm_op_handle))
            return shm_region_handles
        else:
            return []

    # Returns (name, byte size, shm_handle)
    def precreate_register_shape_tensor_regions(self,
                                                value_list,
                                                dtype,
                                                i,
                                                batch_size=1,
                                                tensor_shape=(1,)):
        self.assertFalse(_test_cuda_shared_memory,
                         "Shape tensors does not support CUDA shared memory")
        if _test_system_shared_memory:
            shm_region_handles = []
            for j, (shape_value, value) in enumerate(value_list):
                input_list = list()
                shape_input_list = list()

                for b in range(batch_size):
                    if dtype == np.object_:
                        in0 = np.full(tensor_shape, value, dtype=np.int32)
                        in0n = np.array([str(x) for x in in0.reshape(in0.size)],
                                        dtype=object)
                        in0 = in0n.reshape(tensor_shape)
                    else:
                        in0 = np.full(tensor_shape, value, dtype=dtype)
                    input_list.append(in0)

                # Only one shape tensor input per batch
                shape_input_list.append(
                    np.full(tensor_shape, shape_value, dtype=np.int32))

                if dtype == np.object_:
                    input_list_tmp = iu.serialize_byte_tensor_list(input_list)
                    input_byte_size = sum(
                        [serialized_byte_size(i0) for i0 in input_list_tmp])
                else:
                    input_list_tmp = input_list
                    input_byte_size = sum([i0.nbytes for i0 in input_list_tmp])

                shape_input_byte_size = sum(
                    [i0.nbytes for i0 in shape_input_list])
                shape_output_byte_size = shape_input_byte_size
                output_byte_size = np.dtype(dtype).itemsize + 2
                resized_output_byte_size = 32 * shape_value

                # create shared memory regions and copy data for input values
                ip_name = 'ip{}{}'.format(i, j)
                shape_ip_name = 'shape_ip{}{}'.format(i, j)
                shape_op_name = 'shape_op{}{}'.format(i, j)
                op_name = 'op{}{}'.format(i, j)
                resized_op_name = 'resized_op{}{}'.format(i, j)

                shm_ip_handle = shm.create_shared_memory_region(
                    ip_name, '/' + ip_name, input_byte_size)
                shm_shape_ip_handle = shm.create_shared_memory_region(
                    shape_ip_name, '/' + shape_ip_name, shape_input_byte_size)
                shm_shape_op_handle = shm.create_shared_memory_region(
                    shape_op_name, '/' + shape_op_name, shape_output_byte_size)
                shm_op_handle = shm.create_shared_memory_region(
                    op_name, '/' + op_name, output_byte_size)
                shm_resized_op_handle = shm.create_shared_memory_region(
                    resized_op_name, '/' + resized_op_name,
                    resized_output_byte_size)
                shm.set_shared_memory_region(shm_ip_handle, input_list_tmp)
                shm.set_shared_memory_region(shm_shape_ip_handle,
                                             shape_input_list)
                self.triton_client_.register_system_shared_memory(
                    ip_name, '/' + ip_name, input_byte_size)
                self.triton_client_.register_system_shared_memory(
                    shape_ip_name, '/' + shape_ip_name, shape_input_byte_size)
                self.triton_client_.register_system_shared_memory(
                    shape_op_name, '/' + shape_op_name, shape_output_byte_size)
                self.triton_client_.register_system_shared_memory(
                    op_name, '/' + op_name, output_byte_size)
                self.triton_client_.register_system_shared_memory(
                    resized_op_name, '/' + resized_op_name,
                    resized_output_byte_size)

                shm_region_handles.append(
                    (ip_name, input_byte_size, shm_ip_handle))
                shm_region_handles.append(
                    (shape_ip_name, shape_input_byte_size, shm_shape_ip_handle))
                shm_region_handles.append(
                    (shape_op_name, shape_output_byte_size,
                     shm_shape_op_handle))
                shm_region_handles.append(
                    (op_name, output_byte_size, shm_op_handle))
                shm_region_handles.append(
                    (resized_op_name, resized_output_byte_size,
                     shm_resized_op_handle))
            return shm_region_handles
        else:
            return []

    # Returns (name, byte size, shm_handle)
    def precreate_register_dynaseq_shape_tensor_regions(self,
                                                        value_list,
                                                        dtype,
                                                        i,
                                                        batch_size=1,
                                                        tensor_shape=(1,)):
        self.assertFalse(_test_cuda_shared_memory,
                         "Shape tensors does not support CUDA shared memory")
        if _test_system_shared_memory:
            shm_region_handles = []
            for j, (shape_value, value) in enumerate(value_list):
                input_list = list()
                shape_input_list = list()
                dummy_input_list = list()

                for b in range(batch_size):
                    if dtype == np.object_:
                        dummy_in0 = np.full(tensor_shape, value, dtype=np.int32)
                        dummy_in0n = np.array(
                            [str(x) for x in dummy_in0.reshape(in0.size)],
                            dtype=object)
                        dummy_in0 = dummy_in0n.reshape(tensor_shape)
                    else:
                        dummy_in0 = np.full(tensor_shape, value, dtype=dtype)
                    dummy_input_list.append(dummy_in0)
                    in0 = np.full(tensor_shape, value, dtype=np.int32)
                    input_list.append(in0)

                # Only one shape tensor input per batch
                shape_input_list.append(
                    np.full(tensor_shape, shape_value, dtype=np.int32))

                if dtype == np.object_:
                    input_list_tmp = iu.serialize_byte_tensor_list(input_list)
                    input_byte_size = sum(
                        [serialized_byte_size(i0) for i0 in input_list_tmp])
                else:
                    input_list_tmp = input_list
                    input_byte_size = sum([i0.nbytes for i0 in input_list_tmp])

                dummy_input_byte_size = sum(
                    [i0.nbytes for i0 in dummy_input_list])

                shape_input_byte_size = sum(
                    [i0.nbytes for i0 in shape_input_list])
                shape_output_byte_size = shape_input_byte_size
                output_byte_size = np.dtype(np.int32).itemsize + 2
                resized_output_byte_size = 32 * shape_value

                # create shared memory regions and copy data for input values
                ip_name = 'ip{}{}'.format(i, j)
                shape_ip_name = 'shape_ip{}{}'.format(i, j)
                dummy_ip_name = 'dummy_ip{}{}'.format(i, j)
                shape_op_name = 'shape_op{}{}'.format(i, j)
                op_name = 'op{}{}'.format(i, j)
                resized_op_name = 'resized_op{}{}'.format(i, j)

                shm_ip_handle = shm.create_shared_memory_region(
                    ip_name, '/' + ip_name, input_byte_size)
                shm_shape_ip_handle = shm.create_shared_memory_region(
                    shape_ip_name, '/' + shape_ip_name, shape_input_byte_size)
                shm_dummy_ip_handle = shm.create_shared_memory_region(
                    dummy_ip_name, '/' + dummy_ip_name, dummy_input_byte_size)
                shm_shape_op_handle = shm.create_shared_memory_region(
                    shape_op_name, '/' + shape_op_name, shape_output_byte_size)
                shm_op_handle = shm.create_shared_memory_region(
                    op_name, '/' + op_name, output_byte_size)
                shm_resized_op_handle = shm.create_shared_memory_region(
                    resized_op_name, '/' + resized_op_name,
                    resized_output_byte_size)
                shm.set_shared_memory_region(shm_ip_handle, input_list_tmp)
                shm.set_shared_memory_region(shm_shape_ip_handle,
                                             shape_input_list)
                shm.set_shared_memory_region(shm_dummy_ip_handle,
                                             dummy_input_list)
                self.triton_client_.register_system_shared_memory(
                    ip_name, '/' + ip_name, input_byte_size)
                self.triton_client_.register_system_shared_memory(
                    shape_ip_name, '/' + shape_ip_name, shape_input_byte_size)
                self.triton_client_.register_system_shared_memory(
                    dummy_ip_name, '/' + dummy_ip_name, dummy_input_byte_size)
                self.triton_client_.register_system_shared_memory(
                    shape_op_name, '/' + shape_op_name, shape_output_byte_size)
                self.triton_client_.register_system_shared_memory(
                    op_name, '/' + op_name, output_byte_size)
                self.triton_client_.register_system_shared_memory(
                    resized_op_name, '/' + resized_op_name,
                    resized_output_byte_size)

                shm_region_handles.append(
                    (ip_name, input_byte_size, shm_ip_handle))
                shm_region_handles.append(
                    (shape_ip_name, shape_input_byte_size, shm_shape_ip_handle))
                shm_region_handles.append(
                    (dummy_ip_name, dummy_input_byte_size, shm_dummy_ip_handle))
                shm_region_handles.append(
                    (shape_op_name, shape_output_byte_size,
                     shm_shape_op_handle))
                shm_region_handles.append(
                    (op_name, output_byte_size, shm_op_handle))
                shm_region_handles.append(
                    (resized_op_name, resized_output_byte_size,
                     shm_resized_op_handle))
            return shm_region_handles
        else:
            return []

    def cleanup_shm_regions(self, shm_handles):
        # Make sure unregister is before shared memory destruction
        if _test_system_shared_memory:
            self.triton_client_.unregister_system_shared_memory()
        if _test_cuda_shared_memory:
            self.triton_client_.unregister_cuda_shared_memory()
        for shm_tmp_handle in shm_handles:
            if _test_system_shared_memory:
                shm.destroy_shared_memory_region(shm_tmp_handle[2])
            elif _test_cuda_shared_memory:
                cudashm.destroy_shared_memory_region(shm_tmp_handle[2])

    def check_sequence(self,
                       trial,
                       model_name,
                       input_dtype,
                       correlation_id,
                       sequence_thresholds,
                       values,
                       expected_result,
                       protocol,
                       batch_size=1,
                       sequence_name="<unknown>",
                       tensor_shape=(1,)):
        """Perform sequence of inferences. The 'values' holds a list of
        tuples, one for each inference with format:

        (flag_str, value, (ls_ms, gt_ms), (pre_delay_ms, post_delay_ms)

        """
        if (("savedmodel" not in trial) and ("graphdef" not in trial) and
            ("custom" not in trial) and ("onnx" not in trial) and
            ("libtorch" not in trial) and ("plan" not in trial) and
            ("python" not in trial)):
            self.assertFalse(True, "unknown trial type: " + trial)

        # Can only send the request exactly once since it is a
        # sequence model with state, so can have only a single config.
        configs = []
        if protocol == "http":
            configs.append((f"{_tritonserver_ipaddr}:8000", "http", False))
        if protocol == "grpc":
            configs.append((f"{_tritonserver_ipaddr}:8001", "grpc", False))
        if protocol == "streaming":
            configs.append((f"{_tritonserver_ipaddr}:8001", "grpc", True))

        self.assertFalse(
            _test_system_shared_memory and _test_cuda_shared_memory,
            "Cannot set both System and CUDA shared memory flags to 1")

        self.assertEqual(len(configs), 1)

        full_shape = tensor_shape if "nobatch" in trial else (
            batch_size,) + tensor_shape

        # create and register shared memory output region in advance,
        # knowing that this function will not be called concurrently.
        if _test_system_shared_memory or _test_cuda_shared_memory:
            self.triton_client_.unregister_system_shared_memory()
            self.triton_client_.unregister_cuda_shared_memory()
            output_byte_size = 512
            if _test_system_shared_memory:
                shm_op_handle = shm.create_shared_memory_region(
                    "output_data", "/output", output_byte_size)
                self.triton_client_.register_system_shared_memory(
                    "output_data", "/output", output_byte_size)
            elif _test_cuda_shared_memory:
                shm_op_handle = cudashm.create_shared_memory_region(
                    "output_data", output_byte_size, 0)
                self.triton_client_.register_cuda_shared_memory(
                    "output_data", cudashm.get_raw_handle(shm_op_handle), 0,
                    output_byte_size)
            shm_ip_handles = []

        for config in configs:
            client_utils = grpcclient if config[1] == "grpc" else httpclient

            triton_client = client_utils.InferenceServerClient(config[0],
                                                               verbose=True)
            if config[2]:
                user_data = UserData()
                triton_client.start_stream(
                    partial(completion_callback, user_data))
            # Execute the sequence of inference...
            try:
                seq_start_ms = int(round(time.time() * 1000))

                INPUT = "INPUT__0" if trial.startswith("libtorch") else "INPUT"
                OUTPUT = "OUTPUT__0" if trial.startswith(
                    "libtorch") else "OUTPUT"
                for flag_str, value, thresholds, delay_ms in values:
                    if _test_valgrind or _test_jetson:
                        if delay_ms is not None:
                            delay_ms[0] = max(_valgrind_delay_ms, delay_ms[0])
                            delay_ms[1] = max(_valgrind_delay_ms, delay_ms[1])
                        else:
                            delay_ms = (_valgrind_delay_ms, _valgrind_delay_ms)

                    if delay_ms is not None:
                        time.sleep(delay_ms[0] / 1000.0)

                    seq_start = False
                    seq_end = False
                    if flag_str is not None:
                        seq_start = ("start" in flag_str)
                        seq_end = ("end" in flag_str)

                    # Construct request IOs
                    inputs = []
                    outputs = []
                    inputs.append(
                        client_utils.InferInput(
                            INPUT, full_shape, np_to_triton_dtype(input_dtype)))
                    outputs.append(client_utils.InferRequestedOutput(OUTPUT))
                    if input_dtype == np.object_:
                        in0 = np.full(full_shape, value, dtype=np.int32)
                        in0n = np.array([str(x) for x in in0.reshape(in0.size)],
                                        dtype=object)
                        in0 = in0n.reshape(full_shape)
                    else:
                        in0 = np.full(full_shape, value, dtype=input_dtype)

                    # create input shared memory and copy input data values into it
                    if _test_system_shared_memory or _test_cuda_shared_memory:
                        if input_dtype == np.object_:
                            input_list_tmp = iu.serialize_byte_tensor_list(
                                [in0])
                            input_byte_size = sum([
                                serialized_byte_size(i0)
                                for i0 in input_list_tmp
                            ])
                        else:
                            input_list_tmp = [in0]
                            input_byte_size = sum(
                                [i0.nbytes for i0 in input_list_tmp])
                        ip_name = "ip{}".format(len(shm_ip_handles))
                        if _test_system_shared_memory:
                            shm_ip_handles.append(
                                shm.create_shared_memory_region(
                                    ip_name, "/" + ip_name, input_byte_size))
                            shm.set_shared_memory_region(
                                shm_ip_handles[-1], input_list_tmp)
                            triton_client.register_system_shared_memory(
                                ip_name, "/" + ip_name, input_byte_size)
                        elif _test_cuda_shared_memory:
                            shm_ip_handles.append(
                                cudashm.create_shared_memory_region(
                                    ip_name, input_byte_size, 0))
                            cudashm.set_shared_memory_region(
                                shm_ip_handles[-1], input_list_tmp)
                            triton_client.register_cuda_shared_memory(
                                ip_name,
                                cudashm.get_raw_handle(shm_ip_handles[-1]), 0,
                                input_byte_size)

                        inputs[0].set_shared_memory(ip_name, input_byte_size)
                        outputs[0].set_shared_memory("output_data",
                                                     output_byte_size)
                    else:
                        inputs[0].set_data_from_numpy(in0)

                    start_ms = int(round(time.time() * 1000))

                    if config[2]:
                        triton_client.async_stream_infer(
                            model_name,
                            inputs,
                            outputs=outputs,
                            sequence_id=correlation_id,
                            sequence_start=seq_start,
                            sequence_end=seq_end)
                        (results, error) = user_data._completed_requests.get()
                        if error is not None:
                            raise error
                    else:
                        results = triton_client.infer(
                            model_name,
                            inputs,
                            outputs=outputs,
                            sequence_id=correlation_id,
                            sequence_start=seq_start,
                            sequence_end=seq_end)

                    end_ms = int(round(time.time() * 1000))

                    # Get value of "OUTPUT", for shared memory, need to get it via
                    # shared memory utils
                    if (not _test_system_shared_memory) and (
                            not _test_cuda_shared_memory):
                        out = results.as_numpy(OUTPUT)
                    else:
                        output = results.get_output(OUTPUT)
                        if config[1] == "http":
                            output_shape = output["shape"]
                        else:
                            output_shape = output.shape
                        output_type = input_dtype
                        if _test_system_shared_memory:
                            out = shm.get_contents_as_numpy(
                                shm_op_handle, output_type, output_shape)
                        else:
                            out = cudashm.get_contents_as_numpy(
                                shm_op_handle, output_type, output_shape)
                    result = out[0] if "nobatch" in trial else out[0][0]
                    print("{}: {}".format(sequence_name, result))

                    if thresholds is not None:
                        lt_ms = thresholds[0]
                        gt_ms = thresholds[1]
                        if lt_ms is not None:
                            self.assertTrue((end_ms - start_ms) < lt_ms,
                                            "expected less than " + str(lt_ms) +
                                            "ms response time, got " +
                                            str(end_ms - start_ms) + " ms")
                        if gt_ms is not None:
                            self.assertTrue(
                                (end_ms - start_ms) > gt_ms,
                                "expected greater than " + str(gt_ms) +
                                "ms response time, got " +
                                str(end_ms - start_ms) + " ms")
                    if delay_ms is not None:
                        time.sleep(delay_ms[1] / 1000.0)

                seq_end_ms = int(round(time.time() * 1000))

                if input_dtype == np.object_:
                    self.assertEqual(int(result), expected_result)
                else:
                    self.assertEqual(result, expected_result)

                if sequence_thresholds is not None:
                    lt_ms = sequence_thresholds[0]
                    gt_ms = sequence_thresholds[1]
                    if lt_ms is not None:
                        if _test_jetson:
                            lt_ms *= _jetson_slowdown_factor
                        self.assertTrue((seq_end_ms - seq_start_ms) < lt_ms,
                                        "sequence expected less than " +
                                        str(lt_ms) + "ms response time, got " +
                                        str(seq_end_ms - seq_start_ms) + " ms")
                    if gt_ms is not None:
                        self.assertTrue((seq_end_ms - seq_start_ms) > gt_ms,
                                        "sequence expected greater than " +
                                        str(gt_ms) + "ms response time, got " +
                                        str(seq_end_ms - seq_start_ms) + " ms")
            except Exception as ex:
                self.add_deferred_exception(ex)
            if config[2]:
                triton_client.stop_stream()

        if _test_system_shared_memory or _test_cuda_shared_memory:
            self.triton_client_.unregister_system_shared_memory()
            self.triton_client_.unregister_cuda_shared_memory()
            destroy_func = shm.destroy_shared_memory_region if _test_system_shared_memory else cudashm.destroy_shared_memory_region
            destroy_func(shm_op_handle)
            for shm_ip_handle in shm_ip_handles:
                destroy_func(shm_ip_handle)

    def check_sequence_async(self,
                             trial,
                             model_name,
                             input_dtype,
                             correlation_id,
                             sequence_thresholds,
                             values,
                             expected_result,
                             shm_region_handles,
                             batch_size=1,
                             sequence_name="<unknown>",
                             tensor_shape=(1,)):
        """Perform sequence of inferences using stream async run.
        The 'values' holds a list of tuples, one for each inference with format:

        (flag_str, value, pre_delay_ms)

        """
        if (("savedmodel" not in trial) and ("graphdef" not in trial) and
            ("custom" not in trial) and ("onnx" not in trial) and
            ("libtorch" not in trial) and ("plan" not in trial) and
            ("python" not in trial)):
            self.assertFalse(True, "unknown trial type: " + trial)

        self.assertFalse(
            _test_system_shared_memory and _test_cuda_shared_memory,
            "Cannot set both System and CUDA shared memory flags to 1")

        full_shape = tensor_shape if "nobatch" in trial else (
            batch_size,) + tensor_shape

        client_utils = grpcclient
        triton_client = client_utils.InferenceServerClient(
            f"{_tritonserver_ipaddr}:8001", verbose=True)
        user_data = UserData()
        triton_client.start_stream(partial(completion_callback, user_data))
        # Execute the sequence of inference...
        try:
            seq_start_ms = int(round(time.time() * 1000))

            INPUT = "INPUT__0" if trial.startswith("libtorch") else "INPUT"
            OUTPUT = "OUTPUT__0" if trial.startswith("libtorch") else "OUTPUT"
            sent_count = 0
            for flag_str, value, pre_delay_ms in values:
                seq_start = False
                seq_end = False
                if flag_str is not None:
                    seq_start = ("start" in flag_str)
                    seq_end = ("end" in flag_str)

                # Construct request IOs
                inputs = []
                outputs = []
                inputs.append(
                    client_utils.InferInput(INPUT, full_shape,
                                            np_to_triton_dtype(input_dtype)))
                outputs.append(client_utils.InferRequestedOutput(OUTPUT))

                if not (_test_system_shared_memory or _test_cuda_shared_memory):
                    if input_dtype == np.object_:
                        in0 = np.full(full_shape, value, dtype=np.int32)
                        in0n = np.array([str(x) for x in in0.reshape(in0.size)],
                                        dtype=object)
                        in0 = in0n.reshape(full_shape)
                    else:
                        in0 = np.full(full_shape, value, dtype=input_dtype)
                    inputs[0].set_data_from_numpy(in0)
                else:
                    offset = 2 * sent_count
                    inputs[0].set_shared_memory(shm_region_handles[offset][0],
                                                shm_region_handles[offset][1])
                    outputs[0].set_shared_memory(
                        shm_region_handles[offset + 1][0],
                        shm_region_handles[offset + 1][1])

                if pre_delay_ms is not None:
                    time.sleep(pre_delay_ms / 1000.0)

                triton_client.async_stream_infer(model_name,
                                                 inputs,
                                                 outputs=outputs,
                                                 sequence_id=correlation_id,
                                                 sequence_start=seq_start,
                                                 sequence_end=seq_end)
                sent_count += 1

            # Wait for the results in the order sent
            result = None
            processed_count = 0
            while processed_count < sent_count:
                (results, error) = user_data._completed_requests.get()
                if error is not None:
                    raise error
                # Get value of "OUTPUT", for shared memory, need to get it via
                # shared memory utils
                if (not _test_system_shared_memory) and (
                        not _test_cuda_shared_memory):
                    out = results.as_numpy(OUTPUT)
                else:
                    output = results.get_output(OUTPUT)
                    offset = 2 * processed_count + 1
                    output_shape = output.shape
                    output_type = input_dtype
                    if _test_system_shared_memory:
                        out = shm.get_contents_as_numpy(
                            shm_region_handles[offset][2], output_type,
                            output_shape)
                    else:
                        out = cudashm.get_contents_as_numpy(
                            shm_region_handles[offset][2], output_type,
                            output_shape)
                result = out[0] if "nobatch" in trial else out[0][0]
                print("{}: {}".format(sequence_name, result))
                processed_count += 1

            seq_end_ms = int(round(time.time() * 1000))

            if input_dtype == np.object_:
                self.assertEqual(int(result), expected_result)
            else:
                self.assertEqual(result, expected_result)

            if sequence_thresholds is not None:
                lt_ms = sequence_thresholds[0]
                gt_ms = sequence_thresholds[1]
                if lt_ms is not None:
                    if _test_jetson:
                        lt_ms *= _jetson_slowdown_factor
                    self.assertTrue((seq_end_ms - seq_start_ms) < lt_ms,
                                    "sequence expected less than " +
                                    str(lt_ms) + "ms response time, got " +
                                    str(seq_end_ms - seq_start_ms) + " ms")
                if gt_ms is not None:
                    self.assertTrue((seq_end_ms - seq_start_ms) > gt_ms,
                                    "sequence expected greater than " +
                                    str(gt_ms) + "ms response time, got " +
                                    str(seq_end_ms - seq_start_ms) + " ms")
        except Exception as ex:
            self.add_deferred_exception(ex)
        triton_client.stop_stream()

    # This sequence util only sends inference via streaming scenario
    def check_sequence_shape_tensor_io(self,
                                       model_name,
                                       input_dtype,
                                       correlation_id,
                                       sequence_thresholds,
                                       values,
                                       expected_result,
                                       shm_region_handles,
                                       using_dynamic_batcher=False,
                                       sequence_name="<unknown>"):
        """Perform sequence of inferences using async run. The 'values' holds
        a list of tuples, one for each inference with format:

        (flag_str, shape_value, value, pre_delay_ms)

        """
        tensor_shape = (1, 1)
        # shape tensor is 1-D tensor that doesn't contain batch size as first value
        shape_tensor_shape = (1,)
        self.assertFalse(_test_cuda_shared_memory,
                         "Shape tensors does not support CUDA shared memory")

        client_utils = grpcclient
        triton_client = client_utils.InferenceServerClient(
            f"{_tritonserver_ipaddr}:8001", verbose=True)
        user_data = UserData()
        triton_client.start_stream(partial(completion_callback, user_data))
        # Execute the sequence of inference...
        try:
            seq_start_ms = int(round(time.time() * 1000))

            sent_count = 0
            shape_values = list()
            for flag_str, shape_value, value, pre_delay_ms in values:
                seq_start = False
                seq_end = False
                if flag_str is not None:
                    seq_start = ("start" in flag_str)
                    seq_end = ("end" in flag_str)

                # Construct request IOs
                inputs = []
                outputs = []
                # input order: input, shape(, dummy)
                inputs.append(
                    client_utils.InferInput(
                        "INPUT", tensor_shape,
                        np_to_triton_dtype(np.int32 if using_dynamic_batcher
                                           else input_dtype)))
                inputs.append(
                    client_utils.InferInput("SHAPE_INPUT", shape_tensor_shape,
                                            np_to_triton_dtype(np.int32)))
                if using_dynamic_batcher:
                    inputs.append(
                        client_utils.InferInput(
                            "DUMMY_INPUT", tensor_shape,
                            np_to_triton_dtype(input_dtype)))
                # output order: shape, output, resized
                outputs.append(
                    client_utils.InferRequestedOutput("SHAPE_OUTPUT"))
                outputs.append(client_utils.InferRequestedOutput("OUTPUT"))
                outputs.append(
                    client_utils.InferRequestedOutput("RESIZED_OUTPUT"))

                # Set IO values
                shape_values.append(
                    np.full(shape_tensor_shape, shape_value, dtype=np.int32))
                if not _test_system_shared_memory:
                    if using_dynamic_batcher:
                        if input_dtype == np.object_:
                            dummy_in0 = np.full(tensor_shape,
                                                value,
                                                dtype=np.int32)
                            dummy_in0n = np.array(
                                [str(x) for x in in0.reshape(dummy_in0.size)],
                                dtype=object)
                            dummy_in0 = dummy_in0n.reshape(tensor_shape)
                        else:
                            dummy_in0 = np.full(tensor_shape,
                                                value,
                                                dtype=input_dtype)
                        in0 = np.full(tensor_shape, value, dtype=np.int32)
                    else:
                        if input_dtype == np.object_:
                            in0 = np.full(tensor_shape, value, dtype=np.int32)
                            in0n = np.array(
                                [str(x) for x in in0.reshape(in0.size)],
                                dtype=object)
                            in0 = in0n.reshape(tensor_shape)
                        else:
                            in0 = np.full(tensor_shape,
                                          value,
                                          dtype=input_dtype)

                    inputs[0].set_data_from_numpy(in0)
                    inputs[1].set_data_from_numpy(shape_values[-1])
                    if using_dynamic_batcher:
                        inputs[2].set_data_from_numpy(dummy_in0)
                else:
                    if using_dynamic_batcher:
                        input_offset = 6 * sent_count
                        output_offset = 6 * sent_count + 3
                    else:
                        input_offset = 5 * sent_count
                        output_offset = 5 * sent_count + 2
                    for i in range(len(inputs)):
                        inputs[i].set_shared_memory(
                            shm_region_handles[input_offset + i][0],
                            shm_region_handles[input_offset + i][1])
                    for i in range(len(outputs)):
                        outputs[i].set_shared_memory(
                            shm_region_handles[output_offset + i][0],
                            shm_region_handles[output_offset + i][1])

                if pre_delay_ms is not None:
                    time.sleep(pre_delay_ms / 1000.0)

                triton_client.async_stream_infer(model_name,
                                                 inputs,
                                                 outputs=outputs,
                                                 sequence_id=correlation_id,
                                                 sequence_start=seq_start,
                                                 sequence_end=seq_end)

                sent_count += 1

            # Wait for the results in the order sent
            result = None
            processed_count = 0
            while processed_count < sent_count:
                (results, error) = user_data._completed_requests.get()
                if error is not None:
                    raise error
                # Get value of "OUTPUT", for shared memory, need to get it via
                # shared memory utils
                if (not _test_system_shared_memory):
                    out = results.as_numpy("OUTPUT")
                else:
                    output = results.get_output("OUTPUT")
                    output_offset = 6 * processed_count + 4 if using_dynamic_batcher else 5 * processed_count + 3
                    output_shape = output.shape
                    output_type = np.int32 if using_dynamic_batcher else np.float32
                    out = shm.get_contents_as_numpy(
                        shm_region_handles[output_offset][2], output_type,
                        output_shape)
                result = out[0][0]

                # Validate the (debatched) shape of the resized output matches
                # with the shape input values
                resized_shape = results.get_output("RESIZED_OUTPUT").shape[1:]
                self.assertTrue(
                    np.array_equal(resized_shape,
                                   shape_values[processed_count]),
                    "{}, {}, slot {}, expected: {}, got {}".format(
                        model_name, "RESIZED_OUTPUT", processed_count,
                        shape_values[processed_count], resized_shape))
                print("{}: {}".format(sequence_name, result))
                processed_count += 1

            seq_end_ms = int(round(time.time() * 1000))

            if input_dtype == np.object_:
                self.assertEqual(int(result), expected_result)
            else:
                self.assertEqual(result, expected_result)

            if sequence_thresholds is not None:
                lt_ms = sequence_thresholds[0]
                gt_ms = sequence_thresholds[1]
                if lt_ms is not None:
                    if _test_jetson:
                        lt_ms *= _jetson_slowdown_factor
                    self.assertTrue((seq_end_ms - seq_start_ms) < lt_ms,
                                    "sequence expected less than " +
                                    str(lt_ms) + "ms response time, got " +
                                    str(seq_end_ms - seq_start_ms) + " ms")
                if gt_ms is not None:
                    self.assertTrue((seq_end_ms - seq_start_ms) > gt_ms,
                                    "sequence expected greater than " +
                                    str(gt_ms) + "ms response time, got " +
                                    str(seq_end_ms - seq_start_ms) + " ms")
        except Exception as ex:
            self.add_deferred_exception(ex)
        triton_client.stop_stream()

    def check_setup(self, model_name):
        # Make sure test.sh set up the correct batcher settings
        config = self.triton_client_.get_model_config(model_name).config
        # Skip the sequence batching check on ensemble model
        if config.platform != "ensemble":
            bconfig = config.sequence_batching
            self.assertEqual(bconfig.max_sequence_idle_microseconds,
                             _max_sequence_idle_ms * 1000)  # 5 secs

    def check_status(self, model_name, batch_exec, exec_cnt, infer_cnt):
        stats = self.triton_client_.get_inference_statistics(model_name, "1")
        self.assertEqual(len(stats.model_stats), 1, "expect 1 model stats")
        self.assertEqual(stats.model_stats[0].name, model_name,
                         "expect model stats for model {}".format(model_name))
        self.assertEqual(
            stats.model_stats[0].version, "1",
            "expect model stats for model {} version 1".format(model_name))

        if batch_exec is not None:
            batch_stats = stats.model_stats[0].batch_stats
            print(batch_stats)
            self.assertEqual(
                len(batch_stats), len(batch_exec),
                "expected {} different batch-sizes, got {}".format(
                    len(batch_exec), len(batch_stats)))

            for batch_stat in batch_stats:
                bs = batch_stat.batch_size
                bc = batch_stat.compute_infer.count
                self.assertTrue(
                    bs in batch_exec,
                    "did not find expected batch-size {}".format(bs))
                # Get count from one of the stats
                self.assertEqual(
                    bc, batch_exec[bs],
                    "expected model-execution-count {} for batch size {}, got {}"
                    .format(batch_exec[bs], bs, bc))

        actual_exec_cnt = stats.model_stats[0].execution_count
        self.assertEqual(
            actual_exec_cnt, exec_cnt,
            "expected model-exec-count {}, got {}".format(
                exec_cnt, actual_exec_cnt))

        actual_infer_cnt = stats.model_stats[0].inference_count
        self.assertEqual(
            actual_infer_cnt, infer_cnt,
            "expected model-inference-count {}, got {}".format(
                infer_cnt, actual_infer_cnt))


# qa/common/shm_util.py

# Copyright 2018-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import os
from os import listdir
import numpy as np
from ctypes import *

import tritonclient.http as httpclient
from tritonclient.utils import *

# By default, find tritonserver on "localhost", but can be overridden
# with TRITONSERVER_IPADDR envvar
_tritonserver_ipaddr = os.environ.get('TRITONSERVER_IPADDR', 'localhost')


def _range_repr_dtype(dtype):
    if dtype == np.float64:
        return np.int32
    elif dtype == np.float32:
        return np.int16
    elif dtype == np.float16:
        return np.int8
    elif dtype == np.object_:  # TYPE_STRING
        return np.int32
    return dtype


def create_set_shm_regions(input0_list, input1_list, output0_byte_size,
                           output1_byte_size, outputs, shm_region_names,
                           precreated_shm_regions, use_system_shared_memory,
                           use_cuda_shared_memory):
    # Lazy shm imports...
    if use_system_shared_memory:
        import tritonclient.utils.shared_memory as shm
    if use_cuda_shared_memory:
        import tritonclient.utils.cuda_shared_memory as cudashm

    if use_system_shared_memory and use_cuda_shared_memory:
        raise ValueError(
            "Cannot set both System and CUDA shared memory flags to 1")

    if not (use_system_shared_memory or use_cuda_shared_memory):
        return [], []

    if input0_list[0].dtype == np.object_:
        input0_byte_size = sum([serialized_byte_size(i0) for i0 in input0_list])
    else:
        input0_byte_size = sum([i0.nbytes for i0 in input0_list])

    if input1_list[0].dtype == np.object_:
        input1_byte_size = sum([serialized_byte_size(i1) for i1 in input1_list])
    else:
        input1_byte_size = sum([i1.nbytes for i1 in input1_list])

    if shm_region_names is None:
        shm_region_names = ['input0', 'input1', 'output0', 'output1']

    shm_op0_handle = None
    shm_op1_handle = None

    if use_system_shared_memory:
        shm_ip0_handle = shm.create_shared_memory_region(
            shm_region_names[0] + '_data', '/' + shm_region_names[0],
            input0_byte_size)
        shm_ip1_handle = shm.create_shared_memory_region(
            shm_region_names[1] + '_data', '/' + shm_region_names[1],
            input1_byte_size)

        i = 0
        if "OUTPUT0" in outputs:
            if precreated_shm_regions is None:
                shm_op0_handle = shm.create_shared_memory_region(
                    shm_region_names[2] + '_data', '/' + shm_region_names[2],
                    output0_byte_size)
            else:
                shm_op0_handle = precreated_shm_regions[0]
            i += 1
        if "OUTPUT1" in outputs:
            if precreated_shm_regions is None:
                shm_op1_handle = shm.create_shared_memory_region(
                    shm_region_names[2 + i] + '_data',
                    '/' + shm_region_names[2 + i], output1_byte_size)
            else:
                shm_op1_handle = precreated_shm_regions[i]

        shm.set_shared_memory_region(shm_ip0_handle, input0_list)
        shm.set_shared_memory_region(shm_ip1_handle, input1_list)

    if use_cuda_shared_memory:
        shm_ip0_handle = cudashm.create_shared_memory_region(
            shm_region_names[0] + '_data', input0_byte_size, 0)
        shm_ip1_handle = cudashm.create_shared_memory_region(
            shm_region_names[1] + '_data', input1_byte_size, 0)
        i = 0
        if "OUTPUT0" in outputs:
            if precreated_shm_regions is None:
                shm_op0_handle = cudashm.create_shared_memory_region(
                    shm_region_names[2] + '_data', output0_byte_size, 0)
            else:
                shm_op0_handle = precreated_shm_regions[0]
            i += 1
        if "OUTPUT1" in outputs:
            if precreated_shm_regions is None:
                shm_op1_handle = cudashm.create_shared_memory_region(
                    shm_region_names[2 + i] + '_data', output1_byte_size, 0)
            else:
                shm_op1_handle = precreated_shm_regions[i]

        cudashm.set_shared_memory_region(shm_ip0_handle, input0_list)
        cudashm.set_shared_memory_region(shm_ip1_handle, input1_list)

    return shm_region_names, [
        shm_ip0_handle, shm_ip1_handle, shm_op0_handle, shm_op1_handle
    ]


def register_add_shm_regions(inputs, outputs, shm_region_names,
                             precreated_shm_regions, shm_handles,
                             input0_byte_size, input1_byte_size,
                             output0_byte_size, output1_byte_size,
                             use_system_shared_memory, use_cuda_shared_memory,
                             triton_client):
    # Lazy shm imports...
    if use_system_shared_memory:
        import tritonclient.utils.shared_memory as shm
    if use_cuda_shared_memory:
        import tritonclient.utils.cuda_shared_memory as cudashm

    if use_system_shared_memory or use_cuda_shared_memory:
        # Unregister then register required shared memory regions
        if use_system_shared_memory:
            triton_client.unregister_system_shared_memory(shm_region_names[0] +
                                                          '_data')
            triton_client.unregister_system_shared_memory(shm_region_names[1] +
                                                          '_data')
            triton_client.register_system_shared_memory(
                shm_region_names[0] + '_data', '/' + shm_region_names[0],
                input0_byte_size)
            triton_client.register_system_shared_memory(
                shm_region_names[1] + '_data', '/' + shm_region_names[1],
                input1_byte_size)
            i = 0
            if "OUTPUT0" in outputs:
                if precreated_shm_regions is None:
                    triton_client.unregister_system_shared_memory(
                        shm_region_names[2] + '_data')
                    triton_client.register_system_shared_memory(
                        shm_region_names[2] + '_data',
                        '/' + shm_region_names[2], output0_byte_size)
                i += 1
            if "OUTPUT1" in outputs:
                if precreated_shm_regions is None:
                    triton_client.unregister_system_shared_memory(
                        shm_region_names[2 + i] + '_data')
                    triton_client.register_system_shared_memory(
                        shm_region_names[2 + i] + '_data',
                        '/' + shm_region_names[2 + i], output1_byte_size)

        if use_cuda_shared_memory:
            triton_client.unregister_cuda_shared_memory(shm_region_names[0] +
                                                        '_data')
            triton_client.unregister_cuda_shared_memory(shm_region_names[1] +
                                                        '_data')
            triton_client.register_cuda_shared_memory(
                shm_region_names[0] + '_data',
                cudashm.get_raw_handle(shm_handles[0]), 0, input0_byte_size)
            triton_client.register_cuda_shared_memory(
                shm_region_names[1] + '_data',
                cudashm.get_raw_handle(shm_handles[1]), 0, input1_byte_size)
            i = 0
            if "OUTPUT0" in outputs:
                if precreated_shm_regions is None:
                    triton_client.unregister_cuda_shared_memory(
                        shm_region_names[2] + '_data')
                    triton_client.register_cuda_shared_memory(
                        shm_region_names[2] + '_data',
                        cudashm.get_raw_handle(shm_handles[2]), 0,
                        output0_byte_size)
                i += 1
            if "OUTPUT1" in outputs:
                if precreated_shm_regions is None:
                    triton_client.unregister_cuda_shared_memory(
                        shm_region_names[2 + i] + '_data')
                    triton_client.register_cuda_shared_memory(
                        shm_region_names[2 + i] + '_data',
                        cudashm.get_raw_handle(shm_handles[3]), 0,
                        output1_byte_size)

        # Add shared memory regions to inputs
        inputs[0].set_shared_memory(shm_region_names[0] + '_data',
                                    input0_byte_size)
        inputs[1].set_shared_memory(shm_region_names[1] + '_data',
                                    input1_byte_size)


def unregister_cleanup_shm_regions(shm_regions, shm_handles,
                                   precreated_shm_regions, outputs,
                                   use_system_shared_memory,
                                   use_cuda_shared_memory):
    # Lazy shm imports...
    if use_system_shared_memory:
        import tritonclient.utils.shared_memory as shm
    if use_cuda_shared_memory:
        import tritonclient.utils.cuda_shared_memory as cudashm

    if not (use_system_shared_memory or use_cuda_shared_memory):
        return None

    triton_client = httpclient.InferenceServerClient(
        f"{_tritonserver_ipaddr}:8000")

    if use_cuda_shared_memory:
        triton_client.unregister_cuda_shared_memory(shm_regions[0] + '_data')
        triton_client.unregister_cuda_shared_memory(shm_regions[1] + '_data')
        cudashm.destroy_shared_memory_region(shm_handles[0])
        cudashm.destroy_shared_memory_region(shm_handles[1])
    else:
        triton_client.unregister_system_shared_memory(shm_regions[0] + '_data')
        triton_client.unregister_system_shared_memory(shm_regions[1] + '_data')
        shm.destroy_shared_memory_region(shm_handles[0])
        shm.destroy_shared_memory_region(shm_handles[1])

    if precreated_shm_regions is None:
        i = 0
        if "OUTPUT0" in outputs:
            if use_cuda_shared_memory:
                triton_client.unregister_cuda_shared_memory(shm_regions[2] +
                                                            '_data')
                cudashm.destroy_shared_memory_region(shm_handles[2])
            else:
                triton_client.unregister_system_shared_memory(shm_regions[2] +
                                                              '_data')
                shm.destroy_shared_memory_region(shm_handles[2])
            i += 1
        if "OUTPUT1" in outputs:
            if use_cuda_shared_memory:
                triton_client.unregister_cuda_shared_memory(shm_regions[2 + i] +
                                                            '_data')
                cudashm.destroy_shared_memory_region(shm_handles[3])
            else:
                triton_client.unregister_system_shared_memory(shm_regions[2 +
                                                                          i] +
                                                              '_data')
                shm.destroy_shared_memory_region(shm_handles[3])


def create_set_either_shm_region(shm_region_names, input_list, input_byte_size,
                                 output_byte_size, use_system_shared_memory,
                                 use_cuda_shared_memory):
    # Lazy shm imports...
    if use_system_shared_memory:
        import tritonclient.utils.shared_memory as shm
    if use_cuda_shared_memory:
        import tritonclient.utils.cuda_shared_memory as cudashm

    if use_cuda_shared_memory and use_system_shared_memory:
        raise ValueError(
            "Cannot set both System and CUDA shared memory flags to 1")

    if not (use_system_shared_memory or use_cuda_shared_memory):
        return []

    if use_cuda_shared_memory:
        shm_ip_handle = cudashm.create_shared_memory_region(
            shm_region_names[0] + "_data", input_byte_size, 0)
        shm_op_handle = cudashm.create_shared_memory_region(
            shm_region_names[1] + "_data", output_byte_size, 0)
        cudashm.set_shared_memory_region(shm_ip_handle, input_list)
    elif use_system_shared_memory:
        shm_ip_handle = shm.create_shared_memory_region(
            shm_region_names[0] + "_data", "/" + shm_region_names[0],
            input_byte_size)
        shm_op_handle = shm.create_shared_memory_region(
            shm_region_names[1] + "_data", "/" + shm_region_names[1],
            output_byte_size)
        shm.set_shared_memory_region(shm_ip_handle, input_list)

    return [shm_ip_handle, shm_op_handle]


def register_add_either_shm_regions(inputs, outputs, shm_region_prefix,
                                    shm_handles, io_num, input_byte_size,
                                    output_byte_size, use_system_shared_memory,
                                    use_cuda_shared_memory, triton_client):
    # Lazy shm imports...
    if use_system_shared_memory:
        import tritonclient.utils.shared_memory as shm
    if use_cuda_shared_memory:
        import tritonclient.utils.cuda_shared_memory as cudashm

    if use_system_shared_memory or use_cuda_shared_memory:
        # Unregister then register required shared memory regions
        input_shm_name = shm_region_prefix[0] + str(io_num)
        output_shm_name = shm_region_prefix[1] + str(io_num)
        if use_system_shared_memory:
            triton_client.unregister_system_shared_memory(input_shm_name +
                                                          '_data')
            triton_client.unregister_system_shared_memory(output_shm_name +
                                                          '_data')
            triton_client.register_system_shared_memory(
                input_shm_name + '_data', '/' + input_shm_name, input_byte_size)
            triton_client.register_system_shared_memory(
                output_shm_name + '_data', '/' + output_shm_name,
                output_byte_size)

        if use_cuda_shared_memory:
            triton_client.unregister_cuda_shared_memory(input_shm_name +
                                                        '_data')
            triton_client.unregister_cuda_shared_memory(output_shm_name +
                                                        '_data')
            triton_client.register_cuda_shared_memory(
                input_shm_name + '_data',
                cudashm.get_raw_handle(shm_handles[0][io_num]), 0,
                input_byte_size)
            triton_client.register_cuda_shared_memory(
                output_shm_name + '_data',
                cudashm.get_raw_handle(shm_handles[1][io_num]), 0,
                output_byte_size)

        # Add shared memory regions to inputs
        inputs[io_num].set_shared_memory(input_shm_name + '_data',
                                         input_byte_size)
        outputs[io_num].set_shared_memory(output_shm_name + '_data',
                                          output_byte_size)


class ShmLeakDetector:
    """Detect shared memory leaks when testing Python backend."""

    class ShmLeakProbe:

        def __init__(self, shm_monitors):
            self._shm_monitors = shm_monitors

        def __enter__(self):
            self._shm_region_free_sizes = []
            for shm_monitor in self._shm_monitors:
                self._shm_region_free_sizes.append(shm_monitor.free_memory())

            return self

        def __exit__(self, type, value, traceback):
            current_shm_sizes = []
            for shm_monitor in self._shm_monitors:
                current_shm_sizes.append(shm_monitor.free_memory())

            shm_leak_detected = False
            for current_shm_size, prev_shm_size in zip(
                    current_shm_sizes, self._shm_region_free_sizes):
                if current_shm_size != prev_shm_size:
                    shm_leak_detected = True
                    print(
                        f'Shared memory leak detected: {current_shm_size} (current) != {prev_shm_size} (prev).'
                    )
            assert not shm_leak_detected, "Shared memory leak detected."

    def __init__(self, prefix='triton_python_backend_shm_region'):
        import triton_shm_monitor
        self._shm_monitors = []
        shm_regions = listdir('/dev/shm')
        for shm_region in shm_regions:
            if shm_region.startswith(prefix):
                self._shm_monitors.append(
                    triton_shm_monitor.SharedMemoryManager(shm_region))

    def Probe(self):
        return self.ShmLeakProbe(self._shm_monitors)


# qa/common/test_util.py

# Copyright 2018-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import numpy as np
import unittest
import json

_last_request_id = 0


def shape_element_count(shape):
    cnt = 0
    for d in shape:
        if d == -1:
            return -1
        if cnt == 0:
            cnt = d
        else:
            cnt = cnt * d
    return cnt


def shape_is_fixed(shape):
    return shape_element_count(shape) != -1


def shape_to_tf_shape(shape):
    return [None if i == -1 else i for i in shape]


def shape_to_onnx_shape(shape, idx=0, increment_index=True):
    # Onnx use string for variable size dimension, and the same string
    # will be inferred to have same value for the model run.
    # So there is an extra "idx" parameter to make sure the string is
    # unique
    res = []
    for dim in shape:
        if dim == -1:
            res.append("var_" + str(idx))
            if increment_index:
                idx += 1
        else:
            res.append(dim)
    return res, idx


def shape_to_dims_str(shape):
    return ','.join(str(i) for i in shape)


def validate_for_tf_model(input_dtype, output0_dtype, output1_dtype,
                          input_shape, output0_shape, output1_shape):
    """Return True if input and output dtypes are supported by a TF model."""

    # Not extending test to uint8 yet
    if input_dtype == np.uint8 or output0_dtype == np.uint8 or output1_dtype == np.uint8:
        return False

    # If the input type is string the output type must be string or
    # int32. This is because the QA models we generate convert strings
    # internally to int32 for compute.
    if ((input_dtype == np.object_) and
        (((output0_dtype != np.object_) and (output0_dtype != np.int32)) or
         ((output1_dtype != np.object_) and (output1_dtype != np.int32)))):
        return False

    return True


def validate_for_trt_model(input_dtype, output0_dtype, output1_dtype,
                           input_shape, output0_shape, output1_shape):
    """Return True if input and output dtypes are supported by a TRT model."""
    supported_datatypes = [
        bool, np.int8, np.int32, np.uint8, np.float16, np.float32
    ]
    # FIXME: Remove this check when jetson supports TRT 8.5 (DLIS-4256)
    if not support_trt_uint8():
        supported_datatypes.remove(np.uint8)
    if not input_dtype in supported_datatypes:
        return False
    if not output0_dtype in supported_datatypes:
        return False
    if not output1_dtype in supported_datatypes:
        return False

    datatype_set = set([input_dtype, output0_dtype, output1_dtype])

    # Incompatible datatype conversions
    if (np.int32 in datatype_set) and (np.int8 in datatype_set):
        return False
    if (np.float32 in datatype_set) and (np.int32 in datatype_set):
        return False

    return True


def validate_for_ensemble_model(ensemble_type, input_dtype, output0_dtype,
                                output1_dtype, input_shape, output0_shape,
                                output1_shape):
    """Return True if input and output dtypes are supported by the ensemble type."""

    # Not extending test to uint8 yet
    if input_dtype == np.uint8 or output0_dtype == np.uint8 or output1_dtype == np.uint8:
        return False

    # Those ensemble types contains "identity" model which doesn't allow STRING
    # data type
    # Test types that use identity for both input and output
    test_type_involved = ["reshape", "zero", "fan"]
    if input_dtype == np.object_ or output0_dtype == np.object_ or output1_dtype == np.object_:
        for type_str in test_type_involved:
            if type_str in ensemble_type:
                return False

    # Otherwise, check input / output separately
    if input_dtype == np.object_ and "sequence" in ensemble_type:
        return False

    return True


def validate_for_onnx_model(input_dtype, output0_dtype, output1_dtype,
                            input_shape, output0_shape, output1_shape):
    """Return True if input and output dtypes are supported by a Onnx model."""

    # Not extending test to uint8 yet
    if input_dtype == np.uint8 or output0_dtype == np.uint8 or output1_dtype == np.uint8:
        return False

    # If the input type is string the output type must be string or
    # int32. This is because the QA models we generate convert strings
    # internally to int32 for compute.
    if ((input_dtype == np.object_) and
        (((output0_dtype != np.object_) and (output0_dtype != np.int32)) or
         ((output1_dtype != np.object_) and (output1_dtype != np.int32)))):
        return False

    return True


def validate_for_libtorch_model(input_dtype,
                                output0_dtype,
                                output1_dtype,
                                input_shape,
                                output0_shape,
                                output1_shape,
                                max_batch=0,
                                reshape=False):
    """Return True if input and output dtypes are supported by a libtorch model."""

    # Not extending test to uint8 yet
    if input_dtype == np.uint8 or output0_dtype == np.uint8 or output1_dtype == np.uint8:
        return False

    # STRING data type does not support I/O with more than 1 dims. It supports
    # batching when 'reshape' field is set properly to empty shape.
    has_string_type = (input_dtype == np.object_) or (
        output0_dtype == np.object_) or (output1_dtype == np.object_)
    is_more_than_one_dimensional = (len(input_shape) > 1) or (
        len(output0_shape) > 1) or (len(output1_shape) > 1) or (max_batch != 0)

    if has_string_type and is_more_than_one_dimensional and not reshape:
        return False

    # FLOAT16 and UINT16 data types are not supported currently
    if (input_dtype == np.uint16) or (output0_dtype
                                      == np.uint16) or (output1_dtype
                                                        == np.uint16):
        return False
    if (input_dtype == np.float16) or (output0_dtype
                                       == np.float16) or (output1_dtype
                                                          == np.float16):
        return False

    return True


def validate_for_openvino_model(input_dtype, output0_dtype, output1_dtype,
                                input_shape, output0_shape, output1_shape):
    """Return True if input and output dtypes are supported by an OpenVino model."""

    # Not extending test to uint8 yet
    if input_dtype == np.uint8 or output0_dtype == np.uint8 or output1_dtype == np.uint8:
        return False

    # float16 is not supported on CPU by OpenVino
    supported_datatypes = [np.int8, np.int32, np.float32]
    if not input_dtype in supported_datatypes:
        return False
    if not output0_dtype in supported_datatypes:
        return False
    if not output1_dtype in supported_datatypes:
        return False

    # Return false if input dtype != output dtype and shape > 1 dims
    # https://github.com/openvinotoolkit/openvino/issues/7173
    if ((output1_dtype != input_dtype) or
        (output0_dtype != input_dtype)) and len(input_shape) > 1:
        return False

    return True


def get_model_name(pf, input_dtype, output0_dtype, output1_dtype):
    return "{}_{}_{}_{}".format(pf,
                                np.dtype(input_dtype).name,
                                np.dtype(output0_dtype).name,
                                np.dtype(output1_dtype).name)


def get_sequence_model_name(pf, dtype):
    return "{}_sequence_{}".format(pf, np.dtype(dtype).name)


def get_dyna_sequence_model_name(pf, dtype):
    return "{}_dyna_sequence_{}".format(pf, np.dtype(dtype).name)


def get_zero_model_name(pf, io_cnt, dtype):
    return "{}_zero_{}_{}".format(pf, io_cnt, np.dtype(dtype).name)


# FIXME: Remove this def when jetson supports TRT 8.5 (DLIS-4256)
def support_trt_uint8():
    try:
        import tensorrt as trt
    except:
        # tensorrt library is not found, detect from environment
        import os
        return not bool(int(os.environ.get("TEST_JETSON", 0)))
    # tensorrt library is found, return if uint8 is defined
    return hasattr(trt, "uint8")


class TestResultCollector(unittest.TestCase):
    # TestResultCollector stores test result and prints it to stdout. In order
    # to use this class, unit tests must inherit this class. Use
    # `check_test_results` bash function from `common/util.sh` to verify the
    # expected number of tests produced by this class

    @classmethod
    def setResult(cls, total, errors, failures):
        cls.total, cls.errors, cls.failures = \
            total, errors, failures

    @classmethod
    def tearDownClass(cls):
        # this method is called when all the unit tests in a class are
        # finished.
        json_res = {
            'total': cls.total,
            'errors': cls.errors,
            'failures': cls.failures
        }
        with open('test_results.txt', 'w+') as f:
            f.write(json.dumps(json_res))

    def run(self, result=None):
        # result argument stores the accumulative test results
        test_result = super().run(result)
        total = test_result.testsRun
        errors = len(test_result.errors)
        failures = len(test_result.failures)
        self.setResult(total, errors, failures)


# qa/common/trace_summary.py

#!/usr/bin/python

# Copyright 2019-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
import csv
import json
import numpy as np

FLAGS = None


def add_span(span_map, timestamps, span_name, ts_start, ts_end):
    for tag in (ts_start, ts_end):
        if tag not in timestamps:
            raise ValueError('timestamps missing "{}": {}'.format(
                tag, timestamps))
    if timestamps[ts_end] < timestamps[ts_start]:
        raise ValueError('end timestamp "{}" < start timestamp "{}"'.format(
            ts_end, ts_start))
    if span_name not in span_map:
        span_map[span_name] = 0
    span_map[span_name] += timestamps[ts_end] - timestamps[ts_start]


class AbstractFrontend():

    @property
    def filter_timestamp(self):
        return None

    def add_frontend_span(self, span_map, timestamps):
        pass

    def summarize_frontend_span(self, span_map, cnt):
        return None


class HttpFrontend(AbstractFrontend):

    @property
    def filter_timestamp(self):
        return "HTTP_RECV_START"

    def add_frontend_span(self, span_map, timestamps):
        if ("HTTP_RECV_START" in timestamps) and ("HTTP_SEND_END"
                                                  in timestamps):
            add_span(span_map, timestamps, "HTTP_INFER", "HTTP_RECV_START",
                     "HTTP_SEND_END")
            add_span(span_map, timestamps, "HTTP_RECV", "HTTP_RECV_START",
                     "HTTP_RECV_END")
            add_span(span_map, timestamps, "HTTP_SEND", "HTTP_SEND_START",
                     "HTTP_SEND_END")

    def summarize_frontend_span(self, span_map, cnt):
        if "HTTP_INFER" in span_map:
            res = "HTTP infer request (avg): {}us\n".format(
                span_map["HTTP_INFER"] / (cnt * 1000))
            res += "\tReceive (avg): {}us\n".format(span_map["HTTP_RECV"] /
                                                    (cnt * 1000))
            res += "\tSend (avg): {}us\n".format(span_map["HTTP_SEND"] /
                                                 (cnt * 1000))
            res += "\tOverhead (avg): {}us\n".format(
                (span_map["HTTP_INFER"] - span_map["REQUEST"] -
                 span_map["HTTP_RECV"] - span_map["HTTP_SEND"]) / (cnt * 1000))
            return res
        else:
            return None


class GrpcFrontend(AbstractFrontend):

    @property
    def filter_timestamp(self):
        return "GRPC_WAITREAD_START"

    def add_frontend_span(self, span_map, timestamps):
        if ("GRPC_WAITREAD_START" in timestamps) and ("GRPC_SEND_END"
                                                      in timestamps):
            add_span(span_map, timestamps, "GRPC_INFER", "GRPC_WAITREAD_START",
                     "GRPC_SEND_END")
            add_span(span_map, timestamps, "GRPC_WAITREAD",
                     "GRPC_WAITREAD_START", "GRPC_WAITREAD_END")
            add_span(span_map, timestamps, "GRPC_SEND", "GRPC_SEND_START",
                     "GRPC_SEND_END")

    def summarize_frontend_span(self, span_map, cnt):
        if "GRPC_INFER" in span_map:
            res = "GRPC infer request (avg): {}us\n".format(
                span_map["GRPC_INFER"] / (cnt * 1000))
            res += "\tWait/Read (avg): {}us\n".format(
                span_map["GRPC_WAITREAD"] / (cnt * 1000))
            res += "\tSend (avg): {}us\n".format(span_map["GRPC_SEND"] /
                                                 (cnt * 1000))
            res += "\tOverhead (avg): {}us\n".format(
                (span_map["GRPC_INFER"] - span_map["REQUEST"] -
                 span_map["GRPC_WAITREAD"] - span_map["GRPC_SEND"]) /
                (cnt * 1000))
            return res
        else:
            return None


def summarize(frontend, traces):
    # map from (model_name, model_version) to # of traces
    model_count_map = dict()
    # map from (model_name, model_version) to map of span->total time
    model_span_map = dict()

    # Order traces by id to be more intuitive if 'show_trace'
    traces = sorted(traces, key=lambda t: t.get('id', -1))

    # Filter the trace that is not for the requested frontend
    match_frontend_id_set = set()
    for trace in traces:
        if "id" not in trace:
            continue

        # Trace without a parent must contain frontend timestamps
        if "parent_id" not in trace:
            if frontend.filter_timestamp is None:
                continue
            if "timestamps" in trace:
                for ts in trace["timestamps"]:
                    if frontend.filter_timestamp in ts["name"]:
                        match_frontend_id_set.add(trace["id"])
        # Otherwise need to check whether parent is filtered
        elif trace["parent_id"] in match_frontend_id_set:
            match_frontend_id_set.add(trace["id"])

    # Filter the trace that is not meaningful and group them by 'id'
    filtered_traces = dict()
    for trace in traces:
        if "id" not in trace:
            continue
        if trace["id"] in match_frontend_id_set:
            if (trace['id'] in filtered_traces.keys()):
                rep_trace = filtered_traces[trace['id']]
                # Apend the timestamp to the trace representing this 'id'
                if "model_name" in trace:
                    rep_trace["model_name"] = trace["model_name"]
                if "model_version" in trace:
                    rep_trace["model_version"] = trace["model_version"]
                if "timestamps" in trace:
                    rep_trace["timestamps"] += trace["timestamps"]
            else:
                # Use this trace to represent this 'id'
                if "timestamps" not in trace:
                    trace["timestamps"] = []
                filtered_traces[trace['id']] = trace

    for trace_id, trace in filtered_traces.items():
        if trace_id not in match_frontend_id_set:
            filtered_traces.pop(trace_id, None)
            continue
        timestamps = dict()
        for ts in trace["timestamps"]:
            timestamps[ts["name"]] = ts["ns"]
        if ("REQUEST_START" in timestamps) and ("REQUEST_END" in timestamps):
            key = (trace["model_name"], trace["model_version"])
            if key not in model_count_map:
                model_count_map[key] = 0
                model_span_map[key] = dict()

            model_count_map[key] += 1

            frontend.add_frontend_span(model_span_map[key], timestamps)

            add_span(model_span_map[key], timestamps, "REQUEST",
                     "REQUEST_START", "REQUEST_END")

            # The tags below will be missing for ensemble model
            if ("QUEUE_START" in timestamps) and ("COMPUTE_START"
                                                  in timestamps):
                add_span(model_span_map[key], timestamps, "QUEUE",
                         "QUEUE_START", "COMPUTE_START")
            if ("COMPUTE_START" in timestamps) and ("COMPUTE_END"
                                                    in timestamps):
                add_span(model_span_map[key], timestamps, "COMPUTE",
                         "COMPUTE_START", "COMPUTE_END")
            if ("COMPUTE_INPUT_END" in timestamps) and ("COMPUTE_OUTPUT_START"
                                                        in timestamps):
                add_span(model_span_map[key], timestamps, "COMPUTE_INPUT",
                         "COMPUTE_START", "COMPUTE_INPUT_END")
                add_span(model_span_map[key], timestamps, "COMPUTE_INFER",
                         "COMPUTE_INPUT_END", "COMPUTE_OUTPUT_START")
                add_span(model_span_map[key], timestamps, "COMPUTE_OUTPUT",
                         "COMPUTE_OUTPUT_START", "COMPUTE_END")
            if FLAGS.show_trace:
                print("{} ({}):".format(trace["model_name"],
                                        trace["model_version"]))
                print("\tid: {}".format(trace["id"]))
                if "parent_id" in trace:
                    print("\tparent id: {}".format(trace["parent_id"]))
                ordered_timestamps = list()
                for ts in trace["timestamps"]:
                    ordered_timestamps.append((ts["name"], ts["ns"]))
                ordered_timestamps.sort(key=lambda tup: tup[1])

                now = None
                for ts in ordered_timestamps:
                    if now is not None:
                        print("\t\t{}us".format((ts[1] - now) / 1000))
                    print("\t{}".format(ts[0]))
                    now = ts[1]

    for key, cnt in model_count_map.items():
        model_name, model_value = key
        print("Summary for {} ({}): trace count = {}".format(
            model_name, model_value, cnt))

        frontend_summary = frontend.summarize_frontend_span(
            model_span_map[key], cnt)
        if frontend_summary is not None:
            print(frontend_summary)

        # collect handler timeline
        print("\tHandler (avg): {}us".format(model_span_map[key]["REQUEST"] /
                                             (cnt * 1000)))
        if ("QUEUE"
                in model_span_map[key]) and "COMPUTE" in model_span_map[key]:
            print("\t\tOverhead (avg): {}us".format(
                (model_span_map[key]["REQUEST"] - model_span_map[key]["QUEUE"] -
                 model_span_map[key]["COMPUTE"]) / (cnt * 1000)))
            print("\t\tQueue (avg): {}us".format(model_span_map[key]["QUEUE"] /
                                                 (cnt * 1000)))
            print("\t\tCompute (avg): {}us".format(
                model_span_map[key]["COMPUTE"] / (cnt * 1000)))
        if ("COMPUTE_INPUT" in model_span_map[key]
           ) and "COMPUTE_OUTPUT" in model_span_map[key]:
            print("\t\t\tInput (avg): {}us".format(
                model_span_map[key]["COMPUTE_INPUT"] / (cnt * 1000)))
            print("\t\t\tInfer (avg): {}us".format(
                model_span_map[key]["COMPUTE_INFER"] / (cnt * 1000)))
            print("\t\t\tOutput (avg): {}us".format(
                model_span_map[key]["COMPUTE_OUTPUT"] / (cnt * 1000)))


def summarize_dataflow(traces):
    # collect data flow
    # - parent input
    #   - child input
    #     - ...
    #   - child output

    # Order traces by id to be more intuitive if 'show_trace'
    traces = sorted(traces, key=lambda t: t.get('id', -1))

    # {3: [4, 5, 6], 4: [7]}
    dataflow_parent_map = dict()
    for trace in traces:
        if "id" not in trace:
            continue
        if "parent_id" in trace:
            if trace["parent_id"] not in dataflow_parent_map:
                dataflow_parent_map[trace["parent_id"]] = []
            dataflow_parent_map[trace["parent_id"]].append(trace["id"])

    if len(dataflow_parent_map) == 0:
        # print the tensors of model
        first_id = find_first_id_with_tensor(traces)
        if first_id != 0:
            print("Data Flow:")
        print_tensor_by_id(first_id, traces, 0, 0)
        return

    # print the tensors of ensemble
    print("Data Flow:")
    first_parent_id = list(dataflow_parent_map.items())[0][0]

    # {3: {4: {7: None}, 5: None, 6: None}}
    dataflow_tree_map = dict()
    depth = [0]
    append_dataflow_tensor(dataflow_tree_map, first_parent_id,
                           dataflow_parent_map, traces, depth)

    print_dataflow_tensor(dataflow_tree_map, traces, depth[0], step=0)


def append_dataflow_tensor(dataflow_tensor_map, parent_id, dataflow_tree_map,
                           traces, depth):
    if parent_id not in dataflow_tree_map:
        dataflow_tensor_map[parent_id] = None
        return

    child_tensor_map = dict()
    dataflow_tensor_map[parent_id] = child_tensor_map
    depth[0] = depth[0] + 1

    child_ids = dataflow_tree_map[parent_id]
    for child_id in child_ids:
        append_dataflow_tensor(child_tensor_map, child_id, dataflow_tree_map,
                               traces, depth)


def print_dataflow_tensor(dataflow_tree_map, traces, depth, step):
    for parent_id in dataflow_tree_map:
        print_tensor_by_id(parent_id, traces, depth, step)

        if dataflow_tree_map[parent_id] is None:
            continue

        print_dataflow_tensor(dataflow_tree_map[parent_id], traces, depth,
                              step + 1)


def print_tensor_by_id(id, traces, depth, step):
    if id == 0:
        return

    tabs = "\t" * (step + 1)

    print("{0}{1}".format(tabs, "=" * (50 + 8 * (depth - step))))
    for trace in traces:
        # print model name and version
        if "id" in trace and "model_name" in trace and "model_version" in trace and "timestamps" in trace and trace[
                "id"] == id:
            print("{0}Name:   {1}".format(tabs, trace["model_name"]))
            print("{0}Version:{1}".format(tabs, trace["model_version"]))
        # print data
        if "id" in trace and "activity" in trace:
            if trace["id"] == id and trace["activity"] == "TENSOR_QUEUE_INPUT":
                print("{0}{1}:".format(tabs, "QUEUE_INPUT"))
                print("{0}\t{1}: {2}".format(tabs, trace["tensor"]["name"],
                                             get_numpy_array(trace["tensor"])))
            elif trace["id"] == id and trace[
                    "activity"] == "TENSOR_BACKEND_INPUT":
                print("{0}{1}:".format(tabs, "BACKEND_INPUT"))
                print("{0}\t{1}: {2}".format(tabs, trace["tensor"]["name"],
                                             get_numpy_array(trace["tensor"])))
            elif trace["id"] == id and trace[
                    "activity"] == "TENSOR_BACKEND_OUTPUT":
                print("{0}{1}:".format(tabs, "BACKEND_OUTPUT"))
                print("{0}\t{1}: {2}".format(tabs, trace["tensor"]["name"],
                                             get_numpy_array(trace["tensor"])))
    print("{0}{1}".format(tabs, "=" * (50 + 8 * (depth - step))))


def find_first_id_with_tensor(traces):
    for trace in traces:
        if "activity" in trace and (
                trace["activity"] == "TENSOR_QUEUE_INPUT" or
                trace["activity"] == "TENSOR_BACKEND_INPUT" or
                trace["activity"] == "TENSOR_BACKEND_OUTPUT"):
            return trace["id"]
    return 0


TRITON_TYPE_TO_NUMPY = {
    "BOOL": bool,
    "UINT8": np.uint8,
    "UINT16": np.uint16,
    "UINT32": np.uint32,
    "UINT64": np.uint64,
    "INT8": np.int8,
    "INT16": np.int16,
    "INT32": np.int32,
    "INT64": np.int64,
    "FP16": np.float16,
    "FP32": np.float32,
    "FP64": np.float64,
    "BYTES": np.object_
}


def get_numpy_array(tensor):
    dtype = TRITON_TYPE_TO_NUMPY[tensor["dtype"]]
    if dtype == np.object_:
        value = next(csv.reader([tensor["data"]], skipinitialspace=True))
    else:
        value = map(float, tensor["data"].split(","))
    shape = map(int, tensor["shape"].split(","))
    array = np.array(list(value), dtype=dtype)
    array = array.reshape(list(shape))
    return array


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-v',
                        '--verbose',
                        action="store_true",
                        required=False,
                        default=False,
                        help='Enable verbose output')
    parser.add_argument('-t',
                        '--show-trace',
                        action="store_true",
                        required=False,
                        default=False,
                        help='Show timestamps for each individual trace')
    parser.add_argument('file', type=argparse.FileType('r'), nargs='+')
    FLAGS = parser.parse_args()

    for f in FLAGS.file:
        trace_data = json.loads(f.read())
        if FLAGS.verbose:
            print(json.dumps(trace_data, sort_keys=True, indent=2))

        # Must summarize HTTP and GRPC separately since they have
        # different ways of accumulating time.
        print("File: {}".format(f.name))
        summarize(HttpFrontend(), trace_data)
        summarize(GrpcFrontend(), trace_data)
        summarize_dataflow(trace_data)


# build.py

#!/usr/bin/env python3
# Copyright 2020-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
import os
import os.path
import multiprocessing
import pathlib
import platform
import stat
import subprocess
import sys
from inspect import getsourcefile

#
# Build Triton Inference Server.
#

# By default build.py builds the Triton Docker image, but can also be
# used to build without Docker.  See docs/build.md and --help for more
# infomation.
#
# The TRITON_VERSION file indicates the Triton version and
# TRITON_VERSION_MAP is used to determine the corresponding container
# version and upstream container version (upstream containers are
# dependencies required by Triton). These versions may be overridden.

# Map from Triton version to corresponding container and component versions.
#
#   triton version ->
#     (triton container version,
#      upstream container version,
#      ORT version,
#      ORT OpenVINO version (use None to disable OpenVINO in ORT),
#      Standalone OpenVINO version,
#      DCGM version,
#      Conda version
#     )
#
# Currently the OpenVINO versions used in ORT and standalone must
# match because of the way dlopen works with loading the backends. If
# different versions are used then one backend or the other will
# incorrectly load the other version of the openvino libraries.
#
TRITON_VERSION_MAP = {
    '2.31.0dev': (
        '23.02dev',  # triton container
        '23.01',  # upstream container
        '1.13.1',  # ORT
        '2022.1.0',  # ORT OpenVINO
        '2022.1.0',  # Standalone OpenVINO
        '2.2.9',  # DCGM version
        'py38_4.12.0')  # Conda version.
}

CORE_BACKENDS = ['ensemble']

FLAGS = None
EXTRA_CORE_CMAKE_FLAGS = {}
OVERRIDE_CORE_CMAKE_FLAGS = {}
EXTRA_BACKEND_CMAKE_FLAGS = {}
OVERRIDE_BACKEND_CMAKE_FLAGS = {}

THIS_SCRIPT_DIR = os.path.dirname(os.path.abspath(getsourcefile(lambda: 0)))


def log(msg, force=False):
    if force or not FLAGS.quiet:
        try:
            print(msg, file=sys.stderr)
        except Exception:
            print('<failed to log>', file=sys.stderr)


def log_verbose(msg):
    if FLAGS.verbose:
        log(msg, force=True)


def fail(msg):
    fail_if(True, msg)


def fail_if(p, msg):
    if p:
        print('error: {}'.format(msg), file=sys.stderr)
        sys.exit(1)


def target_platform():
    if FLAGS.target_platform is not None:
        return FLAGS.target_platform
    return platform.system().lower()


def target_machine():
    if FLAGS.target_machine is not None:
        return FLAGS.target_machine
    return platform.machine().lower()


def container_versions(version, container_version, upstream_container_version):
    if container_version is None:
        if version not in TRITON_VERSION_MAP:
            fail('container version not known for {}'.format(version))
        container_version = TRITON_VERSION_MAP[version][0]
    if upstream_container_version is None:
        if version not in TRITON_VERSION_MAP:
            fail('upstream container version not known for {}'.format(version))
        upstream_container_version = TRITON_VERSION_MAP[version][1]
    return container_version, upstream_container_version


class BuildScript:
    """Utility class for writing build scripts"""

    def __init__(self, filepath, desc=None, verbose=False):
        self._filepath = filepath
        self._file = open(self._filepath, "w")
        self._verbose = verbose
        self.header(desc)

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        self.close()

    def __del__(self):
        self.close()

    def close(self):
        if self._file is not None:
            if target_platform() == 'windows':
                self.blankln()
                self._file.write('}\n')
                self._file.write('catch {\n')
                self._file.write('    $_;\n')
                self._file.write('    ExitWithCode 1;\n')
                self._file.write('}\n')
            """Close the file"""
            self._file.close()
            self._file = None
            st = os.stat(self._filepath)
            os.chmod(self._filepath, st.st_mode | stat.S_IEXEC)

    def blankln(self):
        self._file.write('\n')

    def commentln(self, cnt):
        self._file.write('#' * cnt + '\n')

    def comment(self, msg=''):
        if not isinstance(msg, str):
            try:
                for m in msg:
                    self._file.write(f'# {msg}\n')
                return
            except TypeError:
                pass
        self._file.write(f'# {msg}\n')

    def comment_verbose(self, msg=''):
        if self._verbose:
            self.comment(msg)

    def header(self, desc=None):
        if target_platform() != 'windows':
            self._file.write('#!/usr/bin/env bash\n\n')

        if desc is not None:
            self.comment()
            self.comment(desc)
            self.comment()
            self.blankln()

        self.comment('Exit script immediately if any command fails')
        if target_platform() == 'windows':
            self._file.write('function ExitWithCode($exitcode) {\n')
            self._file.write('    $host.SetShouldExit($exitcode)\n')
            self._file.write('    exit $exitcode\n')
            self._file.write('}\n')
            self.blankln()
            if self._verbose:
                self._file.write('Set-PSDebug -Trace 1\n')
            self.blankln()
            self._file.write('try {\n')
        else:
            self._file.write('set -e\n')
            if self._verbose:
                self._file.write('set -x\n')
        self.blankln()

    def envvar_ref(self, v):
        if target_platform() == 'windows':
            return f'${{env:{v}}}'
        return f'${{{v}}}'

    def cmd(self, clist, check_exitcode=False):
        if isinstance(clist, str):
            self._file.write(f'{clist}\n')
        else:
            for c in clist:
                self._file.write(f'{c} ')
            self.blankln()

        if check_exitcode:
            if target_platform() == 'windows':
                self._file.write('if ($LASTEXITCODE -ne 0) {\n')
                self._file.write(
                    '  Write-Output "exited with status code $LASTEXITCODE";\n')
                self._file.write('  ExitWithCode 1;\n')
                self._file.write('}\n')

    def cwd(self, path):
        if target_platform() == 'windows':
            self.cmd(f'Set-Location -EV Err -EA Stop {path}')
        else:
            self.cmd(f'cd {path}')

    def cp(self, src, dest):
        if target_platform() == 'windows':
            self.cmd(f'Copy-Item -EV Err -EA Stop {src} -Destination {dest}')
        else:
            self.cmd(f'cp {src} {dest}')

    def mkdir(self, path):
        if target_platform() == 'windows':
            self.cmd(
                f'New-Item -EV Err -EA Stop -ItemType Directory -Force -Path {path}'
            )
        else:
            self.cmd(f'mkdir -p {pathlib.Path(path)}')

    def rmdir(self, path):
        if target_platform() == 'windows':
            self.cmd(f'if (Test-Path -Path {path}) {{')
            self.cmd(f'  Remove-Item -EV Err -EA Stop -Recurse -Force {path}')
            self.cmd('}')
        else:
            self.cmd(f'rm -fr {pathlib.Path(path)}')

    def cpdir(self, src, dest):
        if target_platform() == 'windows':
            self.cmd(
                f'Copy-Item -EV Err -EA Stop -Recurse {src} -Destination {dest}'
            )
        else:
            self.cmd(f'cp -r {src} {dest}')

    def tar(self, subdir, tar_filename):
        if target_platform() == 'windows':
            fail('unsupported operation: tar')
        else:
            self.cmd(f'tar zcf {tar_filename} {subdir}')

    def cmake(self, args):
        # Pass some additional envvars into cmake...
        env_args = []
        for k in ('TRT_VERSION', 'CMAKE_TOOLCHAIN_FILE',
                  'VCPKG_TARGET_TRIPLET'):
            env_args += [f'"-D{k}={self.envvar_ref(k)}"']
        self.cmd(f'cmake {" ".join(env_args)} {" ".join(args)}',
                 check_exitcode=True)

    def makeinstall(self, target='install'):
        if target_platform() == 'windows':
            verbose_flag = '' if self._verbose else '-clp:ErrorsOnly'
            self.cmd(
                f'msbuild.exe -m:{FLAGS.build_parallel} {verbose_flag} -p:Configuration={FLAGS.build_type} {target}.vcxproj',
                check_exitcode=True)
        else:
            verbose_flag = 'VERBOSE=1' if self._verbose else 'VERBOSE=0'
            self.cmd(f'make -j{FLAGS.build_parallel} {verbose_flag} {target}')

    def gitclone(self, repo, tag, subdir, org):
        clone_dir = subdir
        if not FLAGS.no_force_clone:
            self.rmdir(clone_dir)

        if target_platform() == 'windows':
            self.cmd(f'if (-Not (Test-Path -Path {clone_dir})) {{')
        else:
            self.cmd(f'if [[ ! -e {clone_dir} ]]; then')

        # FIXME [DLIS-4045 - Currently the tag starting with "pull/" is not
        # working with "--repo-tag" as the option is not forwared to the
        # individual repo build correctly.]
        # If 'tag' starts with "pull/" then it must be of form
        # "pull/<pr>/head". We just clone at "main" and then fetch the
        # reference onto a new branch we name "tritonbuildref".
        if tag.startswith("pull/"):
            self.cmd(
                f'  git clone --recursive --depth=1 {org}/{repo}.git {subdir};',
                check_exitcode=True)
            self.cmd('}' if target_platform() == 'windows' else 'fi')
            self.cwd(subdir)
            self.cmd(f'git fetch origin {tag}:tritonbuildref',
                     check_exitcode=True)
            self.cmd(f'git checkout tritonbuildref', check_exitcode=True)
        else:
            self.cmd(
                f'  git clone --recursive --single-branch --depth=1 -b {tag} {org}/{repo}.git {subdir};',
                check_exitcode=True)
            self.cmd('}' if target_platform() == 'windows' else 'fi')


def cmake_core_arg(name, type, value):
    # Return cmake -D setting to set name=value for core build. Use
    # command-line specified value if one is given.
    if name in OVERRIDE_CORE_CMAKE_FLAGS:
        value = OVERRIDE_CORE_CMAKE_FLAGS[name]
    if type is None:
        type = ''
    else:
        type = ':{}'.format(type)
    return '"-D{}{}={}"'.format(name, type, value)


def cmake_core_enable(name, flag):
    # Return cmake -D setting to set name=flag?ON:OFF for core
    # build. Use command-line specified value for 'flag' if one is
    # given.
    if name in OVERRIDE_CORE_CMAKE_FLAGS:
        value = OVERRIDE_CORE_CMAKE_FLAGS[name]
    else:
        value = 'ON' if flag else 'OFF'
    return '"-D{}:BOOL={}"'.format(name, value)


def cmake_core_extra_args():
    args = []
    for k, v in EXTRA_CORE_CMAKE_FLAGS.items():
        args.append('"-D{}={}"'.format(k, v))
    return args


def cmake_backend_arg(backend, name, type, value):
    # Return cmake -D setting to set name=value for backend build. Use
    # command-line specified value if one is given.
    if backend in OVERRIDE_BACKEND_CMAKE_FLAGS:
        if name in OVERRIDE_BACKEND_CMAKE_FLAGS[backend]:
            value = OVERRIDE_BACKEND_CMAKE_FLAGS[backend][name]
    if type is None:
        type = ''
    else:
        type = ':{}'.format(type)
    return '"-D{}{}={}"'.format(name, type, value)


def cmake_backend_enable(backend, name, flag):
    # Return cmake -D setting to set name=flag?ON:OFF for backend
    # build. Use command-line specified value for 'flag' if one is
    # given.
    value = None
    if backend in OVERRIDE_BACKEND_CMAKE_FLAGS:
        if name in OVERRIDE_BACKEND_CMAKE_FLAGS[backend]:
            value = OVERRIDE_BACKEND_CMAKE_FLAGS[backend][name]
    if value is None:
        value = 'ON' if flag else 'OFF'
    return '"-D{}:BOOL={}"'.format(name, value)


def cmake_backend_extra_args(backend):
    args = []
    if backend in EXTRA_BACKEND_CMAKE_FLAGS:
        for k, v in EXTRA_BACKEND_CMAKE_FLAGS[backend].items():
            args.append('"-D{}={}"'.format(k, v))
    return args


def cmake_repoagent_arg(name, type, value):
    # For now there is no override for repo-agents
    if type is None:
        type = ''
    else:
        type = ':{}'.format(type)
    return '"-D{}{}={}"'.format(name, type, value)


def cmake_repoagent_enable(name, flag):
    # For now there is no override for repo-agents
    value = 'ON' if flag else 'OFF'
    return '"-D{}:BOOL={}"'.format(name, value)


def cmake_repoagent_extra_args():
    # For now there is no extra args for repo-agents
    args = []
    return args


def core_cmake_args(components, backends, cmake_dir, install_dir):
    cargs = [
        cmake_core_arg('CMAKE_BUILD_TYPE', None, FLAGS.build_type),
        cmake_core_arg('CMAKE_INSTALL_PREFIX', 'PATH', install_dir),
        cmake_core_arg('TRITON_VERSION', 'STRING', FLAGS.version),
        cmake_core_arg('TRITON_COMMON_REPO_TAG', 'STRING',
                       components['common']),
        cmake_core_arg('TRITON_CORE_REPO_TAG', 'STRING', components['core']),
        cmake_core_arg('TRITON_BACKEND_REPO_TAG', 'STRING',
                       components['backend']),
        cmake_core_arg('TRITON_THIRD_PARTY_REPO_TAG', 'STRING',
                       components['thirdparty'])
    ]

    cargs.append(
        cmake_core_enable('TRITON_ENABLE_LOGGING', FLAGS.enable_logging))
    cargs.append(cmake_core_enable('TRITON_ENABLE_STATS', FLAGS.enable_stats))
    cargs.append(
        cmake_core_enable('TRITON_ENABLE_METRICS', FLAGS.enable_metrics))
    cargs.append(
        cmake_core_enable('TRITON_ENABLE_METRICS_GPU',
                          FLAGS.enable_gpu_metrics))
    cargs.append(
        cmake_core_enable('TRITON_ENABLE_METRICS_CPU',
                          FLAGS.enable_cpu_metrics))
    cargs.append(
        cmake_core_enable('TRITON_ENABLE_TRACING', FLAGS.enable_tracing))
    cargs.append(cmake_core_enable('TRITON_ENABLE_NVTX', FLAGS.enable_nvtx))

    cargs.append(cmake_core_enable('TRITON_ENABLE_GPU', FLAGS.enable_gpu))
    cargs.append(
        cmake_core_arg('TRITON_MIN_COMPUTE_CAPABILITY', None,
                       FLAGS.min_compute_capability))

    cargs.append(
        cmake_core_enable('TRITON_ENABLE_MALI_GPU', FLAGS.enable_mali_gpu))

    cargs.append(
        cmake_core_enable('TRITON_ENABLE_GRPC', 'grpc' in FLAGS.endpoint))
    cargs.append(
        cmake_core_enable('TRITON_ENABLE_HTTP', 'http' in FLAGS.endpoint))
    cargs.append(
        cmake_core_enable('TRITON_ENABLE_SAGEMAKER', 'sagemaker'
                          in FLAGS.endpoint))
    cargs.append(
        cmake_core_enable('TRITON_ENABLE_VERTEX_AI', 'vertex-ai'
                          in FLAGS.endpoint))

    cargs.append(
        cmake_core_enable('TRITON_ENABLE_GCS', 'gcs' in FLAGS.filesystem))
    cargs.append(cmake_core_enable('TRITON_ENABLE_S3', 's3'
                                   in FLAGS.filesystem))
    cargs.append(
        cmake_core_enable('TRITON_ENABLE_AZURE_STORAGE', 'azure_storage'
                          in FLAGS.filesystem))

    cargs.append(
        cmake_core_enable('TRITON_ENABLE_ENSEMBLE', 'ensemble' in backends))
    cargs.append(
        cmake_core_enable('TRITON_ENABLE_TENSORRT', 'tensorrt' in backends))

    cargs += cmake_core_extra_args()
    cargs.append(cmake_dir)
    return cargs


def repoagent_repo(ra):
    return '{}_repository_agent'.format(ra)


def repoagent_cmake_args(images, components, ra, install_dir):
    args = []

    cargs = args + [
        cmake_repoagent_arg('CMAKE_BUILD_TYPE', None, FLAGS.build_type),
        cmake_repoagent_arg('CMAKE_INSTALL_PREFIX', 'PATH', install_dir),
        cmake_repoagent_arg('TRITON_COMMON_REPO_TAG', 'STRING',
                            components['common']),
        cmake_repoagent_arg('TRITON_CORE_REPO_TAG', 'STRING',
                            components['core'])
    ]

    cargs.append(cmake_repoagent_enable('TRITON_ENABLE_GPU', FLAGS.enable_gpu))
    cargs += cmake_repoagent_extra_args()
    cargs.append('..')
    return cargs


def backend_repo(be):
    if (be == 'tensorflow1') or (be == 'tensorflow2'):
        return 'tensorflow_backend'
    return '{}_backend'.format(be)


def backend_cmake_args(images, components, be, install_dir, library_paths):
    if be == 'onnxruntime':
        args = onnxruntime_cmake_args(images, library_paths)
    elif be == 'openvino':
        args = openvino_cmake_args()
    elif be == 'tensorflow1':
        args = tensorflow_cmake_args(1, images, library_paths)
    elif be == 'tensorflow2':
        args = tensorflow_cmake_args(2, images, library_paths)
    elif be == 'python':
        args = []
    elif be == 'dali':
        args = dali_cmake_args()
    elif be == 'pytorch':
        args = pytorch_cmake_args(images)
    elif be == 'armnn_tflite':
        args = armnn_tflite_cmake_args()
    elif be == 'fil':
        args = fil_cmake_args(images)
    elif be == 'fastertransformer':
        args = fastertransformer_cmake_args()
    elif be == 'tensorrt':
        args = tensorrt_cmake_args()
    else:
        args = []

    cargs = args + [
        cmake_backend_arg(be, 'CMAKE_BUILD_TYPE', None, FLAGS.build_type),
        cmake_backend_arg(be, 'CMAKE_INSTALL_PREFIX', 'PATH', install_dir),
        cmake_backend_arg(be, 'TRITON_COMMON_REPO_TAG', 'STRING',
                          components['common']),
        cmake_backend_arg(be, 'TRITON_CORE_REPO_TAG', 'STRING',
                          components['core']),
        cmake_backend_arg(be, 'TRITON_BACKEND_REPO_TAG', 'STRING',
                          components['backend'])
    ]

    cargs.append(cmake_backend_enable(be, 'TRITON_ENABLE_GPU',
                                      FLAGS.enable_gpu))
    cargs.append(
        cmake_backend_enable(be, 'TRITON_ENABLE_MALI_GPU',
                             FLAGS.enable_mali_gpu))
    cargs.append(
        cmake_backend_enable(be, 'TRITON_ENABLE_STATS', FLAGS.enable_stats))
    cargs.append(
        cmake_backend_enable(be, 'TRITON_ENABLE_METRICS', FLAGS.enable_metrics))

    cargs += cmake_backend_extra_args(be)
    cargs.append('..')
    return cargs


def pytorch_cmake_args(images):

    # If platform is jetpack do not use docker based build
    if target_platform() == 'jetpack':
        if 'pytorch' not in library_paths:
            raise Exception(
                "Must specify library path for pytorch using --library-paths=pytorch:<path_to_pytorch>"
            )
        pt_lib_path = library_paths['pytorch'] + "/lib"
        pt_include_paths = ""
        for suffix in [
                'include/torch', 'include/torch/torch/csrc/api/include',
                'include/torchvision'
        ]:
            pt_include_paths += library_paths['pytorch'] + '/' + suffix + ';'
        cargs = [
            cmake_backend_arg('pytorch', 'TRITON_PYTORCH_INCLUDE_PATHS', None,
                              pt_include_paths),
            cmake_backend_arg('pytorch', 'TRITON_PYTORCH_LIB_PATHS', None,
                              pt_lib_path),
        ]
    else:
        if "pytorch" in images:
            image = images["pytorch"]
        else:
            image = 'nvcr.io/nvidia/pytorch:{}-py3'.format(
                FLAGS.upstream_container_version)
        cargs = [
            cmake_backend_arg('pytorch', 'TRITON_PYTORCH_DOCKER_IMAGE', None,
                              image),
        ]

        if FLAGS.enable_gpu:
            cargs.append(
                cmake_backend_enable('pytorch',
                                     'TRITON_PYTORCH_ENABLE_TORCHTRT', True))
        cargs.append(
            cmake_backend_enable('pytorch', 'TRITON_ENABLE_NVTX',
                                 FLAGS.enable_nvtx))
    return cargs


def onnxruntime_cmake_args(images, library_paths):
    cargs = [
        cmake_backend_arg('onnxruntime', 'TRITON_BUILD_ONNXRUNTIME_VERSION',
                          None, TRITON_VERSION_MAP[FLAGS.version][2])
    ]

    # TRITON_ENABLE_GPU is already set for all backends in backend_cmake_args()
    if FLAGS.enable_gpu:
        cargs.append(
            cmake_backend_enable('onnxruntime',
                                 'TRITON_ENABLE_ONNXRUNTIME_TENSORRT', True))

    # If platform is jetpack do not use docker based build
    if target_platform() == 'jetpack':
        if 'onnxruntime' not in library_paths:
            raise Exception(
                "Must specify library path for onnxruntime using --library-paths=onnxruntime:<path_to_onnxruntime>"
            )
        ort_lib_path = library_paths['onnxruntime'] + "/lib"
        ort_include_path = library_paths['onnxruntime'] + "/include"
        cargs += [
            cmake_backend_arg('onnxruntime', 'TRITON_ONNXRUNTIME_INCLUDE_PATHS',
                              None, ort_include_path),
            cmake_backend_arg('onnxruntime', 'TRITON_ONNXRUNTIME_LIB_PATHS',
                              None, ort_lib_path),
            cmake_backend_enable('onnxruntime',
                                 'TRITON_ENABLE_ONNXRUNTIME_OPENVINO', False)
        ]
    else:
        if target_platform() == 'windows':
            if 'ort-base' in images:
                cargs.append(
                    cmake_backend_arg('onnxruntime', 'TRITON_BUILD_CONTAINER',
                                      None, images['ort-base']))
        else:
            if 'ort-base' in images:
                cargs.append(
                    cmake_backend_arg('onnxruntime', 'TRITON_BUILD_CONTAINER',
                                      None, images['ort-base']))
            else:
                cargs.append(
                    cmake_backend_arg('onnxruntime',
                                      'TRITON_BUILD_CONTAINER_VERSION', None,
                                      TRITON_VERSION_MAP[FLAGS.version][1]))

            if ((target_machine() != 'aarch64') and
                (TRITON_VERSION_MAP[FLAGS.version][3] is not None)):
                cargs.append(
                    cmake_backend_enable('onnxruntime',
                                         'TRITON_ENABLE_ONNXRUNTIME_OPENVINO',
                                         True))
                cargs.append(
                    cmake_backend_arg(
                        'onnxruntime',
                        'TRITON_BUILD_ONNXRUNTIME_OPENVINO_VERSION', None,
                        TRITON_VERSION_MAP[FLAGS.version][3]))

    return cargs


def openvino_cmake_args():
    cargs = [
        cmake_backend_arg('openvino', 'TRITON_BUILD_OPENVINO_VERSION', None,
                          TRITON_VERSION_MAP[FLAGS.version][4])
    ]
    if target_platform() == 'windows':
        if 'base' in images:
            cargs.append(
                cmake_backend_arg('openvino', 'TRITON_BUILD_CONTAINER', None,
                                  images['base']))
    else:
        if 'base' in images:
            cargs.append(
                cmake_backend_arg('openvino', 'TRITON_BUILD_CONTAINER', None,
                                  images['base']))
        else:
            cargs.append(
                cmake_backend_arg('openvino', 'TRITON_BUILD_CONTAINER_VERSION',
                                  None, TRITON_VERSION_MAP[FLAGS.version][1]))
    return cargs


def tensorrt_cmake_args():
    cargs = [
        cmake_backend_enable('tensorrt', 'TRITON_ENABLE_NVTX',
                             FLAGS.enable_nvtx),
    ]
    if target_platform() == 'windows':
        cargs.append(
            cmake_backend_arg('tensorrt', 'TRITON_TENSORRT_INCLUDE_PATHS', None,
                              'c:/TensorRT/include'))

    return cargs


def tensorflow_cmake_args(ver, images, library_paths):
    backend_name = "tensorflow{}".format(ver)

    # If platform is jetpack do not use docker images
    extra_args = []
    if target_platform() == 'jetpack':
        if backend_name in library_paths:
            extra_args = [
                cmake_backend_arg(backend_name, 'TRITON_TENSORFLOW_LIB_PATHS',
                                  None, library_paths[backend_name])
            ]
        else:
            raise Exception(
                f"Must specify library path for {backend_name} using --library-paths={backend_name}:<path_to_{backend_name}>"
            )
    else:
        # If a specific TF image is specified use it, otherwise pull from NGC.
        if backend_name in images:
            image = images[backend_name]
        else:
            image = 'nvcr.io/nvidia/tensorflow:{}-tf{}-py3'.format(
                FLAGS.upstream_container_version, ver)
        extra_args = [
            cmake_backend_arg(backend_name, 'TRITON_TENSORFLOW_DOCKER_IMAGE',
                              None, image)
        ]
    return [
        cmake_backend_arg(backend_name, 'TRITON_TENSORFLOW_VERSION', None, ver)
    ] + extra_args


def dali_cmake_args():
    return [
        cmake_backend_enable('dali', 'TRITON_DALI_SKIP_DOWNLOAD', False),
    ]


def fil_cmake_args(images):
    cargs = [cmake_backend_enable('fil', 'TRITON_FIL_DOCKER_BUILD', True)]
    if 'base' in images:
        cargs.append(
            cmake_backend_arg('fil', 'TRITON_BUILD_CONTAINER', None,
                              images['base']))
    else:
        cargs.append(
            cmake_backend_arg('fil', 'TRITON_BUILD_CONTAINER_VERSION', None,
                              TRITON_VERSION_MAP[FLAGS.version][1]))

    return cargs


def armnn_tflite_cmake_args():
    return [
        cmake_backend_arg('armnn_tflite', 'JOBS', None,
                          multiprocessing.cpu_count()),
    ]


def fastertransformer_cmake_args():
    print("Warning: FasterTransformer backend is not officially supported.")
    cargs = [
        cmake_backend_arg('fastertransformer', 'CMAKE_EXPORT_COMPILE_COMMANDS',
                          None, 1),
        cmake_backend_arg('fastertransformer', 'ENABLE_FP8', None, 'OFF')
    ]
    return cargs


def install_dcgm_libraries(dcgm_version, target_machine):
    if dcgm_version == '':
        fail(
            'unable to determine default repo-tag, DCGM version not known for {}'
            .format(FLAGS.version))
        return ''
    else:
        if target_machine == 'aarch64':
            return '''
ENV DCGM_VERSION {}
# Install DCGM. Steps from https://developer.nvidia.com/dcgm#Downloads
RUN curl -o /tmp/cuda-keyring.deb \
    https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/sbsa/cuda-keyring_1.0-1_all.deb \
    && apt install /tmp/cuda-keyring.deb && rm /tmp/cuda-keyring.deb && \
    apt-get update && apt-get install -y datacenter-gpu-manager=1:{}
'''.format(dcgm_version, dcgm_version)
        else:
            return '''
ENV DCGM_VERSION {}
# Install DCGM. Steps from https://developer.nvidia.com/dcgm#Downloads
RUN curl -o /tmp/cuda-keyring.deb \
    https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.deb \
    && apt install /tmp/cuda-keyring.deb && rm /tmp/cuda-keyring.deb && \
    apt-get update && apt-get install -y datacenter-gpu-manager=1:{}
'''.format(dcgm_version, dcgm_version)


def install_miniconda(conda_version, target_machine):
    if target_machine == "arm64":
        # This branch used for the case when linux container builds on MacOS with ARM chip
        # macos arm arch names "arm64" when in linux it's names "aarch64".
        # So we just replace the architecture to able find right conda version for Linux
        target_machine = "aarch64"
    if conda_version == '':
        fail(
            'unable to determine default repo-tag, CONDA version not known for {}'
            .format(FLAGS.version))
    miniconda_url = f"https://repo.anaconda.com/miniconda/Miniconda3-{conda_version}-Linux-{target_machine}.sh"
    if target_machine == 'x86_64':
        sha_sum = "3190da6626f86eee8abf1b2fd7a5af492994eb2667357ee4243975cdbb175d7a"
    else:
        sha_sum = "0c20f121dc4c8010032d64f8e9b27d79e52d28355eb8d7972eafc90652387777"
    return f'''
RUN mkdir -p /opt/
RUN wget "{miniconda_url}" -O miniconda.sh -q && \
    echo "{sha_sum}" "miniconda.sh" > shasum && \
    sha256sum -c ./shasum && \
    sh miniconda.sh -b -p /opt/conda && \
    rm miniconda.sh shasum && \
    find /opt/conda/ -follow -type f -name '*.a' -delete && \
    find /opt/conda/ -follow -type f -name '*.js.map' -delete && \
    /opt/conda/bin/conda clean -afy
ENV PATH /opt/conda/bin:${{PATH}}
'''


def create_dockerfile_buildbase(ddir, dockerfile_name, argmap):
    df = '''
ARG TRITON_VERSION={}
ARG TRITON_CONTAINER_VERSION={}
ARG BASE_IMAGE={}
'''.format(argmap['TRITON_VERSION'], argmap['TRITON_CONTAINER_VERSION'],
           argmap['BASE_IMAGE'])

    df += '''
FROM ${BASE_IMAGE}

ARG TRITON_VERSION
ARG TRITON_CONTAINER_VERSION
'''
    # Install the windows- or linux-specific buildbase dependencies
    if target_platform() == 'windows':
        df += '''
SHELL ["cmd", "/S", "/C"]
'''
    else:
        df += '''
# Ensure apt-get won't prompt for selecting options
ENV DEBIAN_FRONTEND=noninteractive

# libcurl4-openSSL-dev is needed for GCS
# python3-dev is needed by Torchvision
# python3-pip and libarchive-dev is needed by python backend
# uuid-dev and pkg-config is needed for Azure Storage
# scons is needed for armnn_tflite backend build dep
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
            ca-certificates \
            autoconf \
            automake \
            build-essential \
            docker.io \
            git \
            gperf \
            libre2-dev \
            libssl-dev \
            libtool \
            libboost-dev \
            libcurl4-openssl-dev \
            libb64-dev \
            libgoogle-perftools-dev \
            patchelf \
            python3-dev \
            python3-pip \
            python3-setuptools \
            rapidjson-dev \
            scons \
            software-properties-common \
            unzip \
            wget \
            zlib1g-dev \
            libarchive-dev \
            pkg-config \
            uuid-dev \
            libnuma-dev && \
    rm -rf /var/lib/apt/lists/*

RUN pip3 install --upgrade pip && \
    pip3 install --upgrade wheel setuptools docker

# Server build requires recent version of CMake (FetchContent required)
RUN wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc 2>/dev/null | \
      gpg --dearmor - |  \
      tee /etc/apt/trusted.gpg.d/kitware.gpg >/dev/null && \
    apt-add-repository 'deb https://apt.kitware.com/ubuntu/ focal main' && \
    apt-get update && \
    apt-get install -y --no-install-recommends \
      cmake-data=3.21.1-0kitware1ubuntu20.04.1 cmake=3.21.1-0kitware1ubuntu20.04.1
'''

        if FLAGS.enable_gpu:
            df += install_dcgm_libraries(argmap['DCGM_VERSION'],
                                         target_machine())

    df += '''
ENV TRITON_SERVER_VERSION ${TRITON_VERSION}
ENV NVIDIA_TRITON_SERVER_VERSION ${TRITON_CONTAINER_VERSION}
'''

    # Copy in the triton source. We remove existing contents first in
    # case the FROM container has something there already.
    if target_platform() == 'windows':
        df += '''
WORKDIR /workspace
RUN rmdir /S/Q * || exit 0
COPY . .
'''
    else:
        df += '''
WORKDIR /workspace
RUN rm -fr *
COPY . .
ENTRYPOINT []
'''

    # Install miniconda required for the DALI backend.
    if target_platform() != 'windows':
        df += install_miniconda(argmap['CONDA_VERSION'], target_machine())

    with open(os.path.join(ddir, dockerfile_name), "w") as dfile:
        dfile.write(df)


def create_dockerfile_cibase(ddir, dockerfile_name, argmap):
    df = '''
ARG TRITON_VERSION={}
ARG TRITON_CONTAINER_VERSION={}
ARG BASE_IMAGE={}
'''.format(argmap['TRITON_VERSION'], argmap['TRITON_CONTAINER_VERSION'],
           argmap['BASE_IMAGE'])

    df += '''
FROM ${BASE_IMAGE}

ARG TRITON_VERSION
ARG TRITON_CONTAINER_VERSION

COPY build/ci /workspace

WORKDIR /workspace

ENV TRITON_SERVER_VERSION ${TRITON_VERSION}
ENV NVIDIA_TRITON_SERVER_VERSION ${TRITON_CONTAINER_VERSION}
'''

    with open(os.path.join(ddir, dockerfile_name), "w") as dfile:
        dfile.write(df)


def create_dockerfile_linux(ddir, dockerfile_name, argmap, backends, repoagents,
                            endpoints):
    df = '''
ARG TRITON_VERSION={}
ARG TRITON_CONTAINER_VERSION={}
ARG BASE_IMAGE={}

'''.format(argmap['TRITON_VERSION'], argmap['TRITON_CONTAINER_VERSION'],
           argmap['BASE_IMAGE'])

    # PyTorch, TensorFlow 1 and TensorFlow 2 backends need extra CUDA and other
    # dependencies during runtime that are missing in the CPU-only base container.
    # These dependencies must be copied from the Triton Min image.
    if not FLAGS.enable_gpu and (('pytorch' in backends) or
                                 ('tensorflow1' in backends) or
                                 ('tensorflow2' in backends)):
        df += '''
############################################################################
##  Triton Min image
############################################################################
FROM {} AS min_container

'''.format(argmap['GPU_BASE_IMAGE'])

    df += '''
############################################################################
##  Production stage: Create container with just inference server executable
############################################################################
FROM ${BASE_IMAGE}
'''

    df += dockerfile_prepare_container_linux(argmap, backends, FLAGS.enable_gpu,
                                             target_machine())

    df += '''
WORKDIR /opt
COPY --chown=1000:1000 build/install tritonserver

WORKDIR /opt/tritonserver
COPY --chown=1000:1000 NVIDIA_Deep_Learning_Container_License.pdf .

'''
    if not FLAGS.no_core_build:
        # Add feature labels for SageMaker endpoint
        if 'sagemaker' in endpoints:
            df += '''
LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
LABEL com.amazonaws.sagemaker.capabilities.multi-models=true
COPY --chown=1000:1000 docker/sagemaker/serve /usr/bin/.
'''

    # This is required since libcublasLt.so is not present during the build
    # stage of the PyTorch backend
    if not FLAGS.enable_gpu and ('pytorch' in backends):
        df += '''
RUN patchelf --add-needed /usr/local/cuda/lib64/stubs/libcublasLt.so.12 backends/pytorch/libtorch_cuda.so
'''

    with open(os.path.join(ddir, dockerfile_name), "w") as dfile:
        dfile.write(df)


def dockerfile_prepare_container_linux(argmap, backends, enable_gpu,
                                       target_machine):
    gpu_enabled = 1 if enable_gpu else 0
    # Common steps to produce docker images shared by build.py and compose.py.
    # Sets enviroment variables, installs dependencies and adds entrypoint
    df = '''
ARG TRITON_VERSION
ARG TRITON_CONTAINER_VERSION

ENV TRITON_SERVER_VERSION ${TRITON_VERSION}
ENV NVIDIA_TRITON_SERVER_VERSION ${TRITON_CONTAINER_VERSION}
LABEL com.nvidia.tritonserver.version="${TRITON_SERVER_VERSION}"

ENV PATH /opt/tritonserver/bin:${PATH}
'''

    # TODO Remove once the ORT-OpenVINO "Exception while Reading network" is fixed
    if 'onnxruntime' in backends:
        df += '''
ENV LD_LIBRARY_PATH /opt/tritonserver/backends/onnxruntime:${LD_LIBRARY_PATH}
ENV LD_LIBRARY_PATH /usr/local/cuda-11.8/lib64:${LD_LIBRARY_PATH}
'''

    backend_dependencies = ""
    # libgomp1 is needed by both onnxruntime and pytorch backends
    if ('onnxruntime' in backends) or ('pytorch' in backends):
        backend_dependencies = "libgomp1"

    # libgfortran5 is needed by pytorch backend on ARM
    if ('pytorch' in backends) and (target_machine == 'aarch64'):
        backend_dependencies += " libgfortran5"
    # openssh-server is needed for fastertransformer
    if ('fastertransformer' in backends):
        backend_dependencies += " openssh-server"

    df += '''
ENV TF_ADJUST_HUE_FUSED         1
ENV TF_ADJUST_SATURATION_FUSED  1
ENV TF_ENABLE_WINOGRAD_NONFUSED 1
ENV TF_AUTOTUNE_THRESHOLD       2
ENV TRITON_SERVER_GPU_ENABLED    {gpu_enabled}

# Create a user that can be used to run triton as
# non-root. Make sure that this user to given ID 1000. All server
# artifacts copied below are assign to this user.
ENV TRITON_SERVER_USER=triton-server
RUN userdel tensorrt-server > /dev/null 2>&1 || true && \
    if ! id -u $TRITON_SERVER_USER > /dev/null 2>&1 ; then \
        useradd $TRITON_SERVER_USER; \
    fi && \
    [ `id -u $TRITON_SERVER_USER` -eq 1000 ] && \
    [ `id -g $TRITON_SERVER_USER` -eq 1000 ]

# Ensure apt-get won't prompt for selecting options
ENV DEBIAN_FRONTEND=noninteractive

# Common dependencies. FIXME (can any of these be conditional? For
# example libcurl only needed for GCS?)
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
            software-properties-common \
            libb64-0d \
            libcurl4-openssl-dev \
            libre2-5 \
            git \
            gperf \
            dirmngr \
            libgoogle-perftools-dev \
            libnuma-dev \
            curl \
            {backend_dependencies} && \
    rm -rf /var/lib/apt/lists/*

# Set TCMALLOC_RELEASE_RATE for users setting LD_PRELOAD with tcmalloc
ENV TCMALLOC_RELEASE_RATE 200
'''.format(gpu_enabled=gpu_enabled, backend_dependencies=backend_dependencies)

    if ('fastertransformer' in backends):
        be = "fastertransformer"
        import importlib.util, requests
        url = 'https://raw.githubusercontent.com/triton-inference-server/fastertransformer_backend/{}/docker/create_dockerfile_and_build.py'.format(backends[be])
        response = requests.get(url)
        spec = importlib.util.spec_from_loader('fastertransformer_buildscript',
                                               loader=None,
                                               origin=url)
        fastertransformer_buildscript = importlib.util.module_from_spec(spec)
        exec(response.content, fastertransformer_buildscript.__dict__)
        df += fastertransformer_buildscript.create_postbuild(
            is_multistage_build=False)

    if enable_gpu:
        df += install_dcgm_libraries(argmap['DCGM_VERSION'], target_machine)
        # This is temporary solution to support 23.01
        df += '''
RUN apt-get update && apt-get install -y libcufft-11-8
'''
        df += '''
# Extra defensive wiring for CUDA Compat lib
RUN ln -sf ${_CUDA_COMPAT_PATH}/lib.real ${_CUDA_COMPAT_PATH}/lib \
 && echo ${_CUDA_COMPAT_PATH}/lib > /etc/ld.so.conf.d/00-cuda-compat.conf \
 && ldconfig \
 && rm -f ${_CUDA_COMPAT_PATH}/lib
'''

    else:
        libs_arch = 'aarch64' if target_machine == 'aarch64' else 'x86_64'
        if ('pytorch' in backends) or ('tensorflow1' in backends):
            # Add extra dependencies for tensorflow1/pytorch backend.
            # Note: Even though the build is CPU-only, the version of tensorflow1/
            # pytorch we are using depend upon libraries like cuda and cudnn. Since
            # these dependencies are not present in the ubuntu base image,
            # we must copy these from the Triton min container ourselves.
            cuda_arch = 'sbsa' if target_machine == 'aarch64' else 'x86_64'
            df += '''
RUN mkdir -p /usr/local/cuda/lib64/stubs
COPY --from=min_container /usr/local/cuda/lib64/stubs/libcusparse.so /usr/local/cuda/lib64/stubs/libcusparse.so.12
COPY --from=min_container /usr/local/cuda/lib64/stubs/libcusolver.so /usr/local/cuda/lib64/stubs/libcusolver.so.11
COPY --from=min_container /usr/local/cuda/lib64/stubs/libcurand.so /usr/local/cuda/lib64/stubs/libcurand.so.10
COPY --from=min_container /usr/local/cuda/lib64/stubs/libcufft.so /usr/local/cuda/lib64/stubs/libcufft.so.11
COPY --from=min_container /usr/local/cuda/lib64/stubs/libcublas.so /usr/local/cuda/lib64/stubs/libcublas.so.12
COPY --from=min_container /usr/local/cuda/lib64/stubs/libcublasLt.so /usr/local/cuda/lib64/stubs/libcublasLt.so.12
COPY --from=min_container /usr/local/cuda/lib64/stubs/libcublasLt.so /usr/local/cuda/lib64/stubs/libcublasLt.so.11

RUN mkdir -p /usr/local/cuda/targets/{cuda_arch}-linux/lib
COPY --from=min_container /usr/local/cuda-12.0/targets/{cuda_arch}-linux/lib/libcudart.so.12 /usr/local/cuda/targets/{cuda_arch}-linux/lib/.
COPY --from=min_container /usr/local/cuda-12.0/targets/{cuda_arch}-linux/lib/libcupti.so.12 /usr/local/cuda/targets/{cuda_arch}-linux/lib/.
COPY --from=min_container /usr/local/cuda-12.0/targets/{cuda_arch}-linux/lib/libnvToolsExt.so.1 /usr/local/cuda/targets/{cuda_arch}-linux/lib/.
COPY --from=min_container /usr/local/cuda-12.0/targets/{cuda_arch}-linux/lib/libnvJitLink.so.12 /usr/local/cuda/targets/{cuda_arch}-linux/lib/.

COPY --from=min_container /usr/lib/{libs_arch}-linux-gnu/libcudnn.so.8 /usr/lib/{libs_arch}-linux-gnu/libcudnn.so.8

# patchelf is needed to add deps of libcublasLt.so.12 to libtorch_cuda.so
RUN apt-get update && \
        apt-get install -y --no-install-recommends openmpi-bin patchelf

ENV LD_LIBRARY_PATH /usr/local/cuda/targets/{cuda_arch}-linux/lib:/usr/local/cuda/lib64/stubs:${{LD_LIBRARY_PATH}}
'''.format(cuda_arch=cuda_arch, libs_arch=libs_arch)

        if ('pytorch' in backends) or ('tensorflow1' in backends) \
                or ('tensorflow2' in backends):
            # Add NCCL dependency for tensorflow1/tensorflow2/pytorch backend.
            # Note: Even though the build is CPU-only, the version of tensorflow1/
            # tensorflow2/pytorch we are using depends upon the NCCL library. Since
            # this dependency is not present in the ubuntu base image, we must
            # copy it from the Triton min container ourselves.
            df += '''
COPY --from=min_container /usr/lib/{libs_arch}-linux-gnu/libnccl.so.2 /usr/lib/{libs_arch}-linux-gnu/libnccl.so.2
'''.format(libs_arch=libs_arch)

    # Add dependencies needed for python backend
    if 'python' in backends:
        df += '''
# python3, python3-pip and some pip installs required for the python backend
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
            python3 libarchive-dev \
            python3-pip \
            libpython3-dev && \
    pip3 install --upgrade pip && \
    pip3 install --upgrade wheel setuptools && \
    pip3 install --upgrade numpy && \
    rm -rf /var/lib/apt/lists/*
'''

    df += '''
WORKDIR /opt/tritonserver
RUN rm -fr /opt/tritonserver/*
ENV NVIDIA_PRODUCT_NAME="Triton Server"
COPY docker/entrypoint.d/ /opt/nvidia/entrypoint.d/
'''

    # The CPU-only build uses ubuntu as the base image, and so the
    # entrypoint files are not available in /opt/nvidia in the base
    # image, so we must provide them ourselves.
    if not enable_gpu:
        df += '''
COPY docker/cpu_only/ /opt/nvidia/
ENTRYPOINT ["/opt/nvidia/nvidia_entrypoint.sh"]
'''

    df += '''
ENV NVIDIA_BUILD_ID {}
LABEL com.nvidia.build.id={}
LABEL com.nvidia.build.ref={}
'''.format(argmap['NVIDIA_BUILD_ID'], argmap['NVIDIA_BUILD_ID'],
           argmap['NVIDIA_BUILD_REF'])

    return df


def create_dockerfile_windows(ddir, dockerfile_name, argmap, backends,
                              repoagents):
    df = '''
ARG TRITON_VERSION={}
ARG TRITON_CONTAINER_VERSION={}
ARG BASE_IMAGE={}

############################################################################
##  Production stage: Create container with just inference server executable
############################################################################
FROM ${{BASE_IMAGE}}

ARG TRITON_VERSION
ARG TRITON_CONTAINER_VERSION

ENV TRITON_SERVER_VERSION ${{TRITON_VERSION}}
ENV NVIDIA_TRITON_SERVER_VERSION ${{TRITON_CONTAINER_VERSION}}
LABEL com.nvidia.tritonserver.version="${{TRITON_SERVER_VERSION}}"

RUN setx path "%path%;C:\opt\tritonserver\bin"

'''.format(argmap['TRITON_VERSION'], argmap['TRITON_CONTAINER_VERSION'],
           argmap['BASE_IMAGE'])
    df += '''
WORKDIR /opt
RUN rmdir /S/Q tritonserver || exit 0
COPY --chown=1000:1000 build/install tritonserver

WORKDIR /opt/tritonserver
COPY --chown=1000:1000 NVIDIA_Deep_Learning_Container_License.pdf .

'''
    df += '''
ENTRYPOINT []
ENV NVIDIA_BUILD_ID {}
LABEL com.nvidia.build.id={}
LABEL com.nvidia.build.ref={}
'''.format(argmap['NVIDIA_BUILD_ID'], argmap['NVIDIA_BUILD_ID'],
           argmap['NVIDIA_BUILD_REF'])

    with open(os.path.join(ddir, dockerfile_name), "w") as dfile:
        dfile.write(df)


def create_build_dockerfiles(container_build_dir, images, backends, repoagents,
                             endpoints):
    if 'base' in images:
        base_image = images['base']
    elif target_platform() == 'windows':
        base_image = 'mcr.microsoft.com/dotnet/framework/sdk:4.8'
    elif FLAGS.enable_gpu:
        base_image = 'nvcr.io/nvidia/tritonserver:{}-py3-min'.format(
            FLAGS.upstream_container_version)
    else:
        base_image = 'ubuntu:20.04'

    dockerfileargmap = {
        'NVIDIA_BUILD_REF':
            '' if FLAGS.build_sha is None else FLAGS.build_sha,
        'NVIDIA_BUILD_ID':
            '<unknown>' if FLAGS.build_id is None else FLAGS.build_id,
        'TRITON_VERSION':
            FLAGS.version,
        'TRITON_CONTAINER_VERSION':
            FLAGS.container_version,
        'BASE_IMAGE':
            base_image,
        'DCGM_VERSION':
            '' if FLAGS.version is None or FLAGS.version
            not in TRITON_VERSION_MAP else TRITON_VERSION_MAP[FLAGS.version][5],
        'CONDA_VERSION':
            '' if FLAGS.version is None or FLAGS.version
            not in TRITON_VERSION_MAP else TRITON_VERSION_MAP[FLAGS.version][6]
    }

    # For CPU-only image we need to copy some cuda libraries and dependencies
    # since we are using PyTorch, TensorFlow 1, TensorFlow 2 containers that
    # are not CPU-only.
    if not FLAGS.enable_gpu and (
        ('pytorch' in backends) or ('tensorflow1' in backends) or
        ('tensorflow2' in backends)) and (target_platform() != 'windows'):
        if 'gpu-base' in images:
            gpu_base_image = images['gpu-base']
        else:
            gpu_base_image = 'nvcr.io/nvidia/tritonserver:{}-py3-min'.format(
                FLAGS.upstream_container_version)
        dockerfileargmap['GPU_BASE_IMAGE'] = gpu_base_image

    create_dockerfile_buildbase(FLAGS.build_dir, 'Dockerfile.buildbase',
                                dockerfileargmap)

    if target_platform() == 'windows':
        create_dockerfile_windows(FLAGS.build_dir, 'Dockerfile',
                                  dockerfileargmap, backends, repoagents)
    else:
        create_dockerfile_linux(FLAGS.build_dir, 'Dockerfile', dockerfileargmap,
                                backends, repoagents, endpoints)

    # Dockerfile used for the creating the CI base image.
    create_dockerfile_cibase(FLAGS.build_dir, 'Dockerfile.cibase',
                             dockerfileargmap)


def create_docker_build_script(script_name, container_install_dir,
                               container_ci_dir):
    with BuildScript(
            os.path.join(FLAGS.build_dir, script_name),
            verbose=FLAGS.verbose,
            desc=('Docker-based build script for Triton Inference Server'
                 )) as docker_script:

        #
        # Build base image... tritonserver_buildbase
        #
        docker_script.commentln(8)
        docker_script.comment('Create Triton base build image')
        docker_script.comment(
            'This image contains all dependencies necessary to build Triton')
        docker_script.comment()

        cachefrommap = [
            'tritonserver_buildbase', 'tritonserver_buildbase_cache0',
            'tritonserver_buildbase_cache1'
        ]

        baseargs = [
            'docker', 'build', '-t', 'tritonserver_buildbase', '-f',
            os.path.join(FLAGS.build_dir, 'Dockerfile.buildbase')
        ]

        if not FLAGS.no_container_pull:
            baseargs += [
                '--pull',
            ]

        # Windows docker runs in a VM and memory needs to be specified
        # explicitly (at least for some configurations of docker).
        if target_platform() == 'windows':
            if FLAGS.container_memory:
                baseargs += ['--memory', FLAGS.container_memory]

        baseargs += ['--cache-from={}'.format(k) for k in cachefrommap]
        baseargs += ['.']

        docker_script.cwd(THIS_SCRIPT_DIR)
        docker_script.cmd(baseargs, check_exitcode=True)

        #
        # Build...
        #
        docker_script.blankln()
        docker_script.commentln(8)
        docker_script.comment('Run build in tritonserver_buildbase container')
        docker_script.comment(
            'Mount a directory into the container where the install')
        docker_script.comment('artifacts will be placed.')
        docker_script.comment()

        # Don't use '-v' to communicate the built artifacts out of the
        # build, because we want this code to work even if run within
        # Docker (i.e. docker-in-docker) and not just if run directly
        # from host.
        runargs = [
            'docker', 'run', '-w', '/workspace/build', '--name',
            'tritonserver_builder'
        ]

        if not FLAGS.no_container_interactive:
            runargs += ['-it']

        if target_platform() == 'windows':
            if FLAGS.container_memory:
                runargs += ['--memory', FLAGS.container_memory]
            runargs += [
                '-v', '\\\\.\pipe\docker_engine:\\\\.\pipe\docker_engine'
            ]
        else:
            runargs += ['-v', '/var/run/docker.sock:/var/run/docker.sock']

        runargs += ['tritonserver_buildbase']

        if target_platform() == 'windows':
            runargs += [
                'powershell.exe', '-noexit', '-File', './cmake_build.ps1'
            ]
        else:
            runargs += ['./cmake_build']

        # Remove existing tritonserver_builder container...
        if target_platform() == 'windows':
            docker_script.cmd(['docker', 'rm', 'tritonserver_builder'])
        else:
            docker_script._file.write(
                'if [ "$(docker ps -a | grep tritonserver_builder)" ]; then  docker rm tritonserver_builder; fi\n'
            )

        docker_script.cmd(runargs, check_exitcode=True)

        docker_script.cmd([
            'docker', 'cp', 'tritonserver_builder:/tmp/tritonbuild/install',
            FLAGS.build_dir
        ],
                          check_exitcode=True)
        docker_script.cmd([
            'docker', 'cp', 'tritonserver_builder:/tmp/tritonbuild/ci',
            FLAGS.build_dir
        ],
                          check_exitcode=True)

        #
        # Final image... tritonserver
        #
        docker_script.blankln()
        docker_script.commentln(8)
        docker_script.comment('Create final tritonserver image')
        docker_script.comment()

        finalargs = [
            'docker', 'build', '-t', 'tritonserver', '-f',
            os.path.join(FLAGS.build_dir, 'Dockerfile'), '.'
        ]

        docker_script.cwd(THIS_SCRIPT_DIR)
        docker_script.cmd(finalargs, check_exitcode=True)

        #
        # CI base image... tritonserver_cibase
        #
        docker_script.blankln()
        docker_script.commentln(8)
        docker_script.comment('Create CI base image')
        docker_script.comment()

        cibaseargs = [
            'docker', 'build', '-t', 'tritonserver_cibase', '-f',
            os.path.join(FLAGS.build_dir, 'Dockerfile.cibase'), '.'
        ]

        docker_script.cwd(THIS_SCRIPT_DIR)
        docker_script.cmd(cibaseargs, check_exitcode=True)


def core_build(cmake_script, repo_dir, cmake_dir, build_dir, install_dir,
               components, backends):
    repo_build_dir = os.path.join(build_dir, 'tritonserver', 'build')
    repo_install_dir = os.path.join(build_dir, 'tritonserver', 'install')

    cmake_script.commentln(8)
    cmake_script.comment('Triton core library and tritonserver executable')
    cmake_script.comment()
    cmake_script.mkdir(repo_build_dir)
    cmake_script.cwd(repo_build_dir)
    cmake_script.cmake(
        core_cmake_args(components, backends, cmake_dir, repo_install_dir))
    cmake_script.makeinstall()

    if target_platform() == 'windows':
        cmake_script.mkdir(os.path.join(install_dir, 'bin'))
        cmake_script.cp(
            os.path.join(repo_install_dir, 'bin', 'tritonserver.exe'),
            os.path.join(install_dir, 'bin'))
        cmake_script.cp(
            os.path.join(repo_install_dir, 'bin', 'tritonserver.dll'),
            os.path.join(install_dir, 'bin'))
    else:
        cmake_script.mkdir(os.path.join(install_dir, 'bin'))
        cmake_script.cp(os.path.join(repo_install_dir, 'bin', 'tritonserver'),
                        os.path.join(install_dir, 'bin'))
        cmake_script.mkdir(os.path.join(install_dir, 'lib'))
        cmake_script.cp(
            os.path.join(repo_install_dir, 'lib', 'libtritonserver.so'),
            os.path.join(install_dir, 'lib'))

    cmake_script.mkdir(os.path.join(install_dir, 'include', 'triton'))
    cmake_script.cpdir(
        os.path.join(repo_install_dir, 'include', 'triton', 'core'),
        os.path.join(install_dir, 'include', 'triton', 'core'))

    cmake_script.cp(os.path.join(repo_dir, 'LICENSE'), install_dir)
    cmake_script.cp(os.path.join(repo_dir, 'TRITON_VERSION'), install_dir)

    # If requested, package the source code for all OSS used to build
    # For windows, Triton is not delivered as a container so skip for
    # windows platform.
    if target_platform() != 'windows':
        if (not FLAGS.no_container_build) and (not FLAGS.no_core_build) and (
                not FLAGS.no_container_source):
            cmake_script.mkdir(os.path.join(install_dir, 'third-party-src'))
            cmake_script.cwd(repo_build_dir)
            cmake_script.tar(
                'third-party-src',
                os.path.join(install_dir, 'third-party-src', 'src.tar.gz'))
            cmake_script.cp(
                os.path.join(repo_dir, 'docker', 'README.third-party-src'),
                os.path.join(install_dir, 'third-party-src', 'README'))

    cmake_script.comment()
    cmake_script.comment('end Triton core library and tritonserver executable')
    cmake_script.commentln(8)
    cmake_script.blankln()


def backend_build(be, cmake_script, tag, build_dir, install_dir,
                  github_organization, images, components, library_paths):
    repo_build_dir = os.path.join(build_dir, be, 'build')
    repo_install_dir = os.path.join(build_dir, be, 'install')

    cmake_script.commentln(8)
    cmake_script.comment(f'\'{be}\' backend')
    cmake_script.comment('Delete this section to remove backend from build')
    cmake_script.comment()
    cmake_script.mkdir(build_dir)
    cmake_script.cwd(build_dir)
    cmake_script.gitclone(backend_repo(be), tag, be, github_organization)

    cmake_script.mkdir(repo_build_dir)
    cmake_script.cwd(repo_build_dir)
    cmake_script.cmake(
        backend_cmake_args(images, components, be, repo_install_dir,
                           library_paths))
    cmake_script.makeinstall()

    cmake_script.mkdir(os.path.join(install_dir, 'backends'))
    cmake_script.rmdir(os.path.join(install_dir, 'backends', be))
    cmake_script.cpdir(os.path.join(repo_install_dir, 'backends', be),
                       os.path.join(install_dir, 'backends'))

    cmake_script.comment()
    cmake_script.comment(f'end \'{be}\' backend')
    cmake_script.commentln(8)
    cmake_script.blankln()


def repo_agent_build(ra, cmake_script, build_dir, install_dir, repoagent_repo,
                     repoagents):
    repo_build_dir = os.path.join(build_dir, ra, 'build')
    repo_install_dir = os.path.join(build_dir, ra, 'install')

    cmake_script.commentln(8)
    cmake_script.comment(f'\'{ra}\' repository agent')
    cmake_script.comment(
        'Delete this section to remove repository agent from build')
    cmake_script.comment()
    cmake_script.mkdir(build_dir)
    cmake_script.cwd(build_dir)
    cmake_script.gitclone(repoagent_repo(ra), repoagents[ra], ra,
                          FLAGS.github_organization)

    cmake_script.mkdir(repo_build_dir)
    cmake_script.cwd(repo_build_dir)
    cmake_script.cmake(
        repoagent_cmake_args(images, components, ra, repo_install_dir))
    cmake_script.makeinstall()

    cmake_script.mkdir(os.path.join(install_dir, 'repoagents'))
    cmake_script.rmdir(os.path.join(install_dir, 'repoagents', ra))
    cmake_script.cpdir(os.path.join(repo_install_dir, 'repoagents', ra),
                       os.path.join(install_dir, 'repoagents'))
    cmake_script.comment()
    cmake_script.comment(f'end \'{ra}\' repository agent')
    cmake_script.commentln(8)
    cmake_script.blankln()


def cibase_build(cmake_script, repo_dir, cmake_dir, build_dir, install_dir,
                 ci_dir, backends):
    repo_install_dir = os.path.join(build_dir, 'tritonserver', 'install')

    cmake_script.commentln(8)
    cmake_script.comment('Collect Triton CI artifacts')
    cmake_script.comment()

    cmake_script.mkdir(ci_dir)

    # On windows we are not yet using a CI/QA docker image for
    # testing, so don't do anything...
    if target_platform() == 'windows':
        return

    # The core build produces some artifacts that are needed for CI
    # testing, so include those in the install.
    cmake_script.cpdir(os.path.join(repo_dir, 'qa'), ci_dir)
    cmake_script.cpdir(os.path.join(repo_dir, 'deploy'), ci_dir)
    cmake_script.mkdir(os.path.join(ci_dir, 'docs'))
    cmake_script.cpdir(os.path.join(repo_dir, 'docs', 'examples'),
                       os.path.join(ci_dir, 'docs'))
    cmake_script.mkdir(os.path.join(ci_dir, 'src', 'test'))
    cmake_script.cpdir(os.path.join(repo_dir, 'src', 'test', 'models'),
                       os.path.join(ci_dir, 'src', 'test'))
    # Skip copying the artifacts in the bin and lib as those directories will
    # be missing when the core build is not enabled.
    if not FLAGS.no_core_build:
        cmake_script.cpdir(os.path.join(repo_install_dir, 'bin'), ci_dir)
        cmake_script.mkdir(os.path.join(ci_dir, 'lib'))
        cmake_script.cp(
            os.path.join(repo_install_dir, 'lib',
                         'libtritonrepoagent_relocation.so'),
            os.path.join(ci_dir, 'lib'))

    # Some of the backends are needed for CI testing
    cmake_script.mkdir(os.path.join(ci_dir, 'backends'))
    for be in ('identity', 'repeat', 'square'):
        be_install_dir = os.path.join(build_dir, be, 'install', 'backends', be)
        if target_platform() == 'windows':
            cmake_script.cmd(f'if (Test-Path -Path {be_install_dir}) {{')
        else:
            cmake_script.cmd(f'if [[ -e {be_install_dir} ]]; then')
        cmake_script.cpdir(be_install_dir, os.path.join(ci_dir, 'backends'))
        cmake_script.cmd('}' if target_platform() == 'windows' else 'fi')

    # Some of the unit-test built backends are needed for CI testing
    cmake_script.mkdir(
        os.path.join(ci_dir, 'tritonbuild', 'tritonserver', 'backends'))
    for be in ('query', 'implicit_state', 'sequence', 'dyna_sequence',
               'distributed_addsub'):
        be_install_dir = os.path.join(repo_install_dir, 'backends', be)
        if target_platform() == 'windows':
            cmake_script.cmd(f'if (Test-Path -Path {be_install_dir}) {{')
        else:
            cmake_script.cmd(f'if [[ -e {be_install_dir} ]]; then')
        cmake_script.cpdir(
            be_install_dir,
            os.path.join(ci_dir, 'tritonbuild', 'tritonserver', 'backends'))
        cmake_script.cmd('}' if target_platform() == 'windows' else 'fi')

    # The onnxruntime_backend build produces some artifacts that
    # are needed for CI testing.
    if 'onnxruntime' in backends:
        ort_install_dir = os.path.join(build_dir, 'onnxruntime', 'install')
        cmake_script.mkdir(os.path.join(ci_dir, 'qa', 'L0_custom_ops'))
        cmake_script.cp(
            os.path.join(ort_install_dir, 'test', 'libcustom_op_library.so'),
            os.path.join(ci_dir, 'qa', 'L0_custom_ops'))
        cmake_script.cp(
            os.path.join(ort_install_dir, 'test', 'custom_op_test.onnx'),
            os.path.join(ci_dir, 'qa', 'L0_custom_ops'))

    # Need the build area for some backends so that they can be
    # rebuilt with specific options.
    cmake_script.mkdir(os.path.join(ci_dir, 'tritonbuild'))
    for be in ('identity', 'python'):
        if be in backends:
            cmake_script.rmdir(os.path.join(build_dir, be, 'build'))
            cmake_script.rmdir(os.path.join(build_dir, be, 'install'))
            cmake_script.cpdir(os.path.join(build_dir, be),
                               os.path.join(ci_dir, 'tritonbuild'))

    cmake_script.comment()
    cmake_script.comment('end Triton CI artifacts')
    cmake_script.commentln(8)
    cmake_script.blankln()


def finalize_build(cmake_script, install_dir, ci_dir):
    cmake_script.cmd(f'chmod -R a+rw {install_dir}')
    cmake_script.cmd(f'chmod -R a+rw {ci_dir}')


def enable_all():
    if target_platform() != 'windows':
        all_backends = [
            'ensemble', 'identity', 'square', 'repeat', 'tensorflow1',
            'tensorflow2', 'onnxruntime', 'python', 'dali', 'pytorch',
            'openvino', 'fil', 'tensorrt'
        ]
        all_repoagents = ['checksum']
        all_filesystems = ['gcs', 's3', 'azure_storage']
        all_endpoints = ['http', 'grpc', 'sagemaker', 'vertex-ai']

        FLAGS.enable_logging = True
        FLAGS.enable_stats = True
        FLAGS.enable_metrics = True
        FLAGS.enable_gpu_metrics = True
        FLAGS.enable_cpu_metrics = True
        FLAGS.enable_tracing = True
        FLAGS.enable_nvtx = True
        FLAGS.enable_gpu = True
    else:
        all_backends = [
            'ensemble', 'identity', 'square', 'repeat', 'onnxruntime',
            'openvino', 'tensorrt'
        ]
        all_repoagents = ['checksum']
        all_filesystems = []
        all_endpoints = ['http', 'grpc']

        FLAGS.enable_logging = True
        FLAGS.enable_stats = True
        FLAGS.enable_tracing = True
        FLAGS.enable_gpu = True

    requested_backends = []
    for be in FLAGS.backend:
        parts = be.split(':')
        requested_backends += [parts[0]]
    for be in all_backends:
        if be not in requested_backends:
            FLAGS.backend += [be]

    requested_repoagents = []
    for ra in FLAGS.repoagent:
        parts = ra.split(':')
        requested_repoagents += [parts[0]]
    for ra in all_repoagents:
        if ra not in requested_repoagents:
            FLAGS.repoagent += [ra]

    for fs in all_filesystems:
        if fs not in FLAGS.filesystem:
            FLAGS.filesystem += [fs]

    for ep in all_endpoints:
        if ep not in FLAGS.endpoint:
            FLAGS.endpoint += [ep]


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    group_qv = parser.add_mutually_exclusive_group()
    group_qv.add_argument('-q',
                          '--quiet',
                          action="store_true",
                          required=False,
                          help='Disable console output.')
    group_qv.add_argument('-v',
                          '--verbose',
                          action="store_true",
                          required=False,
                          help='Enable verbose output.')

    parser.add_argument(
        '--dryrun',
        action="store_true",
        required=False,
        help='Output the build scripts, but do not perform build.')
    parser.add_argument('--no-container-build',
                        action="store_true",
                        required=False,
                        help='Do not use Docker container for build.')
    parser.add_argument(
        '--no-container-interactive',
        action="store_true",
        required=False,
        help=
        'Do not use -it argument to "docker run" when performing container build.'
    )
    parser.add_argument(
        '--no-container-pull',
        action="store_true",
        required=False,
        help='Do not use Docker --pull argument when building container.')
    parser.add_argument(
        '--container-memory',
        default=None,
        required=False,
        help='Value for Docker --memory argument. Used only for windows builds.'
    )
    parser.add_argument(
        '--target-platform',
        required=False,
        default=None,
        help=
        'Target platform for build, can be "linux", "windows" or "jetpack". If not specified, build targets the current platform.'
    )
    parser.add_argument(
        '--target-machine',
        required=False,
        default=None,
        help=
        'Target machine/architecture for build. If not specified, build targets the current machine/architecture.'
    )

    parser.add_argument('--build-id',
                        type=str,
                        required=False,
                        help='Build ID associated with the build.')
    parser.add_argument('--build-sha',
                        type=str,
                        required=False,
                        help='SHA associated with the build.')
    parser.add_argument(
        '--build-dir',
        type=str,
        required=False,
        help=
        'Build directory. All repo clones and builds will be performed in this directory.'
    )
    parser.add_argument(
        '--install-dir',
        type=str,
        required=False,
        default=None,
        help='Install directory, default is <builddir>/opt/tritonserver.')
    parser.add_argument(
        '--cmake-dir',
        type=str,
        required=False,
        help='Directory containing the CMakeLists.txt file for Triton server.')
    parser.add_argument(
        '--tmp-dir',
        type=str,
        required=False,
        default='/tmp',
        help=
        'Temporary directory used for building inside docker. Default is /tmp.')
    parser.add_argument(
        '--library-paths',
        action='append',
        required=False,
        default=None,
        help=
        'Specify library paths for respective backends in build as <backend-name>[:<library_path>].'
    )
    parser.add_argument(
        '--build-type',
        required=False,
        default='Release',
        help=
        'Build type, one of "Release", "Debug", "RelWithDebInfo" or "MinSizeRel". Default is "Release".'
    )
    parser.add_argument(
        '-j',
        '--build-parallel',
        type=int,
        required=False,
        default=None,
        help='Build parallelism. Defaults to 2 * number-of-cores.')

    parser.add_argument(
        '--github-organization',
        type=str,
        required=False,
        default='https://github.com/triton-inference-server',
        help=
        'The GitHub organization containing the repos used for the build. Defaults to "https://github.com/triton-inference-server".'
    )
    parser.add_argument(
        '--version',
        type=str,
        required=False,
        help=
        'The Triton version. If not specified defaults to the value in the TRITON_VERSION file.'
    )
    parser.add_argument(
        '--container-version',
        type=str,
        required=False,
        help=
        'The Triton container version to build. If not specified the container version will be chosen automatically based on --version value.'
    )
    parser.add_argument(
        '--upstream-container-version',
        type=str,
        required=False,
        help=
        'The upstream container version to use for the build. If not specified the upstream container version will be chosen automatically based on --version value.'
    )
    parser.add_argument(
        '--container-prebuild-command',
        type=str,
        required=False,
        help=
        'When performing a container build, this command will be executed within the container just before the build it performed.'
    )
    parser.add_argument(
        '--no-container-source',
        action="store_true",
        required=False,
        help='Do not include OSS source code in Docker container.')
    parser.add_argument(
        '--image',
        action='append',
        required=False,
        help=
        'Use specified Docker image in build as <image-name>,<full-image-name>. <image-name> can be "base", "gpu-base", "tensorflow1", "tensorflow2", or "pytorch".'
    )

    parser.add_argument(
        '--enable-all',
        action="store_true",
        required=False,
        help=
        'Enable all standard released Triton features, backends, repository agents, endpoints and file systems.'
    )
    parser.add_argument('--enable-logging',
                        action="store_true",
                        required=False,
                        help='Enable logging.')
    parser.add_argument('--enable-stats',
                        action="store_true",
                        required=False,
                        help='Enable statistics collection.')
    parser.add_argument('--enable-metrics',
                        action="store_true",
                        required=False,
                        help='Enable metrics reporting.')
    parser.add_argument('--enable-gpu-metrics',
                        action="store_true",
                        required=False,
                        help='Include GPU metrics in reported metrics.')
    parser.add_argument('--enable-cpu-metrics',
                        action="store_true",
                        required=False,
                        help='Include CPU metrics in reported metrics.')
    parser.add_argument('--enable-tracing',
                        action="store_true",
                        required=False,
                        help='Enable tracing.')
    parser.add_argument('--enable-nvtx',
                        action="store_true",
                        required=False,
                        help='Enable NVTX.')
    parser.add_argument('--enable-gpu',
                        action="store_true",
                        required=False,
                        help='Enable GPU support.')
    parser.add_argument('--enable-mali-gpu',
                        action="store_true",
                        required=False,
                        help='Enable ARM MALI GPU support.')
    parser.add_argument(
        '--min-compute-capability',
        type=str,
        required=False,
        default='6.0',
        help='Minimum CUDA compute capability supported by server.')

    parser.add_argument(
        '--endpoint',
        action='append',
        required=False,
        help=
        'Include specified endpoint in build. Allowed values are "grpc", "http", "vertex-ai" and "sagemaker".'
    )
    parser.add_argument(
        '--filesystem',
        action='append',
        required=False,
        help=
        'Include specified filesystem in build. Allowed values are "gcs", "azure_storage" and "s3".'
    )
    parser.add_argument(
        '--no-core-build',
        action="store_true",
        required=False,
        help='Do not build Triton core shared library or executable.')
    parser.add_argument(
        '--backend',
        action='append',
        required=False,
        help=
        'Include specified backend in build as <backend-name>[:<repo-tag>]. If <repo-tag> starts with "pull/" then it refers to a pull-request reference, otherwise <repo-tag> indicates the git tag/branch to use for the build. If the version is non-development then the default <repo-tag> is the release branch matching the container version (e.g. version YY.MM -> branch rYY.MM); otherwise the default <repo-tag> is "main" (e.g. version YY.MMdev -> branch main).'
    )
    parser.add_argument(
        '--repo-tag',
        action='append',
        required=False,
        help=
        'The version of a component to use in the build as <component-name>:<repo-tag>. <component-name> can be "common", "core", "backend" or "thirdparty". <repo-tag> indicates the git tag/branch to use for the build. Currently <repo-tag> does not support pull-request reference. If the version is non-development then the default <repo-tag> is the release branch matching the container version (e.g. version YY.MM -> branch rYY.MM); otherwise the default <repo-tag> is "main" (e.g. version YY.MMdev -> branch main).'
    )
    parser.add_argument(
        '--repoagent',
        action='append',
        required=False,
        help=
        'Include specified repo agent in build as <repoagent-name>[:<repo-tag>]. If <repo-tag> starts with "pull/" then it refers to a pull-request reference, otherwise <repo-tag> indicates the git tag/branch to use for the build. If the version is non-development then the default <repo-tag> is the release branch matching the container version (e.g. version YY.MM -> branch rYY.MM); otherwise the default <repo-tag> is "main" (e.g. version YY.MMdev -> branch main).'
    )
    parser.add_argument(
        '--no-force-clone',
        action="store_true",
        default=False,
        help='Do not create fresh clones of repos that have already been cloned.'
    )
    parser.add_argument(
        '--extra-core-cmake-arg',
        action='append',
        required=False,
        help=
        'Extra CMake argument as <name>=<value>. The argument is passed to CMake as -D<name>=<value> and is included after all CMake arguments added by build.py for the core builds.'
    )
    parser.add_argument(
        '--override-core-cmake-arg',
        action='append',
        required=False,
        help=
        'Override specified CMake argument in the build as <name>=<value>. The argument is passed to CMake as -D<name>=<value>. This flag only impacts CMake arguments that are used by build.py. To unconditionally add a CMake argument to the core build use --extra-core-cmake-arg.'
    )
    parser.add_argument(
        '--extra-backend-cmake-arg',
        action='append',
        required=False,
        help=
        'Extra CMake argument for a backend build as <backend>:<name>=<value>. The argument is passed to CMake as -D<name>=<value> and is included after all CMake arguments added by build.py for the backend.'
    )
    parser.add_argument(
        '--override-backend-cmake-arg',
        action='append',
        required=False,
        help=
        'Override specified backend CMake argument in the build as <backend>:<name>=<value>. The argument is passed to CMake as -D<name>=<value>. This flag only impacts CMake arguments that are used by build.py. To unconditionally add a CMake argument to the backend build use --extra-backend-cmake-arg.'
    )

    FLAGS = parser.parse_args()

    if FLAGS.image is None:
        FLAGS.image = []
    if FLAGS.repo_tag is None:
        FLAGS.repo_tag = []
    if FLAGS.backend is None:
        FLAGS.backend = []
    if FLAGS.endpoint is None:
        FLAGS.endpoint = []
    if FLAGS.filesystem is None:
        FLAGS.filesystem = []
    if FLAGS.repoagent is None:
        FLAGS.repoagent = []
    if FLAGS.library_paths is None:
        FLAGS.library_paths = []
    if FLAGS.extra_core_cmake_arg is None:
        FLAGS.extra_core_cmake_arg = []
    if FLAGS.override_core_cmake_arg is None:
        FLAGS.override_core_cmake_arg = []
    if FLAGS.override_backend_cmake_arg is None:
        FLAGS.override_backend_cmake_arg = []
    if FLAGS.extra_backend_cmake_arg is None:
        FLAGS.extra_backend_cmake_arg = []

    # if --enable-all is specified, then update FLAGS to enable all
    # settings, backends, repo-agents, file systems, endpoints, etc.
    if FLAGS.enable_all:
        enable_all()

    # When doing a docker build, --build-dir, --install-dir and
    # --cmake-dir must not be set. We will use the build/ subdir
    # within the server/ repo that contains this build.py script for
    # --build-dir. If not doing a docker build, --build-dir must be
    # set.
    if FLAGS.no_container_build:
        if FLAGS.build_dir is None:
            fail('--no-container-build requires --build-dir')
        if FLAGS.install_dir is None:
            FLAGS.install_dir = os.path.join(FLAGS.build_dir, "opt",
                                             "tritonserver")
        if FLAGS.cmake_dir is None:
            FLAGS.cmake_dir = THIS_SCRIPT_DIR
    else:
        if FLAGS.build_dir is not None:
            fail('--build-dir must not be set for container-based build')
        if FLAGS.install_dir is not None:
            fail('--install-dir must not be set for container-based build')
        if FLAGS.cmake_dir is not None:
            fail('--cmake-dir must not be set for container-based build')
        FLAGS.build_dir = os.path.join(THIS_SCRIPT_DIR, 'build')

    # Determine the versions. Start with Triton version, if --version
    # is not explicitly specified read from TRITON_VERSION file.
    if FLAGS.version is None:
        with open(os.path.join(THIS_SCRIPT_DIR, 'TRITON_VERSION'),
                  "r") as vfile:
            FLAGS.version = vfile.readline().strip()

    if FLAGS.build_parallel is None:
        FLAGS.build_parallel = multiprocessing.cpu_count() * 2

    log('Building Triton Inference Server')
    log('platform {}'.format(target_platform()))
    log('machine {}'.format(target_machine()))
    log('version {}'.format(FLAGS.version))
    log('build dir {}'.format(FLAGS.build_dir))
    log('install dir {}'.format(FLAGS.install_dir))
    log('cmake dir {}'.format(FLAGS.cmake_dir))

    # Determine the default repo-tag that should be used for images,
    # backends and repo-agents if a repo-tag is not given
    # explicitly. For release branches we use the release branch as
    # the default, otherwise we use 'main'.
    default_repo_tag = 'main'
    cver = FLAGS.container_version
    if cver is None:
        if FLAGS.version not in TRITON_VERSION_MAP:
            fail(
                'unable to determine default repo-tag, container version not known for {}'
                .format(FLAGS.version))
        cver = TRITON_VERSION_MAP[FLAGS.version][0]
    if not cver.endswith('dev'):
        default_repo_tag = 'r' + cver
    log('default repo-tag: {}'.format(default_repo_tag))

    # For other versions use the TRITON_VERSION_MAP unless explicitly
    # given.
    FLAGS.container_version, FLAGS.upstream_container_version = container_versions(
        FLAGS.version, FLAGS.container_version,
        FLAGS.upstream_container_version)

    log('container version {}'.format(FLAGS.container_version))
    log('upstream container version {}'.format(
        FLAGS.upstream_container_version))

    for ep in FLAGS.endpoint:
        log(f'endpoint "{ep}"')
    for fs in FLAGS.filesystem:
        log(f'filesystem "{fs}"')

    # Initialize map of backends to build and repo-tag for each.
    backends = {}
    for be in FLAGS.backend:
        parts = be.split(':')
        if len(parts) == 1:
            parts.append(default_repo_tag)
        log('backend "{}" at tag/branch "{}"'.format(parts[0], parts[1]))
        backends[parts[0]] = parts[1]

    # Initialize map of repo agents to build and repo-tag for each.
    repoagents = {}
    for be in FLAGS.repoagent:
        parts = be.split(':')
        if len(parts) == 1:
            parts.append(default_repo_tag)
        log('repoagent "{}" at tag/branch "{}"'.format(parts[0], parts[1]))
        repoagents[parts[0]] = parts[1]

    # Initialize map of docker images.
    images = {}
    for img in FLAGS.image:
        parts = img.split(',')
        fail_if(
            len(parts) != 2,
            '--image must specify <image-name>,<full-image-registry>')
        # REMOVEME: ONNXRUNTIME 1.13.1 build is failing with cuda 12. Hence, using
        # cuda 11.8 container as build image for ORT.
        fail_if(
            parts[0] not in [
                'base', 'gpu-base', 'pytorch', 'tensorflow1', 'tensorflow2',
                'ort-base'
            ], 'unsupported value for --image')
        log('image "{}": "{}"'.format(parts[0], parts[1]))
        images[parts[0]] = parts[1]

    # Initialize map of library paths for each backend.
    library_paths = {}
    for lpath in FLAGS.library_paths:
        parts = lpath.split(':')
        if len(parts) == 2:
            log('backend "{}" library path "{}"'.format(parts[0], parts[1]))
            library_paths[parts[0]] = parts[1]

    # Parse any explicitly specified cmake arguments
    for cf in FLAGS.extra_core_cmake_arg:
        parts = cf.split('=')
        fail_if(
            len(parts) != 2,
            '--extra-core-cmake-arg must specify <name>=<value>')
        log('CMake core extra "-D{}={}"'.format(parts[0], parts[1]))
        EXTRA_CORE_CMAKE_FLAGS[parts[0]] = parts[1]

    for cf in FLAGS.override_core_cmake_arg:
        parts = cf.split('=')
        fail_if(
            len(parts) != 2,
            '--override-core-cmake-arg must specify <name>=<value>')
        log('CMake core override "-D{}={}"'.format(parts[0], parts[1]))
        OVERRIDE_CORE_CMAKE_FLAGS[parts[0]] = parts[1]

    for cf in FLAGS.extra_backend_cmake_arg:
        parts = cf.split(':', 1)
        fail_if(
            len(parts) != 2,
            '--extra-backend-cmake-arg must specify <backend>:<name>=<value>')
        be = parts[0]
        parts = parts[1].split('=', 1)
        fail_if(
            len(parts) != 2,
            '--extra-backend-cmake-arg must specify <backend>:<name>=<value>')
        fail_if(
            be not in backends,
            '--extra-backend-cmake-arg specifies backend "{}" which is not included in build'
            .format(be))
        log('backend "{}" CMake extra "-D{}={}"'.format(be, parts[0], parts[1]))
        if be not in EXTRA_BACKEND_CMAKE_FLAGS:
            EXTRA_BACKEND_CMAKE_FLAGS[be] = {}
        EXTRA_BACKEND_CMAKE_FLAGS[be][parts[0]] = parts[1]

    for cf in FLAGS.override_backend_cmake_arg:
        parts = cf.split(':', 1)
        fail_if(
            len(parts) != 2,
            '--override-backend-cmake-arg must specify <backend>:<name>=<value>'
        )
        be = parts[0]
        parts = parts[1].split('=', 1)
        fail_if(
            len(parts) != 2,
            '--override-backend-cmake-arg must specify <backend>:<name>=<value>'
        )
        fail_if(
            be not in backends,
            '--override-backend-cmake-arg specifies backend "{}" which is not included in build'
            .format(be))
        log('backend "{}" CMake override "-D{}={}"'.format(
            be, parts[0], parts[1]))
        if be not in OVERRIDE_BACKEND_CMAKE_FLAGS:
            OVERRIDE_BACKEND_CMAKE_FLAGS[be] = {}
        OVERRIDE_BACKEND_CMAKE_FLAGS[be][parts[0]] = parts[1]

    # Initialize map of common components and repo-tag for each.
    components = {
        'common': default_repo_tag,
        'core': default_repo_tag,
        'backend': default_repo_tag,
        'thirdparty': default_repo_tag
    }
    for be in FLAGS.repo_tag:
        parts = be.split(':')
        fail_if(
            len(parts) != 2,
            '--repo-tag must specify <component-name>:<repo-tag>')
        fail_if(
            parts[0] not in components,
            '--repo-tag <component-name> must be "common", "core", "backend", or "thirdparty"'
        )
        components[parts[0]] = parts[1]
    for c in components:
        log('component "{}" at tag/branch "{}"'.format(c, components[c]))

    # Set the build, install, and cmake directories to use for the
    # generated build scripts and Dockerfiles. If building without
    # Docker, these are the directories specified on the cmdline. If
    # building with Docker, we change these to be directories within
    # FLAGS.tmp_dir inside the Docker container.
    script_repo_dir = THIS_SCRIPT_DIR
    script_build_dir = FLAGS.build_dir
    script_install_dir = script_ci_dir = FLAGS.install_dir
    script_cmake_dir = FLAGS.cmake_dir
    if not FLAGS.no_container_build:
        # FLAGS.tmp_dir may be specified with "\" on Windows, adjust
        # to "/" for docker usage.
        script_build_dir = os.path.normpath(
            os.path.join(FLAGS.tmp_dir, 'tritonbuild').replace("\\", "/"))
        script_install_dir = os.path.normpath(
            os.path.join(script_build_dir, 'install'))
        script_ci_dir = os.path.normpath(os.path.join(script_build_dir, 'ci'))
        if target_platform() == 'windows':
            script_repo_dir = script_cmake_dir = os.path.normpath(
                'c:/workspace')
        else:
            script_repo_dir = script_cmake_dir = '/workspace'

    script_name = 'cmake_build'
    if target_platform() == 'windows':
        script_name += '.ps1'

    # Write the build script that invokes cmake for the core, backends, and repo-agents.
    pathlib.Path(FLAGS.build_dir).mkdir(parents=True, exist_ok=True)
    with BuildScript(
            os.path.join(FLAGS.build_dir, script_name),
            verbose=FLAGS.verbose,
            desc=('Build script for Triton Inference Server')) as cmake_script:

        # Run the container pre-build command if the cmake build is
        # being done within the build container.
        if not FLAGS.no_container_build and FLAGS.container_prebuild_command:
            cmake_script.cmd(FLAGS.container_prebuild_command,
                             check_exitcode=True)
            cmake_script.blankln()

        # Commands to build the core shared library and the server executable.
        if not FLAGS.no_core_build:
            core_build(cmake_script, script_repo_dir, script_cmake_dir,
                       script_build_dir, script_install_dir, components,
                       backends)

        # Commands to build each backend...
        for be in backends:
            # Core backends are not built separately from core so skip...
            if (be in CORE_BACKENDS):
                continue

            # If armnn_tflite backend, source from external repo for git clone
            if be == 'armnn_tflite':
                github_organization = 'https://gitlab.com/arm-research/smarter/'
            else:
                github_organization = FLAGS.github_organization

            backend_build(be, cmake_script, backends[be], script_build_dir,
                          script_install_dir, github_organization, images,
                          components, library_paths)

        # Commands to build each repo agent...
        for ra in repoagents:
            repo_agent_build(ra, cmake_script, script_build_dir,
                             script_install_dir, repoagent_repo, repoagents)

        # Commands needed only when building with Docker...
        if not FLAGS.no_container_build:
            # Commands to collect all the build artifacts needed for CI
            # testing.
            cibase_build(cmake_script, script_repo_dir, script_cmake_dir,
                         script_build_dir, script_install_dir, script_ci_dir,
                         backends)

            # When building with Docker the install and ci artifacts
            # written to the build-dir while running the docker container
            # may have root ownership, so give them permissions to be
            # managed by all users on the host system.
            if target_platform() != 'windows':
                finalize_build(cmake_script, script_install_dir, script_ci_dir)

    # If --no-container-build is not specified then we perform the
    # actual build within a docker container and from that create the
    # final tritonserver docker image. For the build we need to
    # generate a few Dockerfiles and a top-level script that drives
    # the build process.
    if not FLAGS.no_container_build:
        script_name = 'docker_build'
        if target_platform() == 'windows':
            script_name += '.ps1'

        create_build_dockerfiles(script_build_dir, images, backends, repoagents,
                                 FLAGS.endpoint)
        create_docker_build_script(script_name, script_install_dir,
                                   script_ci_dir)

    # In not dry-run, execute the script to perform the build...  If a
    # container-based build is requested use 'docker_build' script,
    # otherwise build directly on this system using cmake script.
    if not FLAGS.dryrun:
        if target_platform() == 'windows':
            p = subprocess.Popen(
                ['powershell.exe', '-noexit', '-File', f'./{script_name}'],
                cwd=FLAGS.build_dir)
        else:
            p = subprocess.Popen([f'./{script_name}'], cwd=FLAGS.build_dir)
        p.wait()
        fail_if(p.returncode != 0, 'build failed')


# compose.py

#!/usr/bin/env python3
# Copyright 2021-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import argparse
import os
import platform
import subprocess
import sys

FLAGS = None


#### helper functions
def log(msg, force=False):
    if force or not FLAGS.quiet:
        try:
            print(msg, file=sys.stderr)
        except Exception:
            print('<failed to log>', file=sys.stderr)


def log_verbose(msg):
    if FLAGS.verbose:
        log(msg, force=True)


def fail(msg):
    print('error: {}'.format(msg), file=sys.stderr)
    sys.exit(1)


def fail_if(p, msg):
    if p:
        fail(msg)


def start_dockerfile(ddir, images, argmap, dockerfile_name, backends):
    # Set enviroment variables, set default user and install dependencies
    df = '''
#
# Multistage build.
#
ARG TRITON_VERSION={}
ARG TRITON_CONTAINER_VERSION={}

FROM {} AS full
'''.format(argmap['TRITON_VERSION'], argmap['TRITON_CONTAINER_VERSION'],
           images["full"])

    # PyTorch, TensorFlow 1 and TensorFlow 2 backends need extra CUDA and other
    # dependencies during runtime that are missing in the CPU-only base container.
    # These dependencies must be copied from the Triton Min image.
    if not FLAGS.enable_gpu and (('pytorch' in backends) or
                                 ('tensorflow1' in backends) or
                                 ('tensorflow2' in backends)):
        df += '''
FROM {} AS min_container

'''.format(images["gpu-min"])

    df += '''
FROM {}
'''.format(images["min"])

    import build
    df += build.dockerfile_prepare_container_linux(argmap, backends,
                                                   FLAGS.enable_gpu,
                                                   platform.machine().lower())
    # Copy over files
    df += '''
WORKDIR /opt/tritonserver
COPY --chown=1000:1000 --from=full /opt/tritonserver/LICENSE .
COPY --chown=1000:1000 --from=full /opt/tritonserver/TRITON_VERSION .
COPY --chown=1000:1000 --from=full /opt/tritonserver/NVIDIA_Deep_Learning_Container_License.pdf .
COPY --chown=1000:1000 --from=full /opt/tritonserver/bin bin/
COPY --chown=1000:1000 --from=full /opt/tritonserver/lib lib/
COPY --chown=1000:1000 --from=full /opt/tritonserver/include include/
'''
    with open(os.path.join(ddir, dockerfile_name), "w") as dfile:
        dfile.write(df)


def add_requested_backends(ddir, dockerfile_name, backends):
    df = "# Copying over backends \n"
    for backend in backends:
        df += '''COPY --chown=1000:1000 --from=full /opt/tritonserver/backends/{} /opt/tritonserver/backends/{}
'''.format(backend, backend)
    if len(backends) > 0:
        df += '''
# Top-level /opt/tritonserver/backends not copied so need to explicitly set permissions here
RUN chown triton-server:triton-server /opt/tritonserver/backends
'''
    with open(os.path.join(ddir, dockerfile_name), "a") as dfile:
        dfile.write(df)


def add_requested_repoagents(ddir, dockerfile_name, repoagents):
    df = "#  Copying over repoagents \n"
    for ra in repoagents:
        df += '''COPY --chown=1000:1000 --from=full /opt/tritonserver/repoagents/{} /opt/tritonserver/repoagents/{}
'''.format(ra, ra)
    if len(repoagents) > 0:
        df += '''
# Top-level /opt/tritonserver/repoagents not copied so need to explicitly set permissions here
RUN chown triton-server:triton-server /opt/tritonserver/repoagents
'''
    with open(os.path.join(ddir, dockerfile_name), "a") as dfile:
        dfile.write(df)


def end_dockerfile(ddir, dockerfile_name, argmap):
    # Install additional dependencies
    df = ""
    if argmap['SAGEMAKER_ENDPOINT']:
        df += '''
LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
COPY --chown=1000:1000 --from=full /usr/bin/serve /usr/bin/.
'''
    with open(os.path.join(ddir, dockerfile_name), "a") as dfile:
        dfile.write(df)


def build_docker_image(ddir, dockerfile_name, container_name):
    # Create container with docker build
    p = subprocess.Popen(['docker', 'build', '-t', container_name, '-f', \
        os.path.join(ddir, dockerfile_name), '.'])
    p.wait()
    fail_if(p.returncode != 0, 'docker build {} failed'.format(container_name))


def get_container_version_if_not_specified():
    if FLAGS.container_version is None:
        # Read from TRITON_VERSION file in server repo to determine version
        with open('TRITON_VERSION', "r") as vfile:
            version = vfile.readline().strip()
        import build
        _, FLAGS.container_version = build.container_versions(
            version, None, FLAGS.container_version)
        log('version {}'.format(version))
    log('using container version {}'.format(FLAGS.container_version))


def create_argmap(images, skip_pull):
    # Extract information from upstream build and create map other functions can
    # use
    full_docker_image = images["full"]
    min_docker_image = images["min"]
    enable_gpu = FLAGS.enable_gpu
    # Docker inspect enviroment variables
    base_run_args = ['docker', 'inspect', '-f']
    import re  # parse all PATH enviroment variables

    # first pull docker images
    if not skip_pull:
        log("pulling container:{}".format(full_docker_image))
        p = subprocess.run(['docker', 'pull', full_docker_image])
        fail_if(
            p.returncode != 0, 'docker pull container {} failed, {}'.format(
                full_docker_image, p.stderr))
    if enable_gpu:
        if not skip_pull:
            pm = subprocess.run(['docker', 'pull', min_docker_image])
            fail_if(
                pm.returncode != 0 and not skip_pull,
                'docker pull container {} failed, {}'.format(
                    min_docker_image, pm.stderr))
        pm_path = subprocess.run(base_run_args + [
            '{{range $index, $value := .Config.Env}}{{$value}} {{end}}',
            min_docker_image
        ],
                                 capture_output=True,
                                 text=True)
        fail_if(
            pm_path.returncode != 0,
            'docker inspect to find triton enviroment variables for min container failed, {}'
            .format(pm_path.stderr))
        # min container needs to be GPU-support-enabled if the build is GPU build
        vars = pm_path.stdout
        e = re.search("CUDA_VERSION", vars)
        gpu_enabled = False if e is None else True
        fail_if(
            not gpu_enabled,
            'Composing container with gpu support enabled but min container provided does not have CUDA installed'
        )

    # Check full container enviroment variables
    p_path = subprocess.run(base_run_args + [
        '{{range $index, $value := .Config.Env}}{{$value}} {{end}}',
        full_docker_image
    ],
                            capture_output=True,
                            text=True)
    fail_if(
        p_path.returncode != 0,
        'docker inspect to find enviroment variables for full container failed, {}'
        .format(p_path.stderr))
    vars = p_path.stdout
    log_verbose("inspect args: {}".format(vars))

    e0 = re.search("TRITON_SERVER_GPU_ENABLED=([\S]{1,}) ", vars)
    e1 = re.search("CUDA_VERSION", vars)
    gpu_enabled = False
    if (e0 != None):
        gpu_enabled = e0.group(1) == "1"
    elif (e1 != None):
        gpu_enabled = True
    fail_if(
        gpu_enabled != enable_gpu,
        'Error: full container provided was build with '
        '\'TRITON_SERVER_GPU_ENABLED\' as {} and you are composing container'
        'with \'TRITON_SERVER_GPU_ENABLED\' as {}'.format(
            gpu_enabled, enable_gpu))
    e = re.search("TRITON_SERVER_VERSION=([\S]{6,}) ", vars)
    version = "" if e is None else e.group(1)
    fail_if(
        len(version) == 0,
        'docker inspect to find triton server version failed, {}'.format(
            p_path.stderr))
    e = re.search("NVIDIA_TRITON_SERVER_VERSION=([\S]{5,}) ", vars)
    container_version = "" if e is None else e.group(1)
    fail_if(
        len(container_version) == 0,
        'docker inspect to find triton container version failed, {}'.format(
            vars))
    dcgm_ver = re.search("DCGM_VERSION=([\S]{4,}) ", vars)
    dcgm_version = ""
    if dcgm_ver is None:
        dcgm_version = "2.2.3"
        log("WARNING: DCGM version not found from image, installing the earlierst version {}"
            .format(dcgm_version))
    else:
        dcgm_version = dcgm_ver.group(1)
    fail_if(
        len(dcgm_version) == 0,
        'docker inspect to find DCGM version failed, {}'.format(vars))

    p_sha = subprocess.run(
        base_run_args +
        ['{{ index .Config.Labels "com.nvidia.build.ref"}}', full_docker_image],
        capture_output=True,
        text=True)
    fail_if(
        p_sha.returncode != 0,
        'docker inspect of upstream docker image build sha failed, {}'.format(
            p_sha.stderr))
    p_build = subprocess.run(
        base_run_args +
        ['{{ index .Config.Labels "com.nvidia.build.id"}}', full_docker_image],
        capture_output=True,
        text=True)
    fail_if(
        p_build.returncode != 0,
        'docker inspect of upstream docker image build sha failed, {}'.format(
            p_build.stderr))

    p_find = subprocess.run(
        ['docker', 'run', full_docker_image, 'bash', '-c', 'ls /usr/bin/'],
        capture_output=True,
        text=True)
    f = re.search("serve", p_find.stdout)
    fail_if(p_find.returncode != 0,
            "Cannot search for 'serve' in /usr/bin, {}".format(p_find.stderr))
    argmap = {
        'NVIDIA_BUILD_REF': p_sha.stdout.rstrip(),
        'NVIDIA_BUILD_ID': p_build.stdout.rstrip(),
        'TRITON_VERSION': version,
        'TRITON_CONTAINER_VERSION': container_version,
        'DCGM_VERSION': dcgm_version,
        'SAGEMAKER_ENDPOINT': f is not None,
    }
    return argmap


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    group_qv = parser.add_mutually_exclusive_group()
    group_qv.add_argument('-q',
                          '--quiet',
                          action="store_true",
                          required=False,
                          help='Disable console output.')
    group_qv.add_argument('-v',
                          '--verbose',
                          action="store_true",
                          required=False,
                          help='Enable verbose output.')
    parser.add_argument(
        '--output-name',
        type=str,
        required=False,
        help='Name for the generated Docker image. Default is "tritonserver".')
    parser.add_argument(
        '--work-dir',
        type=str,
        required=False,
        help=
        'Generated dockerfiles are placed here. Default to current directory.')
    parser.add_argument(
        '--container-version',
        type=str,
        required=False,
        help=
        'The version to use for the generated Docker image. If not specified '
        'the container version will be chosen automatically based on the '
        'repository branch.')
    parser.add_argument(
        '--image',
        action='append',
        required=False,
        help='Use specified Docker image to generate Docker image. Specified as '
        '<image-name>,<full-image-name>. <image-name> can be "min", "gpu-min" '
        'or "full". Both "min" and "full" need to be specified at the same time.'
        'This will override "--container-version". "gpu-min" is needed for '
        'CPU-only container to copy TensorFlow and PyTorch deps.')
    parser.add_argument('--enable-gpu',
                        nargs='?',
                        type=lambda x: (str(x).lower() == 'true'),
                        const=True,
                        default=True,
                        required=False,
                        help=argparse.SUPPRESS)
    parser.add_argument(
        '--backend',
        action='append',
        required=False,
        help=
        'Include <backend-name> in the generated Docker image. The flag may be '
        'specified multiple times.')
    parser.add_argument(
        '--repoagent',
        action='append',
        required=False,
        help=
        'Include <repoagent-name> in the generated Docker image. The flag may '
        'be specified multiple times.')
    parser.add_argument(
        '--skip-pull',
        action='store_true',
        required=False,
        help='Do not pull the required docker images. The user is responsible '
        'for pulling the upstream images needed to compose the image.')
    parser.add_argument(
        '--dry-run',
        action="store_true",
        required=False,
        help='Only creates Dockerfile.compose, does not build the Docker image.'
    )

    FLAGS = parser.parse_args()

    if FLAGS.work_dir is None:
        FLAGS.work_dir = "."
    if FLAGS.output_name is None:
        FLAGS.output_name = "tritonserver"

    dockerfile_name = 'Dockerfile.compose'

    if FLAGS.backend is None:
        FLAGS.backend = []
    if FLAGS.repoagent is None:
        FLAGS.repoagent = []

    # Initialize map of docker images.
    images = {}
    if FLAGS.image:
        for img in FLAGS.image:
            parts = img.split(',')
            fail_if(
                len(parts) != 2,
                '--image must specific <image-name>,<full-image-registry>')
            fail_if(
                parts[0] not in ['min', 'full', 'gpu-min'],
                'unsupported image-name \'{}\' for --image'.format(parts[0]))
            log('image "{}": "{}"'.format(parts[0], parts[1]))
            images[parts[0]] = parts[1]
    else:
        get_container_version_if_not_specified()
        if FLAGS.enable_gpu:
            images = {
                "full":
                    "nvcr.io/nvidia/tritonserver:{}-py3".format(
                        FLAGS.container_version),
                "min":
                    "nvcr.io/nvidia/tritonserver:{}-py3-min".format(
                        FLAGS.container_version)
            }
        else:
            images = {
                "full":
                    "nvcr.io/nvidia/tritonserver:{}-cpu-only-py3".format(
                        FLAGS.container_version),
                "min":
                    "ubuntu:20.04"
            }
    fail_if(
        len(images) < 2,
        "Need to specify both 'full' and 'min' images if at all")

    # For CPU-only image we need to copy some cuda libraries and dependencies
    # since we are using PyTorch, TensorFlow 1, TensorFlow 2 containers that
    # are not CPU-only.
    if (('pytorch' in FLAGS.backend) or ('tensorflow1' in FLAGS.backend) or
        ('tensorflow2' in FLAGS.backend)) and ('gpu-min' not in images):
        images["gpu-min"] = "nvcr.io/nvidia/tritonserver:{}-py3-min".format(
            FLAGS.container_version)

    argmap = create_argmap(images, FLAGS.skip_pull)

    start_dockerfile(FLAGS.work_dir, images, argmap, dockerfile_name,
                     FLAGS.backend)
    add_requested_backends(FLAGS.work_dir, dockerfile_name, FLAGS.backend)
    add_requested_repoagents(FLAGS.work_dir, dockerfile_name, FLAGS.repoagent)
    end_dockerfile(FLAGS.work_dir, dockerfile_name, argmap)

    if (not FLAGS.dry_run):
        build_docker_image(FLAGS.work_dir, dockerfile_name, FLAGS.output_name)


# qa/L0_model_namespacing/test.py
#!/usr/bin/env python
# Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys
import os

sys.path.append(os.path.join(os.environ["TRITON_QA_ROOT_DIR"], "common"))

import numpy as np
import unittest
import time
import shutil
import test_util as tu

import tritonclient.http as httpclient
from tritonclient.utils import InferenceServerException

#
# Test utilities
#

# Checker to perform inference on given model, expecting model to have
# [INPUT0, INPUT1] and produce [OUTPUT0, OUTPUT1] where:
# OUTPUT0 = INPUT0 + INPUT1
# OUTPUT1 = INPUT0 - INPUT1
class AddSubChecker:
    # Optional 'checker_client' may be provided to use a different
    # Triton client library, currently it must be either Triton HTTP client
    # library or Triton GRPC client library
    def __init__(self, checker_client=None):
        # client library selection
        if checker_client is None:
            import tritonclient.http as checker_client
        if "http" in checker_client.__name__:
            self.client_ = checker_client.InferenceServerClient("localhost:8000")
        else:
            self.client_ = checker_client.InferenceServerClient("localhost:8001")

        # Create infer input tensors
        self.inputs_ = []
        self.inputs_.append(checker_client.InferInput('INPUT0', [16], "INT32"))
        self.inputs_.append(checker_client.InferInput('INPUT1', [16], "INT32"))

        # Initialize the data and expected output
        input_data = np.arange(start=0, stop=16, dtype=np.int32)
        self.inputs_[0].set_data_from_numpy(input_data)
        self.inputs_[1].set_data_from_numpy(input_data)
        self.expected_outputs_ = {
          "add" : (input_data + input_data),
          "sub" : (input_data - input_data)
        }
        
    def infer(self, model):
        res = self.client_.infer(model, self.inputs_)
        np.testing.assert_allclose(res.as_numpy('OUTPUT0'), self.expected_outputs_["add"])
        np.testing.assert_allclose(res.as_numpy('OUTPUT1'), self.expected_outputs_["sub"])

# Checker to perform inference on given model, expecting model to have
# [INPUT0, INPUT1] and produce [OUTPUT0, OUTPUT1] where:
# OUTPUT0 = INPUT0 - INPUT1
# OUTPUT1 = INPUT0 + INPUT1
class SubAddChecker(AddSubChecker):
    def infer(self, model):
        res = self.client_.infer(model, self.inputs_)
        np.testing.assert_allclose(res.as_numpy('OUTPUT0'), self.expected_outputs_["sub"])
        np.testing.assert_allclose(res.as_numpy('OUTPUT1'), self.expected_outputs_["add"])

#
# Test suites and cases
#

class ModelNamespacePoll(tu.TestResultCollector):
    def setUp(self):
        self.addsub_ = AddSubChecker()
        self.subadd_ = SubAddChecker()
        # For other server interaction
        self.client_ = httpclient.InferenceServerClient("localhost:8000")


    def check_health(self, expect_live=True, expect_ready=True):
        self.assertEqual(self.client_.is_server_live(), expect_live)
        self.assertEqual(self.client_.is_server_ready(), expect_ready)

    def test_no_duplication(self):
        # Enable model namspacing on repositories that is already valid without
        # enabling model namespacing.
        # All models should be visible and can be inferred individually
        self.check_health()

        # infer check
        for model in ["simple_addsub", "composing_addsub"]:
            self.addsub_.infer(model)
        for model in ["simple_subadd", "composing_subadd"]:
            self.subadd_.infer(model)

    def test_duplication(self):
        # Enable model namspacing on repositories that each repo has one
        # ensemble and it requires an composing model ('composing_model') that
        # exists in both repos.
        # Expect all models are visible, the ensemble will pick up the correct
        # model even the composing model can't be inferred individually.
        self.check_health()

        # infer check
        for model in ["simple_addsub",]:
            self.addsub_.infer(model)
        for model in ["simple_subadd",]:
            self.subadd_.infer(model)
            
        # error check
        try:
            self.addsub_.infer("composing_model")
            self.assertTrue(False,
                            "expected error for inferring ambiguous named model")
        except InferenceServerException as ex:
            self.assertIn("ambiguity", ex.message())

    def test_ensemble_duplication(self):
        # Enable model namspacing on repositories that each repo has one
        # ensemble with the same name. Expect the ensemble will pick up the correct
        # model.
        # Expect all models are visible, the ensemble will pick up the correct
        # model even the ensemble itself can't be inferred without providing
        # namespace.
        self.check_health()

        # infer
        for model in ["composing_addsub",]:
            self.addsub_.infer(model)
        for model in ["composing_subadd",]:
            self.subadd_.infer(model)

        # error check
        try:
            self.addsub_.infer("simple_ensemble")
            self.assertTrue(False,
                            "expected error for inferring ambiguous named model")
        except InferenceServerException as ex:
            self.assertIn("ambiguity", ex.message())

    def test_dynamic_resolution(self):
        # Same model setup as 'test_duplication', will remove / add one of the
        # composing model at runtime and expect the ensemble to be properly
        # linked to exisiting composing model at different steps.
        # 1. Remove 'composing_model' in addsub_repo, expect both ensembles use
        #    'composing_model' in subadd_repo and act as subadd
        # 2. Add back 'composing_model' in addsub_repo, expect the ensembles to behave the
        #    same as before the removal.
        self.assertTrue("NAMESPACE_TESTING_DIRCTORY" in os.environ)
        td = os.environ["NAMESPACE_TESTING_DIRCTORY"]
        composing_before_path = os.path.join(td, "addsub_repo", "composing_model")
        composing_after_path = os.path.join(td, "composing_model")

        self.check_health()
        # step 1.
        shutil.move(composing_before_path, composing_after_path)
        time.sleep(5)

        # infer
        for model in ["simple_subadd", "simple_addsub", "composing_model"]:
            self.subadd_.infer(model)
  
        # step 2.
        shutil.move(composing_after_path, composing_before_path)
        time.sleep(5)

        # infer
        for model in ["simple_addsub",]:
            self.addsub_.infer(model)
        for model in ["simple_subadd",]:
            self.subadd_.infer(model)

        # error check
        try:
            self.addsub_.infer("composing_model")
            self.assertTrue(False,
                            "expected error for inferring ambiguous named model")
        except InferenceServerException as ex:
            self.assertIn("ambiguity", ex.message())


class ModelNamespaceExplicit(tu.TestResultCollector):
    def setUp(self):
        self.addsub_ = AddSubChecker()
        self.subadd_ = SubAddChecker()
        # For other server interaction
        self.client_ = httpclient.InferenceServerClient("localhost:8000")

    def check_health(self, expect_live=True, expect_ready=True):
        self.assertEqual(self.client_.is_server_live(), expect_live)
        self.assertEqual(self.client_.is_server_ready(), expect_ready)

    def test_no_duplication(self):
        # Enable model namspacing on repositories that is already valid without
        # enabling model namespacing.
        # All models should be visible and can be inferred individually
        self.check_health()
        # load ensembles, cascadingly load composing model
        for model in ["simple_addsub", "simple_subadd"]:
            self.client_.load_model(model)

        # infer
        for model in ["simple_addsub", "composing_addsub"]:
            self.addsub_.infer(model)
        for model in ["simple_subadd", "composing_subadd"]:
            self.subadd_.infer(model)

    def test_duplication(self):
        # Enable model namspacing on repositories that each repo has one
        # ensemble and it requires an composing model ('composing_model') that
        # exists in both repos.
        # Expect all models are visible, the ensemble will pick up the correct
        # model even the composing model can't be inferred individually.
        self.check_health()
        # load ensembles, cascadingly load composing model
        for model in ["simple_addsub", "simple_subadd"]:
            self.client_.load_model(model)

        # infer
        for model in ["simple_addsub",]:
            self.addsub_.infer(model)
        for model in ["simple_subadd",]:
            self.subadd_.infer(model)

        # error check
        try:
            self.addsub_.infer("composing_model")
            self.assertTrue(False,
                            "expected error for inferring ambiguous named model")
        except InferenceServerException as ex:
            self.assertIn("ambiguity", ex.message())

    def test_ensemble_duplication(self):
        # Enable model namspacing on repositories that each repo has one
        # ensemble with the same name. Expect the ensemble will pick up the correct
        # model.
        # Expect all models are visible, the ensemble will pick up the correct
        # model even the ensemble itself can't be inferred without providing
        # namespace.
        self.check_health()
        # load ensembles, cascadingly load composing model
        for model in ["simple_ensemble"]:
            self.client_.load_model(model)
        
        # infer
        for model in ["composing_addsub",]:
            self.addsub_.infer(model)
        for model in ["composing_subadd",]:
            self.subadd_.infer(model)

        # error check
        try:
            self.addsub_.infer("simple_ensemble")
            self.assertTrue(False,
                            "expected error for inferring ambiguous named model")
        except InferenceServerException as ex:
            self.assertIn("ambiguity", ex.message())

    def test_dynamic_resolution(self):
        # Same model setup as 'test_duplication', will remove / add one of the
        # composing model at runtime and expect the ensemble to be properly
        # linked to exisiting composing model at different steps.
        # 1. Remove 'composing_model' in addsub_repo, expect both ensembles use
        #    'composing_model' in subadd_repo and act as subadd.
        # 2. Add back 'composing_model' in addsub_repo, expect the ensembles to behave the
        #    same as before the removal.
        self.assertTrue("NAMESPACE_TESTING_DIRCTORY" in os.environ)
        td = os.environ["NAMESPACE_TESTING_DIRCTORY"]
        composing_before_path = os.path.join(td, "addsub_repo", "composing_model")
        composing_after_path = os.path.join(td, "composing_model")

        self.check_health()
        # step 1.
        shutil.move(composing_before_path, composing_after_path)
        # load ensembles, cascadingly load composing model
        for model in ["simple_addsub", "simple_subadd"]:
            self.client_.load_model(model)

        # infer
        for model in ["simple_subadd", "simple_addsub", "composing_model"]:
            self.subadd_.infer(model)
  
        # step 2.
        shutil.move(composing_after_path, composing_before_path)
        # Explicitly load one of the ensembel, should still trigger cascading
        # (re-)load
        for model in ["simple_addsub", ]:
            self.client_.load_model(model)

        # infer
        for model in ["simple_addsub",]:
            self.addsub_.infer(model)
        for model in ["simple_subadd",]:
            self.subadd_.infer(model)

        # error check
        try:
            self.addsub_.infer("composing_model")
            self.assertTrue(False,
                            "expected error for inferring ambiguous named model")
        except InferenceServerException as ex:
            self.assertIn("ambiguity", ex.message())


if __name__ == '__main__':
    unittest.main()
