Skip to content

test: PyTorch PT2 Model Generation #8666

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
whoisj merged 4 commits into from
Mar 13, 2026
Merged

test: PyTorch PT2 Model Generation #8666

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d3baf30
feat: PyTorch PT2 Model Generation
whoisj Nov 25, 2025
d716408
fixup pre-commit complaints
whoisj Feb 19, 2026
f9fdaea
Potential fix for code scanning alert no. 1026: Unused import
whoisj Feb 23, 2026
283d1e0
remove batch_size from aoti model generation
whoisj Mar 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
89 changes: 67 additions & 22 deletions qa/common/gen_qa_model_repository
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,10 +50,16 @@
TRITON_MDLS_BASE_SCRIPT_DIR="$(dirname $(readlink -f $0))"
TRITON_MDLS_BASE_SCRIPT_FILE="$(readlink -f $0)"

log_message.info() { local message=$@ ; echo -e "\033[34m$(date +"%Y-%m-%d %H:%M:%S") - [ INFO ] - ${message} \033[0m"; } ;
log_message.status() { local message=$@ ; echo -e "\033[37m$(date +"%Y-%m-%d %H:%M:%S") - [ STATUS ] - ${message} \033[0m"; } ;
log_message.warning() { local message=$@ ; echo -e "\033[33m$(date +"%Y-%m-%d %H:%M:%S") - [ WARNING ] - ${message} \033[0m"; } ;
log_message.error() { local message=$@ ; echo -e "\033[31m$(date +"%Y-%m-%d %H:%M:%S") - [ ERROR ] - ${message} \033[0m"; } ;
COLOR_ERROR="\033[31m"
COLOR_INFO="\033[94m"
COLOR_RESET="\033[0m"
COLOR_STATUS="\033[36m"
COLOR_WARNING="\033[33m"

log_message.info() { local message=$@ ; echo -e "${COLOR_INFO}$(date +"%Y-%m-%d %H:%M:%S") - [ INFO ] - ${message} ${COLOR_RESET}"; } ;
log_message.status() { local message=$@ ; echo -e "${COLOR_STATUS}$(date +"%Y-%m-%d %H:%M:%S") - [ STATUS ] - ${message} ${COLOR_RESET}"; } ;
log_message.warning() { local message=$@ ; echo -e "${COLOR_WARNING}$(date +"%Y-%m-%d %H:%M:%S") - [ WARNING ] - ${message} ${COLOR_RESET}"; } ;
log_message.error() { local message=$@ ; echo -e "${COLOR_ERROR}$(date +"%Y-%m-%d %H:%M:%S") - [ ERROR ] - ${message} ${COLOR_RESET}"; } ;
Comment thread
whoisj marked this conversation as resolved.

log_message.status "Changing working directory to the script directory to: " "${TRITON_MDLS_BASE_SCRIPT_DIR}"

