wheels_build.yml

name: wheels_build

on:
  workflow_call:
    inputs:
      os:
        required: true
        type: string
      python:
        required: true
        type: string
      torch_version:
        required: true
        type: string
        description: "Example: 1.13.1"
      cuda_short_version:
        required: true
        type: string
        description: "Example: 117 for 11.7"
      cudnn_version_major:
        required: false
        default: '8'
        type: string
        description: "Example: 8"
  workflow_dispatch:
    inputs:
      os:
        required: true
        type: string
      python:
        required: true
        type: string
      torch_version:
        required: true
        type: string
        description: "Example: 1.13.1"
      cuda_short_version:
        required: true
        type: string
        description: "Example: 117 for 11.7"
      cudnn_version_major:
        required: false
        default: '8'
        type: string
        description: "Example: 8"

# this yaml file can be cleaned up using yaml anchors, but they're not supported in github actions yet
# https://github.com/actions/runner/issues/1182

env:
  # you need at least cuda 5.0 for some of the stuff compiled here.
  # TORCH_CUDA_ARCH_LIST: "5.0+PTX 6.0 6.1 7.0 7.5 8.0+PTX"
  # Feature 'f16 arithemetic and compare instructions' requires .target sm_53 or higher
  TORCH_CUDA_ARCH_LIST: "6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX"
  MAX_JOBS: 2
  DISTUTILS_USE_SDK: 1 # otherwise distutils will complain on windows about multiple versions of msvc
  SFAST_APPEND_VERSION: 1
  TWINE_USERNAME: __token__

