benchmark.yml

name: Run Benchmarks

on:
  workflow_dispatch:
    inputs:
      instance_families:
        description: 'Instance families to benchmark (comma-separated: c7i,c7a,m7i,m7a)'
        required: true
        type: string
        default: 'r7i.16xlarge'
      cpu_sizes:
        description: 'CPU sizes to benchmark (comma-separated: 8,16,32,64,96)'
        required: true
        type: string
        default: '64'
      docker_tag:
        description: 'Benchmark Docker image tag'
        required: true
        type: string
        default: 'main'
      programs:
        description: 'Comma-separated list of programs (ed25519_commonware,scep256r1_commonware,helios_light_client,fibonacci,rsa,ecdsa,json,regex,sha,tendermint,bubble_sort,iseven)'
        required: true
        type: string
        default: 'ed25519_commonware,scep256r1_commonware,helios_light_client,fibonacci,rsa,ecdsa,json,regex,sha,tendermint,bubble_sort,iseven'
      provers:
        description: 'Comma-separated list of provers (sp1,risc0)'
        required: true
        type: string
        default: 'sp1,risc0'
      benchmark_sizes:
        description: 'Comma-separated list of benchmark sizes (e.g. 10,50,100,200)'
        required: false
        type: string
        default: ''
      use_cache:
        description: 'Use GitHub Actions Rust cache'
        required: true
        type: boolean
        default: true
      enable_gpu:
        description: 'Enable GPU acceleration'
        required: true
        type: boolean
        default: false
      debug_mode:
        description: 'Enable debug logging'
        required: true
        type: boolean
        default: false

