name: 'Setup Linux Environment'
description: 'Install dependencies for Linux builds'

inputs:
  toolkit:
    description: 'Which toolkit to install'
    required: false
    default: 'cpu'
  python-version:
    description: 'Version of python to set up'
    required: false
    default: '3.14'
  use-ccache:
    description: 'Whether to enable ccache'
    required: false
    default: 'true'

runs:
  using: "composite"
  steps:
    - name: Install common dependencies
      shell: bash
      run: |
        echo "::group::Install common dependencies"
        sudo apt-get update
        sudo apt-get install -y --no-install-recommends \
            zip \
            libblas-dev liblapack-dev liblapacke-dev \
            openmpi-bin openmpi-common libopenmpi-dev
        echo "::endgroup::"

    - name: Use ccache
      if: ${{ inputs.use-ccache == 'true' }}
      uses: hendrikmuhs/ccache-action@v1.2
      with:
        key: ccache-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }}
        max-size: 1GB
        # ccache-action bug: running "apt-get update" fails on large arm runner.
        update-package-index: false

    - uses: actions/setup-python@v6
      with:
        python-version: ${{ inputs.python-version }}

    - name: Setup Python venv
      shell: bash
      run: |
        echo "::group::Setup Python venv"
        python -m venv .venv
        source .venv/bin/activate
        pip install setuptools cmake typing_extensions
        echo PATH=$PATH >> $GITHUB_ENV
        # Search python packages in .venv
        echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
        echo "::endgroup::"

    - name: Install CUDA toolkit
      if: ${{ startsWith(inputs.toolkit, 'cuda') }}
      shell: bash
      env:
        # Note: the CI machine does not meet CUDA 13's driver requirement.
        # Compatibility matrix:
        # https://docs.nvidia.com/deeplearning/cudnn/backend/latest/reference/support-matrix.html
        PACKAGES: |
          {
            "cuda-12.6": "libcudnn9-dev-cuda-12 cuda-compiler-12-6 cuda-libraries-dev-12-6",
            "cuda-12.9": "libcudnn9-dev-cuda-12 cuda-compiler-12-9 cuda-libraries-dev-12-9",
            "cuda-13.0": "libcudnn9-dev-cuda-13 cuda-compiler-13-0 cuda-libraries-dev-13-0"
          }
      run: |
        echo "::group::Install CUDA toolkit"
        # The CUDA binaries are hosted in the "sbsa" repo, the "arm64" repo is
        # Jetson specific. SBSA means Arm Server Base System Architecture.
        ARCH=${{ runner.arch == 'arm64' && 'sbsa' || 'x86_64' }}
        wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/$ARCH/cuda-keyring_1.1-1_all.deb
        sudo dpkg -i cuda-keyring_1.1-1_all.deb
        sudo apt-get update
        sudo apt-get install -y --no-install-recommends \
            libnccl2 libnccl-dev \
            ${{ fromJson(env.PACKAGES)[inputs.toolkit] }}
        echo "/usr/local/${{ inputs.toolkit }}/bin" >> $GITHUB_PATH
        echo "::endgroup::"

    - name: CUDA packages and driver report
      if: ${{ startsWith(inputs.toolkit, 'cuda') }}
      shell: bash
      run: |
        echo "::group::Installed NVIDIA and CUDA packages"
        dpkg -l | egrep "cuda|nvidia" -i
        echo "::endgroup::"
        echo "::group::NVIDIA-SMI Status"
        nvidia-smi || true
        echo "::endgroup::"
