name: Build and Release tokenspeed-kernel

on:
  workflow_dispatch:
    inputs:
      tag_name:
        description: "Git tag prefix to create or update. Defaults to tokenspeed-kernel-v<version>."
        required: false
        type: string
      version_date:
        description: "Version date override for development builds, formatted as YYYYMMDD."
        required: false
        type: string
      cuda_variant:
        description: "CUDA variant to build."
        required: false
        default: all
        type: choice
        options:
          - all
          - cu130
          - cu129
      pypi_cuda_variant:
        description: "CUDA variant to publish to PyPI."
        required: false
        default: cu130
        type: choice
        options:
          - cu130
          - cu129
      prerelease:
        description: "Mark the GitHub release as a prerelease."
        required: false
        default: false
        type: boolean
      publish_github:
        description: "Publish distributions to lightseekorg/whl."
        required: false
        default: false
        type: boolean
      publish_pypi:
        description: "Publish the selected PyPI CUDA variant."
        required: false
        default: false
        type: boolean

permissions:
  contents: read

concurrency:
  group: ${{ github.workflow }}
  cancel-in-progress: true

jobs:
  metadata:
    name: Resolve release metadata
    runs-on: ubuntu-24.04-32core-x64
    outputs:
      version: ${{ steps.metadata.outputs.version }}
      version_date: ${{ steps.metadata.outputs.version_date }}
      tag_name: ${{ steps.metadata.outputs.tag_name }}
    steps:
      - name: Validate publish inputs
        if: inputs.publish_pypi && inputs.cuda_variant != 'all' && inputs.cuda_variant != inputs.pypi_cuda_variant
        run: |
          echo "pypi_cuda_variant must be included by cuda_variant." >&2
          exit 1

      - name: Checkout repository
        uses: actions/checkout@v4

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: "3.12"

      - name: Install build metadata dependencies
        run: python -m pip install --upgrade "setuptools<82" wheel packaging

      - name: Resolve release metadata
        id: metadata
        env:
          PUBLISH_PYPI: ${{ inputs.publish_pypi }}
        run: |
          version_date="${{ inputs.version_date }}"
          if [ -z "$version_date" ]; then
            version_date="$(date -u +%Y%m%d)"
          fi

          version=$(
            cd tokenspeed-kernel/python
            TOKENSPEED_KERNEL_BACKEND=cuda \
            TOKENSPEED_KERNEL_VERSION_DATE="$version_date" \
            TOKENSPEED_KERNEL_GIT_SHA="$GITHUB_SHA" \
            TOKENSPEED_KERNEL_GIT_BRANCH="$GITHUB_REF_NAME" \
              python setup.py --version
          )

          if [ "$PUBLISH_PYPI" = "true" ] && [[ "$version" == *+* ]]; then
            echo "Version ${version} contains a local segment and cannot be published to PyPI; dispatch from a release/* branch." >&2
            exit 1
          fi

          tag_name="${{ inputs.tag_name }}"
          if [ -z "$tag_name" ]; then
            tag_name="tokenspeed-kernel-v${version}"
          fi

          echo "version=${version}" >> "$GITHUB_OUTPUT"
          echo "version_date=${version_date}" >> "$GITHUB_OUTPUT"
          echo "tag_name=${tag_name}" >> "$GITHUB_OUTPUT"
          echo "Release version: ${version}"
          echo "Release tag: ${tag_name}"

  build-wheel:
    name: Build ${{ matrix.cuda.variant }} ${{ matrix.arch.name }} wheel for Python ${{ matrix.python.version }}
    needs: metadata
    runs-on: ${{ matrix.arch.runner }}
    strategy:
      fail-fast: false
      matrix:
        arch:
          - name: x64
            runner: ubuntu-24.04-32core-x64
            manylinux_arch: x86_64
            cuda_repo_arch: x86_64
          - name: arm64
            runner: ubuntu-24.04-32core-arm64
            manylinux_arch: aarch64
            cuda_repo_arch: sbsa
        python:
          - version: "3.10"
            tag: cp310-cp310
          - version: "3.11"
            tag: cp311-cp311
          - version: "3.12"
            tag: cp312-cp312
          - version: "3.13"
            tag: cp313-cp313
        cuda:
          - variant: cu130
            package_version: "13-0"
            torch_index_url: https://download.pytorch.org/whl/cu130
          - variant: cu129
            package_version: "12-9"
            torch_index_url: https://download.pytorch.org/whl/cu129
    steps:
      - name: Checkout repository
        if: inputs.cuda_variant == 'all' || inputs.cuda_variant == matrix.cuda.variant
        uses: actions/checkout@v4

      - name: Build manylinux CUDA image
        if: inputs.cuda_variant == 'all' || inputs.cuda_variant == matrix.cuda.variant
        run: |
          cat > Dockerfile.tokenspeed-kernel-release <<'DOCKERFILE'
          ARG MANYLINUX_ARCH=x86_64
          FROM quay.io/pypa/manylinux_2_28_${MANYLINUX_ARCH}

          ARG CUDA_PKG_VERSION=13-0
          ARG CUDA_REPO_ARCH=x86_64

          RUN curl -o /etc/yum.repos.d/cuda-rhel8.repo \
                  https://developer.download.nvidia.com/compute/cuda/repos/rhel8/${CUDA_REPO_ARCH}/cuda-rhel8.repo \
              && yum install -y \
                  cuda-nvcc-${CUDA_PKG_VERSION} \
                  cuda-cccl-${CUDA_PKG_VERSION} \
                  cuda-cudart-devel-${CUDA_PKG_VERSION} \
                  cuda-driver-devel-${CUDA_PKG_VERSION} \
                  libcublas-devel-${CUDA_PKG_VERSION} \
                  libcurand-devel-${CUDA_PKG_VERSION} \
                  libcusolver-devel-${CUDA_PKG_VERSION} \
                  libcusparse-devel-${CUDA_PKG_VERSION} \
                  gcc-toolset-13-gcc \
                  gcc-toolset-13-gcc-c++ \
                  ninja-build \
              && yum clean all \
              && rm -rf /var/cache/yum \
              && ln -s /usr/bin/ninja-build /usr/local/bin/ninja

          ENV CC=/opt/rh/gcc-toolset-13/root/usr/bin/gcc
          ENV CXX=/opt/rh/gcc-toolset-13/root/usr/bin/g++
          ENV CUDA_HOME=/usr/local/cuda
          ENV PATH=${CUDA_HOME}/bin:${PATH}
          ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${CUDA_HOME}/lib64/stubs:${LD_LIBRARY_PATH}
          ENV LIBRARY_PATH=${CUDA_HOME}/lib64/stubs:${LIBRARY_PATH}

          WORKDIR /workspace
          DOCKERFILE

          docker build \
            --build-arg MANYLINUX_ARCH=${{ matrix.arch.manylinux_arch }} \
            --build-arg CUDA_PKG_VERSION=${{ matrix.cuda.package_version }} \
            --build-arg CUDA_REPO_ARCH=${{ matrix.arch.cuda_repo_arch }} \
            -f Dockerfile.tokenspeed-kernel-release \
            -t tokenspeed-kernel-release:${{ matrix.cuda.variant }}-${{ matrix.arch.name }} .

      - name: Build wheel
        if: inputs.cuda_variant == 'all' || inputs.cuda_variant == matrix.cuda.variant
        run: |
          mkdir -p dist
          docker run --rm \
            -e PYTHON_TAG=${{ matrix.python.tag }} \
            -e TORCH_INDEX_URL=${{ matrix.cuda.torch_index_url }} \
            -e TOKENSPEED_KERNEL_BACKEND=cuda \
            -e TOKENSPEED_KERNEL_VERSION_DATE=${{ needs.metadata.outputs.version_date }} \
            -e TOKENSPEED_KERNEL_GIT_SHA=$GITHUB_SHA \
            -e TOKENSPEED_KERNEL_GIT_BRANCH=$GITHUB_REF_NAME \
            -e MAX_JOBS=8 \
            -v "$PWD:/workspace" \
            -v "$PWD/dist:/output" \
            tokenspeed-kernel-release:${{ matrix.cuda.variant }}-${{ matrix.arch.name }} \
            bash -lc '
              set -euo pipefail

              PYBIN="/opt/python/${PYTHON_TAG}/bin"
              cd /workspace/tokenspeed-kernel/python

              "${PYBIN}/python" -m pip install --upgrade \
                pip "setuptools<82" wheel build twine packaging "apache-tvm-ffi>=0.1.5,<=0.1.11"
              "${PYBIN}/python" -m pip install --no-cache-dir "torch==2.11.0" --index-url "${TORCH_INDEX_URL}"

              rm -rf build *.egg-info
              PIP_EXTRA_INDEX_URL="${TORCH_INDEX_URL}" \
                "${PYBIN}/python" -m build --wheel --no-isolation --outdir /output

              python_tag="${PYTHON_TAG%%-*}"
              abi_tag="${PYTHON_TAG#*-}"
              platform_tag="manylinux_2_28_$(uname -m)"
              "${PYBIN}/python" -m wheel tags \
                --python-tag "${python_tag}" \
                --abi-tag "${abi_tag}" \
                --platform-tag "${platform_tag}" \
                --remove /output/*.whl

              "${PYBIN}/python" -m twine check /output/*.whl
            '

      - name: Upload wheel
        if: inputs.cuda_variant == 'all' || inputs.cuda_variant == matrix.cuda.variant
        uses: actions/upload-artifact@v4
        with:
          name: tokenspeed-kernel-wheel-${{ matrix.cuda.variant }}-${{ matrix.arch.name }}-py${{ matrix.python.version }}
          path: dist/*.whl

  build-sdist:
    name: Build source distribution
    needs: metadata
    runs-on: ubuntu-24.04-32core-x64
    steps:
      - name: Checkout repository
        uses: actions/checkout@v4

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: "3.12"

      - name: Build source distribution
        run: |
          python -m pip install --upgrade \
            build twine "setuptools<82" wheel packaging "apache-tvm-ffi>=0.1.5,<=0.1.11"
          cd tokenspeed-kernel/python
          TOKENSPEED_KERNEL_BACKEND=cuda \
          TOKENSPEED_KERNEL_VERSION_DATE=${{ needs.metadata.outputs.version_date }} \
          TOKENSPEED_KERNEL_GIT_SHA=$GITHUB_SHA \
          TOKENSPEED_KERNEL_GIT_BRANCH=$GITHUB_REF_NAME \
            python -m build --sdist --no-isolation --outdir ../../dist
          python -m twine check ../../dist/*.tar.gz

      - name: Upload source distribution
        uses: actions/upload-artifact@v4
        with:
          name: tokenspeed-kernel-sdist
          path: dist/*.tar.gz

  release:
    name: Create wheelhouse release
    if: inputs.publish_github
    needs: [metadata, build-wheel, build-sdist]
    runs-on: ubuntu-24.04-32core-x64
    steps:
      - name: Checkout repository
        uses: actions/checkout@v4

      - name: Download distributions
        uses: actions/download-artifact@v4
        with:
          path: dist
          pattern: tokenspeed-kernel-*
          merge-multiple: false

      - name: Create or update GitHub release
        env:
          GH_TOKEN: ${{ secrets.LIGHTSEEK_BOT_TOKEN }}
          RELEASE_REPOSITORY: lightseekorg/whl
          TAG_NAME: ${{ needs.metadata.outputs.tag_name }}
          VERSION: ${{ needs.metadata.outputs.version }}
          CUDA_VARIANT: ${{ inputs.cuda_variant }}
          PRERELEASE: ${{ inputs.prerelease }}
          SOURCE_SHA: ${{ github.sha }}
        run: |
          prerelease_args=()
          if [ "$PRERELEASE" = "true" ]; then
            prerelease_args+=(--prerelease)
          fi

          upload_release() {
            local cuda_variant="$1"
            local release_tag="${TAG_NAME}-${cuda_variant}"
            local release_title="tokenspeed-kernel ${VERSION} ${cuda_variant}"

            shopt -s nullglob
            local assets=(
              dist/tokenspeed-kernel-wheel-"${cuda_variant}"-*/*.whl
              dist/tokenspeed-kernel-sdist/*.tar.gz
            )
            shopt -u nullglob

            if [ "${#assets[@]}" -eq 0 ]; then
              echo "No distributions found for ${cuda_variant}" >&2
              exit 1
            fi

            if gh release view "$release_tag" --repo "$RELEASE_REPOSITORY" >/dev/null 2>&1; then
              gh release upload "$release_tag" "${assets[@]}" \
                --repo "$RELEASE_REPOSITORY" \
                --clobber
            else
              gh release create "$release_tag" "${assets[@]}" \
                --repo "$RELEASE_REPOSITORY" \
                --title "$release_title" \
                --notes "$release_title built from lightseekorg/tokenspeed@$SOURCE_SHA" \
                "${prerelease_args[@]}"
            fi
          }

          case "$CUDA_VARIANT" in
            all)
              upload_release cu130
              upload_release cu129
              ;;
            cu130|cu129)
              upload_release "$CUDA_VARIANT"
              ;;
            *)
              echo "Unsupported CUDA variant: ${CUDA_VARIANT}" >&2
              exit 1
              ;;
          esac

  publish-pypi:
    name: Publish ${{ inputs.pypi_cuda_variant }} to PyPI
    if: inputs.publish_pypi
    needs: [metadata, build-wheel, build-sdist]
    runs-on: ubuntu-24.04-32core-x64
    environment: pypi
    permissions:
      id-token: write
    steps:
      - name: Download wheels
        uses: actions/download-artifact@v4
        with:
          path: dist
          pattern: tokenspeed-kernel-wheel-${{ inputs.pypi_cuda_variant }}-*
          merge-multiple: true

      - name: Download source distribution
        uses: actions/download-artifact@v4
        with:
          name: tokenspeed-kernel-sdist
          path: dist

      - name: Publish distributions to PyPI
        uses: pypa/gh-action-pypi-publish@release/v1
        with:
          packages-dir: dist