# Run one of the ML benchmarks on an AWS instance and return parsed results to Slab CI bot.
name: Application benchmarks

on:
  workflow_dispatch:
    inputs:
      instance_id:
        description: 'Instance ID'
        type: string
      instance_image_id:
        description: 'Instance AMI ID'
        type: string
      instance_type:
        description: 'Instance product type'
        type: string
      runner_name:
        description: 'Action runner name'
        type: string
      request_id:
        description: 'Slab request ID'
        type: string
      matrix_item:
        description: 'Build matrix item'
        type: string


env:
  CARGO_TERM_COLOR: always
  RESULTS_FILENAME: parsed_benchmark_results_${{ github.sha }}.json

jobs:
  run-ml-benchmarks:
    name: Execute ML benchmarks subset in EC2
    runs-on: ${{ github.event.inputs.runner_name }}
    if: ${{ !cancelled() }}
    steps:
      - name: Instance configuration used
        run: |
          echo "IDs: ${{ inputs.instance_id }}"
          echo "AMI: ${{ inputs.instance_image_id }}"
          echo "Type: ${{ inputs.instance_type }}"
          echo "Request ID: ${{ inputs.request_id }}"
          echo "Matrix item: ${{ inputs.matrix_item }}"

      - name: Get benchmark date
        run: |
          echo "BENCH_DATE=$(date --iso-8601=seconds)" >> "${GITHUB_ENV}"

      # SSH private key is required as some dependencies are from private repos
      - uses: webfactory/ssh-agent@v0.7.0
        with:
          ssh-private-key: ${{ secrets.CONCRETE_CI_SSH_PRIVATE }}

      - name: Fetch submodules
        uses: actions/checkout@v3
        with:
          fetch-depth: 0
          submodules: recursive
          token: ${{ secrets.CONCRETE_ACTIONS_TOKEN }}

      - name: Set up home
        # "Install rust" step require root user to have a HOME directory which is not set.
        run: |
          echo "HOME=/home/ubuntu" >> "${GITHUB_ENV}"

      - name: Setup rust toolchain for concrete-cpu
        uses: ./.github/workflows/setup_rust_toolchain_for_concrete_cpu

      - name: Build compiler and ML benchmarks
        run: |
          set -e
          cd compiler
          make BINDINGS_PYTHON_ENABLED=OFF build-mlbench

      - name: Download KeySetCache
        if: ${{ !contains(github.head_ref, 'newkeysetcache') }}
        continue-on-error: true
        run: |
          cd compiler
          GITHUB_TOKEN=${{ secrets.GITHUB_TOKEN }} make keysetcache_ci_populated

      - name: Run ML benchmarks
        run: |
          set -e
          cd compiler
          make BINDINGS_PYTHON_ENABLED=OFF ML_BENCH_SUBSET_ID=${{ inputs.matrix_item }} run-mlbench-subset

      - name: Upload raw results artifact
        uses: actions/upload-artifact@v4
        with:
          name: ${{ github.sha }}_raw
          path: compiler/benchmarks_results.json

      - name: Parse results
        shell: bash
        run: |
          COMMIT_DATE="$(git --no-pager show -s --format=%cd --date=iso8601-strict ${{ github.sha }})"
          COMMIT_HASH="$(git describe --tags --dirty)"
          python3 ./ci/benchmark_parser.py compiler/benchmarks_results.json ${{ env.RESULTS_FILENAME }} \
          --database compiler_benchmarks \
          --hardware ${{ inputs.instance_type }} \
          --project-version ${COMMIT_HASH} \
          --branch ${{ github.ref_name }} \
          --commit-date ${COMMIT_DATE} \
          --bench-date "${{ env.BENCH_DATE }}"

      - name: Upload parsed results artifact
        uses: actions/upload-artifact@v4
        with:
          name: ${{ github.sha }}
          path: ${{ env.RESULTS_FILENAME }}

      - name: Checkout Slab repo
        uses: actions/checkout@v3
        with:
          repository: zama-ai/slab
          path: slab
          token: ${{ secrets.CONCRETE_ACTIONS_TOKEN }}

      - name: Send data to Slab
        shell: bash
        run: |
          echo "Computing HMac on downloaded artifact"
          SIGNATURE="$(slab/scripts/hmac_calculator.sh ${{ env.RESULTS_FILENAME }} '${{ secrets.JOB_SECRET }}')"
          echo "Sending results to Slab..."
          curl -v -k \
          -H "Content-Type: application/json" \
          -H "X-Slab-Repository: ${{ github.repository }}" \
          -H "X-Slab-Command: store_data" \
          -H "X-Hub-Signature-256: sha256=${SIGNATURE}" \
          -d @${{ env.RESULTS_FILENAME }} \
          ${{ secrets.SLAB_URL }}
