diff --git a/.github/workflows/_pr-test-check-changes.yml b/.github/workflows/_pr-test-check-changes.yml index a657dac69..3da5c824a 100644 --- a/.github/workflows/_pr-test-check-changes.yml +++ b/.github/workflows/_pr-test-check-changes.yml @@ -32,12 +32,8 @@ on: value: ${{ jobs.run.outputs.jit_kernel }} multimodal_gen: value: ${{ jobs.run.outputs.multimodal_gen }} - max_parallel: - value: ${{ jobs.run.outputs.max_parallel }} - max_parallel_small: - value: ${{ jobs.run.outputs.max_parallel_small }} - max_parallel_2gpu: - value: ${{ jobs.run.outputs.max_parallel_2gpu }} + partitions: + value: ${{ jobs.run.outputs.partitions }} b200_runner: value: ${{ jobs.run.outputs.b200_runner }} enable_retry: @@ -63,9 +59,7 @@ jobs: sgl_kernel_raw: ${{ steps.filter-api.outputs.sgl_kernel || steps.filter.outputs.sgl_kernel }} jit_kernel: ${{ steps.filter-api.outputs.jit_kernel || steps.filter.outputs.jit_kernel || steps.run-mode.outputs.run_all_tests }} multimodal_gen: ${{ steps.filter-api.outputs.multimodal_gen || steps.filter.outputs.multimodal_gen || steps.run-mode.outputs.run_all_tests }} - max_parallel: ${{ steps.set-parallel.outputs.max_parallel }} - max_parallel_small: ${{ steps.set-parallel.outputs.max_parallel_small }} - max_parallel_2gpu: ${{ steps.set-parallel.outputs.max_parallel_2gpu }} + partitions: ${{ steps.partitions.outputs.partitions }} b200_runner: ${{ steps.set-runner.outputs.b200_runner }} enable_retry: ${{ steps.set-retry.outputs.enable_retry }} continue_on_error: ${{ steps.set-continue-on-error.outputs.continue_on_error }} @@ -198,22 +192,29 @@ jobs: echo "multimodal_gen=false" >> $GITHUB_OUTPUT fi - - name: Set max-parallel based on run type - id: set-parallel + - name: Determine full-parallel mode + id: parallel-mode env: GH_TOKEN: ${{ github.token }} run: | - # Determine if this run gets full parallelism (scheduled / high priority) + # `full=true` lifts the matrix-fanout throttle so each suite's + # max_parallel = size. Conditions (matching the prior set-parallel + # step exactly): + # 1. Scheduled cron run. + # 2. pull_request event with the `high priority` label. + # 3. workflow_dispatch with target_stage set (i.e. /rerun-stage) + # whose underlying PR carries `high priority`. The labels + # aren't on the dispatch payload, so look them up via API: + # try SHA -> /pulls (works for fork PRs), fall back to + # branch name -> gh pr list (works for non-fork PRs). FULL=false if [[ "${{ github.event_name }}" == "schedule" ]]; then FULL=true - echo "Scheduled run detected, using full parallelism" + echo "Scheduled run -> full parallelism" elif [[ "${{ github.event_name }}" == "pull_request" && "${{ contains(github.event.pull_request.labels.*.name, 'high priority') }}" == "true" ]]; then FULL=true - echo "High priority PR detected, using full parallelism" + echo "high priority PR -> full parallelism" elif [[ -n "${{ inputs.target_stage }}" ]]; then - # /rerun-stage (workflow_dispatch): query PR labels via GitHub API - # Try SHA lookup first (fork PRs), fallback to branch name (non-fork PRs) LABELS="" PR_HEAD_SHA="${{ inputs.pr_head_sha }}" if [[ -n "$PR_HEAD_SHA" ]]; then @@ -224,28 +225,25 @@ jobs: LABELS=$(gh pr list --head "${{ github.ref_name }}" --repo "${{ github.repository }}" \ --json labels --jq '.[0].labels[].name' 2>/dev/null || true) fi - echo "PR labels: ${LABELS:-"(none)"}" + echo "PR labels (via API): ${LABELS:-"(none)"}" if echo "$LABELS" | grep -Fxq "high priority"; then FULL=true - echo "High priority PR detected via API (/rerun-stage), using full parallelism" + echo "high priority PR (via API) -> full parallelism" fi fi + echo "full=$FULL" >> "$GITHUB_OUTPUT" - # Set max-parallel for each runner type - # 1-gpu-h100: 14 partitions, 1-gpu-5090: 8 partitions, 2-gpu-h100: 4 partitions - if [[ "$FULL" == "true" ]]; then - LEVEL=full - echo "max_parallel=14" >> $GITHUB_OUTPUT - echo "max_parallel_small=8" >> $GITHUB_OUTPUT - echo "max_parallel_2gpu=4" >> $GITHUB_OUTPUT - else - LEVEL=low - echo "max_parallel=3" >> $GITHUB_OUTPUT - echo "max_parallel_small=3" >> $GITHUB_OUTPUT - echo "max_parallel_2gpu=2" >> $GITHUB_OUTPUT - fi - echo "parallel_level=$LEVEL" >> $GITHUB_OUTPUT - echo "Parallelism level: $LEVEL" + - name: Compute partitions + id: partitions + run: | + # Emit a single JSON output `partitions` keyed by suite name with + # {size, arr, max_parallel} fields per suite. Replaces the prior + # full/low max-parallel presets; `--full-parallel` keeps the + # `high priority` PR / scheduled cron escape hatch. + # See scripts/ci/utils/compute_partitions.py. + python3 scripts/ci/utils/compute_partitions.py \ + --full-parallel ${{ steps.parallel-mode.outputs.full }} \ + >> "$GITHUB_OUTPUT" - name: Set B200 runner tag id: set-runner @@ -318,7 +316,6 @@ jobs: echo "| multimodal_gen | ${{ steps.filter-api.outputs.multimodal_gen || steps.filter.outputs.multimodal_gen || steps.run-mode.outputs.run_all_tests }} |" echo "| target_stage | ${{ inputs.target_stage || '(none)' }} |" echo "| detection_method | ${{ inputs.target_stage && 'GitHub API' || 'dorny/paths-filter' }} |" - echo "| max_parallel | ${{ steps.set-parallel.outputs.parallel_level }} (h100=${{ steps.set-parallel.outputs.max_parallel }}, 5090=${{ steps.set-parallel.outputs.max_parallel_small }}, 2gpu=${{ steps.set-parallel.outputs.max_parallel_2gpu }}) |" echo "| b200_runner | ${{ steps.set-runner.outputs.b200_runner }} |" echo "| enable_retry | ${{ steps.set-retry.outputs.enable_retry }} |" echo "| continue_on_error | ${{ steps.set-continue-on-error.outputs.continue_on_error }} |" diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index e684c0a72..f0f8f45e3 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -122,7 +122,11 @@ jobs: id: wait with: stage-name: stage-a - jobs: '["stage-a-test-1-gpu-small", {"prefix": "stage-a-test-cpu", "expected_count": 4}]' + jobs: | + [ + {"prefix": "stage-a-test-1-gpu-small", "expected_count": ${{ fromJson(needs.check-changes.outputs.partitions)['stage-a-test-1-gpu-small'].size }}}, + {"prefix": "stage-a-test-cpu", "expected_count": ${{ fromJson(needs.check-changes.outputs.partitions)['stage-a-test-cpu'].size }}} + ] max-wait-minutes: '240' wait-for-stage-b: @@ -150,10 +154,10 @@ jobs: stage-name: stage-b jobs: | [ - {"prefix": "stage-b-test-1-gpu-small", "expected_count": 8}, - {"prefix": "stage-b-test-1-gpu-large", "expected_count": 14}, - {"prefix": "stage-b-test-2-gpu-large", "expected_count": 4}, - {"prefix": "stage-b-test-4-gpu-b200", "expected_count": 1} + {"prefix": "stage-b-test-1-gpu-small", "expected_count": ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-1-gpu-small'].size }}}, + {"prefix": "stage-b-test-1-gpu-large", "expected_count": ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-1-gpu-large'].size }}}, + {"prefix": "stage-b-test-2-gpu-large", "expected_count": ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-2-gpu-large'].size }}}, + {"prefix": "stage-b-test-4-gpu-b200", "expected_count": ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-4-gpu-b200'].size }}} ] max-wait-minutes: '480' @@ -285,6 +289,11 @@ jobs: ) runs-on: 1-gpu-5090 timeout-minutes: 240 + strategy: + fail-fast: false + max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-a-test-1-gpu-small'].max_parallel }} + matrix: + partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-a-test-1-gpu-small'].arr }} steps: - name: Checkout code uses: actions/checkout@v4 @@ -314,7 +323,7 @@ jobs: CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }} run: | cd test/ - python3 run_suite.py --hw cuda --suite stage-a-test-1-gpu-small $CONTINUE_ON_ERROR_FLAG + python3 run_suite.py --hw cuda --suite stage-a-test-1-gpu-small --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-a-test-1-gpu-small'].size }} $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: failure() @@ -339,8 +348,9 @@ jobs: timeout-minutes: 240 strategy: fail-fast: false + max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-a-test-cpu'].max_parallel }} matrix: - partition: [0, 1, 2, 3] + partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-a-test-cpu'].arr }} steps: - name: Free disk space run: | @@ -391,7 +401,7 @@ jobs: CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }} run: | cd test/ - python3 run_suite.py --hw cpu --suite stage-a-test-cpu --auto-partition-id ${{ matrix.partition }} --auto-partition-size 4 $CONTINUE_ON_ERROR_FLAG + python3 run_suite.py --hw cpu --suite stage-a-test-cpu --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-a-test-cpu'].size }} $CONTINUE_ON_ERROR_FLAG # Runs on 5090 (32GB, SM120) stage-b-test-1-gpu-small: @@ -410,9 +420,9 @@ jobs: timeout-minutes: 240 strategy: fail-fast: false - max-parallel: ${{ fromJson(needs.check-changes.outputs.max_parallel_small) }} + max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-1-gpu-small'].max_parallel }} matrix: - partition: [0, 1, 2, 3, 4, 5, 6, 7] + partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-1-gpu-small'].arr }} steps: - name: Checkout code uses: actions/checkout@v4 @@ -442,7 +452,7 @@ jobs: CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }} run: | cd test/ - python3 run_suite.py --hw cuda --suite stage-b-test-1-gpu-small --auto-partition-id ${{ matrix.partition }} --auto-partition-size 8 $CONTINUE_ON_ERROR_FLAG + python3 run_suite.py --hw cuda --suite stage-b-test-1-gpu-small --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-1-gpu-small'].size }} $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: failure() @@ -470,9 +480,9 @@ jobs: timeout-minutes: 240 strategy: fail-fast: false - max-parallel: ${{ fromJson(needs.check-changes.outputs.max_parallel) }} + max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-1-gpu-large'].max_parallel }} matrix: - partition: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13] + partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-1-gpu-large'].arr }} steps: - name: Checkout code uses: actions/checkout@v4 @@ -502,7 +512,7 @@ jobs: CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }} run: | cd test/ - python3 run_suite.py --hw cuda --suite stage-b-test-1-gpu-large --auto-partition-id ${{ matrix.partition }} --auto-partition-size 14 --timeout-per-file 1800 $CONTINUE_ON_ERROR_FLAG + python3 run_suite.py --hw cuda --suite stage-b-test-1-gpu-large --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-1-gpu-large'].size }} --timeout-per-file 1800 $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: failure() @@ -529,9 +539,9 @@ jobs: timeout-minutes: 240 strategy: fail-fast: false - max-parallel: ${{ fromJson(needs.check-changes.outputs.max_parallel_2gpu) }} + max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-2-gpu-large'].max_parallel }} matrix: - partition: [0, 1, 2, 3] + partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-2-gpu-large'].arr }} steps: - name: Checkout code uses: actions/checkout@v4 @@ -561,7 +571,7 @@ jobs: CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }} run: | cd test/ - python3 run_suite.py --hw cuda --suite stage-b-test-2-gpu-large --auto-partition-id ${{ matrix.partition }} --auto-partition-size 4 $CONTINUE_ON_ERROR_FLAG + python3 run_suite.py --hw cuda --suite stage-b-test-2-gpu-large --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-2-gpu-large'].size }} $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: failure() @@ -588,6 +598,9 @@ jobs: timeout-minutes: 240 strategy: fail-fast: false + max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-4-gpu-b200'].max_parallel }} + matrix: + partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-4-gpu-b200'].arr }} steps: - name: Checkout code @@ -618,7 +631,7 @@ jobs: CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }} run: | cd test - python3 run_suite.py --hw cuda --suite stage-b-test-4-gpu-b200 $CONTINUE_ON_ERROR_FLAG + python3 run_suite.py --hw cuda --suite stage-b-test-4-gpu-b200 --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-4-gpu-b200'].size }} $CONTINUE_ON_ERROR_FLAG - name: Run FA4 jit_kernel tests (SM100+) timeout-minutes: 10 @@ -681,8 +694,9 @@ jobs: timeout-minutes: 240 strategy: fail-fast: false + max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-4-gpu-h100'].max_parallel }} matrix: - part: [0, 1, 2] + partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-4-gpu-h100'].arr }} steps: - name: Checkout code uses: actions/checkout@v4 @@ -712,12 +726,12 @@ jobs: CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }} run: | cd test - python3 run_suite.py --hw cuda --suite stage-c-test-4-gpu-h100 --auto-partition-id ${{ matrix.part }} --auto-partition-size 3 $CONTINUE_ON_ERROR_FLAG + python3 run_suite.py --hw cuda --suite stage-c-test-4-gpu-h100 --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-4-gpu-h100'].size }} $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: failure() with: - artifact-suffix: ${{ matrix.part }} + artifact-suffix: ${{ matrix.partition }} - name: Cleanup venv if: always() @@ -739,8 +753,9 @@ jobs: timeout-minutes: 240 strategy: fail-fast: false + max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-8-gpu-h200'].max_parallel }} matrix: - part: [0, 1, 2, 3] + partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-8-gpu-h200'].arr }} steps: - name: Checkout code uses: actions/checkout@v4 @@ -789,12 +804,12 @@ jobs: CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }} run: | cd test - python3 run_suite.py --hw cuda --suite stage-c-test-8-gpu-h200 --auto-partition-id ${{ matrix.part }} --auto-partition-size 4 $CONTINUE_ON_ERROR_FLAG + python3 run_suite.py --hw cuda --suite stage-c-test-8-gpu-h200 --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-8-gpu-h200'].size }} $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: failure() with: - artifact-suffix: ${{ matrix.part }} + artifact-suffix: ${{ matrix.partition }} - name: Cleanup venv if: always() @@ -816,6 +831,11 @@ jobs: timeout-minutes: 240 env: SGLANG_CI_RDMA_ALL_DEVICES: "mlx5_1,mlx5_2,mlx5_3,mlx5_4" + strategy: + fail-fast: false + max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-8-gpu-h20'].max_parallel }} + matrix: + partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-8-gpu-h20'].arr }} steps: - name: Checkout code uses: actions/checkout@v4 @@ -845,7 +865,7 @@ jobs: CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }} run: | cd test - python3 run_suite.py --hw cuda --suite stage-c-test-8-gpu-h20 $CONTINUE_ON_ERROR_FLAG + python3 run_suite.py --hw cuda --suite stage-c-test-8-gpu-h20 --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-8-gpu-h20'].size }} $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: failure() @@ -868,6 +888,11 @@ jobs: ) runs-on: 4-gpu-h100 timeout-minutes: 240 + strategy: + fail-fast: false + max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-deepep-4-gpu-h100'].max_parallel }} + matrix: + partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-deepep-4-gpu-h100'].arr }} steps: - name: Checkout code uses: actions/checkout@v4 @@ -914,7 +939,7 @@ jobs: CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }} run: | cd test - python3 run_suite.py --hw cuda --suite stage-c-test-deepep-4-gpu-h100 $CONTINUE_ON_ERROR_FLAG + python3 run_suite.py --hw cuda --suite stage-c-test-deepep-4-gpu-h100 --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-deepep-4-gpu-h100'].size }} $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: failure() @@ -937,6 +962,11 @@ jobs: ) runs-on: 8-gpu-h200-deepep timeout-minutes: 240 + strategy: + fail-fast: false + max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-deepep-8-gpu-h200'].max_parallel }} + matrix: + partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-deepep-8-gpu-h200'].arr }} steps: - name: Checkout code uses: actions/checkout@v4 @@ -984,7 +1014,7 @@ jobs: CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }} run: | cd test - python3 run_suite.py --hw cuda --suite stage-c-test-deepep-8-gpu-h200 $CONTINUE_ON_ERROR_FLAG + python3 run_suite.py --hw cuda --suite stage-c-test-deepep-8-gpu-h200 --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-deepep-8-gpu-h200'].size }} $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: failure() @@ -1009,8 +1039,9 @@ jobs: timeout-minutes: 240 strategy: fail-fast: false + max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-4-gpu-b200'].max_parallel }} matrix: - part: [0, 1, 2, 3, 4, 5] + partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-4-gpu-b200'].arr }} steps: - name: Checkout code @@ -1041,12 +1072,12 @@ jobs: CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }} run: | cd test - python3 run_suite.py --hw cuda --suite stage-c-test-4-gpu-b200 --auto-partition-id ${{ matrix.part }} --auto-partition-size 6 --timeout-per-file 1800 $CONTINUE_ON_ERROR_FLAG + python3 run_suite.py --hw cuda --suite stage-c-test-4-gpu-b200 --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-4-gpu-b200'].size }} --timeout-per-file 1800 $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: failure() with: - artifact-suffix: ${{ matrix.part }} + artifact-suffix: ${{ matrix.partition }} - name: Cleanup venv if: always() @@ -1066,6 +1097,11 @@ jobs: ) runs-on: ${{ needs.check-changes.outputs.b200_runner }} timeout-minutes: 240 + strategy: + fail-fast: false + max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-dsv4-4-gpu-b200'].max_parallel }} + matrix: + partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-dsv4-4-gpu-b200'].arr }} steps: - name: Checkout code uses: actions/checkout@v4 @@ -1095,7 +1131,7 @@ jobs: CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }} run: | cd test - python3 run_suite.py --hw cuda --suite stage-c-test-dsv4-4-gpu-b200 --timeout-per-file 1800 $CONTINUE_ON_ERROR_FLAG + python3 run_suite.py --hw cuda --suite stage-c-test-dsv4-4-gpu-b200 --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-dsv4-4-gpu-b200'].size }} --timeout-per-file 1800 $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: failure() @@ -1118,6 +1154,11 @@ jobs: ) runs-on: 8-gpu-h200 timeout-minutes: 240 + strategy: + fail-fast: false + max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-dsv4-8-gpu-h200'].max_parallel }} + matrix: + partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-dsv4-8-gpu-h200'].arr }} steps: - name: Checkout code uses: actions/checkout@v4 @@ -1147,7 +1188,7 @@ jobs: CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }} run: | cd test - python3 run_suite.py --hw cuda --suite stage-c-test-dsv4-8-gpu-h200 --timeout-per-file 1800 $CONTINUE_ON_ERROR_FLAG + python3 run_suite.py --hw cuda --suite stage-c-test-dsv4-8-gpu-h200 --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-dsv4-8-gpu-h200'].size }} --timeout-per-file 1800 $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: failure() diff --git a/scripts/ci/utils/compute_partitions.py b/scripts/ci/utils/compute_partitions.py new file mode 100644 index 000000000..c86a6d9c6 --- /dev/null +++ b/scripts/ci/utils/compute_partitions.py @@ -0,0 +1,163 @@ +"""Sum CIRegistry est_time per per-commit suite and emit one $GITHUB_OUTPUT line +keyed by suite name. Consumed by pr-test.yml stage jobs as +`fromJson(needs.check-changes.outputs.partitions)['']`. + + partitions={"stage-b-test-1-gpu-small": {"size": 8, "arr": [0,...,7], "max_parallel": 2}, ...} +""" + +import argparse +import glob +import importlib.util +import json +import math +import os +from collections import defaultdict + +REPO_ROOT = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +) + +# Load ci_register.py directly: `import sglang.test...` pulls torch/orjson via +# sglang.__init__ but check-changes runs on bare ubuntu-latest. ci_register +# itself is stdlib-only (AST). +_CI_REGISTER_PATH = os.path.join( + REPO_ROOT, "python", "sglang", "test", "ci", "ci_register.py" +) +_spec = importlib.util.spec_from_file_location("ci_register", _CI_REGISTER_PATH) +_ci_register = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_ci_register) +collect_tests = _ci_register.collect_tests +HWBackend = _ci_register.HWBackend + +# pr-test-amd.yml / pr-test-npu.yml have their own dispatch. +_TARGET_BACKENDS = {HWBackend.CUDA, HWBackend.CPU} + +# stage-a is the critical-path entry gate; pin its fanout to smoke-coverage +# defaults instead of est_time. max_parallel = size (no throttle). +_STAGE_A_OVERRIDES = { + "stage-a-test-cpu": 4, + "stage-a-test-1-gpu-small": 1, +} + +# Per-partition wall-clock target + ceiling. Single knob for the whole +# pipeline. ~17 min avg under perfect LPT (TARGET / LPT_SLOP), ~22 min under +# worst-case LPT 4/3 imbalance, fail-fast above 30 min. +TARGET_SECONDS = 20 * 60 + +# LPT (Longest Processing Time first) worst case is 4/3 * OPT; pad ~15% so a +# slightly-unlucky LPT result still fits inside MAX_PARTITION_SECONDS. +LPT_SLOP = 1.15 + +# Hard ceiling. Exceeded → raise, forcing the maintainer to split a slow file +# or bump TARGET_SECONDS deliberately. +MAX_PARTITION_SECONDS = 30 * 60 + + +def discover_files(repo_root: str) -> list[str]: + test_dir = os.path.join(repo_root, "test") + files = [ + f + for f in glob.glob( + os.path.join(test_dir, "registered", "**", "*.py"), recursive=True + ) + if not f.endswith("/conftest.py") and not f.endswith("/__init__.py") + ] + jit_kernel_dir = os.path.join(repo_root, "python", "sglang", "jit_kernel") + files += glob.glob( + os.path.join(jit_kernel_dir, "tests", "**", "test_*.py"), recursive=True + ) + files += glob.glob( + os.path.join(jit_kernel_dir, "benchmark", "**", "bench_*.py"), recursive=True + ) + return files + + +def compute_max_parallel(size: int) -> int: + return max(size // 4, 1) + + +def compute_partitions(tests, full_parallel=False): + """Group per-commit tests by suite and emit partition metadata. + + `full_parallel=True` (scheduled cron or `high priority` PR) sets + max_parallel = size, lifting the matrix-fanout throttle. + """ + suite_tests = defaultdict(list) + for t in tests: + if t.backend not in _TARGET_BACKENDS: + continue + if t.nightly or t.disabled is not None: + continue + suite_tests[t.suite].append(t) + + result = {} + for suite, group in suite_tests.items(): + total = sum(t.est_time for t in group) + if suite in _STAGE_A_OVERRIDES: + size = _STAGE_A_OVERRIDES[suite] + max_parallel = size + else: + size = max(1, math.ceil(total * LPT_SLOP / TARGET_SECONDS)) + max_parallel = size if full_parallel else compute_max_parallel(size) + # Check naive average (total/size). LPT can be ~4/3 of that, but the + # ceil + LPT_SLOP padding above absorbs that slack. + if total / size > MAX_PARTITION_SECONDS: + raise RuntimeError( + f"Suite {suite!r}: total est_time {total:.0f}s / size {size} " + f"= {total / size:.0f}s exceeds MAX_PARTITION_SECONDS " + f"({MAX_PARTITION_SECONDS}s). Split a slow file or raise " + f"TARGET_SECONDS deliberately." + ) + result[suite] = { + "size": size, + "arr": list(range(size)), + "max_parallel": max_parallel, + } + return result + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--repo-root", default=REPO_ROOT) + parser.add_argument( + "--output-format", + choices=("gha", "json"), + default="gha", + help="`gha` emits `partitions=` for $GITHUB_OUTPUT; `json` is raw", + ) + parser.add_argument( + "--full-parallel", + choices=("true", "false"), + default="false", + help="Lift the max_parallel throttle (set by schedule / `high priority`)", + ) + args = parser.parse_args() + + files = discover_files(args.repo_root) + # Warn-not-fail on unregistered files: run_suite.py catches this at + # test-execution time with sanity_check=True; dispatch should keep going. + all_tests = collect_tests(files, sanity_check=False) + + result = compute_partitions(all_tests, full_parallel=(args.full_parallel == "true")) + payload = json.dumps(result, separators=(",", ":"), sort_keys=True) + if args.output_format == "gha": + print(f"partitions={payload}") + else: + print(payload) + + summary_path = os.environ.get("GITHUB_STEP_SUMMARY") + if summary_path: + with open(summary_path, "a") as f: + f.write("## Partitions\n\n") + f.write( + f"`full_parallel={args.full_parallel}` " + f"(`size//4` throttle is lifted when true)\n\n" + ) + f.write("| Suite | size | max_parallel |\n") + f.write("|---|---:|---:|\n") + for suite, info in sorted(result.items()): + f.write(f"| `{suite}` | {info['size']} | {info['max_parallel']} |\n") + + +if __name__ == "__main__": + main()