Build only required prek hook environments in CI jobs on a cache miss (#76111)

Signed-off-by: PoAn Yang <payang@apache.org>

.github/actions/install-prek/action.yml

@@ -25,6 +25,14 @@ inputs:
   platform:
     description: 'Platform for the build - linux/amd64 or linux/arm64'
     required: true
+  hooks:
+    description: >-
+      Set this input to the space-separated ids of the hooks that the job runs. When this input is set,
+      the action still restores the shared cache of hook environments, but it prepares only these hooks and
+      never saves the cache. A hook left out of the list still runs, but prek creates its environment when
+      the hook first runs.
+    required: false
+    default: ''
 runs:
   using: "composite"
   steps:
@@ -116,6 +124,8 @@ runs:
     - name: Validate and install prek hooks
       id: install-hooks
       shell: bash
+      env:
+        HOOKS: ${{ inputs.hooks }}
       # A successfully extracted archive can still contain a partial or unusable environment.
       # This is idempotent on a healthy hit; prek validates/reuses it and repairs misses normally.
       # Keep the existing bounded retries for transient toolchain download failures.
@@ -131,10 +141,16 @@ runs:
           change_detection_uncertain=true
         fi
 
+        # A YAML block scalar puts each hook id on its own line, and `read` stops at the first newline.
+        read -ra hook_selection <<< "${HOOKS//$'\n'/ }"
+        if [[ "${#hook_selection[@]}" -eq 0 ]]; then
+          hook_selection=(--skip run-skill-eval --skip run-skill-eval-codex --skip view-skill-eval)
+        fi
+
         started_at=${SECONDS}
         max_attempts=4
         for attempt in $(seq 1 "${max_attempts}"); do
-          if prek install-hooks --skip run-skill-eval --skip run-skill-eval-codex --skip view-skill-eval; then
+          if prek install-hooks "${hook_selection[@]}"; then
             installation_succeeded=true
             break
           fi
@@ -175,6 +191,7 @@ runs:
       env:
         EVENT_NAME: ${{ github.event_name }}
         SAVE_CACHE: ${{ inputs.save-cache }}
+        HOOKS: ${{ inputs.hooks }}
         STASH_HIT: ${{ steps.restore-prek-cache.outputs.stash-hit }}
         TAR_RESTORED: ${{ steps.restore-prek-tar.outputs.tar-restored }}
         CACHE_CHANGED: ${{ steps.install-hooks.outputs.cache-changed }}
@@ -182,7 +199,12 @@ runs:
       run: |
         SAVE=false
         REASON=unchanged
-        if [[ "${SAVE_CACHE}" != "true" ]]; then
+        # If this job saved its archive after a cache miss, the archive would hold only the selected
+        # environments under the cache key of the full archive. A later job that looks up that key could
+        # restore an archive that lacks most environments.
+        if [[ -n "${HOOKS//[[:space:]]/}" ]]; then
+          REASON=selected-hooks
+        elif [[ "${SAVE_CACHE}" != "true" ]]; then
           REASON=reader
         elif [[ "${STASH_HIT}" != "true" ]]; then
           SAVE=true

.github/workflows/ci-amd.yml

@@ -643,6 +643,7 @@ jobs:
         with:
           platform: ${{ needs.build-info.outputs.platform }}
           save-cache: false
+          hooks: mypy-providers
       - name: "MyPy checks for providers"
         run: prek --color always --verbose --stage manual mypy-providers --all-files
         env:
@@ -684,6 +685,7 @@ jobs:
         with:
           platform: ${{ needs.build-info.outputs.platform }}
           save-cache: false
+          hooks: migration-round-trip
       - name: "Migration round-trip check"
         run: prek --color always --verbose --stage manual migration-round-trip --all-files
         env:

.github/workflows/ci-arm.yml

@@ -519,6 +519,7 @@ jobs:
         with:
           platform: ${{ needs.build-info.outputs.platform }}
           save-cache: false
+          hooks: mypy-providers
       - name: "MyPy checks for providers"
         run: prek --color always --verbose --stage manual mypy-providers --all-files
         env:
@@ -560,6 +561,7 @@ jobs:
         with:
           platform: ${{ needs.build-info.outputs.platform }}
           save-cache: false
+          hooks: migration-round-trip
       - name: "Migration round-trip check"
         run: prek --color always --verbose --stage manual migration-round-trip --all-files
         env:

.github/workflows/generate-constraints.yml

@@ -131,6 +131,8 @@ jobs:
         with:
           platform: ${{ inputs.platform }}
           save-cache: false
+          # The hatch_build.py files of the fab, edge3 and common.ai providers run these hooks.
+          hooks: compile-fab-assets compile-edge-assets compile-common-ai-provider-assets
       # Resolves the providers from the sources in this checkout. Skipped for a release, which
       # ships only the PyPI constraints and must not have a source tree answer for a provider.
       - name: "Source constraints"

.github/workflows/prod-image-build.yml

@@ -151,6 +151,8 @@ jobs:
         with:
           platform: ${{ inputs.platform }}
           save-cache: false
+          # The hatch_build.py files of the fab, edge3 and common.ai providers run these hooks.
+          hooks: compile-fab-assets compile-edge-assets compile-common-ai-provider-assets
         if: inputs.upload-package-artifact == 'true'
       - name: "Prepare providers packages - all providers built from sources"
         shell: bash

.github/workflows/test-providers.yml

@@ -105,6 +105,8 @@ jobs:
         with:
           platform: ${{ inputs.platform }}
           save-cache: false
+          # The hatch_build.py files of the fab, edge3 and common.ai providers run these hooks.
+          hooks: compile-fab-assets compile-edge-assets compile-common-ai-provider-assets
       - name: "Cleanup dist files"
         run: rm -fv ./dist/*
       - name: "Set current date as RELEASE_DATE variable"
@@ -214,6 +216,8 @@ jobs:
         with:
           platform: ${{ inputs.platform }}
           save-cache: false
+          # The hatch_build.py files of the fab, edge3 and common.ai providers run these hooks.
+          hooks: compile-fab-assets compile-edge-assets compile-common-ai-provider-assets
       - name: "Cleanup dist files"
         run: rm -fv ./dist/*
       - name: "Prepare provider distributions: wheel"

scripts/tests/ci/test_prek_cache_optimizations.py

@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 import os
+import re
 import shlex
 import shutil
 import subprocess
@@ -65,7 +66,7 @@ def sandbox(tmp_path):
     home.mkdir()
     output = tmp_path / "output"
     output.touch()
-    return {"HOME": str(home), "GITHUB_OUTPUT": str(output)}
+    return {"HOME": str(home), "GITHUB_OUTPUT": str(output), "HOOKS": ""}
 
 
 def read_outputs(env):
@@ -157,6 +158,8 @@ def test_cache_key_keeps_only_environment_inputs():
             "change-detection-uncertain",
         ),
         ({"SAVE_CACHE": "true"}, False, "unchanged"),
+        ({"HOOKS": "mypy-providers"}, False, "selected-hooks"),
+        ({"HOOKS": "mypy-providers", "SAVE_CACHE": "true", "STASH_HIT": "false"}, False, "selected-hooks"),
     ],
 )
 def test_cache_refresh_policy(sandbox, inputs, expected_save, expected_reason):
@@ -198,6 +201,85 @@ def test_restored_hooks_are_always_validated(sandbox, fake_tools):
     assert read_outputs(sandbox)["change-detection-uncertain"] == "false"
 
 
+@pytest.mark.parametrize(
+    ("hooks", "expected_arguments", "expected_reason"),
+    [
+        (
+            "compile-fab-assets compile-edge-assets",
+            ["compile-fab-assets", "compile-edge-assets"],
+            "selected-hooks",
+        ),
+        (
+            "compile-fab-assets\ncompile-edge-assets\n",
+            ["compile-fab-assets", "compile-edge-assets"],
+            "selected-hooks",
+        ),
+        (
+            " \n",
+            ["--skip", "run-skill-eval", "--skip", "run-skill-eval-codex", "--skip", "view-skill-eval"],
+            "cache-miss",
+        ),
+    ],
+    ids=("space-separated", "newline-separated", "whitespace-only"),
+)
+def test_hooks_input_decides_what_is_prepared_and_saved(
+    sandbox, fake_tools, hooks, expected_arguments, expected_reason
+):
+    install_step = find_step(PREK_ACTION, step_id="install-hooks")
+    policy_step = find_step(PREK_ACTION, step_id="cache-policy")
+    assert install_step["env"]["HOOKS"] == policy_step["env"]["HOOKS"] == "${{ inputs.hooks }}"
+    result = run_shell(install_step["run"], {**sandbox, **fake_tools, "HOOKS": hooks})
+    assert result.returncode == 0, result.stderr
+    assert read_commands(fake_tools) == [["prek", "install-hooks", *expected_arguments]]
+    result = run_shell(
+        policy_step["run"], {**sandbox, "HOOKS": hooks, "SAVE_CACHE": "true", "STASH_HIT": "false"}
+    )
+    assert result.returncode == 0, result.stderr
+    assert read_outputs(sandbox)["reason"] == expected_reason
+
+
+def find_hatch_build_hooks():
+    return {
+        hook
+        for path in (*ROOT.glob("providers/*/hatch_build.py"), *ROOT.glob("providers/*/*/hatch_build.py"))
+        for hook in re.findall(r'"prek", "run", "([\w-]+)"', path.read_text())
+    }
+
+
+def find_hooks_run_by_job(job):
+    hooks = set()
+    for step in job["steps"]:
+        command = step.get("run", "")
+        if "prepare-provider-distributions" in command:
+            hooks |= find_hatch_build_hooks()
+        hooks.update(re.findall(r"^prek .* ([\w-]+) --all-files$", command, re.MULTILINE))
+    return hooks
+
+
+def test_selected_hooks_match_the_hooks_each_job_runs():
+    selections = {
+        f"{workflow.name}:{job_id}": (set(step["with"]["hooks"].split()), find_hooks_run_by_job(job))
+        for workflow in (ROOT / ".github/workflows").glob("*.yml")
+        for job_id, job in load_yaml(workflow)["jobs"].items()
+        for step in job.get("steps", [])
+        if step.get("uses") == "./.github/actions/install-prek" and "hooks" in step.get("with", {})
+    }
+    assert sorted(selections) == [
+        "ci-amd.yml:migration-round-trip",
+        "ci-amd.yml:mypy-providers",
+        "ci-arm.yml:migration-round-trip",
+        "ci-arm.yml:mypy-providers",
+        "generate-constraints.yml:generate-constraints-matrix",
+        "prod-image-build.yml:build-prod-packages",
+        "test-providers.yml:prepare-install-verify-provider-distributions",
+        "test-providers.yml:providers-compatibility-tests-matrix",
+    ]
+    for site, (selected, run) in selections.items():
+        assert selected == run, (
+            f"Set the hooks input of the install-prek step in {site} to: {' '.join(sorted(run))}"
+        )
+
+
 @pytest.mark.parametrize("cache_state", ("missing", "empty", "renamed-markers", "log-only"))
 @pytest.mark.parametrize("save_cache", ("true", "false"))
 def test_markerless_cache_refresh_policy(sandbox, fake_tools, cache_state, save_cache):