ci: harden the flash-attn cache and run unit-test on retargeted PRs

.github/workflows/scripts/ensure_flash_attention.sh

@@ -52,20 +52,37 @@ export UV_CACHE_DIR="$SPECFORGE_CI_CACHE/flash-attn/$build_key"
 # flash-attn 2.8.3's setup.py still appends -std=c++17 after torch's default
 # flags, so a plain source build now fails. Patch the flag and cache the built
 # wheel per build_key; uv never caches builds from local sources.
+#
+# Only a wheel that installed and passed the ABI check is published to the
+# cache, by an atomic rename, so an interrupted build or a broken wheel cannot
+# be reused by later runs. A cached wheel that fails is dropped to force a
+# rebuild next time.
 wheel_dir="$SPECFORGE_CI_CACHE/flash-attn/$build_key/wheels"
 wheel=$(ls "$wheel_dir"/flash_attn-*.whl 2>/dev/null | head -1 || true)
+built=
 if [ -z "$wheel" ]; then
     src=$(mktemp -d)
-    sdist_url=$(curl -sSfL "https://pypi.org/pypi/flash-attn/$FLASH_ATTN_VERSION/json" |
-        python -c 'import json, sys; print(next(u["url"] for u in json.load(sys.stdin)["urls"] if u["packagetype"] == "sdist"))')
+    trap 'rm -rf "$src"' EXIT
+    sdist=$(curl -sSfL "https://pypi.org/pypi/flash-attn/$FLASH_ATTN_VERSION/json" |
+        python -c 'import json, sys; u = next(u for u in json.load(sys.stdin)["urls"] if u["packagetype"] == "sdist"); print(u["url"], u["digests"]["sha256"])')
+    sdist_url=${sdist% *}
+    sdist_sha256=${sdist#* }
     curl -sSfL "$sdist_url" -o "$src/sdist.tar.gz"
+    echo "$sdist_sha256  $src/sdist.tar.gz" | sha256sum -c -
     tar xzf "$src/sdist.tar.gz" -C "$src"
     sed -i 's/-std=c++17/-std=c++20/g' "$src"/flash_attn-*/setup.py
-    mkdir -p "$wheel_dir"
-    (cd "$src"/flash_attn-* && FLASH_ATTENTION_FORCE_BUILD=TRUE MAX_JOBS=8 python setup.py bdist_wheel -d "$wheel_dir")
-    rm -rf "$src"
-    wheel=$(ls "$wheel_dir"/flash_attn-*.whl | head -1)
+    (cd "$src"/flash_attn-* && FLASH_ATTENTION_FORCE_BUILD=TRUE MAX_JOBS=8 python setup.py bdist_wheel -d "$src/dist")
+    wheel=$(ls "$src"/dist/flash_attn-*.whl | head -1)
+    built=1
 fi
 uv pip uninstall flash-attn
-uv pip install --no-deps "$wheel"
-check_flash_attention
+if ! { uv pip install --no-deps "$wheel" && check_flash_attention; }; then
+    [ -n "$built" ] || rm -f "$wheel"
+    exit 1
+fi
+if [ -n "$built" ]; then
+    mkdir -p "$wheel_dir"
+    staged=$(mktemp "$wheel_dir/.staging.XXXXXX")
+    cp "$wheel" "$staged"
+    mv "$staged" "$wheel_dir/$(basename "$wheel")"
+fi

.github/workflows/test.yaml

@@ -3,10 +3,15 @@ name: PR Test
 on:
   pull_request:
     branches: [ main ]
+    # edited: a stacked PR retargeted to main after its parent merges gets no
+    # synchronize event, so without this it merges without unit-test.
+    types: [ opened, synchronize, reopened, edited ]
   workflow_dispatch:
 
+# Title/body edits run in their own group, so they neither cancel nor replace
+# a queued test run; their jobs are skipped below.
 concurrency:
-  group: pr-test-${{ github.ref }}
+  group: pr-test-${{ github.ref }}${{ github.event.action == 'edited' && !github.event.changes.base && format('-edit-{0}', github.run_id) || '' }}
   cancel-in-progress: true
 
 permissions:
@@ -15,7 +20,8 @@ permissions:
 jobs:
   cleanup-runner:
     if: (github.repository == 'sgl-project/SpecForge' || github.event_name == 'pull_request') &&
-        github.event.pull_request.draft == false
+        github.event.pull_request.draft == false &&
+        (github.event.action != 'edited' || github.event.changes.base)
     runs-on: [self-hosted]
     concurrency:
       group: specforge-gpu-ci
@@ -32,7 +38,8 @@ jobs:
   unit-test:
     needs: cleanup-runner
     if: (github.repository == 'sgl-project/SpecForge' || github.event_name == 'pull_request') &&
-        github.event.pull_request.draft == false
+        github.event.pull_request.draft == false &&
+        (github.event.action != 'edited' || github.event.changes.base)
     runs-on: [self-hosted]
     concurrency:
       group: specforge-gpu-ci
@@ -60,10 +67,13 @@ jobs:
       # checkout@v4 whitelists the workspace only in its own temporary git
       # config; later steps that build the package with --no-build-isolation
       # invoke git again (setuptools-scm's file finder) and fail with
-      # "detected dubious ownership" on this self-hosted container.
+      # "detected dubious ownership" on this self-hosted container. HOME
+      # persists across jobs here, so add the entry only once.
       - name: Trust the workspace for later git calls
         shell: bash
-        run: git config --global --add safe.directory "$GITHUB_WORKSPACE"
+        run: |
+          git config --global --fixed-value --get safe.directory "$GITHUB_WORKSPACE" >/dev/null ||
+            git config --global --add safe.directory "$GITHUB_WORKSPACE"
 
       # Cancelled or timed-out runs can leave processes behind which keep GPU
       # memory allocated and make this run OOM. --pid=host above is what makes

tests/test_scripts/test_ci_dependencies.py

@@ -101,41 +101,56 @@ def setup_fa(self):
         )
         uv.chmod(0o755)
         # curl serves the PyPI JSON and an sdist tarball whose setup.py builds a
