# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0

name: Fetch mini CTK

description: Fetch (or create) a mini CUDA Toolkit from cache

inputs:
  host-platform:
    required: true
  cuda-version:
    required: true
  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
    default: "cuda_nvcc,cuda_cudart,cuda_crt,libnvvm,cuda_nvrtc,cuda_profiler_api,cuda_cccl,cuda_cupti,libnvjitlink,libcufile,libnvfatbin,libcudla,cuda_gdb"
  cuda-path:
    description: "where the CTK components will be installed to, relative to $PWD"
    required: false
    default: "./cuda_toolkit"
  redist-base-url:
    description: >
      Base URL of a redistrib server (manifest and component archives).
      Defaults to the public NVIDIA developer server.
      Override to install from an alternate redistrib server.
    required: false
    default: "https://developer.download.nvidia.com/compute/cuda/redist"

runs:
  using: composite
  steps:
    - name: Install Unix dependencies
      if: ${{ runner.os == 'Linux' }}
      uses: ./.github/actions/install_unix_deps
      with:
        dependencies: "zstd curl xz-utils"
        dependent_exes: "zstd curl xz"

    - name: Install Windows dependencies
      if: ${{ runner.os == 'Windows' }}
      uses: ./.github/actions/install_windows_deps
      with:
        dependent_exes: "curl unzip tar sha256sum cygpath"

    - name: Set up CTK cache variable
      shell: bash --noprofile --norc -xeuo pipefail {0}
      run: |
        # 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="${{ inputs.redist-base-url }}"
        # A trailing slash would defeat the public-server comparison below and double up in URLs.
        CTK_BASE_URL="${CTK_BASE_URL%/}"
        # RUNNER_TEMP persists across steps, so "Get CUDA components" reuses this file.
        CTK_JSON_FILE="$RUNNER_TEMP/redistrib-${{ inputs.cuda-version }}.json"
        curl -LSs --fail --retry 3 --retry-delay 2 "${CTK_BASE_URL}/redistrib_${{ inputs.cuda-version }}.json" -o "$CTK_JSON_FILE"

        # Pin the metadata JSON: the per-archive sha256 checks below are only
        # trustworthy if the JSON that carries those hashes is itself trusted.
        # The pins describe the public server; alternate servers only warn.
        if [[ "$CTK_BASE_URL" == "https://developer.download.nvidia.com/compute/cuda/redist" ]]; then
          METADATA_SOURCE=public
        else
          METADATA_SOURCE=custom
        fi
        python "$CTK_REDIST_TOOL" verify-metadata \
          --cuda-version "${{ inputs.cuda-version }}" \
          --metadata-path "$CTK_JSON_FILE" \
          --metadata-source "$METADATA_SOURCE"

        # Pre-process the component list to ensure hash uniqueness
        CTK_CACHE_COMPONENTS="$(python "$CTK_REDIST_TOOL" filter-components \
          --host-platform "${{ inputs.host-platform }}" \
          --cuda-version "${{ inputs.cuda-version }}" \
          --components "${{ inputs.cuda-components }}" \
          --metadata-path "$CTK_JSON_FILE")"

        # Alternate redist servers (non-default) publish different manifests for the
        # same CUDA version; keying on the full manifest hash keeps them from restoring
        # each other's toolkit (a truncated hash could be collided on purpose).
        # Hash via stdin: with a backslash in the path (Windows RUNNER_TEMP), sha256sum
        # prefixes its output line with "\", which would end up in the cache key.
        MANIFEST_SHA256="$(sha256sum < "$CTK_JSON_FILE" | awk '{print $1}')"
        HASH=$(echo -n "${CTK_CACHE_COMPONENTS}" | sha256sum | awk '{print $1}')
        {
          echo "CTK_CACHE_KEY=mini-ctk-${{ inputs.cuda-version }}-${{ inputs.host-platform }}-$HASH-m$MANIFEST_SHA256"
          echo "CTK_CACHE_FILENAME=mini-ctk-${{ inputs.cuda-version }}-${{ inputs.host-platform }}-$HASH-m$MANIFEST_SHA256.tar.gz"
          echo "CTK_CACHE_COMPONENTS=${CTK_CACHE_COMPONENTS}"
          echo "CTK_BASE_URL=${CTK_BASE_URL}"
          echo "CTK_JSON_FILE=${CTK_JSON_FILE}"
        } >> "$GITHUB_ENV"

    - name: Download CTK cache
      id: ctk-get-cache
      uses: actions/cache/restore@55cc8345863c7cc4c66a329aec7e433d2d1c52a9  # v6.1.0
      continue-on-error: true
      with:
        key: ${{ env.CTK_CACHE_KEY }}
        path: ./${{ env.CTK_CACHE_FILENAME }}
        fail-on-cache-miss: false

    - name: Get CUDA components
      if: ${{ steps.ctk-get-cache.outputs.cache-hit != 'true' }}
      shell: bash --noprofile --norc -xeuo pipefail {0}
      run: |
        # Everything under this folder is packed and stored in the GitHub Cache space,
        # and unpacked after retrieving from the cache.
        CACHE_TMP_DIR="./cache_tmp_dir"
        rm -rf $CACHE_TMP_DIR
        mkdir $CACHE_TMP_DIR

        # 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 and the verified manifest at CTK_JSON_FILE come from
        # "Set up CTK cache variable".
        if [[ "${{ inputs.host-platform }}" == linux* ]]; then
          function extract() {
            tar -xvf $1 -C $CACHE_TMP_DIR --strip-components=1
          }
        elif [[ "${{ inputs.host-platform }}" == win* ]]; 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 -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
          # Verify the downloaded archive against the sha256 in the (pinned)
          # metadata before extracting; fail closed on any mismatch.
          CTK_COMPONENT_SHA256="$(python "$CTK_REDIST_TOOL" component-sha256 \
            --host-platform "${{ inputs.host-platform }}" \
            --component "$CTK_COMPONENT" \
            --metadata-path "$CTK_JSON_FILE")"
          echo "${CTK_COMPONENT_SHA256}  ${CTK_COMPONENT_COMPONENT_FILENAME}" | sha256sum --check --strict -
          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
        fi
        ls -l $CACHE_TMP_DIR

        # Prepare the cache
        # Note: try to escape | and > ...
        tar -czvf ${CTK_CACHE_FILENAME} ${CACHE_TMP_DIR}

        # 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
        # Git Bash). We do it in the dumb way.
        cp -r $CACHE_TMP_DIR/* $CUDA_PATH
        rm -rf $CACHE_TMP_DIR
        ls -l $CUDA_PATH

    - name: Upload CTK cache
      if: ${{ !cancelled() &&
              steps.ctk-get-cache.outputs.cache-hit != 'true' }}
      uses: actions/cache/save@55cc8345863c7cc4c66a329aec7e433d2d1c52a9  # v6.1.0
      with:
        key: ${{ env.CTK_CACHE_KEY }}
        path: ./${{ env.CTK_CACHE_FILENAME }}

    - name: Restore CTK cache
      if: ${{ steps.ctk-get-cache.outputs.cache-hit == 'true' }}
      shell: bash --noprofile --norc -xeuo pipefail {0}
      run: |
        ls -l
        CACHE_TMP_DIR="./cache_tmp_dir"
        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
        cp -r $CACHE_TMP_DIR/* $CUDA_PATH
        rm -rf $CACHE_TMP_DIR $CTK_CACHE_FILENAME
        ls -l $CUDA_PATH

    - name: Set output environment variables
      shell: bash --noprofile --norc -xeuo pipefail {0}
      run: |
        # mimics actual CTK installation
        if [[ "${{ inputs.host-platform }}" == linux* ]]; then
          CUDA_PATH=$(realpath "${{ inputs.cuda-path }}")
          echo "${CUDA_PATH}/bin" >> $GITHUB_PATH
          echo "LD_LIBRARY_PATH=${CUDA_PATH}/lib64:${LD_LIBRARY_PATH:-}" >> $GITHUB_ENV
        elif [[ "${{ inputs.host-platform }}" == win* ]]; then
          function normpath() {
            echo "$(echo $(cygpath -w $1) | sed 's/\\/\\\\/g')"
          }
          CUDA_PATH=$(normpath $(realpath "${{ inputs.cuda-path }}"))
          echo "$(normpath ${CUDA_PATH}/bin)" >> $GITHUB_PATH
        fi
        echo "CUDA_PATH=${CUDA_PATH}" >> $GITHUB_ENV
        echo "CUDA_HOME=${CUDA_PATH}" >> $GITHUB_ENV