# This workflow will:
# - Create a new Github release
# - Build wheels for supported architectures
# - Deploy the wheels to the Github release
# - Release the static code to PyPi
# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries

name: Build wheels and deploy

on:
  push:
    tags:
      - v*

jobs:
  setup_release:
    name: Create Release
    runs-on: ubuntu-latest
    outputs:
      release-version: ${{ steps.extract_branch.outputs.branch }}
    steps:
      - name: Get the tag version
        id: extract_branch
        run: echo "branch=${GITHUB_REF#refs/tags/}" >> $GITHUB_OUTPUT
        shell: bash

      - name: Create Release
        env:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
        run: gh release create ${{ steps.extract_branch.outputs.branch }} --repo $GITHUB_REPOSITORY --title ${{ steps.extract_branch.outputs.branch }} --generate-notes
        shell: bash

  build_wheels:
    name: Build Wheel
    needs: setup_release

    strategy:
      fail-fast: false
      matrix:
        # Using ubuntu-22.04 instead of 24.04 for more compatibility (glibc). Ideally we'd use the
        # manylinux docker image, but I haven't figured out how to install CUDA on manylinux.
        os: [ubuntu-22.04, ubuntu-22.04-arm]
        python-version: ["3.10", "3.11", "3.12", "3.13"]
        torch-version: ["2.6.0", "2.7.1", "2.8.0", "2.9.1", "2.10.0"]
        cuda-version: ["11.8.0", "12.9.1", "13.0.1"]
        # We need separate wheels that either uses C++11 ABI (-D_GLIBCXX_USE_CXX11_ABI) or not.
        # Pytorch wheels currently don't use it, but nvcr images have Pytorch compiled with C++11 ABI.
        # Without this we get import error (undefined symbol: _ZN3c105ErrorC2ENS_14SourceLocationESs)
        # when building without C++11 ABI and using it on nvcr images.
        cxx11_abi: ["FALSE", "TRUE"]
        exclude:
          # CUDA 11.8 is not supported by PyTorch 2.8+
          - torch-version: "2.8.0"
            cuda-version: "11.8.0"
          - torch-version: "2.9.1"
            cuda-version: "11.8.0"
          - torch-version: "2.10.0"
            cuda-version: "11.8.0"
          # CUDA 13.0 is only supported by PyTorch 2.9+
          - torch-version: "2.6.0"
            cuda-version: "13.0.1"
          - torch-version: "2.7.1"
            cuda-version: "13.0.1"
          - torch-version: "2.8.0"
            cuda-version: "13.0.1"
          # No aarch64 PyTorch wheels for 2.6.0, or 2.7.1+cu118
          - torch-version: "2.6.0"
            os: ubuntu-22.04-arm
          - torch-version: "2.7.1"
            cuda-version: "11.8.0"
            os: ubuntu-22.04-arm
          # PyTorch 2.7+ pip wheels use CXX11_ABI=1 by default, no need for FALSE
          - torch-version: "2.7.1"
            cxx11_abi: "FALSE"
          - torch-version: "2.8.0"
            cxx11_abi: "FALSE"
          - torch-version: "2.9.1"
            cxx11_abi: "FALSE"
          - torch-version: "2.10.0"
            cxx11_abi: "FALSE"
    uses: ./.github/workflows/_build.yml
    with:
      runs-on: ${{ matrix.os }}
      python-version: ${{ matrix.python-version }}
      cuda-version: ${{ matrix.cuda-version }}
      torch-version: ${{ matrix.torch-version }}
      cxx11_abi: ${{ matrix.cxx11_abi }}
      release-version: ${{ needs.setup_release.outputs.release-version }}
      upload-to-release: true

  check_for_ngc_images:
    runs-on: ubuntu-latest
    outputs:
      images: ${{ steps.check_for_ngc_images.outputs.IMAGES }}
    steps:
      - name: Checkout repository
        uses: actions/checkout@v4

      - name: Check for NGC PyTorch images
        id: check_for_ngc_images
        run: |
          bash ./.github/scripts/check_for_ngc_images.sh
          echo "IMAGES=$(cat ngc_images.json| jq -cr)" | tee -a $GITHUB_OUTPUT

  build_ngc_wheels:
    name: Build Wheel for NGC PyTorch
    needs: [setup_release, check_for_ngc_images]
    strategy:
      fail-fast: false
      matrix:
        os: [ubuntu-22.04, ubuntu-22.04-arm]
        container-image: ${{ fromJson(needs.check_for_ngc_images.outputs.images) }}
    uses: ./.github/workflows/_build_in_container.yml
    with:
      runs-on: ${{ matrix.os }}
      container-image: ${{ matrix.container-image }}
      release-version: ${{ needs.setup_release.outputs.release-version }}
      upload-to-release: true

  publish_package:
    name: Publish package
    needs: [build_wheels]
    if: always() && !cancelled()

    runs-on: ubuntu-latest

    steps:
      - uses: actions/checkout@v4

      - uses: actions/setup-python@v5
        with:
          python-version: "3.10"

      - name: Install dependencies
        run: |
          pip install ninja packaging setuptools wheel twine
          # We don't want to download anything CUDA-related here
          pip install torch --index-url https://download.pytorch.org/whl/cpu

      - name: Build core package
        run: |
          python setup.py sdist --dist-dir=dist

      - name: Deploy
        env:
          TWINE_USERNAME: "__token__"
          TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
        run: |
          python -m twine upload dist/*
