Update to model generation script

Spartee · Oct 4, 2022 · d5424fe · d5424fe
1 parent 03dda11
commit d5424fe
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 3 deletions.
diff --git a/qa/common/gen_qa_custom_ops b/qa/common/gen_qa_custom_ops
@@ -41,7 +41,7 @@ TRITON_VERSION=${TRITON_VERSION:=22.09}
 TENSORFLOW_IMAGE=${TENSORFLOW_IMAGE:=nvcr.io/nvidia/tensorflow:$TRITON_VERSION-tf2-py3}
 PYTORCH_IMAGE=${PYTORCH_IMAGE:=nvcr.io/nvidia/pytorch:$TRITON_VERSION-py3}
 
-CUDA_DEVICE=0
+CUDA_DEVICE=${RUNNER_ID:=0}
 
 ###
 HOST_BUILD_DIR=${HOST_BUILD_DIR:=/tmp}
@@ -66,6 +66,7 @@ DESTDIR=/tmp/custom_ops
 # Tensorflow
 cat >$HOST_SRCDIR/$TFSCRIPT <<EOF
 #!/bin/bash -x
+nvidia-smi -L || true
 set -e
 
 TF_CFLAGS=\$(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_compile_flags()))')
@@ -114,6 +115,7 @@ fi
 # PyTorch
 cat >$HOST_SRCDIR/$PYTSCRIPT <<EOF
 #!/bin/bash -x
+nvidia-smi -L || true
 set -e
 python3 $SRCDIR/gen_qa_custom_ops_models.py --libtorch --models_dir=$DESTDIR
 cp /root/.cache/torch_extensions/py38_cu118/custom_modulo/custom_modulo.so $DESTDIR/libtorch_modulo/.

diff --git a/qa/common/gen_qa_model_repository b/qa/common/gen_qa_model_repository
@@ -61,7 +61,7 @@ UBUNTU_IMAGE=${UBUNTU_IMAGE:=ubuntu:20.04}
 PYTORCH_IMAGE=${PYTORCH_IMAGE:=nvcr.io/nvidia/pytorch:$TRITON_VERSION-py3}
 TENSORFLOW_IMAGE=${TENSORFLOW_IMAGE:=nvcr.io/nvidia/tensorflow:$TRITON_VERSION-tf1-py3}
 TENSORRT_IMAGE=${TENSORRT_IMAGE:=nvcr.io/nvidia/tensorrt:$TRITON_VERSION-py3}
-CUDA_DEVICE=0
+CUDA_DEVICE=${RUNNER_ID:=0}
 
 ###
 HOST_BUILD_DIR=${HOST_BUILD_DIR:=/tmp}
@@ -187,6 +187,7 @@ if [[ "aarch64" != $(uname -m) ]] ; then
 
 cat >$HOST_SRCDIR/$OPENVINOSCRIPT <<EOF
 #!/bin/bash
+nvidia-smi -L || true
 set -e
 export DEBIAN_FRONTEND=noninteractive
 apt-get update && \
@@ -252,6 +253,7 @@ fi # [[ "aarch64" != $(uname -m) ]]
 # ONNX
 cat >$HOST_SRCDIR/$ONNXSCRIPT <<EOF
 #!/bin/bash
+nvidia-smi -L || true
 set -e
 export DEBIAN_FRONTEND=noninteractive
 apt-get update && \
@@ -321,6 +323,7 @@ fi
 # PyTorch
 cat >$HOST_SRCDIR/$TORCHSCRIPT <<EOF
 #!/bin/bash
+nvidia-smi -L || true
 set -e
 python3 $SRCDIR/gen_qa_models.py --libtorch --models_dir=$DESTDIR
 chmod -R 777 $DESTDIR
@@ -366,6 +369,7 @@ fi
 # Tensorflow
 cat >$HOST_SRCDIR/$TFSCRIPT <<EOF
 #!/bin/bash
+nvidia-smi -L || true
 set -e
 python3 $SRCDIR/gen_qa_models.py --graphdef --savedmodel --models_dir=$DESTDIR
 chmod -R 777 $DESTDIR
@@ -428,6 +432,7 @@ fi
 # TensorRT
 cat >$HOST_SRCDIR/$TRTSCRIPT <<EOF
 #!/bin/bash
+nvidia-smi -L || true
 set -e
 export TRT_SUPPRESS_DEPRECATION_WARNINGS=1
 # Models using shape tensor i/o

diff --git a/qa/common/gen_xavier_trt_models b/qa/common/gen_xavier_trt_models
@@ -51,7 +51,7 @@
 ############################################################################
 
 TRITON_VERSION=${TRITON_VERSION:=22.09}
-CUDA_DEVICE=0
+CUDA_DEVICE=${RUNNER_ID:=0}
 
 HOST_BUILD_DIR=${HOST_BUILD_DIR:=/tmp}
 TMP_DIR=$HOST_BUILD_DIR/${TRITON_VERSION}