gpu-load-test.yaml

name: GPU Load Test

on:
  workflow_dispatch:
    inputs:
      runners:
        description: 'Runners to test (comma-separated, or "all")'
        type: string
        default: 'all'
      image:
        description: 'Docker image'
        type: string
        default: 'rocm/atom-dev:latest'

jobs:
  parse-runners:
    name: Parse runners
    runs-on: ubuntu-latest
    outputs:
      matrix: ${{ steps.parse.outputs.matrix }}
    steps:
      - name: Parse runner input
        id: parse
        run: |
          INPUT="${{ inputs.runners || 'all' }}"
          if [ "$INPUT" = "all" ]; then
            MATRIX='[{"runner":"mia1-p01-g33","label":"mia1-p01-g33"},{"runner":"mia1-p01-g34","label":"mia1-p01-g34"},{"runner":"mia1-p01-g40","label":"mia1-p01-g40"},{"runner":"mia1-p01-g42","label":"mia1-p01-g42"},{"runner":"mia1-p01-g45","label":"mia1-p01-g45"},{"runner":"mia1-p01-g64","label":"mia1-p01-g64"}]'
          else
            MATRIX="["
            SEP=""
            IFS=',' read -ra RUNNERS <<< "$INPUT"
            for r in "${RUNNERS[@]}"; do
              r=$(echo "$r" | xargs)
              MATRIX="${MATRIX}${SEP}{\"runner\":\"${r}\",\"label\":\"${r}\"}"
              SEP=","
            done
            MATRIX="${MATRIX}]"
          fi
          echo "matrix=${MATRIX}" >> $GITHUB_OUTPUT
          echo "Runner matrix: ${MATRIX}"

  gpu-load-test:
    name: GPU Load Test (${{ matrix.config.label }})
    needs: parse-runners
    strategy:
      fail-fast: false
      matrix:
        config: ${{ fromJson(needs.parse-runners.outputs.matrix) }}
    runs-on: ${{ matrix.config.runner }}

    env:
      CONTAINER_NAME: gpu_load_test
      MODEL_NAME: "deepseek-ai/DeepSeek-R1-0528"
      TENSOR_PARALLEL: 8
      KV_CACHE_DTYPE: "fp8"

    steps:
      - name: Kill all Docker containers and clean up workspace
        run: |
          echo "=== Cleaning up containers on $(hostname) ==="
          containers=$(docker ps -q)
          if [ -n "$containers" ]; then
            docker kill $containers || true
          fi
          docker run --rm -v "${{ github.workspace }}":/workspace -w /workspace --privileged rocm/pytorch:latest bash -lc "ls -la /workspace/ && rm -rf /workspace/*" || true

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

      - name: GPU status
        run: |
          echo "Hostname: $(hostname)"
          echo ""
          if command -v rocm-smi &> /dev/null; then
            rocm-smi --showid || true
            echo ""
            rocm-smi --showtemp 2>&1 | grep "Temperature (Sensor junction)" | head -8 || true
            echo ""
            rocm-smi --showmemuse || true
          fi

      - name: Resolve model path
        id: model
        run: |
          if [ -f "/models/${{ env.MODEL_NAME }}/config.json" ]; then
            echo "path=/models/${{ env.MODEL_NAME }}" >> $GITHUB_OUTPUT
            echo "Found model at /models/${{ env.MODEL_NAME }}"
          elif [ -f "/data/${{ env.MODEL_NAME }}/config.json" ]; then
            echo "path=/data/${{ env.MODEL_NAME }}" >> $GITHUB_OUTPUT
            echo "Found model at /data/${{ env.MODEL_NAME }}"
          else
            echo "path=${{ env.MODEL_NAME }}" >> $GITHUB_OUTPUT
            echo "Model not found locally, will download from HuggingFace"
          fi

      - name: Start container
        run: |
          if [ -f "/etc/podinfo/gha-render-devices" ]; then
            DEVICE_FLAG=$(cat /etc/podinfo/gha-render-devices)
          else
            DEVICE_FLAG="--device /dev/dri"
          fi

          MOUNT_FLAGS=""
          [ -d "/data" ] && MOUNT_FLAGS="$MOUNT_FLAGS -v /data:/data"
          [ -d "/models" ] && MOUNT_FLAGS="$MOUNT_FLAGS -v /models:/models"

          docker run -dt \
            --name ${{ env.CONTAINER_NAME }} \
            --network=host \
            --device=/dev/kfd $DEVICE_FLAG \
            --group-add video \
            --cap-add=SYS_PTRACE \
            --security-opt seccomp=unconfined \
            $MOUNT_FLAGS \
            -v "${{ github.workspace }}:/workspace" \
            -w /workspace \
            -e HF_HOME=/data/huggingface_cache \
            -e NCCL_DEBUG=WARN \
            -e RCCL_DEBUG=WARN \
            --shm-size=16G \
            --ulimit memlock=-1 \
            --ulimit stack=67108864 \
            ${{ inputs.image || 'rocm/atom-dev:latest' }}

      - name: Run GPU load test
        id: test
        timeout-minutes: 45
        run: |
          MODEL_PATH="${{ steps.model.outputs.path }}"

          docker exec ${{ env.CONTAINER_NAME }} bash -c '
            MODEL_RUNNER="/app/ATOM/atom/model_engine/model_runner.py"

            # Add timing instrumentation
            if ! grep -q "^import time$" "$MODEL_RUNNER"; then
                sed -i "1a import time" "$MODEL_RUNNER"
            fi

            sed -i "/load_model(self.model, config.model, config.hf_config, config.load_dummy)/i\\
                load_start_time = time.time()\\
                logger.info(f\"[LOAD START] GPU {self.rank} | Time: {load_start_time:.6f}\")" \
                "$MODEL_RUNNER"

            sed -i "/load_model(self.model, config.model, config.hf_config, config.load_dummy)/a\\
                load_elapsed = time.time() - load_start_time\\
                logger.info(f\"[LOAD DONE] GPU {self.rank} | Duration: {load_elapsed:.2f}s | Time: {time.time():.6f}\")" \
                "$MODEL_RUNNER"

            python3 -m atom.examples.simple_inference \
              --model "'"$MODEL_PATH"'" \
              --kv_cache_dtype '"${{ env.KV_CACHE_DTYPE }}"' \
              -tp '"${{ env.TENSOR_PARALLEL }}"' \
              --temperature 0
          ' 2>&1 | tee gpu_load_test.log

      - name: Analyze results and generate summary
        if: always()
        run: |
          LOG_FILE="gpu_load_test.log"
          HOSTNAME=$(hostname)
          LOAD_COUNT=$(grep -c "\[LOAD DONE\]" "$LOG_FILE" 2>/dev/null || echo 0)

          # Start summary
          {
            echo "## GPU Load Test - ${HOSTNAME} (${{ matrix.config.label }})"
            echo ""
            echo "| Item | Value |"
            echo "|------|-------|"
            echo "| Hostname | \`${HOSTNAME}\` |"
            echo "| Runner | \`${{ matrix.config.runner }}\` |"
            echo "| Date | $(date -u '+%Y-%m-%d %H:%M:%S UTC') |"
            echo "| Model | \`${{ env.MODEL_NAME }}\` |"
            echo "| GPUs Completed | ${LOAD_COUNT} / 8 |"
            echo ""
          } >> $GITHUB_STEP_SUMMARY

          if [ "$LOAD_COUNT" -eq 0 ]; then
            echo "### Result: FAILED" >> $GITHUB_STEP_SUMMARY
            echo "No GPU load completion markers found. Check the logs." >> $GITHUB_STEP_SUMMARY
            exit 1
          fi

          # Extract timing data
          LOAD_DATA=$(grep "\[LOAD DONE\]" "$LOG_FILE" | \
            sed 's/.*\[atom\] //' | \
            grep -oP 'GPU \K\d+.*Duration: [0-9.]+s' | \
            sort -t' ' -k1 -n)

          MIN_TIME=$(echo "$LOAD_DATA" | awk '{print $3}' | sed 's/s$//' | sort -n | head -1)
          MAX_TIME=$(echo "$LOAD_DATA" | awk '{print $3}' | sed 's/s$//' | sort -n | tail -1)
          DELTA=$(awk "BEGIN {printf \"%.2f\", $MAX_TIME - $MIN_TIME}")

          # Determine overall status
          if (( $(echo "$DELTA < 1" | bc -l) )); then
            STATUS_EMOJI="✅"
            STATUS_TEXT="EXCELLENT"
          elif (( $(echo "$DELTA < 5" | bc -l) )); then
            STATUS_EMOJI="✅"
            STATUS_TEXT="GOOD"
          elif (( $(echo "$DELTA < 10" | bc -l) )); then
            STATUS_EMOJI="⚠️"
            STATUS_TEXT="MODERATE"
          else
            STATUS_EMOJI="❌"
            STATUS_TEXT="HIGH VARIANCE"
          fi

          # Per-GPU table
          {
            echo "### Per-GPU Load Times"
            echo ""
            echo "| GPU | Load Time | Delta from Fastest | Status |"
            echo "|-----|-----------|-------------------|--------|"
          } >> $GITHUB_STEP_SUMMARY

          while IFS= read -r line; do
            GPU=$(echo "$line" | awk '{print $1}')
            TIME=$(echo "$line" | awk '{print $3}' | sed 's/s$//')
            D=$(awk "BEGIN {printf \"%.2f\", $TIME - $MIN_TIME}")
            PCT=$(awk "BEGIN {printf \"%.1f\", ($TIME - $MIN_TIME) / $MIN_TIME * 100}")

            if (( $(echo "$D < 1" | bc -l) )); then
              S="✅ Excellent"
            elif (( $(echo "$D < 5" | bc -l) )); then
              S="✅ Good"
            elif (( $(echo "$D < 10" | bc -l) )); then
              S="⚠️ Moderate"
            else
              S="❌ SLOW"
            fi

            echo "| GPU ${GPU} | ${TIME}s | +${D}s (${PCT}%) | ${S} |" >> $GITHUB_STEP_SUMMARY
          done <<< "$LOAD_DATA"

          # Statistics
          STATS=$(grep "\[LOAD DONE\]" "$LOG_FILE" | grep -oP 'Duration: \K[0-9.]+' | awk '{
              times[NR] = $1; sum += $1;
              if(NR==1) { min=max=$1 }
              if($1 < min) min=$1;
              if($1 > max) max=$1;
          } END {
              if(NR > 0) {
                  avg = sum / NR;
                  n = asort(times, sorted);
                  median = (n % 2) ? sorted[(n+1)/2] : (sorted[n/2] + sorted[n/2+1]) / 2;
                  sum_sq = 0;
                  for(i=1; i<=NR; i++) { d = times[i] - avg; sum_sq += d*d; }
                  stddev = sqrt(sum_sq / NR);
                  printf "%.2f %.2f %.2f %.2f %.2f %.2f", min, max, avg, median, stddev, max-min;
              }
          }')
          read S_MIN S_MAX S_AVG S_MED S_STD S_DELTA <<< "$STATS"

          {
            echo ""
            echo "### Statistics"
            echo ""
            echo "| Metric | Value |"
            echo "|--------|-------|"
            echo "| Min Load Time | ${S_MIN}s |"
            echo "| Max Load Time | ${S_MAX}s |"
            echo "| Average | ${S_AVG}s |"
            echo "| Median | ${S_MED}s |"
            echo "| Std Deviation | ${S_STD}s |"
            echo "| Delta (Max-Min) | **${S_DELTA}s** |"
            echo ""
            echo "### Overall: ${STATUS_EMOJI} ${STATUS_TEXT} (delta: ${S_DELTA}s)"
            echo ""
          } >> $GITHUB_STEP_SUMMARY

          # Detailed diagnostics for high variance
          if (( $(echo "$DELTA >= 10" | bc -l) )); then
            {
              echo "<details>"
              echo "<summary>Diagnostic Details (High Variance Detected)</summary>"
              echo ""
              echo "**Likely causes (in order of probability):**"
              echo "1. Storage I/O bottleneck - check if \`/data\` is NFS/network mounted"
              echo "2. Model shard distribution - different shard sizes for different GPUs"
              echo "3. NUMA/memory issues - cross-NUMA memory access"
              echo "4. PCIe link degradation - reduced bandwidth on specific GPUs"
              echo ""
              echo "**Recommended actions:**"
              echo "- Run test again to check if same GPUs are consistently slow"
              echo "- Check storage mount: \`mount | grep /data\`"
              echo "- Check NUMA topology: \`rocm-smi --showtopo\`"
              echo "- Check PCIe link: \`lspci -vv | grep -A 10 'AMD/ATI' | grep LnkSta\`"
              echo "</details>"
              echo ""
            } >> $GITHUB_STEP_SUMMARY
          fi

      - name: Upload test log
        if: always()
        uses: actions/upload-artifact@v4
        with:
          name: gpu-load-test-${{ matrix.config.label }}-${{ github.run_id }}
          path: gpu_load_test.log

      - name: Clean up
        if: always()
        run: |
          docker stop ${{ env.CONTAINER_NAME }} || true
          docker rm ${{ env.CONTAINER_NAME }} || true