Mooncake/.github/workflows/ci_efa.yml

name: 'Build Wheel (AWS EFA)'

on:
  workflow_call: {}

# Builds the AWS EFA (libfabric) wheel variants on a stock ubuntu runner.
# No EFA hardware is required to *build*: USE_EFA only needs the libfabric
# headers/lib to compile and link. auditwheel later excludes libfabric/libefa
# from the wheel so they resolve to the user's system EFA install
# (/opt/amazon/efa/lib) at runtime. The distro libfabric (1.x) is ABI-forward-
# compatible with the AWS EFA libfabric (2.x) that loads at runtime; the EFA
# transport only uses long-stable fi_* core APIs.
#
# Two variants, since the EFA transport's memory path is CUDA-aware
# (FI_HMEM_CUDA / GPUDirect under USE_CUDA=ON, FI_HMEM=system otherwise):
#   efa            USE_CUDA=ON   (GPU)
#   efa-non-cuda   USE_CUDA=OFF  (CPU/DRAM)
# PR validation builds one python version per variant to keep CI cheap;
# the release workflow builds the full python matrix.
jobs:
  build-wheel-efa:
    runs-on: ubuntu-22.04
    strategy:
      matrix:
        include:
          - variant: cuda
            use_cuda: "ON"
            build_env: "EFA_BUILD"
            python-version: "3.12"
          - variant: non-cuda
            use_cuda: "OFF"
            build_env: "EFA_NON_CUDA_BUILD"
            python-version: "3.10"
    env:
      TORCH_CUDA_ARCH_LIST: "8.0;9.0"
      SCCACHE_GHA_ENABLED: "true"

    steps:
    - uses: actions/checkout@v4
      with:
        persist-credentials: false

    - 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
      if: matrix.use_cuda == 'ON'
      uses: Jimver/cuda-toolkit@v0.2.24
      with:
        cuda: '12.8.1'
        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: Install dependencies
      run: |
        sudo apt update -y
        sudo apt install -y ninja-build libfabric-dev libfabric1
        sudo bash -x dependencies.sh -y
        df -h        
      shell: bash

    - name: Configure project
      run: |
        mkdir build
        cd build
        EXTRA_FLAGS=""
        if [ "${{ matrix.use_cuda }}" = "ON" ]; then
          EXTRA_FLAGS="-DCMAKE_EXE_LINKER_FLAGS=-L/usr/local/cuda/lib64/stubs"
        fi
        cmake -G Ninja .. \
          -DUSE_ETCD=ON \
          -DUSE_HTTP=ON \
          -DWITH_STORE=ON \
          -DWITH_METRICS=ON \
          -DBUILD_UNIT_TESTS=OFF \
          -DBUILD_EXAMPLES=ON \
          -DENABLE_SCCACHE=ON \
          -DBUILD_BENCHMARK=ON \
          -DUSE_EFA=ON \
          -DUSE_CUDA=${{ matrix.use_cuda }} \
          -DLIBFABRIC_INCLUDE_DIR=/usr/include \
          -DLIBFABRIC_LIBRARY=/usr/lib/x86_64-linux-gnu/libfabric.so \
          -DCMAKE_BUILD_TYPE=Release \
          -DENABLE_DEBUG_SYMBOLS=OFF \
          $EXTRA_FLAGS        
      shell: bash

    - name: Build project
      run: |
        if [ "${{ matrix.use_cuda }}" = "ON" ]; then
          export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
          export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
        fi
        cd build
        cmake --build .
        sudo cmake --install .
        df -h        
      shell: bash

    - name: Run sccache stat for check
      if: ${{ env.SCCACHE_PATH != '' }}
      shell: bash
      run: ${SCCACHE_PATH} --show-stats

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

    - name: Build Python wheel
      run: |
        export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
        export ${{ matrix.build_env }}=1
        PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag.outputs.python_version_tag }} ./scripts/build_wheel.sh        
      shell: bash

    - name: Verify libfabric is excluded from the wheel
      run: |
        WHL=$(ls mooncake-wheel/dist-py${{ steps.generate_tag.outputs.python_version_tag }}/*.whl | head -1)
        echo "Inspecting $WHL"
        if unzip -l "$WHL" | grep -iE 'libfabric|libefa'; then
          echo "::error::libfabric/libefa must NOT be bundled in the EFA wheel"
          exit 1
        fi
        echo "OK: libfabric/libefa correctly excluded (resolve to system EFA at runtime)"        
      shell: bash

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