From 14ac14287cf3d38ec3b7a7e6e0345fdfa5e1e0f4 Mon Sep 17 00:00:00 2001 From: Kangyan-Zhou Date: Wed, 22 Apr 2026 11:28:06 -0700 Subject: [PATCH] [CI] /rerun-stage: auto-include wheel build when PR modifies sgl-kernel/ (#23492) Co-authored-by: Claude Opus 4.7 (1M context) --- .github/workflows/pr-test.yml | 91 ++++++++++++++++------- scripts/ci/utils/slash_command_handler.py | 40 ++++++---- 2 files changed, 92 insertions(+), 39 deletions(-) diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index bf0b0fb6d..1342dbc42 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -25,6 +25,11 @@ on: required: false type: string default: "" + include_wheel_build: + description: "When set with target_stage, also run sgl-kernel-build-wheels so the target stage uses the freshly-built kernel (for /rerun-stage on PRs that modify sgl-kernel/)" + required: false + type: boolean + default: false test_parallel_dispatch: description: "Test parallel dispatch behavior (simulates scheduled run)" required: false @@ -80,10 +85,14 @@ jobs: outputs: # Use API-based detection for target_stage mode (filter-api), otherwise use dorny/paths-filter (filter) main_package: ${{ steps.filter-api.outputs.main_package || steps.filter.outputs.main_package || steps.run-mode.outputs.run_all_tests }} - # sgl_kernel is forced to false when target_stage is set, since sgl-kernel-build-wheels won't run - # This prevents CUSTOM_BUILD_SGL_KERNEL=true when the wheel artifacts aren't available - # Note: If PR has kernel changes AND target_stage is set, the validate-target-stage step will fail - sgl_kernel: ${{ !inputs.target_stage && (steps.filter-api.outputs.sgl_kernel || steps.filter.outputs.sgl_kernel) }} + # sgl_kernel is forced to false when target_stage is set AND include_wheel_build is NOT set, + # since sgl-kernel-build-wheels normally skips in target_stage mode. When include_wheel_build + # is true, keep the real value so the wheel build runs and the target stage downloads its + # artifact (used by /rerun-stage on PRs that modify sgl-kernel/). + # This prevents CUSTOM_BUILD_SGL_KERNEL=true when the wheel artifacts aren't available. + # Note: If PR has kernel changes AND target_stage is set AND include_wheel_build is NOT set, + # the validate-target-stage step will fail. + sgl_kernel: ${{ (!inputs.target_stage || inputs.include_wheel_build) && (steps.filter-api.outputs.sgl_kernel || steps.filter.outputs.sgl_kernel) }} # Raw sgl_kernel value before target_stage override (used for validation) 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 }} @@ -298,20 +307,26 @@ jobs: fi - name: Validate target_stage with kernel changes - # Use API-based detection (filter-api) for target_stage mode, otherwise use dorny/paths-filter (filter) - if: inputs.target_stage && (steps.filter-api.outputs.sgl_kernel == 'true' || steps.filter.outputs.sgl_kernel == 'true') + # Fail only when PR has sgl-kernel changes AND the caller didn't opt into include_wheel_build. + # include_wheel_build=true means sgl-kernel-build-wheels will run alongside the target stage + # (see the sgl_kernel output and sgl-kernel-build-wheels if-conditions above/below), so it's + # safe to proceed. + if: inputs.target_stage && !inputs.include_wheel_build && (steps.filter-api.outputs.sgl_kernel == 'true' || steps.filter.outputs.sgl_kernel == 'true') run: | - echo "::error::Cannot use /rerun-stage when PR has sgl-kernel changes." - echo "::error::The sgl-kernel-build-wheels job is skipped in target_stage mode, but this PR modifies sgl-kernel/ files." - echo "::error::Please use /tag-and-rerun-ci to run the full workflow including kernel builds." + echo "::error::Cannot use /rerun-stage when PR has sgl-kernel changes without include_wheel_build." + echo "::error::The sgl-kernel-build-wheels job is skipped in target_stage mode by default, but this PR modifies sgl-kernel/ files." + echo "::error::The slash-command handler should have set include_wheel_build=true automatically; falling back to /tag-and-rerun-ci." echo "" - echo "ERROR: Cannot use /rerun-stage when PR has sgl-kernel changes." + echo "ERROR: Cannot use /rerun-stage when PR has sgl-kernel changes without include_wheel_build." echo "" echo "This PR modifies files in sgl-kernel/, which requires building custom kernel wheels." - echo "The /rerun-stage command skips the wheel build job, so the test would run against" - echo "the wrong (PyPI) version of sgl-kernel instead of your changes." + echo "Running the target stage without rebuilding the kernel would use the wrong (PyPI)" + echo "version of sgl-kernel instead of your changes." echo "" - echo "To properly test your kernel changes, use one of these commands instead:" + echo "The /rerun-stage handler sets include_wheel_build=true automatically when it detects" + echo "sgl-kernel/ changes on the PR. If you see this error, the handler may be outdated." + echo "" + echo "Alternatives:" echo " /tag-and-rerun-ci - Re-run the full workflow including kernel builds" echo " /rerun-ci - Re-run the full workflow" echo "" @@ -326,7 +341,7 @@ jobs: echo "|-------------------|---------|" echo "| main_package | ${{ steps.filter-api.outputs.main_package || steps.filter.outputs.main_package || steps.run-mode.outputs.run_all_tests }} |" echo "| sgl_kernel (raw) | ${{ steps.filter-api.outputs.sgl_kernel || steps.filter.outputs.sgl_kernel }} |" - echo "| sgl_kernel (used) | ${{ !inputs.target_stage && (steps.filter-api.outputs.sgl_kernel || steps.filter.outputs.sgl_kernel) }} |" + echo "| sgl_kernel (used) | ${{ (!inputs.target_stage || inputs.include_wheel_build) && (steps.filter-api.outputs.sgl_kernel || steps.filter.outputs.sgl_kernel) }} |" echo "| jit_kernel | ${{ steps.filter-api.outputs.jit_kernel || steps.filter.outputs.jit_kernel || steps.run-mode.outputs.run_all_tests }} |" 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)' }} |" @@ -421,8 +436,22 @@ jobs: sgl-kernel-build-wheels: needs: [check-changes, call-gate] - # Skip for scheduled runs (they run stages independently) and when target_stage is set - if: github.event_name != 'schedule' && inputs.test_parallel_dispatch != true && !inputs.target_stage && needs.check-changes.outputs.sgl_kernel == 'true' + # Skip for scheduled runs (they run stages independently). Runs in target_stage mode only when + # include_wheel_build is true (i.e. /rerun-stage on a PR with sgl-kernel changes), so the + # target stage can download the freshly-built wheel. + # + # `always()` lets us run when call-gate is skipped (which it always is in target_stage mode by + # design). The explicit needs..result checks preserve old gating for the normal PR path. + if: | + always() && + github.event_name != 'schedule' && + inputs.test_parallel_dispatch != true && + needs.check-changes.result == 'success' && + needs.check-changes.outputs.sgl_kernel == 'true' && + ( + (!inputs.target_stage && needs.call-gate.result == 'success') || + (inputs.target_stage && inputs.include_wheel_build) + ) runs-on: x64-kernel-build-node timeout-minutes: 240 strategy: @@ -469,8 +498,20 @@ jobs: sgl-kernel-build-wheels-arm: needs: [check-changes, call-gate] - # Skip for scheduled runs (they run stages independently) and when target_stage is set - if: github.event_name != 'schedule' && inputs.test_parallel_dispatch != true && !inputs.target_stage && needs.check-changes.outputs.sgl_kernel == 'true' + # Skip for scheduled runs (they run stages independently). Runs in target_stage mode only when + # include_wheel_build is true (i.e. /rerun-stage on a PR with sgl-kernel changes). + # + # See sgl-kernel-build-wheels above for the always() + result-check rationale. + if: | + always() && + github.event_name != 'schedule' && + inputs.test_parallel_dispatch != true && + needs.check-changes.result == 'success' && + needs.check-changes.outputs.sgl_kernel == 'true' && + ( + (!inputs.target_stage && needs.call-gate.result == 'success') || + (inputs.target_stage && inputs.include_wheel_build) + ) runs-on: arm-kernel-build-node timeout-minutes: 240 strategy: @@ -946,7 +987,7 @@ jobs: secrets: inherit stage-c-test-4-gpu-h100: - needs: [check-changes, call-gate, wait-for-stage-b] + needs: [check-changes, call-gate, wait-for-stage-b, sgl-kernel-build-wheels] if: | always() && ( @@ -1004,7 +1045,7 @@ jobs: run: bash scripts/ci/cuda/ci_cleanup_venv.sh stage-c-test-8-gpu-h200: - needs: [check-changes, call-gate, wait-for-stage-b] + needs: [check-changes, call-gate, wait-for-stage-b, sgl-kernel-build-wheels] if: | always() && ( @@ -1081,7 +1122,7 @@ jobs: run: bash scripts/ci/cuda/ci_cleanup_venv.sh stage-c-test-8-gpu-h20: - needs: [check-changes, call-gate, wait-for-stage-b] + needs: [check-changes, call-gate, wait-for-stage-b, sgl-kernel-build-wheels] if: | always() && ( @@ -1142,7 +1183,7 @@ jobs: run: bash scripts/ci/cuda/ci_cleanup_venv.sh stage-c-test-deepep-4-gpu-h100: - needs: [check-changes, call-gate, wait-for-stage-b] + needs: [check-changes, call-gate, wait-for-stage-b, sgl-kernel-build-wheels] if: | always() && ( @@ -1211,7 +1252,7 @@ jobs: run: bash scripts/ci/cuda/ci_cleanup_venv.sh stage-c-test-deepep-8-gpu-h200: - needs: [check-changes, call-gate, wait-for-stage-b] + needs: [check-changes, call-gate, wait-for-stage-b, sgl-kernel-build-wheels] if: | always() && ( @@ -1281,7 +1322,7 @@ jobs: run: bash scripts/ci/cuda/ci_cleanup_venv.sh stage-c-test-4-gpu-b200: - needs: [check-changes, call-gate, wait-for-stage-b] + needs: [check-changes, call-gate, wait-for-stage-b, sgl-kernel-build-wheels] if: | always() && ( @@ -1340,7 +1381,7 @@ jobs: run: bash scripts/ci/cuda/ci_cleanup_venv.sh stage-c-test-4-gpu-b200-small: - needs: [check-changes, call-gate, wait-for-stage-b] + needs: [check-changes, call-gate, wait-for-stage-b, sgl-kernel-build-wheels] if: | always() && ( diff --git a/scripts/ci/utils/slash_command_handler.py b/scripts/ci/utils/slash_command_handler.py index f779aa279..59d597551 100644 --- a/scripts/ci/utils/slash_command_handler.py +++ b/scripts/ci/utils/slash_command_handler.py @@ -367,6 +367,18 @@ def handle_rerun_stage( ) print(f"PR is from fork: {is_fork}") + # If the PR modifies sgl-kernel/, the target stage would otherwise use the + # PyPI sgl-kernel wheel instead of the PR's changes (sgl-kernel-build-wheels + # skips in target_stage mode by default). Set include_wheel_build=true so the + # workflow runs sgl-kernel-build-wheels alongside the target stage; the target + # stage waits for the build via its needs list. + kernel_changes = has_sgl_kernel_changes(pr) + if kernel_changes: + print( + "PR modifies sgl-kernel/ - setting include_wheel_build=true so the " + "target stage gets the freshly-built wheel instead of the PyPI one." + ) + # pr_head_sha is used for fork PRs (passed to workflow and used for URL lookup) pr_head_sha = None @@ -378,25 +390,25 @@ def handle_rerun_stage( print( f"Triggering {workflow_name} workflow on ref: {ref}, PR head SHA: {pr_head_sha}" ) - if is_amd_stage: - inputs = { - "target_stage": stage_name, - "pr_head_sha": pr_head_sha, - } - else: - inputs = { - "target_stage": stage_name, - "pr_head_sha": pr_head_sha, - } + inputs = { + "target_stage": stage_name, + "pr_head_sha": pr_head_sha, + } else: # For non-fork PRs: dispatch on the PR branch directly # This allows testing workflow changes before merge ref = pr.head.ref print(f"Triggering {workflow_name} workflow on branch: {ref}") - if is_amd_stage: - inputs = {"target_stage": stage_name} - else: - inputs = {"target_stage": stage_name} + inputs = {"target_stage": stage_name} + + # For NVIDIA stages, honor the sgl-kernel / include_wheel_build flow. AMD is + # a separate workflow that doesn't share the same wheel-build pipeline. + if kernel_changes and not is_amd_stage: + inputs["include_wheel_build"] = "true" + # include_wheel_build relies on filter-api detecting kernel changes, which + # requires pr_head_sha. Ensure it's set even for non-fork PRs. + if not is_fork: + inputs["pr_head_sha"] = pr.head.sha # Record dispatch time before triggering dispatch_time = time.time()