Get CTK prereleases installing in CI on Linux x86

.github/actions/fetch_ctk/action.yml

@@ -11,6 +11,10 @@ inputs:
     required: true
   cuda-version:
     required: true
+  cuda-channel:
+    description: "CUDA package channel: stable redistributables or prerelease packages"
+    required: false
+    default: "stable"
   cuda-components:
     description: "A list of the CTK components to install as a comma-separated list. e.g. 'cuda_nvcc,cuda_nvrtc,cuda_cudart'"
     required: false
@@ -30,17 +34,34 @@ runs:
         # Use the runtime workspace mount so this also works inside container jobs.
         CTK_REDIST_TOOL="${GITHUB_WORKSPACE}/ci/tools/fetch_ctk_redistrib.py"
         CTK_CACHE_COMPONENTS=${{ inputs.cuda-components }}
-        CTK_JSON_URL="https://developer.download.nvidia.com/compute/cuda/redist/redistrib_${{ inputs.cuda-version }}.json"
-        CTK_CACHE_COMPONENTS="$(python "$CTK_REDIST_TOOL" filter-components \
-          --host-platform "${{ inputs.host-platform }}" \
-          --cuda-version "${{ inputs.cuda-version }}" \
-          --components "$CTK_CACHE_COMPONENTS" \
-          --metadata-url "$CTK_JSON_URL")"
+        CTK_PREVIEW_PACKAGES=
+        if [[ "${{ inputs.cuda-channel }}" == "stable" ]]; then
+          CTK_JSON_URL="https://developer.download.nvidia.com/compute/cuda/redist/redistrib_${{ inputs.cuda-version }}.json"
+          CTK_CACHE_COMPONENTS="$(python "$CTK_REDIST_TOOL" filter-components \
+            --host-platform "${{ inputs.host-platform }}" \
+            --cuda-version "${{ inputs.cuda-version }}" \
+            --components "$CTK_CACHE_COMPONENTS" \
+            --metadata-url "$CTK_JSON_URL")"
+        elif [[ "${{ inputs.cuda-channel }}" == "prerelease" ]]; then
+          CTK_PREVIEW_PACKAGES="$(python "$CTK_REDIST_TOOL" preview-packages \
+            --host-platform "${{ inputs.host-platform }}" \
+            --cuda-version "${{ inputs.cuda-version }}" \
+            --components "$CTK_CACHE_COMPONENTS")"
+          CTK_CACHE_COMPONENTS="$CTK_PREVIEW_PACKAGES"
+        else
+          echo "Unsupported CUDA package channel: ${{ inputs.cuda-channel }}" >&2
+          exit 1
+        fi
 
         HASH=$(echo -n "${CTK_CACHE_COMPONENTS}" | sha256sum | awk '{print $1}')
-        echo "CTK_CACHE_KEY=mini-ctk-${{ inputs.cuda-version }}-${{ inputs.host-platform }}-$HASH" >> $GITHUB_ENV
+        CHANNEL_CACHE_SEGMENT=
+        if [[ "${{ inputs.cuda-channel }}" != "stable" ]]; then
+          CHANNEL_CACHE_SEGMENT="-${{ inputs.cuda-channel }}"
+        fi
+        echo "CTK_CACHE_KEY=mini-ctk${CHANNEL_CACHE_SEGMENT}-${{ inputs.cuda-version }}-${{ inputs.host-platform }}-$HASH" >> $GITHUB_ENV
         echo "CTK_CACHE_FILENAME=mini-ctk-${{ inputs.cuda-version }}-${{ inputs.host-platform }}-$HASH.tar.gz" >> $GITHUB_ENV
         echo "CTK_CACHE_COMPONENTS=${CTK_CACHE_COMPONENTS}" >> $GITHUB_ENV
+        echo "CTK_PREVIEW_PACKAGES=${CTK_PREVIEW_PACKAGES}" >> $GITHUB_ENV
 
     - name: Install dependencies
       uses: ./.github/actions/install_unix_deps
