name: Release

on:
  push:
    tags:
      - 'v*'

env:
  SCCACHE_GHA_ENABLED: "true"
jobs:
  build:
    runs-on: ubuntu-22.04
    permissions:
      contents: write
    strategy:
      matrix:
        python-version: ['3.10', '3.11', '3.12', '3.13']
    env:
      BUILD_WITH_EP: "1"
      EP_TORCH_VERSIONS: "2.9.0;2.9.1;2.10.0"
      TORCH_CUDA_ARCH_LIST: "8.0;9.0"
    steps:
      - name: Checkout source
        uses: actions/checkout@v4

      - name: Set up Python ${{ matrix.python-version }}
        uses: actions/setup-python@v5
        with:
          python-version: ${{ matrix.python-version }}

      - name: Free up disk space
        run: |
          sudo rm -rf /usr/share/dotnet
          sudo rm -rf /opt/ghc
          sudo rm -rf /opt/hostedtoolcache/CodeQL
          sudo rm -rf /usr/local/lib/android
          df -h

      - name: Install CUDA Toolkit
        uses: Jimver/cuda-toolkit@v0.2.24
        with:
          cuda: '12.8.1'
          linux-local-args: '["--toolkit"]'
          method: 'network'
          sub-packages: '["nvcc", "nvrtc-dev"]'
          non-cuda-sub-packages: '["libcusparse-dev", "libcublas-dev", "libcusolver-dev"]'

      - name: Run sccache-cache
        uses: mozilla-actions/sccache-action@v0.0.9

      - name: Configure sccache
        uses: actions/github-script@v7
        with:
          script: |
            core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
            core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');

      - name: Run sccache stat for check
        shell: bash
        run: ${SCCACHE_PATH} --show-stats

      - name: Configure project
        run: |
          sudo apt update -y
          sudo bash -x dependencies.sh -y
          mkdir build
          cd build
          cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
        shell: bash

      - name: Build project
        run: |
          export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:LIBRARY_PATH
          cd build
          make -j
          sudo make install
        shell: bash

      - name: Build nvlink_allocator.so
        run: |
          export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
          export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
          export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
          mkdir -p build/mooncake-transfer-engine/nvlink-allocator
          cd mooncake-transfer-engine/nvlink-allocator
          bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/
        shell: bash

      - name: Generate Python version tag
        id: generate_tag_release
        run: |
          echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
        shell: bash

      - name: Build Python wheel
        run: |
          # Set LD_LIBRARY_PATH for wheel building
          export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
          PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh
        env:
          VERSION: ${{ env.VERSION }}

      - name: Upload Python wheel artifact
        uses: actions/upload-artifact@v4
        with:
          name: mooncake-wheel-py${{ steps.generate_tag_release.outputs.python_version_tag }}
          path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl

  publish-release:
    needs: build
    runs-on: ubuntu-22.04
    permissions:
      contents: write
      id-token: write
    steps:
      - name: Checkout source
        uses: actions/checkout@v4

      - name: Download all wheel artifacts
        uses: actions/download-artifact@v4
        with:
          path: mooncake-wheel/dist-all

      - name: Prepare wheels for release
        run: |
          # Move all wheels to a single directory
          mkdir -p mooncake-wheel/dist-release
          find mooncake-wheel/dist-all -name "*.whl" -exec cp {} mooncake-wheel/dist-release/ \;
          ls -la mooncake-wheel/dist-release/
          # List all collected wheels
          echo "Collected wheels for release:"
          ls -la mooncake-wheel/dist-release/

      - name: Upload wheels to GitHub Release
        uses: softprops/action-gh-release@v1
        with:
          files: mooncake-wheel/dist-release/*.whl

      - name: Publish package to PyPI
        if: github.repository == 'kvcache-ai/Mooncake'
        uses: pypa/gh-action-pypi-publish@release/v1
        with:
          packages-dir: mooncake-wheel/dist-release/
          password: ${{ secrets.PYPI_API_TOKEN }}
