[CI] Add smoke test for wheel release

Author: lanluo-nvidia

.github/workflows/mlir-tensorrt-build-test.yml

@@ -181,7 +181,7 @@ jobs:
 
       # Run NCCL long tests
       - name: Run Long Tests With CUDA:${{ matrix.cuda }} + TensorRT:${{ matrix.trt }}
-        id: long-tests
+        id: nccl-long-tests
         if: ${{ inputs.channel == 'nightly' || (github.event_name == 'push' && github.ref == 'refs/heads/main') }}
         env:
           LONG_TESTS: "ON"

.github/workflows/mlir-tensorrt-release.yml

@@ -27,7 +27,6 @@ jobs:
       CPM_SOURCE_CACHE: ${{ github.workspace }}/mlir-tensorrt/.cache.cpm
       CMAKE_PRESET: distribution-wheels
       CCACHE_RESTORE_KEY: mlir-tensorrt-ccache-v1-${{ matrix.arch }}-distribution-wheels
-      CPM_KEY: mlir-tensorrt-cpm-v1
       CPM_RESTORE_KEY: mlir-tensorrt-cpm-v1
     runs-on: ${{ matrix.github_runner }}
     timeout-minutes: 120
@@ -68,6 +67,13 @@ jobs:
                  | sort | sha256sum | cut -d' ' -f1)
           echo "key=${{ env.CCACHE_RESTORE_KEY }}-${hash}" >> $GITHUB_OUTPUT
 
+      # Compute cpm key after checkout
+      - name: Compute CPM Key
+        id: cpm-key
+        run: |
+          hash=$(sha256sum mlir-tensorrt/DependencyProvider.cmake | cut -d' ' -f1)
+          echo "key=${{ env.CPM_RESTORE_KEY }}-${hash}" >> $GITHUB_OUTPUT
+
       # Restore cache, if exists.
       - name: Restore CCache
         id: restore-ccache
@@ -83,7 +89,7 @@ jobs:
         id: restore-cpm
         uses: actions/cache/restore@v4
         with:
-          key: ${{ env.CPM_KEY }}
+          key: ${{ steps.cpm-key.outputs.key }}
           enableCrossOsArchive: true
           restore-keys: |
             ${{ env.CPM_RESTORE_KEY }}
@@ -102,11 +108,9 @@ jobs:
           set -euo pipefail
           set -x
           cd mlir-tensorrt
-          export MLIR_TRT_DOWNLOAD_TENSORRT_VERSION=${{ matrix.trt }}
           # TODO: remove this, once patch is in our rockylinux prebuilt image
           dnf install -yq patch
-          # build wheels for mlir-tensorrt-compiler, mlir-tensorrt-runtime
-          PACKAGES="mlir_tensorrt_compiler mlir_tensorrt_runtime" \
+          # build wheels for all available packages
           ./build_tools/scripts/cicd-build-wheels.sh
 
       # Save ccache when cache is not hit or cache was a fallback(cache-matched is not the same as the cache key)
@@ -122,22 +126,42 @@ jobs:
       - name: Save CPM Cache
         # cpm cache is shared across x86_64 and aarch64, we let only x86_64 to save cpm cache when in cache miss case
         # this is to avoid both x86_64 and aarch64 to save cpm cache when in cache miss case
-        if: ${{ matrix.arch == 'x86_64' &&
-              (
-                steps.restore-cpm.outputs.cache-hit != 'true' ||
-                steps.restore-cpm.outputs['cache-matched-key'] != env.CPM_KEY
-              )
-            }}
+        if: ${{ matrix.arch == 'x86_64' && steps.restore-cpm.outputs.cache-hit != 'true' }}
         uses: actions/cache/save@v4
         with:
-          key: ${{ env.CPM_KEY }}
+          key: ${{ steps.cpm-key.outputs.key }}
           enableCrossOsArchive: true
           # exclude only works for the relative path pattern
           path: |
             mlir-tensorrt/.cache.cpm/*
             !mlir-tensorrt/.cache.cpm/tensorrt
             !mlir-tensorrt/.cache.cpm/tensorrt/**
 
+      - name: Smoke Test Wheels With CUDA:${{ matrix.cuda }} + TensorRT:${{ matrix.trt }}
+        env:
+          ENABLE_ASAN: "OFF"
+        run: |
+          set -euo pipefail
+          set -x
+          cd mlir-tensorrt
+          uv sync --python 3.10 --extra cu13
+          source .venv/bin/activate
+          export JAX_PLATFORMS="mlir_tensorrt,cpu"
+
+          # Compute TensorRT library path from CPM cache by locating libnvinfer.so
+          echo "CPM_SOURCE_CACHE: ${{ env.CPM_SOURCE_CACHE }}"
+          CPM_TRT_DIR="${{ env.CPM_SOURCE_CACHE }}/tensorrt"
+          TensorRT_LIB_PATH="$(find "${CPM_TRT_DIR}" -type f -name 'libnvinfer.so*' 2>/dev/null | head -n1 | xargs -r dirname)"
+          if [[ -z "${TensorRT_LIB_PATH}" ]]; then
+            echo "Error: Could not locate libnvinfer.so under ${CPM_TRT_DIR}" >&2
+            exit 1
+          fi
+          echo "TensorRT_LIB_PATH: ${TensorRT_LIB_PATH}"
+          export LD_LIBRARY_PATH="${TensorRT_LIB_PATH}:$LD_LIBRARY_PATH"
+          uv pip install dist/mlir_tensorrt_jax-*-cp310-cp310*.whl
+          # smoke test wheels
+          python3 -m pytest integrations/PJRT/test/JAX/ -v
+
       # Upload wheels to GitHub Actions artifact
       - name: Upload Wheels
         uses: actions/upload-artifact@v4

mlir-tensorrt/build_tools/scripts/cicd-build-wheels.sh

@@ -62,8 +62,13 @@ _raw_py_versions="${PYTHON_VERSIONS:-${python_versions:-"3.10 3.11 3.12 3.13"}}"
 _raw_py_versions="${_raw_py_versions//,/ }"
 read -r -a PY_VERSIONS <<< "${_raw_py_versions}"
 
-# Resolve packages
-_raw_packages="${PACKAGES:-"mlir_tensorrt_compiler mlir_tensorrt_runtime"}"
+# Resolve packages: if PACKAGES is set, only build those; otherwise build all defaults
+if [[ -n "${PACKAGES:-}" ]]; then
+  _raw_packages="${PACKAGES}"
+else
+  # default: build all available packages
+  _raw_packages="mlir_tensorrt_compiler mlir_tensorrt_runtime pjrt"
+fi
 _raw_packages="${_raw_packages//,/ }"
 read -r -a PACKAGES <<< "${_raw_packages}"
 
@@ -76,9 +81,23 @@ function build_package() {
   done
 }
 
-build_package "integrations/python/mlir_tensorrt_compiler"
-build_package "integrations/python/mlir_tensorrt_runtime"
-build_package "integrations/PJRT/python"
+# Map package keys to paths and build
+for pkg_key in "${PACKAGES[@]}"; do
+  case "${pkg_key}" in
+    mlir_tensorrt_compiler)
+      build_package "integrations/python/mlir_tensorrt_compiler"
+      ;;
+    mlir_tensorrt_runtime)
+      build_package "integrations/python/mlir_tensorrt_runtime"
+      ;;
+    pjrt|PJRT)
+      build_package "integrations/PJRT/python"
+      ;;
+    *)
+      echo "WARNING: Unknown package key '${pkg_key}'. Supported: mlir_tensorrt_compiler, mlir_tensorrt_runtime, pjrt" >&2
+      ;;
+  esac
+done
 
 # print the size of the ccache and cpm source cache
 echo "🔨🧪 Printing the size of the ccache"