@@ -68,48 +89,92 @@ runs:
         rm -rf $CACHE_TMP_DIR
         mkdir $CACHE_TMP_DIR
 
-        # The binary archives (redist) are guaranteed to be updated as part of the release posting.
-        # Use the runtime workspace mount so this also works inside container jobs.
-        CTK_REDIST_TOOL="${GITHUB_WORKSPACE}/ci/tools/fetch_ctk_redistrib.py"
-        CTK_BASE_URL="https://developer.download.nvidia.com/compute/cuda/redist/"
-        CTK_JSON_URL="$CTK_BASE_URL/redistrib_${{ inputs.cuda-version }}.json"
-        CTK_JSON_FILE="$CACHE_TMP_DIR/redistrib.json"
-        curl -LSs "$CTK_JSON_URL" -o "$CTK_JSON_FILE"
-        if [[ "${{ inputs.host-platform }}" == linux* ]]; then
-          function extract() {
-            tar -xvf $1 -C $CACHE_TMP_DIR --strip-components=1
-          }
-        elif [[ "${{ inputs.host-platform }}" == "win-64" ]]; then
-          function extract() {
-            _TEMP_DIR_=$(mktemp -d)
-            unzip $1 -d $_TEMP_DIR_
-            cp -r $_TEMP_DIR_/*/* $CACHE_TMP_DIR
-            rm -rf $_TEMP_DIR_
-            # see commit NVIDIA/cuda-python@69410f1d9228e775845ef6c8b4a9c7f37ffc68a5
-            chmod 644 $CACHE_TMP_DIR/LICENSE
+        if [[ "${{ inputs.cuda-channel }}" == "prerelease" ]]; then
+          if [[ "${{ inputs.host-platform }}" != linux* ]]; then
+            echo "CUDA prerelease package extraction currently supports Linux only" >&2
+            exit 1
+          fi
+
+          source /etc/os-release
+          DISTRO_CODENAME="${VERSION_CODENAME:-}"
+          case "$DISTRO_CODENAME" in
+            bookworm|jammy|noble|resolute|trixie) ;;
+            *)
+              echo "Unsupported distribution for CUDA prerelease packages: ${DISTRO_CODENAME:-unknown}" >&2
+              exit 1
+              ;;
+          esac
+
+          KEYRING_DEB="$CACHE_TMP_DIR/nvidia-preview-keyring.deb"
+          curl -fLSs "https://packages.nvidia.com/${DISTRO_CODENAME}/nvidia-preview-keyring.deb" -o "$KEYRING_DEB"
+          sudo dpkg -i "$KEYRING_DEB"
+          sudo apt-get update
+
+          DEB_DIR="$CACHE_TMP_DIR/debs"
+          DEB_ROOT="$CACHE_TMP_DIR/deb-root"
+          mkdir -p "$DEB_DIR/partial" "$DEB_ROOT"
+          DEB_DIR="$(realpath "$DEB_DIR")"
+          IFS=, read -ra PREVIEW_PACKAGES <<< "$CTK_PREVIEW_PACKAGES"
+          sudo apt-get install --yes --download-only --no-install-recommends \
+            -o "Dir::Cache::archives=$DEB_DIR" \
+            "${PREVIEW_PACKAGES[@]}"
+          for package in "$DEB_DIR"/*.deb; do
+            dpkg-deb -x "$package" "$DEB_ROOT"
+          done
+
+          CUDA_SHORT_VERSION="${{ inputs.cuda-version }}"
+          CUDA_SHORT_VERSION="${CUDA_SHORT_VERSION%.*}"
+          CUDA_PACKAGE_ROOT="$DEB_ROOT/usr/local/cuda-${CUDA_SHORT_VERSION}"
+          if [[ ! -d "$CUDA_PACKAGE_ROOT/include" ]]; then
+            echo "CUDA prerelease packages did not provide $CUDA_PACKAGE_ROOT/include" >&2
+            exit 1
+          fi
+          cp -a "$CUDA_PACKAGE_ROOT/." "$CACHE_TMP_DIR/"
+          rm -rf "$DEB_DIR" "$DEB_ROOT" "$KEYRING_DEB"
+        else
+          # The binary archives (redist) are guaranteed to be updated as part of the release posting.
+          # Use the runtime workspace mount so this also works inside container jobs.
+          CTK_REDIST_TOOL="${GITHUB_WORKSPACE}/ci/tools/fetch_ctk_redistrib.py"
+          CTK_BASE_URL="https://developer.download.nvidia.com/compute/cuda/redist/"
+          CTK_JSON_URL="$CTK_BASE_URL/redistrib_${{ inputs.cuda-version }}.json"
+          CTK_JSON_FILE="$CACHE_TMP_DIR/redistrib.json"
+          curl -fLSs "$CTK_JSON_URL" -o "$CTK_JSON_FILE"
+          if [[ "${{ inputs.host-platform }}" == linux* ]]; then
+            function extract() {
+              tar -xvf $1 -C $CACHE_TMP_DIR --strip-components=1
+            }
+          elif [[ "${{ inputs.host-platform }}" == "win-64" ]]; then
+            function extract() {
+              _TEMP_DIR_=$(mktemp -d)
+              unzip $1 -d $_TEMP_DIR_
+              cp -r $_TEMP_DIR_/*/* $CACHE_TMP_DIR
+              rm -rf $_TEMP_DIR_
+              # see commit NVIDIA/cuda-python@69410f1d9228e775845ef6c8b4a9c7f37ffc68a5
+              chmod 644 $CACHE_TMP_DIR/LICENSE
+            }
+          fi
+          function populate_cuda_path() {
+            # take the component name as a argument
+            function download() {
+              curl -fLSs $1 -o $2
+            }
+            CTK_COMPONENT=$1
+            CTK_COMPONENT_REL_PATH="$(python "$CTK_REDIST_TOOL" component-relative-path \
+              --host-platform "${{ inputs.host-platform }}" \
+              --component "$CTK_COMPONENT" \
+              --metadata-path "$CTK_JSON_FILE")"
+            CTK_COMPONENT_URL="${CTK_BASE_URL}/${CTK_COMPONENT_REL_PATH}"
+            CTK_COMPONENT_COMPONENT_FILENAME="$(basename $CTK_COMPONENT_REL_PATH)"
+            download $CTK_COMPONENT_URL $CTK_COMPONENT_COMPONENT_FILENAME
+            extract $CTK_COMPONENT_COMPONENT_FILENAME
+            rm $CTK_COMPONENT_COMPONENT_FILENAME
           }
