import_time.yml

name: "⏰ Import Time vs main"

on:
  pull_request:
      branches:
      - main

permissions:
  contents: read

jobs:
  benchmark-comparison:
    name: benchmark
    runs-on: ubuntu-latest
    strategy:
      matrix:
        include:
          # Ubuntu tests
          - extra: dataset,denoisers,test,training
            name: "all optional dependencies"
          - extra: ""
            name: "no optional dependencies"

    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with:
          python-version: 3.11
          cache: 'pip' # caching pip dependencies

      - name: "Set up Hyperfine"
        run: |
          sudo apt-get update && sudo apt-get install -y hyperfine

      - name: Install deepinv and its dependencies
        shell: bash
        run: |
          # if the extra is empty, install without optional dependencies
          if [ "${{ matrix.extra }}" = "" ]; then
            pip install .
          else
            pip install .[${{ matrix.extra }}]
          fi

      - name: "Benchmark torch+torchvision"
        run: |
          echo "Benchmarking torch..."
          hyperfine --warmup 3 --show-output --export-json torch.json 'python -c "import torch, torchvision"'

      - name: "Benchmark PR"
        run: |
          echo "Benchmarking PR branch..."
          hyperfine --warmup 3 --show-output --export-json head.json 'python -c "import deepinv"'

      - name: "Checkout main branch"
        run: |
          pip uninstall -y deepinv # Clean up for the next step
          git fetch origin main --depth=1
          git checkout main
          pip install .

      - name: "Benchmark main"
        run: |
          echo "Benchmarking main branch..."
          hyperfine --warmup 3 --show-output --export-json base.json 'python -c "import deepinv"'

      - name: "Compare results and fail on regression"
        uses: actions/github-script@v7
        with:
          script: |
            const fs = require('fs');

            const torchResult = JSON.parse(fs.readFileSync('torch.json', 'utf8')).results[0];
            const headResult = JSON.parse(fs.readFileSync('head.json', 'utf8')).results[0];
            const baseResult = JSON.parse(fs.readFileSync('base.json', 'utf8')).results[0];

            const torchMean = torchResult.mean;
            const headMean = headResult.mean;
            const baseMean = baseResult.mean;

            const rel_diff = (headMean - baseMean) / baseMean * 100;
            const threshold = 5; // 5% regression threshold

            core.info(`Torch mean time: ${(torchMean * 1000).toFixed(2)} ms`);
            core.info(`Main mean time: ${(baseMean * 1000).toFixed(2)} ms`);
            core.info(`PR mean time:   ${(headMean * 1000).toFixed(2)} ms`);
            core.info(`Rel difference (PR-main)/main: ${rel_diff.toFixed(2)} %`);
            core.info(`Failure threshold >+${threshold} %`);

            if (rel_diff > threshold) {
              core.setFailed(`❌ Performance regression detected! Import time increased by ${rel_diff}%, which is over the 5% threshold.`);
            } else {
              core.info('✅ No significant performance regression detected.');
            }