Model-Optimizer/.github/workflows/example_tests.yml

name: Example tests

on:
  push:
    branches: ["pull-request/[0-9]+"]
    # NOTE: paths cannot be used since push happens to copied PR and only latest commit to PR is used
  schedule:
    - cron: "0 0 * * *" # Nightly
  workflow_dispatch:
    # On-demand

concurrency:
  # Cancel previous runs if new commit is pushed to the same PR
  group: ${{ github.workflow }}-${{ startsWith(github.ref, 'refs/heads/pull-request/') && github.ref || github.sha }}
  cancel-in-progress: true

# Each job's `allow_failure` reads the repo variable ALLOW_FAILURE_EXAMPLE_TESTS:
# a comma-separated list of example names whose test failures should be non-blocking, e.g. "torch_trt,llm_qat"

jobs:
  # One changed-files pass decides which lanes run. Lane granularity, not per-example: a job's
  # `if` cannot read `matrix`, so gating an individual example would need a job per example.
  pr-gate:
    runs-on: ubuntu-latest
    permissions:
      checks: read
      contents: read
    outputs:
      any_changed: ${{ steps.lanes.outputs.any_changed }}
      torch: ${{ steps.lanes.outputs.torch }}
      trtllm: ${{ steps.lanes.outputs.trtllm }}
      megatron: ${{ steps.lanes.outputs.megatron }}
      onnx: ${{ steps.lanes.outputs.onnx }}
    steps:
      - uses: actions/checkout@v6
        with:
          fetch-depth: 0
      - if: startsWith(github.ref, 'refs/heads/pull-request/')
        id: base
        uses: ./.github/actions/pr-merge-base
      - if: startsWith(github.ref, 'refs/heads/pull-request/')
        id: changed
        uses: step-security/changed-files@v47.0.5
        with:
          base_sha: ${{ steps.base.outputs.merge_base }}
          sha: ${{ steps.base.outputs.head_sha }}
          fail_on_initial_diff_error: true
          # `common` runs every lane. Each lane lists the example directories it covers, plus any
          # extra directory those tests reach into (e.g. hf_ptq's script runs lm_eval from
          # ../llm_eval, and the speculative_decoding test drives hf_ptq).
          files_yaml: |
            common:
              - .github/actions/**
              - .github/workflows/_example_tests_runner.yml
              - .github/workflows/example_tests.yml
              - modelopt/**
              - modelopt_recipes/**
              - pyproject.toml
              - tests/_test_utils/**
              - tests/conftest.py
              - tests/examples/conftest.py
              - "!**/*.{md,rst,png,ipynb}"
            torch:
              - "{examples,tests/examples}/{llm_distill,llm_qat,llm_sparsity,specdec_bench,speculative_decoding}/**"
              - examples/dataset/** # data prep for llm_qat and speculative_decoding
              - examples/hf_ptq/** # the speculative_decoding test drives hf_ptq
              - "!**/*.{md,rst,png,ipynb}"
            trtllm:
              - "{examples,tests/examples}/{gpt-oss,hf_ptq,llm_eval}/**"
              - examples/dataset/** # data prep for gpt-oss
              - "!**/*.{md,rst,png,ipynb}"
            megatron:
              - "{examples,tests/examples}/megatron_bridge/**"
              - "!**/*.{md,rst,png,ipynb}"
            onnx:
              - "{examples,tests/examples}/{diffusers,torch_onnx,torch_trt}/**"
              - examples/onnx_ptq/** # torch_trt reuses onnx_ptq
              - "!**/*.{md,rst,png,ipynb}"
      - id: lanes
        env:
          # Nightly and on-demand runs have no diff to inspect, so they run everything.
          NON_PR: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') }}
          COMMON: ${{ steps.changed.outputs.common_any_modified }}
          TORCH: ${{ steps.changed.outputs.torch_any_modified }}
          TRTLLM: ${{ steps.changed.outputs.trtllm_any_modified }}
          MEGATRON: ${{ steps.changed.outputs.megatron_any_modified }}
          ONNX: ${{ steps.changed.outputs.onnx_any_modified }}
        run: |
          any=false
          for lane in TORCH TRTLLM MEGATRON ONNX; do
            if [[ "$NON_PR" == "true" || "$COMMON" == "true" || "${!lane}" == "true" ]]; then
              run=true
              any=true
            else
              run=false
            fi
            echo "${lane,,}=$run" | tee --append "$GITHUB_OUTPUT"
          done
          echo "any_changed=$any" | tee --append "$GITHUB_OUTPUT"
      - name: Wait for unit tests before spending GPU runners
        if: ${{ steps.lanes.outputs.any_changed == 'true' && startsWith(github.ref, 'refs/heads/pull-request/') }}
        uses: poseidon/wait-for-status-checks@v0.7.0
        with:
          token: ${{ secrets.GITHUB_TOKEN }}
          match_pattern: "^linux$" # Wait for Unit tests / linux
          delay: 300s

  ##### PyTorch Example Tests #####
  torch:
    needs: [pr-gate]
    if: needs.pr-gate.outputs.torch == 'true'
    strategy:
      fail-fast: false
      matrix:
        example: [llm_distill, llm_qat, llm_sparsity, specdec_bench, speculative_decoding]
    uses: ./.github/workflows/_example_tests_runner.yml
    permissions:
      contents: read
    secrets: inherit
    with:
      docker_image: "nvcr.io/nvidia/pytorch:26.07-py3"
      example: ${{ matrix.example }}
      timeout_minutes: 30
      pip_install_extras: "[hf,dev-test]"
      runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
      allow_failure: ${{ contains(format(',{0},', vars.ALLOW_FAILURE_EXAMPLE_TESTS), format(',{0},', matrix.example)) }}

  ##### TensorRT-LLM Example Tests #####
  trtllm:
    needs: [pr-gate]
    if: needs.pr-gate.outputs.trtllm == 'true'
    strategy:
      fail-fast: false
      matrix:
        example: [gpt-oss, hf_ptq, llm_eval]
    uses: ./.github/workflows/_example_tests_runner.yml
    permissions:
      contents: read
    secrets: inherit
    with:
      docker_image: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20"
      example: ${{ matrix.example }}
      pip_install_extras: "[hf,dev-test]"
      runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
      allow_failure: ${{ contains(format(',{0},', vars.ALLOW_FAILURE_EXAMPLE_TESTS), format(',{0},', matrix.example)) }}

  ##### Megatron Example Tests #####
  megatron:
    needs: [pr-gate]
    if: needs.pr-gate.outputs.megatron == 'true'
    uses: ./.github/workflows/_example_tests_runner.yml
    permissions:
      contents: read
    secrets: inherit
    with:
      docker_image: "nvcr.io/nvidia/nemo:26.08"
      example: megatron_bridge
      timeout_minutes: 60
      pip_install_extras: "[hf,puzzletron,dev-test]"
      runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
      allow_failure: ${{ contains(format(',{0},', vars.ALLOW_FAILURE_EXAMPLE_TESTS), ',megatron_bridge,') }}

  ##### ONNX/TensorRT Example Tests #####
  onnx:
    needs: [pr-gate]
    if: needs.pr-gate.outputs.onnx == 'true'
    strategy:
      fail-fast: false
      matrix:
        example: [diffusers, torch_onnx, torch_trt]
    uses: ./.github/workflows/_example_tests_runner.yml
    permissions:
      contents: read
    secrets: inherit
    with:
      # Pinned to 26.05 (TensorRT 10): torch-tensorrt is capped at <2.13 (== 2.12.1),
      # which needs libnvinfer.so.10; newer tensorrt containers drop it. Bump only once
      # a torch-tensorrt build for the newer TensorRT is available.
      docker_image: "nvcr.io/nvidia/tensorrt:26.05-py3"
      example: ${{ matrix.example }}
      timeout_minutes: 45
      pip_install_extras: "[onnx,hf,dev-test]"
      runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
      allow_failure: ${{ contains(format(',{0},', vars.ALLOW_FAILURE_EXAMPLE_TESTS), format(',{0},', matrix.example)) }}

  ##### Required Check for PR #####
  example-pr-required-check:
    # Run even if example tests are skipped
    if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }}
    needs: [pr-gate, torch, trtllm, megatron, onnx]
    runs-on: ubuntu-latest
    steps:
      - name: Required example tests did not succeed
        if: |
          needs.pr-gate.result != 'success' ||
          (needs.pr-gate.outputs.any_changed == 'true' && (
            !contains(fromJSON('["success", "skipped"]'), needs.torch.result) ||
            !contains(fromJSON('["success", "skipped"]'), needs.trtllm.result) ||
            !contains(fromJSON('["success", "skipped"]'), needs.megatron.result) ||
            !contains(fromJSON('["success", "skipped"]'), needs.onnx.result)
          ))
        run: exit 1