# 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:
  create:
    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 ::set-output name=branch::${GITHUB_REF#refs/tags/}
        shell: bash
      - name: Create Release
        id: create_release
        uses: actions/create-release@v1
        env:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
        with:
          tag_name: ${{ steps.extract_branch.outputs.branch }}
          release_name: ${{ steps.extract_branch.outputs.branch }}

  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]
        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
        torch-version: ["2.5.1", "2.6.0", "2.7.1", "2.8.0"]
        cuda-version: ["12.9.1", "13.0.0"]
        # 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:
          # see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix
          - torch-version: "2.4.0"
            python-version: "3.13"
        include:
          - os: "ubuntu-22.04"
            cxx11_abi: "FALSE"
            cuda-version: "12.1.0"
            python-version: "3.8"
            torch-version: "2.1.0"
    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
      use-local-version: false
  
  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]
        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]
    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 wheel twine
          # Install latest setuptools with support for pypi metadata 2.2 (improved compat w/ uv)
          pip install setuptools==75.8.0
          # 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
        env:
          DG_USE_LOCAL_VERSION: "0"
          DG_SKIP_CUDA_BUILD: "1"
        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/*