From 6ab10e204570488966c2ee8bc1617dd27d7d8dbb Mon Sep 17 00:00:00 2001
From: Won-Kyu Park <wkpark@gmail.com>
Date: Wed, 7 Feb 2024 20:47:42 +0900
Subject: [PATCH] reduce CAPABILITY for pull_requests

---
 .github/workflows/python-package.yml | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml
index b2a8be65c..16ed0cf5f 100644
--- a/.github/workflows/python-package.yml
+++ b/.github/workflows/python-package.yml
@@ -212,6 +212,14 @@ jobs:
 
         nvcc --version
 
+    - name: Select COMPUTE_CAPABILITY
+      shell: bash -el {0}
+      run: |
+        # set COMPUTE_CAPABILITY
+        COMPUTE_CAPABILITY="61;75;86;89"
+        [[ "${{ env.GITHUB_REF }}" = refs/tags/* ]] && COMPUTE_CAPABILITY="50;52;60;61;62;70;72;75;80;86;87;89;90"
+        echo "COMPUTE_CAPABILITY=$COMPUTE_CAPABILITY" >> "$GITHUB_ENV"
+
     - name: Update environment
       if: startsWith(matrix.os, 'windows')
       run: mamba env update -n bnb-env -f environment-bnb.yml
@@ -228,7 +236,7 @@ jobs:
         -G Ninja ${{ env.DCMAKE_CUDA_COMPILER }} \
         ${{ env.CXX_COMPILER }} \
         -DCMAKE_BUILD_TYPE=${{ matrix.build_type }} \
-        -DCOMPUTE_CAPABILITY="50;52;60;61;62;70;72;75;80;86;87;89;90" \
+        -DCOMPUTE_CAPABILITY="${{ env.COMPUTE_CAPABILITY }}" \
         -DCOMPUTE_BACKEND=cuda \
         -S .
 
@@ -239,7 +247,7 @@ jobs:
       uses: addnab/docker-run-action@v3
       with:
         image: ${{ format('nvidia/cuda:{0}-{1}', matrix.cuda-version, 'devel-ubuntu22.04') }}
-        options: --platform linux/${{ matrix.arch }} -w /src -v ${{ github.workspace }}:/src
+        options: --platform linux/${{ matrix.arch }} -w /src -v ${{ github.workspace }}:/src -e "COMPUTE_CAPABILITY"
         run: |
           apt-get update
           DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends cmake python3 python3-pip
@@ -250,7 +258,7 @@ jobs:
             cmake -B build \
             -G Ninja \
             -DCMAKE_BUILD_TYPE=${{ matrix.build_type }} \
-            -DCOMPUTE_CAPABILITY="50;52;60;61;62;70;72;75;80;86;87;89;90" \
+            -DCOMPUTE_CAPABILITY="${{ env.COMPUTE_CAPABILITY }}" \
             -DCOMPUTE_BACKEND=cuda \
             -DNO_CUBLASLT=$NO_CUBLAST \
             -S .