+
+          # Get headers and shared libraries in place
+          for item in $(echo $CTK_CACHE_COMPONENTS | tr ',' ' '); do
+              populate_cuda_path "$item"
+          done
         fi
-        function populate_cuda_path() {
-          # take the component name as a argument
-          function download() {
-            curl -LSs $1 -o $2
-          }
-          CTK_COMPONENT=$1
-          CTK_COMPONENT_REL_PATH="$(python "$CTK_REDIST_TOOL" component-relative-path \
-            --host-platform "${{ inputs.host-platform }}" \
-            --component "$CTK_COMPONENT" \
-            --metadata-path "$CTK_JSON_FILE")"
-          CTK_COMPONENT_URL="${CTK_BASE_URL}/${CTK_COMPONENT_REL_PATH}"
-          CTK_COMPONENT_COMPONENT_FILENAME="$(basename $CTK_COMPONENT_REL_PATH)"
-          download $CTK_COMPONENT_URL $CTK_COMPONENT_COMPONENT_FILENAME
-          extract $CTK_COMPONENT_COMPONENT_FILENAME
-          rm $CTK_COMPONENT_COMPONENT_FILENAME
-        }
-
-        # Get headers and shared libraries in place
-        for item in $(echo $CTK_CACHE_COMPONENTS | tr ',' ' '); do
-            populate_cuda_path "$item"
-        done
         # TODO: check Windows
         if [[ "${{ inputs.host-platform }}" == linux* && -d "${CACHE_TMP_DIR}/lib" ]]; then
           mv $CACHE_TMP_DIR/lib $CACHE_TMP_DIR/lib64
