feat: Adding Artifact reuse and stage skip functionality to multi-arch CI

.github/workflows/multi_arch_ci.yml

@@ -42,6 +42,22 @@ on:
         type: string
         default: ""
         description: "Workflow run ID to copy prebuilt stage artifacts from; required when prebuilt_stages is set"
+      stage_reuse_mode:
+        type: string
+        default: "dry-run"
+        # This drives the AUTOMATIC stage-reuse detection layer, which is
+        # separate from the manual `prebuilt_stages` input above. dry-run only
+        # reports which unaffected stages could be reused; reuse-stage actually
+        # reuses them (skips the build and copies artifacts from a baseline run).
+        description: "dry-run (report only) | reuse-stage (auto-reuse unaffected stages)"
+      stage_reuse_max_age_hours:
+        type: string
+        default: "72"
+        description: "Automatic stage reuse: reject baseline runs older than this many hours"
+      stage_reuse_commit_history:
+        type: string
+        default: "50"
+        description: "Automatic stage reuse: number of recent branch commits to fetch for ancestry"
       changed_projects:
         type: string
         default: ""
@@ -82,6 +98,9 @@ jobs:
       windows_test_labels: ${{ inputs.windows_test_labels || '' }}
       prebuilt_stages: ${{ inputs.prebuilt_stages || '' }}
       baseline_run_id: ${{ inputs.baseline_run_id || '' }}
+      stage_reuse_mode: ${{ inputs.stage_reuse_mode || 'dry-run' }}
+      stage_reuse_max_age_hours: ${{ inputs.stage_reuse_max_age_hours || '72' }}
+      stage_reuse_commit_history: ${{ inputs.stage_reuse_commit_history || '50' }}
       build_pytorch: ${{ github.event_name != 'workflow_dispatch' || inputs.build_pytorch }}
       build_jax: false
 

.github/workflows/setup_multi_arch.yml

@@ -43,6 +43,23 @@ on:
         description: "Workflow run ID to copy prebuilt stage artifacts from."
         type: string
         default: ""
+      stage_reuse_mode:
+        description: >-
+          Automatic stage-reuse mode (separate from the manual prebuilt_stages
+          input): "dry-run" (default) only reports which unaffected stages could
+          be reused; "reuse-stage" actually reuses them.
+        type: string
+        default: "dry-run"
+      stage_reuse_max_age_hours:
+        description: >-
+         Automatic stage reuse: reject baseline runs older than this many hours. default: "72"
+        type: string
+        default: "72"
+      stage_reuse_commit_history:
+        description: >-
+          Automatic stage reuse: number of recent branch commits to fetch for ancestry. default: "50"
+        type: string
+        default: "50"
       build_pytorch:
         description: "Build PyTorch wheels."
         type: boolean
@@ -159,6 +176,10 @@ jobs:
           WINDOWS_TEST_LABELS: ${{ inputs.windows_test_labels }}
           PREBUILT_STAGES: ${{ inputs.prebuilt_stages }}
           BASELINE_RUN_ID: ${{ inputs.baseline_run_id }}
+          STAGE_REUSE_MODE: ${{ inputs.stage_reuse_mode }}
+          STAGE_REUSE_CURRENT_SHA: ${{ steps.checkout.outputs.commit }}
+          STAGE_REUSE_MAX_AGE_HOURS: ${{ inputs.stage_reuse_max_age_hours }}
+          STAGE_REUSE_COMMIT_HISTORY: ${{ inputs.stage_reuse_commit_history }}
           # Skip path filtering for external repos (git sha won't exist in TheRock)
           SKIP_PATH_FILTERS: ${{ inputs.external_repo != '' && 'true' || '' }}
           CI_CONFIG_PATH: ci-config

build_tools/github_actions/configure_multi_arch_ci.py

