name: Setup Python & Install
description: Sets up Python and installs project dependencies.
inputs:
  python-version:
    description: 'Python version to use'
    required: false
    default: '3.14'
  key:
    description: 'Key for the python cache'
    required: false
    default: ''  # if you don't set a key, it doesn't cache
  deps:
    description: 'Extra dependency groups (space separated)'
    required: false
    default: ''
  pydeps:
    description: 'Extra Python dependency groups (space separated)'
    required: false
    default: ''
  opencl:
    description: "Install OpenCL?"
    required: false
    default: 'false'
  amd:
    description: "Install AMD?"
    required: false
    default: 'false'
  cuda:
    description: "Install CUDA?"
    required: false
    default: 'false'
  ocelot:
    description: "Install gpuocelot?"
    required: false
    default: 'false'
  webgpu:
    description: "Install webgpu?"
    required: false
    default: 'false'
  llvm:
    description: "Install LLVM?"
    required: false
    default: 'false'
  qemu:
    description: "Install qemu?"
    required: false
    default: 'false'
  ninja:
    description: "Install ninja?"
    required: false
    default: 'false'
runs:
  using: "composite"
  steps:
    - name: Setup environment
      shell: bash
      run: |
        echo "UV_CACHE_DIR=/tmp/.uv-cache" >> "$GITHUB_ENV"
        echo "OMP_NUM_THREADS=1" >> "$GITHUB_ENV"
        # no buffers should be over 300MB in CI
        echo "MAX_BUFFER_SIZE=300000000" >> "$GITHUB_ENV"
        if [[ "$RUNNER_OS" == "Linux" ]]; then
          echo "VIRTUAL_ENV=/opt/venv/${{ inputs.python-version }}" >> "$GITHUB_ENV"
          echo "UV_PYTHON_INSTALL_DIR=/opt/python" >> "$GITHUB_ENV"
        else
          echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
        fi

    - name: Set up uv
      uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b
      with:
        enable-cache: 'false' # see below for manual caching

    # **** Caching packages ****

    - name: Cache Python packages (PR)
      if: github.event_name == 'pull_request'
      id: restore-venv-pr
      uses: actions/cache/restore@v5
      with:
        path: /tmp/.uv-cache
        key: uv-${{ runner.os }}-${{ runner.arch }}-python-${{ inputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
    - name: Cache Python packages
      if: github.event_name != 'pull_request'
      id: restore-venv
      uses: actions/cache@v5
      with:
        path: /tmp/.uv-cache
        key: uv-${{ runner.os }}-${{ runner.arch }}-python-${{ inputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}

    # **** Caching downloads ****

    - name: Cache downloads (PR)
      if: inputs.key != '' && github.event_name == 'pull_request'
      uses: actions/cache/restore@v5
      with:
        path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
        key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
    - name: Cache downloads
      if: inputs.key != '' && github.event_name != 'pull_request'
      uses: actions/cache@v5
      with:
        path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
        key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}

    # **** Python deps ****

    - name: Install dependencies in venv (with extra)
      if: inputs.deps != ''
      shell: bash
      run: |
        uv venv --allow-existing --python ${{ inputs.python-version }} "$VIRTUAL_ENV"
        DEPS="${{ inputs.deps }}"
        uv pip install --python "$VIRTUAL_ENV" -e ".[${DEPS// /,}]" ${{ inputs.pydeps }} --torch-backend cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
    - name: Install dependencies in venv (without extra)
      if: inputs.deps == ''
      shell: bash
      run: |
        uv venv --allow-existing --python ${{ inputs.python-version }} "$VIRTUAL_ENV"
        uv pip install --python "$VIRTUAL_ENV" -e . ${{ inputs.pydeps }}
    - name: Prune uv cache
      if: github.event_name != 'pull_request'
      shell: bash
      run: uv cache prune --ci
    - name: Configure venv
      shell: bash
      run: |
        if [[ "$RUNNER_OS" == "Windows" ]]; then
          echo "$VIRTUAL_ENV/Scripts" >> "$GITHUB_PATH"
        else
          echo "$VIRTUAL_ENV/bin" >> "$GITHUB_PATH"
        fi

    # ******************* apt *******************
    - name: Setup apt
      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
      shell: bash
      run: |
        sudo mkdir -p /var/cache/apt/archives
        sudo chown -R $USER:$USER /var/cache/apt/archives

        echo 'Acquire::GzipIndexes "true";' | sudo tee /etc/apt/apt.conf.d/gzip
        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
        echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' | sudo tee -a /etc/apt/apt.conf.d/99keep-debs

    - name: Add AMD Repo (Linux)
      if: inputs.amd == 'true' && runner.os == 'Linux'
      shell: bash
      run: |
        wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
        sudo tee /etc/apt/sources.list.d/rocm.list <<EOF
        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/7.1 $(lsb_release -cs) main
        EOF
        echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600

    - name: Add LLVM Repo (Linux)
      if: inputs.llvm == 'true' && runner.os == 'Linux'
      shell: bash
      run: |
        wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
        echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-20 main" | sudo tee /etc/apt/sources.list.d/llvm.list

    - name: Compute Package List + Hash
      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
      id: apt-pkgs
      shell: bash
      run: |
        pkgs=""
        # **** OpenCL ****
        if [[ "${{ inputs.opencl }}" == "true" ]]; then
          pkgs+=" ocl-icd-opencl-dev"
        fi
        # **** AMD ****
        if [[ "${{ inputs.amd }}" == "true" ]]; then
          pkgs+=" comgr"
        fi
        # **** WebGPU (dependencies for software-based vulkan) ****
        if [[ "${{ inputs.webgpu }}" == "true" ]]; then
          pkgs+=" mesa-vulkan-drivers"
        fi
        # **** LLVM ****
        if [[ "${{ inputs.llvm }}" == "true" ]]; then
          pkgs+=" libllvm20 clang-20 lld-20"
        fi
        # **** QEMU ****
        if [[ "${{ inputs.qemu }}" == "true" ]]; then
          pkgs+=" qemu-user-static"
        fi
        # **** ninja ****
        if [[ "${{ inputs.ninja }}" == "true" ]]; then
          pkgs+=" ninja-build"
        fi

        echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
        echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"

        installed=true
        for pkg in $pkgs; do
          info=$(dpkg-query -W -f='${db:Status-Abbrev} ${Version}' "$pkg" 2> /dev/null || true)
          echo "${pkg}: ${info:-not in dpkg database}"
          [[ "$info" == ii* ]] || installed=false
        done
        echo "installed=$installed" >> "$GITHUB_OUTPUT"

    - name: Cache apt (PR)
      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name == 'pull_request' && steps.apt-pkgs.outputs.installed == 'false'
      uses: actions/cache/restore@v5
      with:
        path: /var/cache/apt/archives/
        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
    - name: Cache apt
      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name != 'pull_request' && steps.apt-pkgs.outputs.installed == 'false'
      uses: actions/cache@v5
      with:
        path: /var/cache/apt/archives/
        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}

    - name: Run apt Update + Install
      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && steps.apt-pkgs.outputs.installed == 'false'
      shell: bash
      run: |
        sudo apt -qq update || true

        # ******** do install ********
        if [[ -n "${{ steps.apt-pkgs.outputs.pkgs }}" ]]; then
          sudo apt-get -y --allow-unauthenticated --no-install-recommends install ${{ steps.apt-pkgs.outputs.pkgs }}
        fi

        sudo mkdir -p /var/cache/apt/archives
        sudo chown -R $USER:$USER /var/cache/apt/archives/

    - name: Add clang to PATH (Linux)
      if: inputs.llvm == 'true' && runner.os == 'Linux'
      shell: bash
      run: echo "/usr/lib/llvm-20/bin" >> "$GITHUB_PATH"

    # **** AMD ****
    - name: Setup AMD (Linux)
      if: inputs.amd == 'true' && runner.os == 'Linux'
      shell: bash
      run: |
        sudo tee --append /etc/ld.so.conf.d/rocm.conf <<'EOF'
          /opt/rocm/lib
          /opt/rocm/lib64
        EOF
        sudo ldconfig
    - name: Setup AMD comgr (macOS)
      if: inputs.amd == 'true' && runner.os == 'macOS'
      shell: bash
      run: |
        sudo "$VIRTUAL_ENV/bin/python" -c "
          from tinygrad.helpers import fetch
          fetch('https://github.com/tinygrad/amdcomgr_dylib/releases/download/v7.2.0/libamd_comgr.dylib', name='/usr/local/lib/libamd_comgr.dylib',
                sha256='7712fbe4fcb9fcdea49aeac989876448df975ce0a8ce7c9b15b55c15e7a05935').chmod(0o644)"

    # **** CUDA ****
    - name: Install CUDA
      if: inputs.cuda == 'true'
      shell: bash
      run: |
        sudo mkdir -p /usr/local/cuda/targets/x86_64-linux
        curl -fL https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvrtc/linux-x86_64/cuda_nvrtc-linux-x86_64-11.5.119-archive.tar.xz \
          | sudo tar -xJ -C /usr/local/cuda/targets/x86_64-linux --strip-components=1
        echo /usr/local/cuda/targets/x86_64-linux/lib | sudo tee /etc/ld.so.conf.d/cuda-nvrtc.conf
        sudo ldconfig

    # **** gpuocelot ****
    - name: Install gpuocelot
      if: inputs.ocelot == 'true'
      shell: bash
      run: |
        sudo "$VIRTUAL_ENV/bin/python" -c "
          from tinygrad.helpers import fetch
          fetch('https://github.com/tinygrad/gpuocelot/releases/download/v0.1.0/libgpuocelot.${{ runner.os == 'Linux' && 'so' || 'dylib' }}',
                name='/usr/local/lib/libgpuocelot.${{ runner.os == 'Linux' && 'so' || 'dylib' }}',
                sha256='${{ runner.os == 'Linux' && 'a24705276a9a187111371465987b3258f8836ef512a34266e3075bc4714e125a' || '5106c998c795a36dec79eb7b2aae324a93d1338236d36eeaae232649ec457663' }}').chmod(0o644)"

    # **** WebGPU ****

    - name: Install WebGPU dawn
      if: inputs.webgpu == 'true'
      shell: bash
      run: |
        sudo "$VIRTUAL_ENV/bin/python" -c "
          from tinygrad.helpers import fetch
          fetch('https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.${{ runner.os == 'Linux' && 'so' || 'dylib' }}',
                name='/usr/local/lib/libwebgpu_dawn.${{ runner.os == 'Linux' && 'so' || 'dylib' }}',
                sha256='${{ runner.os == 'Linux' && 'cf36091d266a32c9d5080f14662de44cece241987939713282ea0ff558db81c6' || '7e87c7acefda8b6af1a1c5debfedcf62958311284b8fd8d9bcf93e312e6636e3' }}').chmod(0o644)"

    # **** LLVM ****

    - name: Install LLVM (macOS)
      if: inputs.llvm == 'true' && runner.os == 'macOS'
      shell: bash
      run: brew install llvm@20

    # *** OpenCL ***
    - name: Install rusticl
      if: inputs.opencl == 'true'
      shell: bash
      run: |
        sudo "$VIRTUAL_ENV/bin/python" -c "
          from tinygrad.helpers import fetch
          fetch('https://github.com/sirhcm/tinymesa/releases/download/rusticl-v1/libRusticlOpenCL.so.1.0.0', name='/usr/lib/libRusticlOpenCL.so',
                sha256='d4f48566d8fd33f6cdd8ef6de35a71966e8a8517e6f68ff3c52dbb43765a2513').chmod(0o644)"
        sudo mkdir -p /etc/OpenCL/vendors
        echo "/usr/lib/libRusticlOpenCL.so" | sudo tee /etc/OpenCL/vendors/rusticl.icd
        echo "RUSTICL_ENABLE=llvmpipe" >> "$GITHUB_ENV"