Merging py-linux-gpu-cuda1.yml with py-linux-gpu.yml

microsoft · Nov 8, 2023 · bf64dc2 · bf64dc2
1 parent 646a8b6
commit bf64dc2
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 4 deletions.
diff --git a/tools/ci_build/github/azure-pipelines/templates/py-linux-gpu.yml b/tools/ci_build/github/azure-pipelines/templates/py-linux-gpu.yml
@@ -17,7 +17,18 @@ parameters:
    - Release
    - RelWithDebInfo
    - MinSizeRel
-
+- name: docker_base_image
+  type: string
+  default: nvidia/cuda:11.8.0-cudnn8-devel-ubi8
+  values:
+   - nvidia/cuda:11.8.0-cudnn8-devel-ubi8
+   - nvidia/cuda:12.2.2-cudnn8-devel-ubi8
+- name: trt_version
+  type: string
+  default: 8.6.1.6-1.cuda11.8
+  values:
+    - 8.6.1.6-1.cuda11.8
+    - 8.6.1.6-1.cuda12.0
 jobs:
 - job: Linux_py_GPU_Wheels_${{ parameters.arch }}
   timeoutInMinutes: 240
@@ -40,8 +51,8 @@ jobs:
         Context: tools/ci_build/github/linux/docker
         DockerBuildArgs: "
         --network=host 
-        --build-arg BASEIMAGE=nvidia/cuda:11.8.0-cudnn8-devel-ubi8 
-        --build-arg TRT_VERSION=8.6.1.6-1.cuda11.8 
+        --build-arg BASEIMAGE=${{ parameters.docker_base_image }}
+        --build-arg TRT_VERSION=${{ parameters.trt_version }}
         --build-arg BUILD_UID=$( id -u )
         --build-arg PLATFORM=${{ parameters.arch }}
         "

diff --git a/tools/ci_build/github/azure-pipelines/templates/py-packaging-stage-cuda12.yml b/tools/ci_build/github/azure-pipelines/templates/py-packaging-stage-cuda12.yml
@@ -66,9 +66,11 @@ stages:
 
 
   - ${{ if eq(parameters.enable_linux_gpu, true) }}:
-      - template: py-linux-gpu-cuda12.yml
+      - template: py-linux-gpu.yml
         parameters:
           arch: 'x86_64'
           machine_pool: 'onnxruntime-Ubuntu2004-AMD-CPU'
           extra_build_arg: ${{ parameters.build_py_parameters }}
           cmake_build_type: ${{ parameters.cmake_build_type }}
+          docker_base_image: nvidia/cuda:12.2.2-cudnn8-devel-ubi8
+          trt_version: 8.6.1.6-1.cuda12.0