publish_wheel.yml

name: Publish PyPI
on:
  workflow_dispatch:
    inputs:
      tag:
        description: 'Tag to build and update whl index'
        required: true
        type: string

  workflow_call:
    inputs:
      tag:
        description: 'Tag to build and update whl index'
        required: true
        type: string
env:
  # Tells where to store caches.
  CI_CACHE_DIR: ${{ github.workspace }}/../../ci_cache

jobs:
  build_wheel:
    strategy:
      matrix:
        python: ["3.9", "3.10", "3.11", "3.12"]
        cuda: ["12.4"]
        torch: ["2.6.0"]
    runs-on: [self-hosted, linux, release]
    env:
      PYTHON_VERSION: ${{ matrix.python }}
      CUDA_VERSION: ${{ matrix.cuda }}
      TORCH_VERSION: ${{ matrix.torch }}
      SCALELLM_VERSION: ${{ inputs.tag }}
    steps:
      - name: Checkout repository
        uses: actions/checkout@v4
        with:
          submodules: recursive

      - name: Create cache directory
        run: |
          mkdir -p $CI_CACHE_DIR/.vcpkg/bincache
          mkdir -p $CI_CACHE_DIR/.ccache
          mkdir -p $CI_CACHE_DIR/.pip

      - name: Build wheel
        timeout-minutes: 60
        run: |
          docker pull vectorchai/scalellm_manylinux:cuda${CUDA_VERSION}
          docker run --rm -t \
              -v "$CI_CACHE_DIR":/ci_cache \
              -v "$GITHUB_WORKSPACE":/ScaleLLM \
              -e PYTHON_VERSION=${PYTHON_VERSION} \
              -e CUDA_VERSION=${CUDA_VERSION} \
              -e TORCH_VERSION=${TORCH_VERSION} \
              -e SCALELLM_VERSION=${SCALELLM_VERSION} \
              -e VCPKG_DEFAULT_BINARY_CACHE=/ci_cache/.vcpkg/bincache \
              -e CCACHE_DIR=/ci_cache/.ccache \
              -e PIP_CACHE_DIR=/ci_cache/.pip \
              -u $(id -u):$(id -g) \
              vectorchai/scalellm_manylinux:cuda${CUDA_VERSION} \
              bash /ScaleLLM/scripts/build_wheel.sh

      - name: show wheels
        run: ls -lh dist

      - name: rename wheel to manylinux
        run: |
          for whl in dist/scalellm-*.whl; do
            new_whl=${whl//"-linux_"/"-manylinux1_"}
            if [ "$whl" != "$new_whl" ]; then
              mv $whl $new_whl
            fi
          done

      - name: show wheels
        run: ls -lh dist

      - uses: actions/upload-artifact@v4
        with:
          name: manylinux1_wheel-cuda${{ matrix.cuda }}-torch${{ matrix.torch }}-python${{ matrix.python }}
          path: dist/*

  publish_wheel:
    needs: build_wheel
    runs-on: ubuntu-latest
    permissions:
      contents: write
    steps:
      - name: Download wheel
        uses: actions/download-artifact@v4
        with:
          path: dist
          merge-multiple: true
          pattern: manylinux1_wheel-*

      - name: Show wheels
        run: ls -lh dist

      - name: Publish package to PyPI
        uses: pypa/gh-action-pypi-publish@release/v1.8
        with:
          packages-dir: dist/
          password: ${{ secrets.PYPI_API_TOKEN }}
          skip-existing: true
          verbose: true