-        # real (trivial) wheel once the CI script patches its -std flag.
+        # real (trivial) wheel once the CI script patches its -std flag. The
+        # tarball is byte-stable so the JSON can carry its sha256.
         curl = self.bin / "curl"
         curl.write_text(
             f"#!{sys.executable}\n"
             + textwrap.dedent(
                 """\
+                import gzip
+                import hashlib
                 import io
                 import json
                 import os
                 import sys
                 import tarfile
+                from pathlib import Path
+
+                def sdist():
+                    version = os.environ["FLASH_ATTN_VERSION"]
+                    payload = (
+                        "import os, sys\\n"
+                        "# flags: -std=c++17\\n"
+                        "if os.environ.get('TEST_INSTALL_FAIL'):\\n"
+                        "    sys.exit(1)\\n"
+                        "assert os.environ.get('FLASH_ATTENTION_FORCE_BUILD') == 'TRUE'\\n"
+                        "src = open(__file__).read()\\n"
+                        "assert '-std=' + 'c++17' not in src, \\"setup.py not patched\\"\\n"
+                        "from setuptools import setup\\n"
+                        "setup(name='flash-attn', version=os.environ['FLASH_ATTN_VERSION'])\\n"
+                    ).encode()
+                    raw = io.BytesIO()
+                    with tarfile.open(fileobj=raw, mode="w") as tar:
+                        info = tarfile.TarInfo(f"flash_attn-{version}/setup.py")
+                        info.size = len(payload)
+                        tar.addfile(info, io.BytesIO(payload))
+                    return gzip.compress(raw.getvalue(), mtime=0)
 
                 args = sys.argv[1:]
+                with (Path(os.environ["TEST_ROOT"]) / "downloads").open("a") as log:
+                    log.write(" ".join(args) + "\\n")
                 if "-o" not in args:
+                    digest = hashlib.sha256(sdist()).hexdigest()
+                    if os.environ.get("TEST_BAD_SHA256"):
+                        digest = "0" * 64
                     print(json.dumps({"urls": [{
                         "packagetype": "sdist",
                         "url": "https://files.example.invalid/flash_attn.tar.gz",
+                        "digests": {"sha256": digest},
                     }]}))
                     sys.exit(0)
-                out = args[args.index("-o") + 1]
-                version = os.environ["FLASH_ATTN_VERSION"]
-                payload = (
-                    "import os, sys\\n"
-                    "# flags: -std=c++17\\n"
-                    "if os.environ.get('TEST_INSTALL_FAIL'):\\n"
-                    "    sys.exit(1)\\n"
-                    "src = open(__file__).read()\\n"
-                    "assert '-std=' + 'c++17' not in src, \\"setup.py not patched\\"\\n"
-                    "from setuptools import setup\\n"
-                    "setup(name='flash-attn', version=os.environ['FLASH_ATTN_VERSION'])\\n"
-                ).encode()
-                with tarfile.open(out, "w:gz") as tar:
-                    info = tarfile.TarInfo(f"flash_attn-{version}/setup.py")
-                    info.size = len(payload)
-                    tar.addfile(info, io.BytesIO(payload))
+                Path(args[args.index("-o") + 1]).write_bytes(sdist())
                 """
             )
         )
