benchmarks_pr.yml

name: Continuous Benchmark (PR)
# Opt-in only: add the "benchmarks/upload" label to run the comparison. The
# baseline and PR legs run once on separate pinned runners for CPU and GPU. The
# downstream job compares those four artifacts. Both revisions use benchmark
# definitions archived from the PR base SHA.
on:
  pull_request:
    types: [opened, synchronize, reopened, labeled]

permissions:
  actions: read
  contents: read

concurrency:
  group: ${{ github.workflow }}-${{ github.ref }}
  cancel-in-progress: true

jobs:
  prepare-definitions:
    name: Pin benchmark definitions
    if: contains(github.event.pull_request.labels.*.name, 'benchmarks/upload')
    runs-on: ubuntu-22.04
    steps:
      - name: Checkout baseline benchmark definitions
        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4
        with:
          fetch-depth: 1
          repository: ${{ github.event.pull_request.base.repo.full_name }}
          ref: ${{ github.event.pull_request.base.sha }}
      - name: Archive benchmark definitions
        env:
          PR_BASE_SHA: ${{ github.event.pull_request.base.sha }}
        run: |
          set -euxo pipefail
          mkdir -p benchmark-definitions/payload
          cp -a benchmarks benchmark-definitions/payload/
          (
            cd benchmark-definitions/payload
            find benchmarks -type f -print0 \
              | sort -z \
              | xargs -0 sha256sum > ../definitions.sha256
            tar -czf ../benchmark-definitions.tar.gz benchmarks
          )
          DEFINITIONS_SHA="$(sha256sum benchmark-definitions/definitions.sha256 | cut -d' ' -f1)"
          export DEFINITIONS_SHA
          python3 - <<'PY'
          import json
          import os
          from pathlib import Path

          Path("benchmark-definitions/manifest.json").write_text(
              json.dumps(
                  {
                      "source_sha": os.environ["PR_BASE_SHA"],
                      "definitions_sha256": os.environ["DEFINITIONS_SHA"],
                  },
                  indent=2,
              )
              + "\n",
              encoding="utf-8",
          )
          PY
          rm -rf benchmark-definitions/payload
      - name: Upload pinned benchmark definitions
        uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4
        with:
          name: benchmark-definitions
          path: benchmark-definitions
          retention-days: 3

  prepare-environment:
    name: Prepare pinned benchmark environment
    if: contains(github.event.pull_request.labels.*.name, 'benchmarks/upload')
    runs-on: mt-l-x86aavx2-11-41-a10g
    container:
      image: nvidia/cuda:12.6.3-cudnn-devel-ubuntu22.04@sha256:b3e7fba84d169f46939f00c25be7d016f712a8d651f4756d6a55e693d84d94f2
      options: --gpus all --shm-size=8g
    defaults:
      run:
        shell: bash -l {0}
    steps:
      - name: Install checkout dependencies
        run: |
          apt-get update -y
          apt-get install -y --no-install-recommends ca-certificates git
          rm -rf /var/lib/apt/lists/*
      - name: Checkout environment lock
        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4
        with:
          fetch-depth: 1
          repository: ${{ github.event.pull_request.head.repo.full_name }}
          ref: ${{ github.event.pull_request.head.sha }}
      - name: Set up pinned Python
        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5
        with:
          python-version: '3.10.20'
      - name: Restore or create benchmark environment
        id: benchmark-env-cache
        uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4
        with:
          path: .benchmark-venv
          key: pr-benchmark-venv-v1-${{ runner.os }}-${{ hashFiles('.github/benchmark-versions.env', '.github/benchmark-requirements.txt') }}
      - name: Install pinned benchmark dependencies
        if: steps.benchmark-env-cache.outputs.cache-hit != 'true'
        run: |
          set -euxo pipefail
          source .github/benchmark-versions.env
          python -m venv .benchmark-venv
          source .benchmark-venv/bin/activate
          python -m pip install \
            "pip==${PIP_VERSION}" \
            "setuptools==${SETUPTOOLS_VERSION}" \
            "wheel==${WHEEL_VERSION}"
          grep -vE '^(autorom-accept-rom-license==|tensordict @ )' \
            .github/benchmark-requirements.txt \
            > /tmp/benchmark-bootstrap-requirements.txt
          grep -E '^(autorom-accept-rom-license==|tensordict @ )' \
            .github/benchmark-requirements.txt \
            > /tmp/benchmark-source-requirements.txt
          python -m pip install --pre --no-build-isolation \
            -r /tmp/benchmark-bootstrap-requirements.txt \
            --index-url https://pypi.org/simple \
            --extra-index-url https://download.pytorch.org/whl/nightly/cu126
          python -m pip install --pre --no-build-isolation --no-deps \
            -r /tmp/benchmark-source-requirements.txt \
            --index-url https://pypi.org/simple \
            --extra-index-url https://download.pytorch.org/whl/nightly/cu126
          python -m pip check
          python -m pip freeze --all > .benchmark-venv/benchmark-freeze.txt
      - name: Validate pinned benchmark environment
        run: |
          set -euxo pipefail
          source .github/benchmark-versions.env
          source .benchmark-venv/bin/activate
          python - <<'PY'
          import importlib.metadata
          import os

          import torch
          import torchvision

          assert torch.__version__ == os.environ["TORCH_VERSION"]
          assert torchvision.__version__ == os.environ["TORCHVISION_VERSION"]
          direct_url = importlib.metadata.distribution("tensordict").read_text(
              "direct_url.json"
          )
          assert os.environ["TENSORDICT_COMMIT"] in (direct_url or "")
          PY
          bash .github/unittest/helpers/assert_torch_version.sh nightly
          bash .github/unittest/helpers/assert_torch_tensordict_versions.sh nightly

  benchmark:
    name: ${{ matrix.device }} ${{ matrix.revision }} benchmark
    if: contains(github.event.pull_request.labels.*.name, 'benchmarks/upload')
    needs: [prepare-definitions, prepare-environment]
    runs-on: mt-l-x86aavx2-11-41-a10g
    strategy:
      fail-fast: false
      max-parallel: 4
      matrix:
        include:
          - device: CPU
            revision: baseline
            repository: ${{ github.event.pull_request.base.repo.full_name }}
            sha: ${{ github.event.pull_request.base.sha }}
            image: nvidia/cuda:12.6.3-cudnn-runtime-ubuntu22.04@sha256:5f0d2d827f6436b3cb7468fd8acbdc8c1d41261614e579ae49afe6141da51133
          - device: CPU
            revision: contender
            repository: ${{ github.event.pull_request.head.repo.full_name }}
            sha: ${{ github.event.pull_request.head.sha }}
            image: nvidia/cuda:12.6.3-cudnn-runtime-ubuntu22.04@sha256:5f0d2d827f6436b3cb7468fd8acbdc8c1d41261614e579ae49afe6141da51133
          - device: GPU
            revision: baseline
            repository: ${{ github.event.pull_request.base.repo.full_name }}
            sha: ${{ github.event.pull_request.base.sha }}
            image: nvidia/cuda:12.6.3-cudnn-devel-ubuntu22.04@sha256:b3e7fba84d169f46939f00c25be7d016f712a8d651f4756d6a55e693d84d94f2
          - device: GPU
            revision: contender
            repository: ${{ github.event.pull_request.head.repo.full_name }}
            sha: ${{ github.event.pull_request.head.sha }}
            image: nvidia/cuda:12.6.3-cudnn-devel-ubuntu22.04@sha256:b3e7fba84d169f46939f00c25be7d016f712a8d651f4756d6a55e693d84d94f2
    container:
      image: ${{ matrix.image }}
      options: --gpus all --shm-size=8g
    defaults:
      run:
        shell: bash -l {0}
    env:
      COMPOSITE_LP_AGGREGATE: '0'
      LANG: C.UTF-8
      LC_ALL: C.UTF-8
      MKL_NUM_THREADS: '1'
      OMP_NUM_THREADS: '1'
      PIP_DISABLE_PIP_VERSION_CHECK: '1'
      PYTHONHASHSEED: '0'
      PYTHONNOUSERSITE: '1'
      TD_GET_DEFAULTS_TO_NONE: '1'
      TORCHDYNAMO_INLINE_INBUILT_NN_MODULES: '1'
      TZ: Europe/London
    steps:
      - name: Install pinned system environment
        run: |
          apt-get update -y
          apt-get install -y --no-install-recommends \
            ca-certificates g++ gcc git libgl1-mesa-glx libgles2-mesa-dev \
            libglew-dev libglfw3-dev libglu1-mesa libosmesa6 libz-dev
          rm -rf /var/lib/apt/lists/*
      - name: Checkout workflow inputs
        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4
        with:
          fetch-depth: 1
          repository: ${{ github.event.pull_request.head.repo.full_name }}
          ref: ${{ github.event.pull_request.head.sha }}
      - name: Checkout ${{ matrix.revision }} source
        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4
        with:
          fetch-depth: 1
          path: source
          repository: ${{ matrix.repository }}
          ref: ${{ matrix.sha }}
      - name: Set up pinned Python
        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5
        with:
          python-version: '3.10.20'
      - name: Restore pinned benchmark environment
        id: benchmark-env-cache
        uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4
        with:
          fail-on-cache-miss: true
          path: .benchmark-venv
          key: pr-benchmark-venv-v1-${{ runner.os }}-${{ hashFiles('.github/benchmark-versions.env', '.github/benchmark-requirements.txt') }}
      - name: Download pinned benchmark definitions
        uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4
        with:
          name: benchmark-definitions
          path: benchmark-definitions
      - name: Install ${{ matrix.revision }} source
        env:
          TORCHRL_BENCHMARK_DEVICE: ${{ matrix.device }}
        run: |
          set -euxo pipefail
          source .benchmark-venv/bin/activate
          if [ "${TORCHRL_BENCHMARK_DEVICE}" = CPU ]; then
            export CUDA_VISIBLE_DEVICES=
          else
            export CUDA_VISIBLE_DEVICES=0
          fi
          rm -rf source/build
          python -m pip install -e source --no-build-isolation --no-deps
          if [ "${TORCHRL_BENCHMARK_DEVICE}" = GPU ]; then
            nvcc --version
            python -c "import torch; assert torch.cuda.device_count()"
            python -c "import torchrl._torchrl as ext; assert hasattr(ext, 'CudaSumSegmentTreeFp32')"
          fi
      - name: Run ${{ matrix.revision }} benchmarks
        env:
          TORCHRL_BENCHMARK_DEVICE: ${{ matrix.device }}
        run: |
          set -euxo pipefail
          source .benchmark-venv/bin/activate
          if [ "${TORCHRL_BENCHMARK_DEVICE}" = CPU ]; then
            export CUDA_VISIBLE_DEVICES=
          else
            export CUDA_VISIBLE_DEVICES=0
          fi
          mkdir -p .benchmark-site "benchmark-result/${TORCHRL_BENCHMARK_DEVICE}"
          cat > .benchmark-site/sitecustomize.py <<'PY'
          import warnings

          try:
              import torch

              torch._dynamo.config.reorderable_logging_functions.add(warnings.warn)
          except (AttributeError, ImportError):
              pass
          PY
          export PYTHONPATH="${GITHUB_WORKSPACE}/.benchmark-site"
          tar -xzf benchmark-definitions/benchmark-definitions.tar.gz \
            -C benchmark-definitions
          RESULT_JSON="${GITHUB_WORKSPACE}/benchmark-result/${TORCHRL_BENCHMARK_DEVICE}/${{ matrix.revision }}.json"
          cd benchmark-definitions/benchmarks
          python -m pytest -vvv --rank 0 --timeout=240 --benchmark-only \
            --benchmark-json "${RESULT_JSON}" --ignore test_llm.py \
            -k 'not async_collection_pixels_64_envs and not 64-processes' .
      - name: Record benchmark metadata
        if: always()
        env:
          BASE_SHA: ${{ github.event.pull_request.base.sha }}
          BENCHMARK_COMMAND: python -m pytest -vvv --rank 0 --timeout=240 --benchmark-only --ignore test_llm.py -k 'not async_collection_pixels_64_envs and not 64-processes' .
          DEVICE: ${{ matrix.device }}
          HEAD_SHA: ${{ github.event.pull_request.head.sha }}
          IMAGE: ${{ matrix.image }}
          PR_NUMBER: ${{ github.event.pull_request.number }}
          REVISION: ${{ matrix.revision }}
          REVISION_SHA: ${{ matrix.sha }}
        run: |
          set -euxo pipefail
          source .benchmark-venv/bin/activate
          DEPENDENCY_LOCK_SHA="$(cat .github/benchmark-versions.env .github/benchmark-requirements.txt | sha256sum | cut -d' ' -f1)"
          DEPENDENCY_ENVIRONMENT_SHA="$(sha256sum .benchmark-venv/benchmark-freeze.txt | cut -d' ' -f1)"
          DEFINITIONS_SHA="$(python -c 'import json; print(json.load(open("benchmark-definitions/manifest.json"))["definitions_sha256"])')"
          DEFINITIONS_SOURCE_SHA="$(python -c 'import json; print(json.load(open("benchmark-definitions/manifest.json"))["source_sha"])')"
          PYTHON_VERSION="$(python -c 'import platform; print(platform.python_version())')"
          SYSTEM_ENVIRONMENT_SHA="$(dpkg-query -W -f='${Package}=${Version}\n' | sort | sha256sum | cut -d' ' -f1)"
          export DEPENDENCY_ENVIRONMENT_SHA DEPENDENCY_LOCK_SHA DEFINITIONS_SHA DEFINITIONS_SOURCE_SHA PYTHON_VERSION SYSTEM_ENVIRONMENT_SHA
          mkdir -p "benchmark-result/${DEVICE}"
          python - <<'PY'
          import json
          import os
          from pathlib import Path

          metadata = {
              "device": os.environ["DEVICE"],
              "revision": os.environ["REVISION"],
              "sha": os.environ["REVISION_SHA"],
              "pr_number": int(os.environ["PR_NUMBER"]),
              "base_sha": os.environ["BASE_SHA"],
              "head_sha": os.environ["HEAD_SHA"],
              "runner": "mt-l-x86aavx2-11-41-a10g",
              "image": os.environ["IMAGE"],
              "python_version": os.environ["PYTHON_VERSION"],
              "system_environment_sha256": os.environ["SYSTEM_ENVIRONMENT_SHA"],
              "dependency_lock_sha256": os.environ["DEPENDENCY_LOCK_SHA"],
              "dependency_environment_sha256": os.environ[
                  "DEPENDENCY_ENVIRONMENT_SHA"
              ],
              "benchmark_definitions_source_sha": os.environ[
                  "DEFINITIONS_SOURCE_SHA"
              ],
              "benchmark_definitions_sha256": os.environ["DEFINITIONS_SHA"],
              "benchmark_command": os.environ["BENCHMARK_COMMAND"],
              "benchmark_environment": {
                  "COMPOSITE_LP_AGGREGATE": "0",
                  "CUDA_VISIBLE_DEVICES": "" if os.environ["DEVICE"] == "CPU" else "0",
                  "LANG": "C.UTF-8",
                  "LC_ALL": "C.UTF-8",
                  "MKL_NUM_THREADS": "1",
                  "OMP_NUM_THREADS": "1",
                  "PYTHONHASHSEED": "0",
                  "PYTHONNOUSERSITE": "1",
                  "TD_GET_DEFAULTS_TO_NONE": "1",
                  "TORCHDYNAMO_INLINE_INBUILT_NN_MODULES": "1",
                  "TZ": "Europe/London",
              },
              "run_id": os.environ["GITHUB_RUN_ID"],
          }
          output = Path("benchmark-result") / os.environ["DEVICE"]
          (output / f"{os.environ['REVISION']}-metadata.json").write_text(
              json.dumps(metadata, indent=2) + "\n", encoding="utf-8"
          )
          PY
      - name: Upload ${{ matrix.device }} ${{ matrix.revision }} JSON
        if: always()
        uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4
        with:
          name: ${{ matrix.device }}-${{ matrix.revision }}-benchmark-json
          path: benchmark-result
          retention-days: 3

  compare:
    name: Compare benchmark results
    if: contains(github.event.pull_request.labels.*.name, 'benchmarks/upload')
    needs: benchmark
    runs-on: ubuntu-22.04
    steps:
      - name: Checkout comparison tooling
        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4
        with:
          fetch-depth: 1
          repository: ${{ github.event.pull_request.head.repo.full_name }}
          ref: ${{ github.event.pull_request.head.sha }}
      - name: Download baseline and PR JSON artifacts
        uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4
        with:
          merge-multiple: true
          path: raw-results
          pattern: '*-benchmark-json'
      - name: Compare benchmark results
        run: |
          set -euxo pipefail
          python .github/scripts/compare_pr_benchmarks.py \
            --raw-root raw-results \
            --output-root final-results \
            --summary final-summary.md \
            --reporting-threshold 5 \
            --image 'CPU=nvidia/cuda:12.6.3-cudnn-runtime-ubuntu22.04@sha256:5f0d2d827f6436b3cb7468fd8acbdc8c1d41261614e579ae49afe6141da51133' \
            --image 'GPU=nvidia/cuda:12.6.3-cudnn-devel-ubuntu22.04@sha256:b3e7fba84d169f46939f00c25be7d016f712a8d651f4756d6a55e693d84d94f2'
          cat final-summary.md >> "${GITHUB_STEP_SUMMARY}"
      - name: Upload final CPU benchmark comparison
        uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4
        with:
          name: CPU-benchmark-pr-results
          path: final-results/CPU
          retention-days: 30
      - name: Upload final GPU benchmark comparison
        uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4
        with:
          name: GPU-benchmark-pr-results
          path: final-results/GPU
          retention-days: 30