@@ -681,18 +681,61 @@ def _determine_test_type(
 def decide_jobs(
     ci_inputs: CIInputs,
     git_context: GitContext,
+    targets: "TargetSelection | None" = None,
 ) -> JobDecisions:
-    """Determine which job groups to run, skip, or satisfy with prebuilt files."""
+    """Determine which job groups to run, skip, or satisfy with prebuilt files.
+    ``targets`` (the per-platform family selection) is used to scope automatic
+    stage reuse to the platforms actually being built, so a stage is only
+    reused when its artifacts exist for every one of those platforms.
+    """
 
     # Build ROCm.
-    # TODO(#3399): Use changed files and build_topology.py to:
-    #   1. set per-stage prebuilt decisions
-    #   2. skip job groups that aren't reachable from the changed files
-    # Parse explicit prebuilt stages from workflow_dispatch input.
+    # Parse explicit prebuilt stages from workflow_dispatch input. These are
+    # the MANUAL inputs and are always honored, unchanged, in every mode.
+
     stage_decisions: dict[str, JobAction] = {}
     if ci_inputs.prebuilt_stages:
         for stage in _parse_prebuilt_stages(ci_inputs.prebuilt_stages):
             stage_decisions[stage] = JobAction.PREBUILT
+
+    # Automatic stage reuse, behind STAGE_REUSE_MODE: in dry-run we only report
+    # which stages WOULD be reused and apply nothing; in "reuse-stage" the
+    # eligible stages are merged into stage_decisions so the orchestrator skips
+    # their builds and copies artifacts instead.
+    # These imports are intentionally local. stage_reuse_decision imports from
+    # sibling CI modules (baseline_runs, stage_impact, github_actions_api) that
+    # in turn pull in heavier deps; importing it lazily here keeps module import
+    # of configure_multi_arch_ci cheap and avoids an import cycle between the two
+    # modules (stage_reuse plumbing conceptually depends on this pipeline).
+
+    from stage_reuse_decision import (
+        StageReuseMode,
+        compute_auto_stage_reuse,
+        render_step_summary,
+        log_report,
+        stage_reuse_target_families,
+    )
+
+    build_platforms = _stage_reuse_build_platforms(ci_inputs, targets)
+    target_families = stage_reuse_target_families(
+        ci_inputs.linux_amdgpu_families,
+        ci_inputs.windows_amdgpu_families,
+    )
+    auto = compute_auto_stage_reuse(
+        changed_files=git_context.changed_files,
+        mode=StageReuseMode.from_environ(),
+        platforms=build_platforms,
+        target_families=target_families,
+    )
+    log_report(auto)
+
+    try:
+        gha_append_step_summary(render_step_summary(auto))
+    except Exception as exc:  # never fail CI config on a reporting error
+        print(f"[STAGE-REUSE] could not write step summary: {exc}")
+    # Only reuse-stage mode returns non-empty applied_reuse_stages.
+    for stage in auto.applied_reuse_stages:
+        stage_decisions.setdefault(stage, JobAction.PREBUILT)
     build_rocm = BuildRocmDecision(
         action=JobAction.RUN,
         stage_decisions=stage_decisions,
@@ -742,6 +785,33 @@ def decide_jobs(
 # ---------------------------------------------------------------------------
 
 
+def _stage_reuse_build_platforms(
+    ci_inputs: "CIInputs",
+    targets: "TargetSelection | None",
+) -> list[str]:
+    """Platforms whose baselines must contain a stage's artifacts to reuse it.
+    Automatic reuse merges into the prebuilt_stages that flow to BOTH the Linux
+    and Windows build configs, so a stage may only be reused when its artifacts
+    are available for every platform actually being built. When target selection
+    is known we use it directly; otherwise we fall back to the requested family
+    inputs, defaulting to Linux only.
+    """
+    if targets is not None:
+        platforms = []
+        if targets.linux_families:
+            platforms.append("linux")
+        if targets.windows_families:
+            platforms.append("windows")
+        return platforms or ["linux"]
+
+    platforms = []
+    if ci_inputs.linux_amdgpu_families:
+        platforms.append("linux")
+    if ci_inputs.windows_amdgpu_families:
+        platforms.append("windows")
+    return platforms or ["linux"]
+
+
 def _validate_family_names(
     names: list[str],
     known: dict[str, dict],