benchmark-multinode-tmpl.yml

name: Template - Multi-Node Benchmark

on:
  workflow_call:
    inputs:
      runner:
        required: true
        type: string
      image:
        required: true
        type: string
      model:
        required: true
        type: string
      framework:
        required: true
        type: string
      precision:
        required: true
        type: string
      exp-name:
        required: true
        type: string
      isl:
        required: true
        type: string
      osl:
        required: true
        type: string
      max-model-len:
        required: true
        type: string
      random-range-ratio:
        required: false
        type: string
        default: '0.8'
      mtp-mode:
        required: true
        type: string

env:
  EXP_NAME: ${{ inputs.exp-name }}
  IMAGE: ${{ inputs.image }}
  FRAMEWORK: ${{ inputs.framework }}
  PRECISION: ${{ inputs.precision }}
  ISL: ${{ inputs.isl }}
  OSL: ${{ inputs.osl }}
  MAX_MODEL_LEN: ${{ inputs.max-model-len }}
  RANDOM_RANGE_RATIO: ${{ inputs.random-range-ratio }}
  MTP_MODE: ${{ inputs.mtp-mode }}

permissions:
  contents: read

jobs:
  benchmark:
    runs-on: ${{ inputs.runner }}
    timeout-minutes: 480
    name: '${{ inputs.exp-name }} ${{ inputs.runner }} ${{ inputs.precision }} mtp-${{ inputs.mtp-mode }}'

    steps:
      - name: Resource cleanup
        run: |
          echo "[Slurm] Cleaning up resources ..."
          scancel -u $USER
          while [ -n "$(squeue -u $USER --noheader --format='%i')" ]; do
            squeue -u $USER
            sleep 5
          done

      - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
        with:
          token: ${{ secrets.REPO_PAT }}
          fetch-depth: 0

      - name: Launch multi-node job script
        env:
          RUNNER_NAME: ${{ runner.name }}
          RESULT_FILENAME: ${{ env.EXP_NAME }}_${{ env.PRECISION }}_${{ env.FRAMEWORK }}_mtp-${{ env.MTP_MODE }}_${{ runner.name }}
        run: |
          bash ./runners/launch_${RUNNER_NAME%%_*}.sh
          # Check if at least one result file was created
          if ls ${RESULT_FILENAME}_*.json 1> /dev/null 2>&1; then
            echo "RESULT_FILENAME=${RESULT_FILENAME}" >> $GITHUB_ENV
            echo "Found result files: $(ls ${RESULT_FILENAME}_*.json)"
          else
            echo "Run failed: No benchmark result files found for ${RESULT_FILENAME}_*.json" >&2
            exit 1
          fi

      - name: Process results
        env:
          RUNNER_TYPE: ${{ inputs.runner }}
        run: |
          # Process each result file
          for result_file in ${RESULT_FILENAME}_*.json; do
            if [ -f "$result_file" ]; then
              echo "Processing $result_file"
              # Extract GPU count, prefill_gpus and decode_gpus from filename for tp_size calculation
              gpus=$(echo "$result_file" | sed -n "s/.*_gpus_\([0-9]*\).*\.json/\1/p")
              prefill_gpus=$(echo "$result_file" | sed -n "s/.*_ctx_\([0-9]*\).*\.json/\1/p")
              decode_gpus=$(echo "$result_file" | sed -n "s/.*_gen_\([0-9]*\).*\.json/\1/p")
              
              if [ -n "$gpus" ]; then
                echo "Extracted: gpus=$gpus, prefill_gpus=$prefill_gpus, decode_gpus=$decode_gpus"
                TP=$gpus RESULT_FILENAME=${result_file%.json} EP_SIZE=1 DP_ATTENTION=false PREFILL_GPUS="$prefill_gpus" DECODE_GPUS="$decode_gpus" python3 utils/process_result.py
              fi
            fi
          done

      - name: Upload results
        uses: actions/upload-artifact@330a01c490aca151604b8cf639adc76d48f6c5d4 # v5.0.0
        with:
          name: ${{ env.RESULT_FILENAME }}
          path: agg_${{ env.RESULT_FILENAME }}_*.json