[OMNIML-6180] Run the small example lanes of a job together

.github/actions/run-example/action.yml

@@ -0,0 +1,113 @@
+name: Run example tests
+description: Install the requirements of one example, run its tests with coverage and upload the coverage under the flag `examples-<example>`. Used by _example_tests_runner.yml,
+  once per example, so that examples sharing a job still get their own step, requirements and Codecov flag. A failing install or failing tests do not fail
+  this action; it reports them in `failed`, so that the caller can run the remaining examples first.
+
+inputs:
+  example:
+    description: Example directory name, e.g. `llm_distill` (tests/examples/<example>, examples/<example>).
+    required: true
+  allow-failure-examples:
+    description: Comma-separated example names (whitespace is ignored) whose failures are reported as a warning instead of as a failed example.
+    required: false
+    default: ""
+  codecov-token:
+    description: Codecov upload token (composite actions cannot read secrets).
+    required: false
+    default: ""
+
+outputs:
+  failed:
+    description: "'true' if the example's install or tests failed and the example is not in the allow-failure list."
+    value: ${{ steps.verdict.outputs.failed }}
+
+runs:
+  using: composite
+  steps:
+    - name: Run tests
+      id: run_tests
+      shell: bash
+      env:
+        # Absolute paths so subprocesses running from different working directories
+        # all find the config and write .coverage.* files to the same location.
+        COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml
+        COVERAGE_FILE: ${{ github.workspace }}/.coverage
+        # sys.monitoring records the same lines as the default C tracer at a fraction of the
+        # overhead; the tracer slowed each hf_ptq example run by ~60%. Needs Python 3.12+.
+        COVERAGE_CORE: sysmon
+      run: |
+        # Nothing here may fail the step: a failed step ends the whole job at once, whatever continue-on-error
+        # the caller puts on this action's step, and the later examples would be skipped. The exit status goes
+        # to the verdict below instead.
+        status=0
+
+        # Examples of one job share the GPU and the environment one after the other: log both, so a failure
+        # caused by an earlier example is easy to trace.
+        echo "::group::GPU state before ${{ inputs.example }}"
+        nvidia-smi --query-gpu=index,memory.used,memory.total --format=csv || true
+        nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv || true
+        echo "::endgroup::"
+
+        if [[ "${{ inputs.example }}" == "diffusers" ]]; then
+          echo "Uninstalling apex for diffusers: T5 Int8 (PixArt) + Apex is not supported as per https://github.com/huggingface/transformers/issues/21391"
+          python3 -m pip uninstall -y apex || true
+        fi
+
+        # Installed right before this example's tests, not up front for all examples of the job: an example
+        # that pins or upgrades packages (e.g. gpt-oss upgrades datasets, torch_trt downgrades torch) must not
+        # change the environment of the examples that run before it.
+        packages_before="${RUNNER_TEMP:-/tmp}/packages-before-${{ inputs.example }}.txt"
+        python3 -m pip freeze 2>/dev/null > "$packages_before" || true
+        echo "::group::Install requirements of ${{ inputs.example }}"
+        find examples/${{ inputs.example }} -name "requirements.txt" | while read req_file; do python3 -m pip install -r "$req_file" || exit 1; done || status=$?
+        echo "::endgroup::"
+        echo "::group::Packages changed by the requirements of ${{ inputs.example }} (< before, > after)"
+        { python3 -m pip freeze 2>/dev/null | diff "$packages_before" - || true; } | grep "^[<>]" || echo "none"
+        echo "::endgroup::"
+
+        if (( status == 0 )); then
+          echo "Running tests for: ${{ inputs.example }}"
+          pytest_args=()
+          if [[ "${{ inputs.example }}" == "diffusers" ]]; then
+            # Keep legacy ONNX/TRT tests available manually during deprecation.
+            pytest_args+=(--ignore=tests/examples/diffusers/test_diffusers_onnx_trt.py)
+          fi
+          if [[ "${{ inputs.example }}" == "hf_ptq" && "${{ github.ref }}" == refs/heads/pull-request/* ]]; then
+            # One GPU on PR runners: small-model tests share it and the rest run alone on it
+            # (tests/examples/hf_ptq/conftest.py). The 2-GPU nightly stays serial.
+            pytest_args+=(-n 4)
+          fi
+          python3 -m pytest tests/examples/${{ inputs.example }} --cov "${pytest_args[@]}" || status=$?
+        fi
+        echo "status=$status" >> "$GITHUB_OUTPUT"
+    - name: Verdict
+      id: verdict
+      shell: bash
+      env:
+        STATUS: ${{ steps.run_tests.outputs.status }}
+        EXAMPLE: ${{ inputs.example }}
+        ALLOW_FAILURE_EXAMPLES: ${{ inputs.allow-failure-examples }}
+      run: |
+        failed=false
+        if [[ "$STATUS" != "0" ]]; then
+          if [[ ",${ALLOW_FAILURE_EXAMPLES//[[:space:]]/}," == *",$EXAMPLE,"* ]]; then
+            echo "::warning title=Allowed example failure::'$EXAMPLE' failed but is in the allow-failure list (vars.ALLOW_FAILURE_EXAMPLE_TESTS); not blocking. Remove it from the variable once fixed."
+          else
+            echo "::error title=Example failed::The requirements install or the tests of '$EXAMPLE' failed (see the log of its Run tests step). The job fails after the remaining examples have run."
+            failed=true
+          fi
+        fi
+        echo "failed=$failed" >> "$GITHUB_OUTPUT"
+    - name: Upload coverage to Codecov
+      # A failed run, allowed or not, only has partial coverage
+      if: ${{ steps.run_tests.outputs.status == '0' }}
+      uses: codecov/codecov-action@303a32d7a59b442fa8d48b6a1cc6825c09c847a5 # v7.1.1
+      with:
+        token: ${{ inputs.codecov-token }}
+        files: coverage.xml
+        # One flag per example, not a shared `examples`: carryforward only applies to a flag
+        # with no upload, so a shared flag would replace every lane's coverage with the subset
+        # that ran once lanes are gated independently.
+        flags: examples-${{ inputs.example }}
+        fail_ci_if_error: false # test may be skipped if relevant file changes are not detected
+        verbose: true

.github/workflows/_example_tests_runner.yml

@@ -9,7 +9,7 @@ on:
         required: true
         type: string
       example:
-        description: "Example name to test (e.g. 'hf_ptq')"
+        description: "Example name to test (e.g. 'hf_ptq'), or up to 3 names separated by spaces: they share the job (one container start and one base install) and run one after the other, each in its own step, in the order given. List an example that pins or upgrades packages last."
         required: true
         type: string
       timeout_minutes:
@@ -27,11 +27,11 @@ on:
         required: false
         type: string
         default: "linux-amd64-gpu-rtxpro6000-latest-1"
-      allow_failure:
-        description: "If true, test failures are reported as a warning and do not fail the job (used to keep a known-broken example non-blocking)"
+      allow_failure_examples:
+        description: "Comma-separated example names whose test failures are reported as a warning and do not fail the job (used to keep a known-broken example non-blocking)"
         required: false
-        type: boolean
-        default: false
+        type: string
+        default: ""
 
 jobs:
   run-test:
@@ -50,6 +50,18 @@ jobs:
         # Build CUDA kernels only for the runner's RTX PRO 6000 (sm_120), not the image's ~6 archs.
         TORCH_CUDA_ARCH_LIST: "12.0"
     steps:
+      - name: Split the examples of this job
+        id: examples
+        shell: bash
+        env:
+          EXAMPLES: ${{ inputs.example }}
+        run: |
+          max=3  # the cap; the "Example N" steps and the final failure check below must match it
+          read -ra examples <<< "$EXAMPLES"
+          (( ${#examples[@]} >= 1 && ${#examples[@]} <= max )) || { echo "::error::A job runs 1 to $max examples, got: '$EXAMPLES'"; exit 1; }
+          for (( i = 0; i < max; i++ )); do echo "e$i=${examples[$i]:-}" >> "$GITHUB_OUTPUT"; done
+          # Expressions have no replace(), and a cache key is better without spaces
+          echo "key=$(IFS=+; echo "${examples[*]}")" >> "$GITHUB_OUTPUT"
       - name: Install git for vLLM image
         if: startsWith(inputs.docker_image, 'vllm/') || startsWith(inputs.docker_image, 'docker.io/vllm/')
         run: apt-get update && apt-get install -y git
@@ -58,8 +70,8 @@ jobs:
       - uses: ./.github/actions/cache-extensions
         id: ext-cache
         with:
-          # Per example: jobs sharing an image build different extensions, and the first to save wins the key.
-          cache-key: rtxpro6000-example-${{ inputs.example }}-${{ inputs.docker_image }}
+          # Per job (its examples): jobs sharing an image build different extensions, and the first to save wins the key.
+          cache-key: rtxpro6000-example-${{ steps.examples.outputs.key }}-${{ inputs.docker_image }}
       - name: Setup environment variables
         run: |
           echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu:/usr/local/tensorrt/targets/x86_64-linux-gnu/lib" >> $GITHUB_ENV
@@ -75,61 +87,47 @@ jobs:
           if [[ "${{ inputs.docker_image }}" == *"/tensorrt:"* ]]; then
             echo "torch<2.14" > /tmp/pip-constraints.txt
             export PIP_CONSTRAINT=/tmp/pip-constraints.txt
+            # The requirements of each example are installed later, in their own step
+            echo "PIP_CONSTRAINT=/tmp/pip-constraints.txt" >> "$GITHUB_ENV"
           fi
 
           # Use `python3 -m pip` instead of `pip` to avoid conflicts with system pip for nemo containers
           # Editable install so example scripts launched as subprocesses resolve modelopt to the same source path as the test process
           python3 -m pip install -e ".${{ inputs.pip_install_extras }}"
 
-          if [[ "${{ inputs.example }}" == *"diffusers"* ]]; then
-            echo "Uninstalling apex for diffusers: T5 Int8 (PixArt) + Apex is not supported as per https://github.com/huggingface/transformers/issues/21391"
-            python3 -m pip uninstall -y apex || true
-          fi
-
-          find examples/${{ inputs.example }} -name "requirements.txt" | while read req_file; do python3 -m pip install -r "$req_file" || exit 1; done
-      - name: Run tests
-        id: run_tests
-        shell: bash
-        continue-on-error: ${{ inputs.allow_failure }}
-        env:
-          # Absolute paths so subprocesses running from different working directories
-          # all find the config and write .coverage.* files to the same location.
-          COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml
-          COVERAGE_FILE: ${{ github.workspace }}/.coverage
-          # sys.monitoring records the same lines as the default C tracer at a fraction of the
-          # overhead; the tracer slowed each hf_ptq example run by ~60%. Needs Python 3.12+.
-          COVERAGE_CORE: sysmon
-        run: |
-          echo "Running tests for: ${{ inputs.example }}"
-          pytest_args=()
-          if [[ "${{ inputs.example }}" == "diffusers" ]]; then
-            # Keep legacy ONNX/TRT tests available manually during deprecation.
-            pytest_args+=(--ignore=tests/examples/diffusers/test_diffusers_onnx_trt.py)
-          fi
-          if [[ "${{ inputs.example }}" == "hf_ptq" && "${{ github.ref }}" == refs/heads/pull-request/* ]]; then
-            # One GPU on PR runners: small-model tests share it and the rest run alone on it
-            # (tests/examples/hf_ptq/conftest.py). The 2-GPU nightly stays serial.
-            pytest_args+=(-n 4)
-          fi
-          python3 -m pytest tests/examples/${{ inputs.example }} --cov "${pytest_args[@]}"
+      # One step per example. A failing example does not fail its step (the action reports `failed`
+      # instead, because a failed step would end the job at once), so the examples after it still run, as
+      # when each had its own job, and the last step of the job fails it. These three steps and that final
+      # check hold the job's cap of 3 (`max` in the first step).
+      - name: "Example 1: ${{ steps.examples.outputs.e0 }}"
+        id: example0
+        uses: ./.github/actions/run-example
+        with:
+          example: ${{ steps.examples.outputs.e0 }}
+          allow-failure-examples: ${{ inputs.allow_failure_examples }}
+          codecov-token: ${{ secrets.CODECOV_TOKEN }}
+      - name: "Example 2: ${{ steps.examples.outputs.e1 || '-' }}"
+        id: example1
+        if: steps.examples.outputs.e1 != ''
+        uses: ./.github/actions/run-example
+        with:
+          example: ${{ steps.examples.outputs.e1 }}
+          allow-failure-examples: ${{ inputs.allow_failure_examples }}
+          codecov-token: ${{ secrets.CODECOV_TOKEN }}
+      - name: "Example 3: ${{ steps.examples.outputs.e2 || '-' }}"
+        id: example2
+        if: steps.examples.outputs.e2 != ''
+        uses: ./.github/actions/run-example
+        with:
+          example: ${{ steps.examples.outputs.e2 }}
+          allow-failure-examples: ${{ inputs.allow_failure_examples }}
+          codecov-token: ${{ secrets.CODECOV_TOKEN }}
       - name: Save extensions
         # Not tied to the tests passing, see the action description
         if: always() && steps.ext-cache.outputs.cache-hit != 'true'
         uses: ./.github/actions/save-extensions
         with:
           cache-key: ${{ steps.ext-cache.outputs.cache-key }}
-      - name: Flag allowed failure
-        if: ${{ inputs.allow_failure && steps.run_tests.outcome == 'failure' }}
-        run: |
-          echo "::warning title=Allowed example failure::'${{ inputs.example }}' failed but is in the allow-failure list (vars.ALLOW_FAILURE_EXAMPLE_TESTS); not blocking. Remove it from the variable once fixed."
-      - name: Upload coverage to Codecov
-        uses: codecov/codecov-action@303a32d7a59b442fa8d48b6a1cc6825c09c847a5 # v7.1.1
-        with:
-          token: ${{ secrets.CODECOV_TOKEN }}
-          files: coverage.xml
-          # One flag per example, not a shared `examples`: carryforward only applies to a flag
-          # with no upload, so a shared flag would replace every lane's coverage with the subset
-          # that ran once lanes are gated independently.
-          flags: examples-${{ inputs.example }}
-          fail_ci_if_error: false # test may be skipped if relevant file changes are not detected
-          verbose: true
+      - name: Fail the job for an example that failed
+        if: steps.example0.outputs.failed == 'true' || steps.example1.outputs.failed == 'true' || steps.example2.outputs.failed == 'true'
+        run: exit 1

.github/workflows/example_tests.yml

@@ -14,8 +14,16 @@ concurrency:
   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:
+# Every job passes the repo variable ALLOW_FAILURE_EXAMPLE_TESTS to the runner as `allow_failure_examples`:
 # a comma-separated list of example names whose test failures should be non-blocking, e.g. "torch_trt,llm_qat"
+#
+# An `example` of several names separated by spaces runs them one after the other in one job, each in its own
+# step, which saves the container start and the base install (3-5 min each) of every job after the first. Only
+# small examples with the same image and extras are grouped; each example installs its own requirements just
+# before its tests, so a group lists first the examples that leave the environment alone and last the ones that
+# pin or upgrade packages. An example that changes torch (torch_trt downgrades it) must not be grouped: the
+# extension cache key is the group name plus the image, with no torch version, and a cache built against another
+# torch fails with undefined symbols.
 
 jobs:
   # One changed-files pass decides which lanes run. Lane granularity, not per-example: a job's
@@ -126,7 +134,9 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        example: [llm_distill, llm_qat, llm_sparsity, specdec_bench, speculative_decoding]
+        # llm_qat and speculative_decoding take 9-16 min and stay on their own jobs; the three small ones share one.
+        # llm_distill pins trl and torchao, so it goes last.
+        example: ["specdec_bench llm_sparsity llm_distill", llm_qat, speculative_decoding]
     uses: ./.github/workflows/_example_tests_runner.yml
     permissions:
       contents: read
@@ -137,7 +147,7 @@ jobs:
       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)) }}
+      allow_failure_examples: ${{ vars.ALLOW_FAILURE_EXAMPLE_TESTS }}
 
   ##### vLLM Example Tests #####
   vllm:
@@ -157,7 +167,7 @@ jobs:
       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)) }}
+      allow_failure_examples: ${{ vars.ALLOW_FAILURE_EXAMPLE_TESTS }}
 
   ##### TensorRT-LLM Example Tests #####
   trtllm:
@@ -166,7 +176,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        example: [gpt-oss, hf_ptq, llm_eval]
+        # hf_ptq is the longest and stays alone. gpt-oss upgrades datasets and trl, so it goes after llm_eval.
+        example: ["llm_eval gpt-oss", hf_ptq]
     uses: ./.github/workflows/_example_tests_runner.yml
     permissions:
       contents: read
@@ -178,7 +189,7 @@ jobs:
       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)) }}
+      allow_failure_examples: ${{ vars.ALLOW_FAILURE_EXAMPLE_TESTS }}
 
   ##### Megatron Example Tests #####
   megatron:
@@ -195,7 +206,7 @@ jobs:
       timeout_minutes: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 30 || 75 }}
       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,') }}
+      allow_failure_examples: ${{ vars.ALLOW_FAILURE_EXAMPLE_TESTS }}
 
   ##### Diffusers Example Tests #####
   diffusers:
@@ -212,7 +223,7 @@ jobs:
       timeout_minutes: 60
       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), ',diffusers,') }}
+      allow_failure_examples: ${{ vars.ALLOW_FAILURE_EXAMPLE_TESTS }}
 
   ##### ONNX/TensorRT Example Tests #####
   onnx:
@@ -235,7 +246,7 @@ jobs:
       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)) }}
+      allow_failure_examples: ${{ vars.ALLOW_FAILURE_EXAMPLE_TESTS }}
 
   ##### Required Check for PR #####
   example-pr-required-check: