_build_test_publish_wheel.yml

on:
  workflow_call:
    inputs:
      ref:
        required: false
        description: Ref (SHA or branch) to release
        type: string
        default: ${{ github.sha }}
      dry-run:
        required: false
        description: Upload to PyPy Test instance
        type: boolean
        default: true
      no-publish:
        required: false
        description: Do not publish the wheel
        type: boolean
        default: true
    secrets:
      TWINE_PASSWORD:
        required: true

jobs:
  build-and-test-wheels:
    strategy:
      fail-fast: false
      matrix:
        include:
          - PACKAGE: megatron-core
            PLATFORM: arm64
            IMAGE: quay.io/pypa/manylinux_2_28_aarch64
          - PACKAGE: megatron-core
            PLATFORM: amd64
            IMAGE: quay.io/pypa/manylinux_2_28_x86_64
          - PACKAGE: megatron-fsdp
            IMAGE: quay.io/pypa/manylinux_2_28_x86_64
            PLATFORM: amd64
    runs-on: ${{ matrix.PLATFORM == 'amd64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }}
    env:
      PACKAGE: ${{ matrix.PACKAGE }}
      IMAGE: ${{ matrix.IMAGE }}
      PLATFORM: ${{ matrix.PLATFORM }}
      PUBLISH_DRYRUN: ${{ inputs.dry-run }}
    steps:
      - name: Checkout repository
        uses: actions/checkout@v6
        with:
          ref: ${{ inputs.ref }}

      - name: Build wheel
        id: build-wheel
        run: |
          set -x

          if [ "$PACKAGE" = "megatron-core" ]; then
            ROOTDIR="megatron/core"
            BUILD_DIR="."
          elif [ "$PACKAGE" = "megatron-fsdp" ]; then
            ROOTDIR="megatron/core/distributed/fsdp/src/megatron_fsdp"
            BUILD_DIR="megatron/core/distributed/fsdp/src"
          else
            echo Unknown package: $PACKAGE
            exit 1
          fi

          if [ "$PUBLISH_DRYRUN" = "true" ]; then
            PRE_RELEASE=$(sed -n "s/.*PRE_RELEASE = '\(.*\)'/\1/p" $ROOTDIR/package_info.py)
            sed -i "/^PRE_RELEASE/c\PRE_RELEASE = '${PRE_RELEASE}.dev$((RANDOM % 900000 + 100000))'" $ROOTDIR/package_info.py
          fi

          pushd $BUILD_DIR
            rm LICENSE || true
            docker run --rm -v $(pwd):/workspace -w /workspace $IMAGE bash -c '\
              for python_version in cp311 cp312 cp313; do \
                /opt/python/${python_version}-${python_version}/bin/pip install --upgrade "setuptools<80.0.0,>=77.0.0" build; \
              done && \
              for python_version in cp311 cp312 cp313; do \
                /opt/python/${python_version}-${python_version}/bin/python -m build; \
              done \
            '

            PLATFORM_WHEELS=$(find dist -name "*.whl" -not -name "*-none-any.whl")
            if [ -n "$PLATFORM_WHEELS" ]; then
                echo "Found platform wheels to repair: $PLATFORM_WHEELS"
                docker run --rm -v $(pwd):/workspace -w /workspace $IMAGE auditwheel repair $PLATFORM_WHEELS
                docker run --rm -v $(pwd):/workspace -w /workspace $IMAGE rm -rf dist/*.whl
                docker run --rm -v $(pwd):/workspace -w /workspace $IMAGE cp -a wheelhouse/* dist/
            fi
          popd

          pushd $ROOTDIR
            EXPECTED_RELEASE_NUMBER=$(python -c "import package_info; print(package_info.__version__)")
          popd

          echo "expected-release-number=$EXPECTED_RELEASE_NUMBER" | tee -a "${GITHUB_OUTPUT}"

          if [ "$PACKAGE" = "megatron-fsdp" ]; then
            mkdir -p dist/
            cp -a megatron/core/distributed/fsdp/src/dist/* dist/
          fi

          ls -al dist/

      - name: Test wheels
        run: |
          ls -al dist/

          if [ "$PACKAGE" = "megatron-core" ]; then
            ROOTPATH="megatron.core"
            WHEEL_PREFIX="megatron_core"
          elif [ "$PACKAGE" = "megatron-fsdp" ]; then
            ROOTPATH="megatron_fsdp"
            WHEEL_PREFIX="megatron_fsdp"
          else
            echo Unknown package: $PACKAGE
            exit 1
          fi

          if [ "$PACKAGE" = "megatron-core" ]; then
            if [[ "$PLATFORM" == "arm64" ]]; then
              WHEEL_GLOB="dist/${WHEEL_PREFIX}*cp312*aarch64.whl"
            else
              WHEEL_GLOB="dist/${WHEEL_PREFIX}*cp312*x86_64.whl"
            fi
          else
            WHEEL_GLOB="dist/${WHEEL_PREFIX}*.whl"
          fi

          docker run --rm -v $(pwd):/workspace -w /workspace $IMAGE bash -c "\
            /opt/python/cp312-cp312/bin/pip install --no-cache-dir $WHEEL_GLOB && \
            rm -rf megatron/ && \
            RELEASE_NUMBER=\$(/opt/python/cp312-cp312/bin/python -c 'import $ROOTPATH; print($ROOTPATH.__version__)') && \
            test '${{ steps.build-wheel.outputs.expected-release-number }}' == \"\$RELEASE_NUMBER\" \
          "

      - name: Upload wheels
        uses: actions/upload-artifact@v6
        with:
          name: wheels-${{ matrix.PACKAGE }}-${{ matrix.PLATFORM }}-${{ inputs.dry-run && 'dry-run' || 'release' }}
          path: dist/

  publish-wheels:
    needs: [build-and-test-wheels]
    runs-on: ubuntu-latest
    if: inputs.no-publish == false
    strategy:
      fail-fast: false
      matrix:
        include:
          - PACKAGE: megatron-core
            PLATFORM: arm64
          - PACKAGE: megatron-core
            PLATFORM: amd64
          - PACKAGE: megatron-fsdp
            PLATFORM: amd64
    env:
      PACKAGE: ${{ matrix.PACKAGE }}
    steps:
      - name: Download wheels
        uses: actions/download-artifact@v7
        with:
          name: wheels-${{ matrix.PACKAGE }}-${{ matrix.PLATFORM }}-${{ inputs.dry-run && 'dry-run' || 'release' }}
          path: dist/
          merge-multiple: true

      - name: Publish wheels
        env:
          TWINE_USERNAME: __token__
          TWINE_PASSWORD: ${{ secrets.TWINE_PASSWORD }}
          TWINE_REPOSITORY: ${{ (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/heads/r')) && 'pypi' || 'testpypi' }}
          PLATFORM: ${{ matrix.PLATFORM }}
        run: |

          # Delete sdist for arm64 since we already upload it with amd64.
          if [ "$PLATFORM" == "arm64" ]; then
            rm dist/*.tar.gz
          fi

          ls -al dist/
          pip install twine
          twine upload \
            --verbose \
            -r $TWINE_REPOSITORY \
            -u $TWINE_USERNAME \
            -p $TWINE_PASSWORD \
            dist/*