@@ -120,8 +185,12 @@ runs:
         # Note: try to escape | and > ...
         tar -czvf ${CTK_CACHE_FILENAME} ${CACHE_TMP_DIR}
 
-        # "Move" files from temp dir to CUDA_PATH
-        CUDA_PATH="./cuda_toolkit"
+        # Populate the final CUDA_PATH directly (never stage through ./cuda_toolkit
+        # unconditionally — a second call with a different cuda-path would leave
+        # symlinks from a prerelease CTK at ./cuda_toolkit/include etc., causing
+        # "cp: cannot overwrite non-directory" failures on the next restore).
+        CUDA_PATH="${{ inputs.cuda-path }}"
+        rm -rf $CUDA_PATH
         mkdir -p $CUDA_PATH
         # Unfortunately we cannot use "rsync -av $CACHE_TMP_DIR/ $CUDA_PATH" because
         # not all runners have rsync pre-installed (or even installable, such as
@@ -144,7 +213,8 @@ runs:
       run: |
         ls -l
         CACHE_TMP_DIR="./cache_tmp_dir"
-        CUDA_PATH="./cuda_toolkit"
+        CUDA_PATH="${{ inputs.cuda-path }}"
+        rm -rf $CUDA_PATH
         mkdir -p $CUDA_PATH
         tar -xzvf $CTK_CACHE_FILENAME
         # Can't use rsync here, see above
@@ -155,12 +225,6 @@ runs:
           exit 1
         fi
 
-    - name: Move CTK to the specified location
-      if: ${{ inputs.cuda-path != './cuda_toolkit' }}
-      shell: bash --noprofile --norc -xeuo pipefail {0}
-      run: |
-        mv ./cuda_toolkit ${{ inputs.cuda-path }}
-
     - name: Set output environment variables
       shell: bash --noprofile --norc -xeuo pipefail {0}
       run: |

.github/workflows/build-docs.yml

@@ -64,8 +64,10 @@ jobs:
         run: |
           if [[ -f ci/versions.yml ]]; then
             BUILD_CTK_VER=$(yq '.cuda.build.version' ci/versions.yml)
+            BUILD_CTK_CHANNEL=$(yq '.cuda.build.channel // "stable"' ci/versions.yml)
           elif [[ -f ci/versions.json ]]; then
             BUILD_CTK_VER=$(jq -r '.cuda.build.version' ci/versions.json)
+            BUILD_CTK_CHANNEL=$(jq -r '.cuda.build.channel // "stable"' ci/versions.json)
           else
             echo "error: cannot find ci/versions.yml or ci/versions.json" >&2
             exit 1
@@ -75,6 +77,7 @@ jobs:
             exit 1
           fi
           echo "BUILD_CTK_VER=${BUILD_CTK_VER}" >> "$GITHUB_ENV"
+          echo "BUILD_CTK_CHANNEL=${BUILD_CTK_CHANNEL}" >> "$GITHUB_ENV"
 
       # TODO: This workflow runs on GH-hosted runner and cannot use the proxy cache
 
@@ -100,6 +103,7 @@ jobs:
         with:
           host-platform: linux-64
           cuda-version: ${{ env.BUILD_CTK_VER }}
+          cuda-channel: ${{ env.BUILD_CTK_CHANNEL }}
 
       - name: Set environment variables
         run: |

.github/workflows/build-wheel.yml

@@ -11,6 +11,10 @@ on:
       cuda-version:
         required: true
         type: string
+      cuda-channel:
+        required: false
+        type: string
+        default: stable
       prev-cuda-version:
         required: true
         type: string
@@ -164,6 +168,7 @@ jobs:
         with:
           host-platform: ${{ inputs.host-platform }}
           cuda-version: ${{ inputs.cuda-version }}
