fix: Fix L0_backend_python test (#8178)

Author: krishung5

qa/L0_backend_python/common.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2021-2025, 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
@@ -85,7 +85,11 @@ create_conda_env_with_specified_path() {
 create_python_backend_stub() {
   rm -rf python_backend
   git clone ${TRITON_REPO_ORGANIZATION}/python_backend -b $PYTHON_BACKEND_REPO_TAG
+  CUDA_PATH=$(readlink -f /usr/local/cuda)
   (cd python_backend/ && mkdir builddir && cd builddir && \
-  cmake -DTRITON_ENABLE_GPU=ON -DTRITON_REPO_ORGANIZATION:STRING=${TRITON_REPO_ORGANIZATION} -DTRITON_BACKEND_REPO_TAG=$TRITON_BACKEND_REPO_TAG -DTRITON_COMMON_REPO_TAG=$TRITON_COMMON_REPO_TAG -DTRITON_CORE_REPO_TAG=$TRITON_CORE_REPO_TAG -DPYBIND11_PYTHON_VERSION=$PY_VERSION ../ && \
+  cmake -DTRITON_ENABLE_GPU=ON -DCMAKE_CUDA_COMPILER=$CUDA_PATH/bin/nvcc \
+     -DCUDAToolkit_ROOT=$CUDA_PATH -DTRITON_REPO_ORGANIZATION:STRING=${TRITON_REPO_ORGANIZATION} \
+    -DTRITON_BACKEND_REPO_TAG=$TRITON_BACKEND_REPO_TAG -DTRITON_COMMON_REPO_TAG=$TRITON_COMMON_REPO_TAG \
+    -DTRITON_CORE_REPO_TAG=$TRITON_CORE_REPO_TAG -DPYBIND11_PYTHON_VERSION=$PY_VERSION ../ && \
   make -j18 triton-python-backend-stub)
 }

qa/L0_backend_python/env/test.sh

@@ -53,8 +53,8 @@ conda install numpy=1.26.4 -y
 if [ $TRITON_RHEL -eq 1 ]; then
     TORCH_VERISON="2.17.0"
 fi
-conda install torch=${TORCH_VERSION} -y
-PY312_VERSION_STRING="Python version is 3.12, NumPy version is 1.26.4, and PyTorch version is ${TORCH_VERISON}"
+conda install pytorch=${TORCH_VERSION} -y
+PY312_VERSION_STRING="Python version is 3.12, NumPy version is 1.26.4, and PyTorch version is ${TORCH_VERSION}"
 conda pack -o python3.12.tar.gz
 mkdir -p models/python_3_12/1/
 cp ../../python_models/python_version/config.pbtxt ./models/python_3_12
@@ -122,7 +122,7 @@ fi
 kill_server
 
 set +e
-grep "Locale is ('en_US', 'UTF-8')" $SERVER_LOG
+grep "Locale is ('C', 'UTF-8')" $SERVER_LOG
     if [ $? -ne 0 ]; then
         cat $SERVER_LOG
         echo -e "\n***\n*** Locale UTF-8 was not found in Triton logs. \n***"
@@ -182,10 +182,6 @@ aws s3 mb "${BUCKET_URL}"
 BUCKET_URL=${BUCKET_URL%/}
 BUCKET_URL_SLASH="${BUCKET_URL}/"
 
-# Remove Python 3.7 model because it contains absolute paths and cannot be used
-# with S3.
-rm -rf models/python_3_7
-
 # Test with the bucket url as model repository
 aws s3 cp models/ "${BUCKET_URL_SLASH}" --recursive --include "*"
 
@@ -205,10 +201,10 @@ fi
 kill_server
 
 set +e
-grep "$PY36_VERSION_STRING" $SERVER_LOG
+grep "$PY312_VERSION_STRING" $SERVER_LOG
 if [ $? -ne 0 ]; then
     cat $SERVER_LOG
-    echo -e "\n***\n*** $PY36_VERSION_STRING was not found in Triton logs. \n***"
+    echo -e "\n***\n*** $PY312_VERSION_STRING was not found in Triton logs. \n***"
     RET=1
 fi
 set -e
@@ -217,8 +213,6 @@ set -e
 aws s3 rm "${BUCKET_URL_SLASH}" --recursive --include "*"
 
 # Test with EXECUTION_ENV_PATH outside the model directory
-sed -i "s/TRITON_MODEL_DIRECTORY\/python_3_6_environment/TRITON_MODEL_DIRECTORY\/..\/python_3_6_environment/" models/python_3_6/config.pbtxt
-mv models/python_3_6/python_3_6_environment.tar.gz models
 sed -i "s/\$\$TRITON_MODEL_DIRECTORY\/python_3_12_environment/s3:\/\/triton-bucket-${CI_JOB_ID}\/python_3_12_environment/" models/python_3_12/config.pbtxt
 mv models/python_3_12/python_3_12_environment.tar.gz models
 
@@ -238,7 +232,7 @@ fi
 kill_server
 
 set +e
-for EXPECTED_VERSION_STRING in "$PY36_VERSION_STRING" "$PY312_VERSION_STRING"; do
+for EXPECTED_VERSION_STRING in "$PY312_VERSION_STRING"; do
     grep "$EXPECTED_VERSION_STRING" $SERVER_LOG
     if [ $? -ne 0 ]; then
         cat $SERVER_LOG

qa/L0_backend_python/examples/test.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2021-2025, 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
@@ -69,7 +69,7 @@ pip3 install validators
 # Install JAX
 # Jax has dropped the support for Python 3.8. See https://jax.readthedocs.io/en/latest/changelog.html
 if [ "$TEST_JETSON" == "0" ] && [ ${PYTHON_ENV_VERSION} != "8" ]; then
-    pip3 install --upgrade "jax[cuda12_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+    pip install -U "jax[cuda12]"
 fi
 
 git clone ${TRITON_REPO_ORGANIZATION}/python_backend -b $PYTHON_BACKEND_REPO_TAG

qa/L0_backend_python/setup_python_enviroment.sh

@@ -56,7 +56,7 @@ conda update -n base -c defaults conda -y
 # been setup correctly.
 if [ ${PYTHON_ENV_VERSION} = "11" ]; then
     create_conda_env "3.11" "python-3-11"
-    conda install torch=2.6.0 -y
+    conda install pytorch=2.6.0 -y
     conda install -c conda-forge libstdcxx-ng=14 -y
     conda install numpy=1.23.5 -y
     EXPECTED_VERSION_STRING="Python version is 3.11, NumPy version is 1.23.5, and PyTorch version is 2.6.0"