fix major version of CUDNN when build in CI

chengzeyi · Nov 14, 2023 · 655ae7c · 655ae7c
1 parent f359a6e
commit 655ae7c
Showing 1 changed file with 15 additions and 4 deletions.
diff --git a/.github/workflows/wheels_build.yml b/.github/workflows/wheels_build.yml
@@ -33,6 +33,11 @@ on:
         required: true
         type: string
         description: "Example: 117 for 11.7"
+      cudnn_version_major:
+        required: false
+        default: '8'
+        type: string
+        description: "Example: 8
 
 # this yaml file can be cleaned up using yaml anchors, but they're not supported in github actions yet
 # https://github.com/actions/runner/issues/1182
@@ -141,9 +146,9 @@ jobs:
           # Set build version to x.x.x.devYYYYMMDD+torchxxxcu111
           torch_version_suffix=torch$(echo ${{ inputs.torch_version }} | sed 's/\.//g')
           cuda_version_suffix=${{ steps.cuda_info.outputs.CUDA_VERSION_SUFFIX }}
-          nightly_tag=$([[ $VERSION_SOURCE == 'tag' ]] && echo '' || echo '.dev'`date +%Y%m%d`)
-          echo "BUILD_VERSION=$version$nightly_tag+$torch_version_suffix$cuda_version_suffix" >> ${GITHUB_ENV}
-          echo "BUILD_VERSION=$version$nightly_tag+$torch_version_suffix$cuda_version_suffix" >> ${GITHUB_OUTPUT}
+          nightly_tag=$([[ ${VERSION_SOURCE} == 'tag' ]] && echo '' || echo '.dev'`date +%Y%m%d`)
+          echo "BUILD_VERSION=${version}${nightly_tag}+${torch_version_suffix}${cuda_version_suffix}" >> ${GITHUB_ENV}
+          echo "BUILD_VERSION=${version}${nightly_tag}+${torch_version_suffix}${cuda_version_suffix}" >> ${GITHUB_OUTPUT}
           cat ${GITHUB_ENV}
       - run: echo "sfast-${BUILD_VERSION}"
       - run: echo "release version"
@@ -163,7 +168,13 @@ jobs:
           python: ${{ inputs.python }}
 
       - name: Install dependencies
-        run: $PY -m pip install wheel setuptools ninja twine ${{ steps.cuda_info.outputs.CUDNN_PYPI_PACKAGE }} -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu${{ inputs.cuda_short_version }} --no-cache-dir
+        env:
+          CUDNN_VERSION_MAJOR: ${{ inputs.cudnn_version_major }}
+          CUDNN_PYPI_PACKAGE: ${{ steps.cuda_info.outputs.CUDNN_PYPI_PACKAGE }}
+        run: |
+          cudnn_next_version_major=$((${CUDNN_VERSION_MAJOR} + 1))
+          cudnn_package_name="${CUDNN_PYPI_PACKAGE}>=${CUDNN_VERSION_MAJOR}.0.0.0,<${cudnn_next_version_major}.0.0.0"
+          $PY -m pip install wheel setuptools ninja twine "${cudnn_package_name}" -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu${{ inputs.cuda_short_version }} --no-cache-dir
 
       - name: Build wheel
         run: |