#!/bin/bash
# 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.

############################################################################
## This script generates the model repository needed for TensorRT testing
## testing on the Jetson Xavier. Generating these models requires having
## TensorRT on the system.
##
## 1. Update Jetpack version to match what is being used by the
## tritonserver release being tested.
##
## 2. Set CUDA_DEVICE to the ID of the CUDA device present on the
## system that you want to target for the generated models.
##
## 3. Run this script to create a '/tmp/<TRITON_VERSION>_xavier' directory
## containing all the TensorRT models needed for CI testing.
##
## 4. Copy all non-TensorRT models from the dlcluster directory
## '/mnt/dldata/inferenceserver/<TRITON_VERSION>' to
## '/mnt/dldata/inferenceserver/<TRITON_VERSION>'_xavier.
##
## 5. Copy the generated models from '/tmp/<TRITON_VERSION>_xavier' to the
## dlcluster directory '/mnt/dldata/inferenceserver/<TRITON_VERSION>'_xavier.
##
## 6. Remember to delete the '/tmp/<TRITON_VERSION>_xavier' directory after.
##
############################################################################

TRITON_VERSION=20.02
CUDA_DEVICE=0

REPO_ARCH="xavier"
TMP_DIR=/tmp/${TRITON_VERSION}_${REPO_ARCH}

SRCDIR=/tmp/src
DESTDIR=$TMP_DIR/models
VARDESTDIR=$TMP_DIR/varmodels
IDENTITYDESTDIR=$TMP_DIR/zeromodels
IDENTITYBIGDESTDIR=$TMP_DIR/zerobigmodels
SHAPEDESTDIR=$TMP_DIR/shapetensormodels
RESHAPEDESTDIR=$TMP_DIR/reshapemodels
SEQDESTDIR=$TMP_DIR/seqmodels
DYNASEQDESTDIR=$TMP_DIR/dynaseqmodels
VARSEQDESTDIR=$TMP_DIR/varseqmodels
ENSEMBLEDESTDIR=$TMP_DIR/ensemblemodels
NOSHAPEDESTDIR=$TMP_DIR/noshapemodels

cp ./gen_qa_models.py $SRCDIR/.
cp ./gen_qa_identity_models.py $SRCDIR/.
cp ./gen_qa_reshape_models.py $SRCDIR/.
cp ./gen_qa_noshape_models.py $SRCDIR/.
cp ./gen_qa_sequence_models.py $SRCDIR/.
cp ./gen_qa_dyna_sequence_models.py $SRCDIR/.
cp ./gen_ensemble_model_utils.py $SRCDIR/.
cp ./gen_qa_trt_plugin_models.py $SRCDIR/.
cp ./test_util.py $SRCDIR/.

# TensorRT
set -e
export TRT_SUPPRESS_DEPRECATION_WARNINGS=1
cd $SRCDIR
# Models using shape tensor i/o
python3 $SRCDIR/gen_qa_identity_models.py --tensorrt-shape-io --models_dir=$SHAPEDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --tensorrt-shape-io --models_dir=$SHAPEDESTDIR
python3 $SRCDIR/gen_qa_dyna_sequence_models.py --tensorrt-shape-io --models_dir=$SHAPEDESTDIR
chown -R $(id -u):$(id -g) $SHAPEDESTDIR
python3 $SRCDIR/gen_qa_models.py --tensorrt --models_dir=$DESTDIR
chown -R $(id -u):$(id -g) $DESTDIR
python3 $SRCDIR/gen_qa_models.py --tensorrt --variable --models_dir=$VARDESTDIR
chown -R $(id -u):$(id -g) $VARDESTDIR
python3 $SRCDIR/gen_qa_identity_models.py --tensorrt --models_dir=$IDENTITYDESTDIR
chown -R $(id -u):$(id -g) $IDENTITYDESTDIR
python3 $SRCDIR/gen_qa_identity_models.py --tensorrt-big --models_dir=$IDENTITYBIGDESTDIR
chown -R $(id -u):$(id -g) $IDENTITYBIGDESTDIR
python3 $SRCDIR/gen_qa_reshape_models.py --tensorrt --variable --models_dir=$RESHAPEDESTDIR
chown -R $(id -u):$(id -g) $RESHAPEDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --tensorrt --models_dir=$SEQDESTDIR
chown -R $(id -u):$(id -g) $SEQDESTDIR
python3 $SRCDIR/gen_qa_dyna_sequence_models.py --tensorrt --models_dir=$DYNASEQDESTDIR
chown -R $(id -u):$(id -g) $DYNASEQDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --tensorrt --variable --models_dir=$VARSEQDESTDIR
chown -R $(id -u):$(id -g) $VARSEQDESTDIR
