|
| 1 | +name: Fetch mini CTK |
| 2 | + |
| 3 | +description: Fetch (or create) a mini CUDA Toolkit from cache |
| 4 | + |
| 5 | +inputs: |
| 6 | + host-platform: |
| 7 | + required: true |
| 8 | + type: string |
| 9 | + cuda-version: |
| 10 | + required: true |
| 11 | + type: string |
| 12 | + fail-on-ctk-cache-miss: |
| 13 | + required: true |
| 14 | + type: boolean |
| 15 | + |
| 16 | +runs: |
| 17 | + using: composite |
| 18 | + steps: |
| 19 | + - name: Set up CTK cache variable |
| 20 | + shell: bash --noprofile --norc -xeuo pipefail {0} |
| 21 | + run: | |
| 22 | + echo "CTK_CACHE_KEY=mini-ctk-${{ inputs.cuda-version }}-${{ inputs.host-platform }}" >> $GITHUB_ENV |
| 23 | + echo "CTK_CACHE_FILENAME=mini-ctk-${{ inputs.cuda-version }}-${{ inputs.host-platform }}.tar.gz" >> $GITHUB_ENV |
| 24 | +
|
| 25 | + - name: Download CTK cache |
| 26 | + id: ctk-get-cache |
| 27 | + uses: actions/cache/restore@v4 |
| 28 | + continue-on-error: true |
| 29 | + with: |
| 30 | + key: ${{ env.CTK_CACHE_KEY }} |
| 31 | + path: ./${{ env.CTK_CACHE_FILENAME }} |
| 32 | + fail-on-cache-miss: ${{ inputs.fail-on-ctk-cache-miss }} |
| 33 | + |
| 34 | + - name: Get CUDA components |
| 35 | + if: ${{ steps.ctk-get-cache.outputs.cache-hit != 'true' }} |
| 36 | + shell: bash --noprofile --norc -xeuo pipefail {0} |
| 37 | + run: | |
| 38 | + CUDA_PATH="$(pwd)/cuda_toolkit" |
| 39 | + mkdir $CUDA_PATH |
| 40 | +
|
| 41 | + # The binary archives (redist) are guaranteed to be updated as part of the release posting. |
| 42 | + CTK_BASE_URL="https://developer.download.nvidia.com/compute/cuda/redist/" |
| 43 | + CTK_JSON_URL="$CTK_BASE_URL/redistrib_${{ inputs.cuda-version }}.json" |
| 44 | + if [[ "${{ inputs.host-platform }}" == linux* ]]; then |
| 45 | + if [[ "${{ inputs.host-platform }}" == "linux-x64" ]]; then |
| 46 | + CTK_SUBDIR="linux-x86_64" |
| 47 | + elif [[ "${{ inputs.host-platform }}" == "linux-aarch64" ]]; then |
| 48 | + CTK_SUBDIR="linux-sbsa" |
| 49 | + fi |
| 50 | + function extract() { |
| 51 | + tar -xvf $1 -C $CUDA_PATH --strip-components=1 |
| 52 | + } |
| 53 | + elif [[ "${{ inputs.host-platform }}" == "win-x64" ]]; then |
| 54 | + CTK_SUBDIR="windows-x86_64" |
| 55 | + function extract() { |
| 56 | + _TEMP_DIR_=$(mktemp -d) |
| 57 | + unzip $1 -d $_TEMP_DIR_ |
| 58 | + cp -r $_TEMP_DIR_/*/* $CUDA_PATH |
| 59 | + rm -rf $_TEMP_DIR_ |
| 60 | + } |
| 61 | + fi |
| 62 | + function populate_cuda_path() { |
| 63 | + # take the component name as a argument |
| 64 | + function download() { |
| 65 | + curl -kLSs $1 -o $2 |
| 66 | + } |
| 67 | + CTK_COMPONENT=$1 |
| 68 | + CTK_COMPONENT_REL_PATH="$(curl -s $CTK_JSON_URL | |
| 69 | + python -c "import sys, json; print(json.load(sys.stdin)['${CTK_COMPONENT}']['${CTK_SUBDIR}']['relative_path'])")" |
| 70 | + CTK_COMPONENT_URL="${CTK_BASE_URL}/${CTK_COMPONENT_REL_PATH}" |
| 71 | + CTK_COMPONENT_COMPONENT_FILENAME="$(basename $CTK_COMPONENT_REL_PATH)" |
| 72 | + download $CTK_COMPONENT_URL $CTK_COMPONENT_COMPONENT_FILENAME |
| 73 | + extract $CTK_COMPONENT_COMPONENT_FILENAME |
| 74 | + rm $CTK_COMPONENT_COMPONENT_FILENAME |
| 75 | + } |
| 76 | +
|
| 77 | + # Get headers and shared libraries in place |
| 78 | + # Note: the existing artifact would need to be manually deleted (ex: through web UI) |
| 79 | + # if this list is changed, as the artifact actions do not offer any option for us to |
| 80 | + # invalidate the artifact. |
| 81 | + populate_cuda_path cuda_nvcc |
| 82 | + populate_cuda_path cuda_cudart |
| 83 | + populate_cuda_path cuda_nvrtc |
| 84 | + populate_cuda_path cuda_profiler_api |
| 85 | + populate_cuda_path libnvjitlink |
| 86 | + ls -l $CUDA_PATH |
| 87 | +
|
| 88 | + # Prepare the cache |
| 89 | + # Note: try to escape | and > ... |
| 90 | + tar -czvf ${CTK_CACHE_FILENAME} ${CUDA_PATH} |
| 91 | +
|
| 92 | + # Note: the headers will be copied into the cibuildwheel manylinux container, |
| 93 | + # so setting the CUDA_PATH env var here is meaningless. |
| 94 | +
|
| 95 | + - name: Upload CTK cache |
| 96 | + if: ${{ always() && |
| 97 | + steps.ctk-get-cache.outputs.cache-hit != 'true' }} |
| 98 | + uses: actions/cache/save@v4 |
| 99 | + with: |
| 100 | + key: ${{ env.CTK_CACHE_KEY }} |
| 101 | + path: ./${{ env.CTK_CACHE_FILENAME }} |
| 102 | + |
| 103 | + - name: Restore CTK cache |
| 104 | + if: ${{ steps.ctk-get-cache.outputs.cache-hit == 'true' }} |
| 105 | + shell: bash --noprofile --norc -xeuo pipefail {0} |
| 106 | + run: | |
| 107 | + ls -l |
| 108 | + CUDA_PATH="$(pwd)/cuda_toolkit" |
| 109 | + tar -xzvf $CTK_CACHE_FILENAME |
| 110 | + ls -l $CUDA_PATH |
| 111 | + if [ ! -d "$CUDA_PATH/include" ]; then |
| 112 | + exit 1 |
| 113 | + fi |
| 114 | +
|
| 115 | + echo "CUDA_PATH=${CUDA_PATH}" >> $GITHUB_ENV |
| 116 | + echo "PATH=${PATH}:${CUDA_PATH}/bin" >> $GITHUB_ENV |
| 117 | + echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${CUDA_PATH}/lib" >> $GITHUB_ENV |
0 commit comments