jobs:
  matrix-prep:
    runs-on: ubuntu-latest
    outputs:
      matrix: ${{ steps.set-matrix.outputs.matrix }}
    steps:
      - id: set-matrix
        run: |
          # Convert comma-separated values to arrays
          IFS=',' read -ra FAMILIES <<< "${{ inputs.instance_families }}"
          IFS=',' read -ra CPUS <<< "${{ inputs.cpu_sizes }}"
          IFS=',' read -ra PROVERS <<< "${{ inputs.provers }}"
          
          # Create matrix with specified CPU sizes and provers for each family
          matrix="{\"include\":["
          for family in "${FAMILIES[@]}"; do
            for cpu in "${CPUS[@]}"; do
              for prover in "${PROVERS[@]}"; do
                [[ $matrix != "{\"include\":[" ]] && matrix="$matrix,"
                matrix="$matrix{\"family\":\"$family\",\"cpu\":$cpu,\"prover\":\"$prover\"}"
              done
            done
          done
          matrix="$matrix]}"
          
          echo "matrix=$matrix" >> $GITHUB_OUTPUT

  run-benchmark:
    needs: matrix-prep
    name: Run ${{ matrix.family }} ${{ matrix.cpu }}CPU ${{ matrix.prover }} Benchmark
    strategy:
      matrix: ${{fromJson(needs.matrix-prep.outputs.matrix)}}
      fail-fast: false
    runs-on:
      - runs-on=${{ github.run_id }}
      - cpu=${{ matrix.cpu }}
      - family=${{ matrix.family }}
      - disk=large
      - spot=true
      - image=${{ inputs.enable_gpu && 'prooflab-cuda' || 'ubuntu22-full-x64' }}
    
    steps:
      - name: Checkout repository
        uses: actions/checkout@v4

      - name: Setup Rust cache
        if: inputs.use_cache
        uses: actions/cache@v4
        with:
          path: |
            ~/.cargo/registry
            ~/.cargo/git
          key: ${{ runner.os }}-${{ matrix.family }}-${{ matrix.cpu }}-${{ matrix.prover }}-cargo-${{ hashFiles('**/Cargo.lock') }}
          restore-keys: |
            ${{ runner.os }}-${{ matrix.family }}-${{ matrix.cpu }}-${{ matrix.prover }}-cargo-

      - name: Login to GitHub Container Registry
        uses: docker/login-action@v3
        with:
          registry: ghcr.io
          username: ${{ github.actor }}
          password: ${{ secrets.GITHUB_TOKEN }}

      - name: Pull benchmark image
        run: |
          # TODO: Figure out how to fix this in the GPU AMI
          sudo systemctl restart docker
          # Append -gpu or -cpu suffix to the tag based on GPU toggle
          IMAGE_TAG="${{ inputs.docker_tag }}${{ inputs.enable_gpu && '-gpu' || '-cpu' }}"
          docker pull ghcr.io/prooflabdev/prooflab-rs:${IMAGE_TAG}

      - name: Run benchmarks
        run: |
          # Set the image tag with appropriate suffix
          IMAGE_TAG="${{ inputs.docker_tag }}${{ inputs.enable_gpu && '-gpu' || '-cpu' }}"
          
          # Split the comma-separated programs into an array
          IFS=',' read -ra PROGRAMS <<< "${{ inputs.programs }}"
          
          # Create results directory with instance info
          mkdir -p "telemetry/${{ matrix.family }}_${{ matrix.cpu }}cpu_${{ matrix.prover }}"
          
          # Create cargo cache directories if they don't exist
          mkdir -p ~/.cargo/registry
          mkdir -p ~/.cargo/git

          # Check for AVX support
          if grep -q avx512 /proc/cpuinfo; then
            echo "CPU supports AVX512, enabling AVX512 acceleration"
            echo "RUSTFLAGS=-C target-cpu=native -C target-feature=+avx512f" >> $GITHUB_ENV
          elif grep -q avx2 /proc/cpuinfo; then
            echo "CPU supports AVX2, enabling AVX256 acceleration"
            echo "RUSTFLAGS=-C target-cpu=native" >> $GITHUB_ENV
          else
            echo "No AVX support detected"
          fi
          
          # Run benchmark for each program
          for program in "${PROGRAMS[@]}"; do
            # Determine benchmark sizes command
            BENCHMARK_SIZES_CMD=""
            if [ ! -z "${{ inputs.benchmark_sizes }}" ]; then
              BENCHMARK_SIZES_CMD="--benchmark-sizes ${{ inputs.benchmark_sizes }}"
            fi
            
            docker run \
              --network host \
              -v "$(pwd)/telemetry/${{ matrix.family }}_${{ matrix.cpu }}cpu_${{ matrix.prover }}:/prooflab-rs/telemetry" \
              ${{ inputs.use_cache && '-v ~/.cargo/registry:/root/.cargo/registry -v ~/.cargo/git:/root/.cargo/git' || '' }} \
              ${{ inputs.enable_gpu && '-v /var/run/docker.sock:/var/run/docker.sock --gpus all' || '' }} \
              ${{ inputs.enable_gpu && '-e PROOFLAB_GPU=true' || '' }} \
              -e RUST_LOG=${{ inputs.debug_mode && 'debug' || 'info' }} \
              -e RUST_BACKTRACE=${{ inputs.debug_mode && 'full' || '1' }} \
              ghcr.io/prooflabdev/prooflab-rs:${IMAGE_TAG} \
              bash -c "cargo run --release -p prooflab -- prove-${{ matrix.prover }} examples/${program} --enable-telemetry ${BENCHMARK_SIZES_CMD}"
          done

      # Fix cache permissions before saving
      - name: Fix cache permissions
        if: inputs.use_cache
        run: |
          sudo chown -R $(id -u):$(id -g) ~/.cargo/registry
          sudo chown -R $(id -u):$(id -g) ~/.cargo/git

      - name: Upload benchmark results
        uses: actions/upload-artifact@v4
        with:
          name: benchmark-results-${{ matrix.family }}-${{ matrix.cpu }}cpu-${{ matrix.prover }}
          path: telemetry/${{ matrix.family }}_${{ matrix.cpu }}cpu_${{ matrix.prover }}/