Expand Down Expand Up @@ -135,8 +141,8 @@ function define_models_generation_scripts() {
umask 0000
nvidia-smi --query-gpu=compute_cap,compute_mode,driver_version,name,index --format=csv || true
nvidia-smi || true
echo -e "${COLOR_INFO}Generating OpenVINO models${COLOR_RESET}"
set -e
set -x
export DEBIAN_FRONTEND=noninteractive
apt-get update && \
apt-get install -y --no-install-recommends \
Expand Down Expand Up @@ -170,6 +176,7 @@ chmod -R 777 $TRITON_MDLS_QA_RESHAPE_MODEL
# chmod -R 777 $SVOLUME_EQDESTDIR
# python3 $TRITON_MDLS_SRC_DIR/gen_qa_dyna_sequence_models.py --openvino --models_dir=$TRITON_MDLS_QA_DYNA_SEQUENCE_MODEL
# chmod -R 777 $TRITON_MDLS_QA_DYNA_SEQUENCE_MODEL
exit 0
EOF

log_message.status "run: chmod a+x ${OPENVINOSCRIPT}"
Expand All @@ -187,8 +194,8 @@ EOF
umask 0000
nvidia-smi --query-gpu=compute_cap,compute_mode,driver_version,name,index --format=csv || true
nvidia-smi || true
echo -e "${COLOR_INFO}Generating ONNX models${COLOR_RESET}"
set -e
set -x
export DEBIAN_FRONTEND=noninteractive
apt-get update && \
apt-get install -y --no-install-recommends build-essential cmake libprotobuf-dev \
Expand Down Expand Up @@ -233,6 +240,7 @@ python3 $TRITON_MDLS_SRC_DIR/gen_qa_identity_models.py --ensemble --models_dir=$
python3 $TRITON_MDLS_SRC_DIR/gen_qa_sequence_models.py --ensemble --models_dir=$TRITON_MDLS_QA_ENSEMBLE_MODEL/qa_sequence_model_repository
python3 $TRITON_MDLS_SRC_DIR/gen_qa_sequence_models.py --ensemble --variable --models_dir=$TRITON_MDLS_QA_ENSEMBLE_MODEL/qa_variable_sequence_model_repository
chmod -R 777 $TRITON_MDLS_QA_ENSEMBLE_MODEL
exit 0
EOF

log_message.status "run: chmod a+x ${ONNXSCRIPT}"
Expand All @@ -249,12 +257,15 @@ EOF
umask 0000
nvidia-smi --query-gpu=compute_cap,compute_mode,driver_version,name,index --format=csv || true
nvidia-smi || true
echo -e "${COLOR_INFO}Generating PyTorch models${COLOR_RESET}"
pip3 install onnxscript
set -e
set -x
PATH=$PATH:/usr/local/cuda-13.0/bin
python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --libtorch --models_dir=$TRITON_MDLS_QA_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --torch-aoti --models_dir=$TRITON_MDLS_QA_MODEL
chmod -R 777 $TRITON_MDLS_QA_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --libtorch --variable --models_dir=$TRITON_MDLS_QA_VARIABLE_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --torch-aoti --variable --models_dir=$TRITON_MDLS_QA_VARIABLE_MODEL
chmod -R 777 $TRITON_MDLS_QA_VARIABLE_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_identity_models.py --libtorch --models_dir=$TRITON_MDLS_QA_IDENTITY_MODEL
chmod -R 777 $TRITON_MDLS_QA_IDENTITY_MODEL
Expand All @@ -271,7 +282,7 @@ chmod -R 777 $TRITON_MDLS_QA_VARIABLE_SEQUENCE_IMPLICIT_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_dyna_sequence_models.py --libtorch --models_dir=$TRITON_MDLS_QA_DYNA_SEQUENCE_MODEL
chmod -R 777 $TRITON_MDLS_QA_DYNA_SEQUENCE_MODEL
if [ -z "$MODEL_TYPE" ] || [ "$MODEL_TYPE" != "igpu" ]; then
nvidia-smi --query-gpu=compute_cap | grep -qz 12.1 && echo -e '\033[33m[WARNING]\033[0m Skipping model generation for Torch TensorRT' || python3 $TRITON_MDLS_SRC_DIR/gen_qa_torchtrt_models.py --models_dir=$TRITON_MDLS_QA_TORCHTRT_MODEL
nvidia-smi --query-gpu=compute_cap | grep -qz 12.1 && echo -e '${COLOR_WARNING}[WARNING]${COLOR_RESET} Skipping model generation for Torch TensorRT${COLOR_RESET}' || python3 $TRITON_MDLS_SRC_DIR/gen_qa_torchtrt_models.py --models_dir=$TRITON_MDLS_QA_TORCHTRT_MODEL
chmod -R 777 $TRITON_MDLS_QA_TORCHTRT_MODEL
Comment thread
whoisj marked this conversation as resolved.
fi
python3 $TRITON_MDLS_SRC_DIR/gen_qa_ragged_models.py --libtorch --models_dir=$TRITON_MDLS_QA_RAGGED_MODEL
Expand All @@ -286,6 +297,7 @@ python3 $TRITON_MDLS_SRC_DIR/gen_qa_custom_ops_models.py --libtorch --models_dir
mkdir -p $TRITON_MDLS_QA_CUSTOM_OPS/libtorch_modulo/
cp \${TORCH_EXTENSIONS_DIR}/custom_modulo/custom_modulo.so $TRITON_MDLS_QA_CUSTOM_OPS/libtorch_modulo/.
chmod -R 777 $TRITON_MDLS_QA_CUSTOM_OPS
exit 0
EOF

log_message.status "run: chmod a+x ${TORCHSCRIPT}"
Expand All @@ -302,8 +314,8 @@ EOF
umask 0000
nvidia-smi --query-gpu=compute_cap,compute_mode,driver_version,name,index --format=csv || true
nvidia-smi || true
echo -e "${COLOR_INFO}Generating TensorRT models${COLOR_RESET}"
set -e
set -x
dpkg -l | grep TensorRT
export TRT_SUPPRESS_DEPRECATION_WARNINGS=1
# Models using shape tensor i/o
Expand Down Expand Up @@ -338,7 +350,7 @@ python3 $TRITON_MDLS_SRC_DIR/gen_qa_ragged_models.py --tensorrt --models_dir=$TR
chmod -R 777 $TRITON_MDLS_QA_RAGGED_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_trt_format_models.py --models_dir=$TRITON_MDLS_QA_TRT_FORMAT_MODEL
chmod -R 777 $TRITON_MDLS_QA_TRT_FORMAT_MODEL
nvidia-smi --query-gpu=compute_cap | grep -qz 11.0 && echo -e '\033[33m[WARNING]\033[0m Skipping model generation for data dependent shape' || python3 $TRITON_MDLS_SRC_DIR/gen_qa_trt_data_dependent_shape.py --models_dir=$TRITON_MDLS_QA_TRT_DATA_DEPENDENT_MODEL
nvidia-smi --query-gpu=compute_cap | grep -qz 11.0 && echo -e '${COLOR_WARNING}[WARNING]${COLOR_RESET} Skipping model generation for data dependent shape${COLOR_RESET}' || python3 $TRITON_MDLS_SRC_DIR/gen_qa_trt_data_dependent_shape.py --models_dir=$TRITON_MDLS_QA_TRT_DATA_DEPENDENT_MODEL
chmod -R 777 $TRITON_MDLS_QA_TRT_DATA_DEPENDENT_MODEL
Comment thread
whoisj marked this conversation as resolved.
# Make shared library for custom Hardmax plugin.
if [ -d "/usr/src/tensorrt/samples/python/onnx_custom_plugin" ]; then
Expand All @@ -362,6 +374,7 @@ rm -rf build && mkdir build && \
cd build && cmake .. && make -j && cp libcustomHardmaxPlugin.so $TRITON_MDLS_QA_TRT_PLUGIN_MODEL/.
LD_PRELOAD=$TRITON_MDLS_QA_TRT_PLUGIN_MODEL/libcustomHardmaxPlugin.so python3 $TRITON_MDLS_SRC_DIR/gen_qa_trt_plugin_models.py --models_dir=$TRITON_MDLS_QA_TRT_PLUGIN_MODEL
chmod -R 777 $TRITON_MDLS_QA_TRT_PLUGIN_MODEL
exit 0
EOF

log_message.status "run: chmod a+x ${TRTSCRIPT}"
Expand Down Expand Up @@ -395,6 +408,7 @@ if [ "$TRITON_MODELS_USE_DOCKER" -eq 1 ] && which docker ; then
docker pull $UBUNTU_IMAGE

log_message.status "docker volume: create destination directory on volume"
log_message.info "docker run -v $DOCKER_VOLUME:/mnt -w /mnt/$CI_JOB_ID $UBUNTU_IMAGE mkdir -p gen_srcdir ${TRITON_VERSION}"
docker run \
--rm \
--label RUNNER_ID=$RUNNER_ID \
Expand All @@ -404,6 +418,7 @@ if [ "$TRITON_MODELS_USE_DOCKER" -eq 1 ] && which docker ; then
$UBUNTU_IMAGE \
mkdir -p gen_srcdir ${TRITON_VERSION}

log_message.status "docker volume: create model directories on volume"
docker run \
--rm \
--label RUNNER_ID=$RUNNER_ID \
Expand Down Expand Up @@ -439,6 +454,7 @@ if [ "$TRITON_MODELS_USE_DOCKER" -eq 1 ] && which docker ; then
$TRITON_MDLS_QA_DYNAMIC_BATCH_IMAGE_MODEL

log_message.status "docker container: create container $DOCKER_VOLUME_CONTAINER"
log_message.info "docker create --name $DOCKER_VOLUME_CONTAINER -v $DOCKER_VOLUME:/mnt -w /mnt/$CI_JOB_ID $UBUNTU_IMAGE"
docker create \
--label RUNNER_ID=$RUNNER_ID \
--label PROJECT_NAME=$PROJECT_NAME \
Expand All @@ -452,74 +468,103 @@ if [ "$TRITON_MODELS_USE_DOCKER" -eq 1 ] && which docker ; then

if [[ "aarch64" != $(uname -m) ]] ; then

log_message.info "docker run: $OPENVINOSCRIPT"
log_message.status "docker run: $OPENVINOSCRIPT"
log_message.info "docker run $DOCKER_GPU_ARGS -v $DOCKER_VOLUME:/mnt $UBUNTU_IMAGE bash -e $TRITON_MDLS_SRC_DIR/$OPENVINOSCRIPT"
docker run \
--rm \
-e TRITON_GENSRCDIR=$TRITON_MDLS_SRC_DIR \
--label RUNNER_ID=$RUNNER_ID \
--label PROJECT_NAME=$PROJECT_NAME \
$DOCKER_GPU_ARGS \
-v $DOCKER_VOLUME:/mnt \
-t \
$UBUNTU_IMAGE \
bash -xe $TRITON_MDLS_SRC_DIR/$OPENVINOSCRIPT
bash -e $TRITON_MDLS_SRC_DIR/$OPENVINOSCRIPT

exit_code=$?

if [ $? -ne 0 ]; then
if [ $exit_code -ne 0 ]; then
log_message.error "docker run: ${OPENVINOSCRIPT} failed"
exit 1
fi

rm $OPENVINOSCRIPT
fi # [[ "aarch64" != $(uname -m) ]]

log_message.info "docker run: $ONNXSCRIPT"
log_message.status "docker run: $ONNXSCRIPT"
log_message.info "docker run $DOCKER_GPU_ARGS -v $DOCKER_VOLUME:/mnt $UBUNTU_IMAGE bash -e $TRITON_MDLS_SRC_DIR/$ONNXSCRIPT"
docker run \
--rm \
-e TRITON_GENSRCDIR=$TRITON_MDLS_SRC_DIR \
--label RUNNER_ID=$RUNNER_ID \
--label PROJECT_NAME=$PROJECT_NAME \
$DOCKER_GPU_ARGS \
-v $DOCKER_VOLUME:/mnt \
-t \
$UBUNTU_IMAGE \
bash -xe $TRITON_MDLS_SRC_DIR/$ONNXSCRIPT
bash -e $TRITON_MDLS_SRC_DIR/$ONNXSCRIPT

if [ $? -ne 0 ]; then
exit_code=$?

if [ $exit_code -ne 0 ]; then
log_message.error "docker run: ${ONNXSCRIPT} failed"
exit 1
fi

rm $ONNXSCRIPT

log_message.status "docker pull: $PYTORCH_IMAGE"
log_message.info "docker pull $PYTORCH_IMAGE"
docker pull $PYTORCH_IMAGE

log_message.info "docker run: $TORCHSCRIPT"
log_message.status "docker run: $TORCHSCRIPT"
log_message.info "docker run $DOCKER_GPU_ARGS -v $DOCKER_VOLUME:/mnt $PYTORCH_IMAGE bash -e $TRITON_MDLS_SRC_DIR/$TORCHSCRIPT"
docker run \
--rm \
-e TRITON_GENSRCDIR=$TRITON_MDLS_SRC_DIR \
--label RUNNER_ID=$RUNNER_ID \
--label PROJECT_NAME=$PROJECT_NAME \
$DOCKER_GPU_ARGS \
-v $DOCKER_VOLUME:/mnt \
-t \
$PYTORCH_IMAGE \
bash -xe $TRITON_MDLS_SRC_DIR/$TORCHSCRIPT
bash -e $TRITON_MDLS_SRC_DIR/$TORCHSCRIPT

if [ $? -ne 0 ]; then
exit_code=$?

if [ $exit_code -ne 0 ]; then
log_message.error "docker run: ${TORCHSCRIPT} failed"
exit 1
fi

rm $TORCHSCRIPT

if [ "$MODEL_TYPE" != "igpu" ] ; then
log_message.status "docker pull: $TENSORRT_IMAGE"
docker pull $TENSORRT_IMAGE

log_message.info "docker run: $TRTSCRIPT"
log_message.status "docker run: $TRTSCRIPT"
log_message.info "docker run $DOCKER_GPU_ARGS -v $DOCKER_VOLUME:/mnt $TENSORRT_IMAGE bash -e $TRITON_MDLS_SRC_DIR/$TRTSCRIPT"
docker run \
--rm \
-e TRITON_GENSRCDIR=$TRITON_MDLS_SRC_DIR \
--label RUNNER_ID=$RUNNER_ID \
--label PROJECT_NAME=$PROJECT_NAME \
$DOCKER_GPU_ARGS \
-v $DOCKER_VOLUME:/mnt \
-t \
-e TRT_VERBOSE \
$TENSORRT_IMAGE \
bash -xe $TRITON_MDLS_SRC_DIR/$TRTSCRIPT
bash -e $TRITON_MDLS_SRC_DIR/$TRTSCRIPT

if [ $? -ne 0 ]; then
exit_code=$?

if [ $exit_code -ne 0 ]; then
log_message.error "docker run: ${TRTSCRIPT} failed"
exit 1
fi

rm $TRTSCRIPT
fi # [ "$MODEL_TYPE" != "igpu" ]

if [ -z $CI ] ; then
Expand Down
Loading
Loading