jobs:
  build_internal:
    name: ${{ inputs.os }}-py${{ inputs.python }}-torch${{ inputs.torch_version }}+cu${{ inputs.cuda_short_version }}
    runs-on: ${{ inputs.os }}
    env:
      # alias for the current python version
      # windows does not have per version binary, it is just 'python3'
      PY: python${{ contains(inputs.os, 'ubuntu') && inputs.python || '3' }}

    # container: ${{ contains(inputs.os, 'ubuntu') && 'quay.io/pypa/manylinux2014_x86_64' || null }}
    timeout-minutes: 360
    defaults:
      run:
        shell: bash
    steps:
      - id: cuda_info
        shell: python
        run: |
          import os
          import sys
          print(sys.version)
          cushort = "${{ inputs.cuda_short_version }}"
          # https://github.com/Jimver/cuda-toolkit/blob/master/src/links/linux-links.ts
          full_version, install_script = {
            "121": ("12.1.0", "https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run"),
            "118": ("11.8.0", "https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run"),
            "117": ("11.7.1", "https://developer.download.nvidia.com/compute/cuda/11.7.1/local_installers/cuda_11.7.1_515.65.01_linux.run"),
            "116": ("11.6.2", "https://developer.download.nvidia.com/compute/cuda/11.6.2/local_installers/cuda_11.6.2_510.47.03_linux.run"),
          }[cushort]
          cudnn_pypi_package = {
            "121": "nvidia-cudnn-cu12",
            "118": "nvidia-cudnn-cu11",
            "117": "nvidia-cudnn-cu11",
            "116": "nvidia-cudnn-cu11",
          }[cushort]
          with open(os.environ['GITHUB_OUTPUT'], "r+") as fp:
            fp.write("CUDA_VERSION=" + full_version + "\n")
            fp.write("CUDA_VERSION_SUFFIX=cu" + cushort + "\n")
            # fp.write("TORCH_ORG_S3_PATH=s3://pytorch/whl/" + cushort + "\n")
            # fp.write("PUBLISH_PYPI=0\n")
            fp.write("CUDA_INSTALL_SCRIPT=" + install_script + "\n")
            fp.write("CUDNN_PYPI_PACKAGE=" + cudnn_pypi_package + "\n")
      - run: echo "CUDA_VERSION_SUFFIX=${{ steps.cuda_info.outputs.CUDA_VERSION_SUFFIX }}"
      # - run: echo "TORCH_ORG_S3_PATH=${{ steps.cuda_info.outputs.TORCH_ORG_S3_PATH }}"
      # - run: echo "PUBLISH_PYPI=${{ steps.cuda_info.outputs.PUBLISH_PYPI }}"

      - name: Add H100 if nvcc 11.08+
        shell: python
        run: |
          import os
          import sys
          print(sys.version)
          cuda_short_version = "${{ inputs.cuda_short_version }}"
          arch_list = os.environ["TORCH_CUDA_ARCH_LIST"]
          if cuda_short_version not in ["116", "117"]:
            arch_list += " 9.0"
          with open(os.environ['GITHUB_ENV'], "r+") as fp:
            fp.write("TORCH_CUDA_ARCH_LIST=" + arch_list + "\n")
      - run: echo "${TORCH_CUDA_ARCH_LIST}"

      - if: contains(inputs.os, 'ubuntu')
        name: Free Disk Space (Ubuntu)
        uses: jlumbroso/free-disk-space@v1.3.1
        with:
          # this might remove tools that are actually needed,
          # if set to "true" but frees about 6 GB
          tool-cache: false
          
          # all of these default to true, but feel free to set to
          # "false" if necessary for your workflow
          android: true
          dotnet: true
          haskell: true
          large-packages: true
          docker-images: true
          swap-storage: true

      - if: runner.os == 'Linux'
        name: (Linux) install cuda
        run: |
          # yum install wget git prename -y
          # yum clean all --verbose
          sudo apt update
          sudo apt install -y wget git rename
          sudo apt clean -y
          sudo apt autoremove -y
          wget -q "${{ steps.cuda_info.outputs.CUDA_INSTALL_SCRIPT }}" -O cuda.run
          sudo sh ./cuda.run --silent --toolkit --toolkitpath=/usr/local/cuda || cat /tmp/cuda-installer.log
          rm ./cuda.run
          echo "CUDA_HOME=/usr/local/cuda" >> ${GITHUB_ENV}
          echo "PATH=/usr/local/cuda/bin:$PATH" >> ${GITHUB_ENV}

      - if: runner.os == 'Linux'
        name: (Linux) install python
        run: |
          sudo add-apt-repository ppa:deadsnakes/ppa -y
          sudo apt update
          sudo apt install -y python${{ inputs.python }} python${{ inputs.python }}-dev python${{ inputs.python }}-venv
          sudo apt clean -y
          sudo apt autoremove -y

      - name: Recursive checkout
        uses: actions/checkout@v3
        with:
          submodules: recursive
          path: "."
          fetch-depth: 0 # for tags

      - if: runner.os != 'Windows'
        name: (Linux) Setup venv for linux
        run: |
          $PY -m venv venv
          . ./venv/bin/activate
          which pip
          echo "PY=$(which python)" >> ${GITHUB_ENV}
          echo "PATH=$PATH" >> ${GITHUB_ENV}
          # echo "MAX_JOBS=2" >> ${GITHUB_ENV}

      - name: Define version
        id: sfast_version
        env:
          VERSION_SOURCE: ${{ github.ref_type == 'tag' && 'tag' || 'dev' }}
        run: |
          set -Eeuo pipefail
          git config --global --add safe.directory "*"
          version=`cat version.txt`
          # Set build version to x.x.x.devYYYYMMDD+torchxxxcu111
          torch_version_suffix=torch$(echo ${{ inputs.torch_version }} | sed 's/\.//g')
          cuda_version_suffix=${{ steps.cuda_info.outputs.CUDA_VERSION_SUFFIX }}
          nightly_tag=$([[ ${VERSION_SOURCE} == 'tag' ]] && echo '' || echo '.dev'`date +%Y%m%d`)
          echo "BUILD_VERSION=${version}${nightly_tag}+${torch_version_suffix}${cuda_version_suffix}" >> ${GITHUB_ENV}
          echo "BUILD_VERSION=${version}${nightly_tag}+${torch_version_suffix}${cuda_version_suffix}" >> ${GITHUB_OUTPUT}
          cat ${GITHUB_ENV}
      - run: echo "sfast-${BUILD_VERSION}"
      - run: echo "release version"
        if: ${{ !contains(steps.sfast_version.outputs.BUILD_VERSION, '.dev') }}

      # - name: Setup proper pytorch dependency in "requirements.txt"
      #   run: |
      #     sed -i '/torch/d' ./requirements.txt
      #     echo "torch==${{ inputs.torch_version }}" >> ./requirements.txt
      #     cat ./requirements.txt

      - if: runner.os == 'Windows'
        name: (Windows) Setup Runner
        uses: ./.github/actions/setup-windows-runner
        with:
          cuda: ${{ steps.cuda_info.outputs.CUDA_VERSION }}
          python: ${{ inputs.python }}

      - name: Install dependencies
        env:
          CUDNN_VERSION_MAJOR: ${{ inputs.cudnn_version_major }}
          CUDNN_PYPI_PACKAGE: ${{ steps.cuda_info.outputs.CUDNN_PYPI_PACKAGE }}
        run: |
          cudnn_next_version_major=$((${CUDNN_VERSION_MAJOR} + 1))
          cudnn_package_name="${CUDNN_PYPI_PACKAGE}>=${CUDNN_VERSION_MAJOR}.0.0.0,<${cudnn_next_version_major}.0.0.0"
          $PY -m pip install --upgrade pip
          $PY -m pip install wheel setuptools ninja twine "torch==${{ inputs.torch_version }}" "${cudnn_package_name}" -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu${{ inputs.cuda_short_version }} --no-cache-dir

      - name: Build wheel
        run: |
          $PY setup.py bdist_wheel -d dist/ -k $PLAT_ARG
        env:
          PLAT_ARG: ${{ contains(inputs.os, 'ubuntu') && '--plat-name manylinux2014_x86_64' || '' }}

      - run: du -h dist/*
      - uses: actions/upload-artifact@v3
        with:
          name: ${{ inputs.os }}-py${{ inputs.python }}-torch${{ inputs.torch_version }}+cu${{ inputs.cuda_short_version }}
          path: dist/*.whl
# Note: it might be helpful to have additional steps that test if the built wheels actually work