+          cuda-channel: ${{ inputs.cuda-channel }}
 
       - name: Build cuda.bindings wheel
         uses: pypa/cibuildwheel@294735312765b09d24a2fbec22660ce817587d55  # v4.1.0

.github/workflows/ci.yml

@@ -31,6 +31,7 @@ jobs:
     runs-on: ubuntu-latest
     outputs:
       CUDA_BUILD_VER: ${{ steps.get-vars.outputs.cuda_build_ver }}
+      CUDA_BUILD_CHANNEL: ${{ steps.get-vars.outputs.cuda_build_channel }}
       CUDA_PREV_BUILD_VER: ${{ steps.get-vars.outputs.cuda_prev_build_ver }}
     steps:
       - name: Checkout repository
@@ -43,6 +44,9 @@ jobs:
           cuda_build_ver=$(yq '.cuda.build.version' ci/versions.yml)
           echo "cuda_build_ver=$cuda_build_ver" >> $GITHUB_OUTPUT
 
+          cuda_build_channel=$(yq '.cuda.build.channel // "stable"' ci/versions.yml)
+          echo "cuda_build_channel=$cuda_build_channel" >> $GITHUB_OUTPUT
+
           cuda_prev_build_ver=$(yq '.cuda.prev_build.version' ci/versions.yml)
           echo "cuda_prev_build_ver=$cuda_prev_build_ver" >> $GITHUB_OUTPUT
 
@@ -250,6 +254,7 @@ jobs:
     with:
       host-platform: ${{ matrix.host-platform }}
       cuda-version: ${{ needs.ci-vars.outputs.CUDA_BUILD_VER }}
+      cuda-channel: ${{ needs.ci-vars.outputs.CUDA_BUILD_CHANNEL }}
       prev-cuda-version: ${{ needs.ci-vars.outputs.CUDA_PREV_BUILD_VER }}
 
   # See build-linux-64 for why build jobs are split by platform.
@@ -269,6 +274,7 @@ jobs:
     with:
       host-platform: ${{ matrix.host-platform }}
       cuda-version: ${{ needs.ci-vars.outputs.CUDA_BUILD_VER }}
+      cuda-channel: ${{ needs.ci-vars.outputs.CUDA_BUILD_CHANNEL }}
       prev-cuda-version: ${{ needs.ci-vars.outputs.CUDA_PREV_BUILD_VER }}
 
   # See build-linux-64 for why build jobs are split by platform.
@@ -282,12 +288,13 @@ jobs:
         host-platform:
           - win-64
     name: Build ${{ matrix.host-platform }}, CUDA ${{ needs.ci-vars.outputs.CUDA_BUILD_VER }}
-    if: ${{ github.repository_owner == 'nvidia' && !fromJSON(needs.should-skip.outputs.skip) && !fromJSON(needs.should-skip.outputs.doc-only) }}
+    if: ${{ github.repository_owner == 'nvidia' && !fromJSON(needs.should-skip.outputs.skip) && !fromJSON(needs.should-skip.outputs.doc-only) && needs.ci-vars.outputs.CUDA_BUILD_CHANNEL != 'prerelease' }}
     secrets: inherit
     uses: ./.github/workflows/build-wheel.yml
     with:
       host-platform: ${{ matrix.host-platform }}
       cuda-version: ${{ needs.ci-vars.outputs.CUDA_BUILD_VER }}
+      cuda-channel: ${{ needs.ci-vars.outputs.CUDA_BUILD_CHANNEL }}
       prev-cuda-version: ${{ needs.ci-vars.outputs.CUDA_PREV_BUILD_VER }}
 
   # NOTE: test-sdist jobs are split by platform (mirroring build-* and test-wheel-*)
@@ -307,14 +314,15 @@ jobs:
     with:
       host-platform: linux-64
       cuda-version: ${{ needs.ci-vars.outputs.CUDA_BUILD_VER }}