@@ -156,6 +171,15 @@ def calls(self):
             json.loads(line) for line in (self.root / "calls").read_text().splitlines()
         ]
 
+    def sdist_downloads(self):
+        downloads = self.root / "downloads"
+        if not downloads.exists():
+            return 0
+        return downloads.read_text().count("flash_attn.tar.gz")
+
+    def cached_wheels(self):
+        return sorted(p.name for p in self.cache.glob("flash-attn/*/wheels/*"))
+
     def test_save_and_restore_remove_stale_versions(self):
         for environment in ("sf", "sf-capture"):
             (self.cache / environment / "old.dist-info").mkdir(parents=True)
@@ -248,6 +272,56 @@ def test_install_or_final_abi_failure_stops_ci(self):
         self.env["TEST_ABI_FAIL"] = "1"
         self.run_fa(success=False)
 
+    def test_only_a_wheel_that_passes_the_abi_check_is_cached(self):
+        self.setup_fa()
+        self.env["TEST_ABI_FAIL"] = "1"
+        self.run_fa(success=False)
+        self.assertEqual(self.cached_wheels(), [])
+        del self.env["TEST_ABI_FAIL"]
+        self.run_fa()
+        self.assertEqual(self.sdist_downloads(), 2)
+        (wheel,) = self.cached_wheels()
+        self.assertTrue(wheel.startswith("flash_attn-") and wheel.endswith(".whl"))
+
+    def test_cached_wheel_that_fails_is_dropped_and_rebuilt(self):
+        self.setup_fa()
+        self.run_fa()
+        self.assertEqual(len(self.cached_wheels()), 1)
+        (self.root / "ready").unlink()
+        self.env["TEST_ABI_FAIL"] = "1"
+        self.run_fa(success=False)
+        self.assertEqual(self.sdist_downloads(), 1)
+        self.assertEqual(self.cached_wheels(), [])
+        del self.env["TEST_ABI_FAIL"]
+        self.run_fa()
+        self.assertEqual(self.sdist_downloads(), 2)
+        self.assertEqual(len(self.cached_wheels()), 1)
+
+    def test_sdist_checksum_mismatch_stops_before_build(self):
+        self.setup_fa()
+        self.env["TEST_BAD_SHA256"] = "1"
+        self.run_fa(success=False)
+        self.assertEqual(self.cached_wheels(), [])
+        self.assertFalse((self.root / "calls").exists())
+
+    def test_trusting_the_workspace_is_idempotent(self):
+        self.env["GIT_CONFIG_GLOBAL"] = str(self.root / "gitconfig")
+        self.env["GITHUB_WORKSPACE"] = str(self.root / "workspace")
+        for _ in range(2):
+            self.run_step("Trust the workspace for later git calls")
+        result = self.run_shell("git config --global --get-all safe.directory")
+        self.assertEqual(result.stdout.splitlines(), [self.env["GITHUB_WORKSPACE"]])
+
+    def test_retargeted_prs_run_but_description_edits_do_not(self):
+        workflow = yaml.safe_load(WORKFLOW.read_text())
+        # PyYAML reads the bare `on` key as boolean True.
+        types = workflow[True]["pull_request"]["types"]
+        self.assertEqual(set(types), {"opened", "synchronize", "reopened", "edited"})
+        self.assertIn("github.run_id", workflow["concurrency"]["group"])
+        for name in ("cleanup-runner", "unit-test"):
+            with self.subTest(job=name):
+                self.assertIn("github.event.changes.base", workflow["jobs"][name]["if"])
+
     def test_workflow_orders_install_check_and_snapshot_before_tests(self):
         names = list(self.steps)
         ordered = [