+      cuda-channel: ${{ needs.ci-vars.outputs.CUDA_BUILD_CHANNEL }}
 
   # See test-sdist-linux for why sdist test jobs are split by platform.
   test-sdist-windows:
     needs:
       - ci-vars
       - should-skip
     name: Test sdist win-64
-    if: ${{ github.repository_owner == 'nvidia' && !fromJSON(needs.should-skip.outputs.skip) && !fromJSON(needs.should-skip.outputs.doc-only) }}
+    if: ${{ github.repository_owner == 'nvidia' && !fromJSON(needs.should-skip.outputs.skip) && !fromJSON(needs.should-skip.outputs.doc-only) && needs.ci-vars.outputs.CUDA_BUILD_CHANNEL != 'prerelease' }}
     secrets: inherit
     uses: ./.github/workflows/test-sdist-windows.yml
     with:
@@ -383,7 +391,7 @@ jobs:
         host-platform:
           - win-64
     name: Test ${{ matrix.host-platform }}
-    if: ${{ github.repository_owner == 'nvidia' && !fromJSON(needs.should-skip.outputs.doc-only) }}
+    if: ${{ github.repository_owner == 'nvidia' && !fromJSON(needs.should-skip.outputs.doc-only) && needs.ci-vars.outputs.CUDA_BUILD_CHANNEL != 'prerelease' }}
     permissions:
       contents: read  # This is required for actions/checkout
     needs:
@@ -421,6 +429,7 @@ jobs:
     if: always()
     runs-on: ubuntu-latest
     needs:
+      - ci-vars
       - should-skip
       - detect-changes
       - test-sdist-linux
@@ -447,6 +456,7 @@ jobs:
           fi
 
           doc_only="${{ needs.should-skip.outputs.doc-only }}"
+          cuda_build_channel="${{ needs.ci-vars.outputs.CUDA_BUILD_CHANNEL }}"
           status="success"
           check_result() {
             name=$1; expected=$2; result=$3
@@ -458,16 +468,28 @@ jobs:
           }
 
           # always expected to succeed (even in [doc-only] mode)
+          check_result "ci-vars"        "success" "${{ needs.ci-vars.result }}"
           check_result "should-skip"    "success" "${{ needs.should-skip.result }}"
           check_result "detect-changes" "success" "${{ needs.detect-changes.result }}"
           check_result "doc"            "success" "${{ needs.doc.result }}"
 
-          # [doc-only] flips these from 'success' to 'skipped'
-          if [[ "$doc_only" == "true" ]]; then expected="skipped"; else expected="success"; fi
-          check_result "test-sdist-linux"   "$expected" "${{ needs.test-sdist-linux.result }}"
-          check_result "test-sdist-windows" "$expected" "${{ needs.test-sdist-windows.result }}"
-          check_result "test-linux-64"      "$expected" "${{ needs.test-linux-64.result }}"
-          check_result "test-linux-aarch64" "$expected" "${{ needs.test-linux-aarch64.result }}"
-          check_result "test-windows"       "$expected" "${{ needs.test-windows.result }}"
+          # [doc-only] skips all tests. Prerelease CUDA packages currently support
+          # Linux only, so Windows tests are also skipped for that channel.
+          if [[ "$doc_only" == "true" ]]; then
+            linux_expected="skipped"
+            windows_expected="skipped"
+          else
+            linux_expected="success"
+            if [[ "$cuda_build_channel" == "prerelease" ]]; then
+              windows_expected="skipped"
+            else
+              windows_expected="success"
+            fi
+          fi
+          check_result "test-sdist-linux"   "$linux_expected" "${{ needs.test-sdist-linux.result }}"
+          check_result "test-sdist-windows" "$windows_expected" "${{ needs.test-sdist-windows.result }}"
+          check_result "test-linux-64"      "$linux_expected" "${{ needs.test-linux-64.result }}"
+          check_result "test-linux-aarch64" "$linux_expected" "${{ needs.test-linux-aarch64.result }}"
+          check_result "test-windows"       "$windows_expected" "${{ needs.test-windows.result }}"
 
           [[ "$status" == "success" ]]