ci: registry-free GPU job (runner-local SIF) + verify the provisioned overlay from a fresh session (#2894)

* ci: registry-free mode — run the GPU job from a runner-local SIF

The GPU job logged in to Docker Hub on every run and pulled the container
from a private repo. With the credentials being rotated the unconditional
login step fails before any test runs, although the SIF is cached on the
runner and the pull is skipped anyway.

- gpu-test action: new input `fa4_local_sif`. When set, use that SIF, skip
  the pull and the stale-SIF prune, fail loudly if the file is missing.
- ci.yml: Docker Hub login only when `vars.FA4_LOCAL_SIF` is empty; pass
  the variable through. Unset variable == old behaviour.
- tools/ci/docker/build_local_sif.sh: docker-build tools/ci/docker/Dockerfile
  on the runner and apptainer-build a SIF into CI_WORK_DIR/local-sif/ (the
  registry-mode prune deletes *.sif directly under CI_WORK_DIR).
- tools/ci/README.md: document the mode.

* ci: verify the provisioned overlay from a fresh session and pin the DSL runtime

All GPU jobs on main failed on 2026-09-16 with
  JIT session error: Symbols not found: [ CuteDSLRT_TVMFFISetRaisedCudaError ]
right after "DSL floor check OK: nvidia-cutlass-dsl=4.8.0.dev0".

The test session saw the SIF-baked cutlass-dsl 4.6.0 runtime (lib/) next
to the installed 4.8.0 (cu13/lib/): apptainer returns before the fuse2fs
that owns the overlay image has finished flushing, so the next session can
read a stale view, and DSL discovery tries lib/ first. Intermittent and
load-dependent; latent since #2685, fatal once #2798 moved the pin off the
SIF's layout. The floor check runs in the same session as the install and
cannot see this.

- wait_for_overlay_release(): after every session, wait until no process
  holds the overlay image (60 s cap). `sync` at the end of provisioning.
- verify_overlay(): fresh session; one visible distribution per runtime
  package, flash_attn.cute imported from the checkout, the runtime that
  get_shared_libs() resolves owned by the installed cutlass-dsl.
  Re-provision once on failure, then fail with "OVERLAY VERIFY FAILED".
- Pin the verified runtime into the test sessions via CUTE_DSL_LIBS.

* ci: fail on overlay hand-off timeout; rerun the floor check in the fresh session

Review follow-ups:
- wait_for_overlay_release() raises after the timeout instead of warning
  and continuing — opening another session on a still-held image is the
  stale-view hazard the wait exists to prevent.
- verify_overlay() runs assert_dsl_floor.py in the fresh session before
  the consistency checks, so the version floors are checked against what
  the test sessions will actually see, not only the provisioning session.

.github/actions/gpu-test/action.yml

@@ -16,6 +16,12 @@ inputs:
   fa4_image_cu130:
     description: Docker image for CUDA 13.0 (used when driver supports CUDA 13.0)
     required: true
+  fa4_local_sif:
+    description: >-
+      Absolute path of a SIF on the runner. When set, the image inputs are ignored and nothing is
+      pulled (build one with tools/ci/docker/build_local_sif.sh). Missing file = job fails.
+    required: false
+    default: ""
 
 runs:
   using: composite
@@ -37,7 +43,21 @@ runs:
 
     - name: Pull FA4 SIF
       shell: bash
+      env:
+        FA4_LOCAL_SIF: ${{ inputs.fa4_local_sif }}
       run: |
+        if [ -n "$FA4_LOCAL_SIF" ]; then
+          # Registry-free mode: the runner keeps a SIF built from tools/ci/docker/Dockerfile.
+          if [ ! -f "$FA4_LOCAL_SIF" ]; then
+            echo "::error::fa4_local_sif is set but $FA4_LOCAL_SIF does not exist on this runner." \
+                 "Rebuild it with tools/ci/docker/build_local_sif.sh or unset the FA4_LOCAL_SIF variable."
+            exit 1
+          fi
+          echo "Using local SIF: $FA4_LOCAL_SIF"
+          apptainer inspect "$FA4_LOCAL_SIF" 2>/dev/null | sed 's/^/  /' || true
+          echo "FA4_SIF=$FA4_LOCAL_SIF" >> "$GITHUB_ENV"
+          exit 0
+        fi
         CI_WORK_DIR="${CI_WORK_DIR:-/scratch/user/$USER}"
         TAG=$(echo "$FA4_IMAGE" | tr '/: ' '---')
         SIF="$CI_WORK_DIR/${TAG}.sif"

.github/workflows/ci.yml

@@ -45,6 +45,8 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - name: Log in to Docker Hub
+        # Only needed to pull the image; skipped when the runner uses a locally built SIF.
+        if: ${{ vars.FA4_LOCAL_SIF == '' }}
         uses: docker/login-action@v3
         with:
           username: ${{ secrets.DOCKERHUB_USERNAME }}
@@ -52,5 +54,6 @@ jobs:
       - uses: ./.github/actions/gpu-test
         with:
           test-filter: ${{ env.FA4_TEST_FILTER }}
+          fa4_local_sif: ${{ vars.FA4_LOCAL_SIF }}
           fa4_image_cu129: "togethercomputer/training-performance:flash-attn-cu12.9-26.03.25@sha256:304a5c3d2b3a75b151cd2a964cd26d444e0d8b5686d63943df13378c9705f943"
           fa4_image_cu130: "togethercomputer/training-performance:flash-attn-cu13.0-26.06.27@sha256:2c31843c7137cbe909611eb70a9b6940581a2e3018d2c8b629f84dfd04b9bfd7"

tools/ci/README.md

@@ -1,7 +1,8 @@
 # FA4 CI
 
-CI runs on a self-hosted GPU runner using an Apptainer (SIF) container pulled from Docker Hub.
-Triggered on every push to `main`.
+CI runs on a self-hosted GPU runner inside an Apptainer (SIF) container. The container is either
+pulled from Docker Hub (default) or, when the runner has one, a SIF built locally from
+`tools/ci/docker/Dockerfile` (registry-free mode, see below). Triggered on every push to `main`.
 
 ## Two-pass test strategy
 
@@ -10,6 +11,13 @@ Triggered on every push to `main`.
 
 See `run_fa4_ci.py` for the shared logic used by both CI and `test_ci_local.sh`.
 
+Before Pass 1 the driver provisions a disk-backed Apptainer overlay (replaces the SIF-baked cutlass-dsl /
+quack / FA4 with the versions pinned in `flash_attn/cute/pyproject.toml` plus the checked-out FA4) and
+**verifies it from a fresh session**: one distribution per package, `flash_attn.cute` from the checkout,
+and the CuTeDSL runtime owned by the installed cutlass-dsl. On failure it re-provisions once, then fails
+with `OVERLAY VERIFY FAILED`. The verified runtime is pinned via `CUTE_DSL_LIBS`. This guards against
+an apptainer session hand-off race where the SIF-baked packages reappear next to the installed ones.
+
 ## Required GitHub secrets / variables
 
 | Name | Kind | Value |
@@ -17,10 +25,23 @@ See `run_fa4_ci.py` for the shared logic used by both CI and `test_ci_local.sh`.
 | `DOCKERHUB_USERNAME` | Secret | Docker Hub username |
 | `DOCKERHUB_TOKEN` | Secret | Docker Hub access token |
 | `CI_WORK_DIR` | Variable | Large-disk path on runner, e.g. `/scratch/user/johnson` |
+| `FA4_LOCAL_SIF` | Variable | Optional. Absolute path of a runner-local SIF; enables registry-free mode |
 
 `CI_WORK_DIR` is used for SIF caching and Apptainer temp files. Falls back to `/scratch/user/<github-actor>` if unset.
 
-## Updating the container image
+## Registry-free mode (runner-local SIF)
+
+When `FA4_LOCAL_SIF` is set, the GPU job skips the Docker Hub login and pull and runs in that SIF, so
+CI needs no registry credentials. The Docker Hub secrets are then unused.
+
+1. On the runner: `tools/ci/docker/build_local_sif.sh` builds `tools/ci/docker/Dockerfile` and writes a SIF
+   to `CI_WORK_DIR/local-sif/` (keep it there: registry mode prunes `*.sif` directly under `CI_WORK_DIR`).
+2. Set the repo variable `FA4_LOCAL_SIF` to the printed path. A missing file fails the job explicitly.
+
+The image only provides OS, Python and torch; `run_fa4_ci.py` installs cutlass-dsl, quack and FA4 at job
+time, so DSL pin bumps never need a rebuild. Rebuild when the Dockerfile changes.
+
+## Updating the container image (registry mode)
 
 1. Build and push a new image via `tools/ci/docker/build.sh` + `tag_and_push.sh`.
 2. Update `FA4_IMAGE` in `.github/workflows/ci.yml` with the new tag and `sha256` digest.

tools/ci/docker/build_local_sif.sh

@@ -0,0 +1,38 @@
+#!/usr/bin/env bash
+# Build the FA4 CI container on the runner and convert it to a SIF — no registry involved.
+#
+#   tools/ci/docker/build_local_sif.sh [OUTPUT_DIR]      # default OUTPUT_DIR = $CI_WORK_DIR
+#
+# Output: OUTPUT_DIR/local-sif/flash-attn-4-cu13.0-<date>-<sha>.sif — publish that path as the repo
+# variable FA4_LOCAL_SIF. Keep it in local-sif/: registry mode prunes *.sif directly under CI_WORK_DIR.
+# Rebuild only when the Dockerfile changes; the DSL stack is installed at job time by run_fa4_ci.py.
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(git -C "$SCRIPT_DIR" rev-parse --show-toplevel)"
+OUT_DIR="${1:-${CI_WORK_DIR:-/scratch/user/$USER}}"
+DATE=$(date +%y.%m.%d)
+GIT_SHA=$(git -C "$REPO_ROOT" rev-parse --short HEAD)
+IMAGE="flash-attn-4:flash-attn-cu13.0-${DATE}"
+SIF_DIR="${OUT_DIR}/local-sif"
+SIF="${SIF_DIR}/flash-attn-4-cu13.0-${DATE}-${GIT_SHA}.sif"
+TMP_DIR="${APPTAINER_TMPDIR:-/scratch/apptainer_tmp}"
+
+mkdir -p "$SIF_DIR"
+sudo mkdir -p "$TMP_DIR"
+
+echo "=== docker build $IMAGE (repo @ $GIT_SHA) ==="
+sudo docker build -t "$IMAGE" --label "org.flash_attn.ci.git_sha=$GIT_SHA" \
+  -f "$SCRIPT_DIR/Dockerfile" "$REPO_ROOT"
+
+echo "=== apptainer build $SIF ==="
+rm -f "$SIF"
+sudo APPTAINER_TMPDIR="$TMP_DIR" apptainer build "$SIF" "docker-daemon://$IMAGE"
+sudo chown "$(id -u):$(id -g)" "$SIF"
+
+echo
+echo "SIF ready: $SIF ($(du -sh "$SIF" | cut -f1))"
+echo "Publish it to CI (repo Settings → Variables, or):"
+echo "  gh variable set FA4_LOCAL_SIF -R Dao-AILab/flash-attention --body '$SIF'"
+echo "Local run against it:"
+echo "  FA4_SIF=$SIF CI_WORK_DIR=$OUT_DIR python3 tools/ci/run_fa4_ci.py --repo-root $REPO_ROOT --test-filter '<filter>'"

tools/ci/run_fa4_ci.py

@@ -11,6 +11,7 @@ import os
 import re
 import shlex
 import subprocess
+import time
 from dataclasses import dataclass
 from pathlib import Path
 
@@ -178,12 +179,157 @@ def prepare_overlay(
         f"python3 {shlex.quote(str(repo_root / 'tools/ci/assert_dsl_floor.py'))} "
         f"{shlex.quote(str(repo_root / 'flash_attn/cute/pyproject.toml'))}"
     )
-    parts = [uv_cache_export, site_packages, nuke_baked_dsl, dsl_install_cmd, fa4_install_cmd, floor_check_cmd]
+    # sync: flush the overlay image before this session tears down (see wait_for_overlay_release).
+    parts = [uv_cache_export, site_packages, nuke_baked_dsl, dsl_install_cmd, fa4_install_cmd, floor_check_cmd, "sync"]
     cmd = ["apptainer", "exec", "--nv", "--overlay", overlay, "--bind", work_dir, sif, "bash", "-c", " && ".join(parts)]
     subprocess.run(cmd, check=True, cwd=repo_root, env=base_env)
 
 
-def run_step(step: Step, repo_root: Path, base_env: dict[str, str], sif: str, work_dir: str, overlay: str) -> None:
+# ── Overlay hand-off between apptainer sessions ────────────────────────────────────────────────
+# `apptainer exec` returns before its fuse2fs child has finished flushing the ext3 overlay image, so a
+# session opened right after may see the image without the previous session's whiteouts: the SIF-baked
+# packages that prepare_overlay() deleted reappear next to the installed ones. Seen 2026-09-16: the
+# baked cutlass-dsl 4.6.0 `lib/` runtime was linked instead of 4.8.0 `cu13/lib/` ("Symbols not found").
+# Defence: wait for the image to be released, verify the stack from a fresh session (re-provision
+# once), and pin the verified runtime via CUTE_DSL_LIBS.
+
+def _pids_holding(path: str) -> list[int]:
+    """PIDs with `path` open (the fuse2fs of a live session)."""
+    holders = []
+    for pid in os.listdir("/proc"):
+        if not pid.isdigit():
+            continue
+        fd_dir = f"/proc/{pid}/fd"
+        try:
+            fds = os.listdir(fd_dir)
+        except OSError:
+            continue
+        for fd in fds:
+            try:
+                target = os.readlink(f"{fd_dir}/{fd}")
+            except OSError:
+                continue
+            if target == path or target == f"{path} (deleted)":
+                holders.append(int(pid))
+                break
+    return holders
+
+
+def wait_for_overlay_release(overlay: str, timeout_s: float = 60.0) -> float:
+    """Block until no process holds the overlay image open; return the seconds waited.
+
+    Fails after `timeout_s`: opening another session on a still-held image is exactly the
+    stale-view hazard this guards against, so continuing would defeat the check.
+    """
+    real = os.path.realpath(overlay)
+    start = time.monotonic()
+    while True:
+        holders = _pids_holding(real)
+        waited = time.monotonic() - start
+        if not holders:
+            if waited > 0.5:
+                print(f"(overlay hand-off: waited {waited:.1f}s for the previous session to release {overlay})")
+            return waited
+        if waited > timeout_s:
+            raise RuntimeError(f"{overlay} still held by pid(s) {holders} after {timeout_s:.0f}s")
+        time.sleep(0.2)
+
+
+# Runs in a fresh session after provisioning; prints "CUTE_DSL_LIBS=<paths>" on success.
+_VERIFY_OVERLAY_PY = r"""
+import importlib.metadata as md, os, sys
+repo = os.path.realpath(sys.argv[1])
+problems = []
+
+def dists():
+    for d in md.distributions():
+        yield (d.metadata.get("Name") or "").lower().replace("_", "-"), d
+
+# 1. One visible distribution per runtime package (two = the SIF-baked copy leaked through).
+seen = {}
+for name, d in dists():
+    if name.startswith("nvidia-cutlass-dsl") or name in ("quack-kernels", "flash-attn-4"):
+        seen.setdefault(name, []).append(f"{d.version} @ {d.locate_file('')}")
+for name, versions in seen.items():
+    if len(versions) > 1:
+        problems.append(f"{name}: {len(versions)} distributions visible: {versions}")
+
+# 2. FA4 imports from the checkout, not from a baked copy.
+import flash_attn.cute.interface as iface
+if not os.path.realpath(iface.__file__).startswith(repo):
+    problems.append(f"flash_attn.cute imports from {iface.__file__}, expected under {repo}")
+
+# 3. The runtime the JIT will link belongs to the installed cutlass-dsl distribution.
+import cutlass.cutlass_dsl as cd
+libs = cd.CuTeDSL._get_dsl().get_shared_libs()
+runtimes = [l for l in libs if os.path.basename(l).startswith("libcute_dsl_runtime")]
+owned = set()
+for name, d in dists():
+    if name.startswith("nvidia-cutlass-dsl"):
+        for f in d.files or []:
+            owned.add(os.path.realpath(str(d.locate_file(f))))
+if not runtimes:
+    problems.append(f"no libcute_dsl_runtime in get_shared_libs(): {libs}")
+for l in runtimes:
+    if os.path.realpath(l) not in owned:
+        problems.append(f"runtime {l} is not part of the installed cutlass-dsl distribution (stale SIF-baked copy?)")
+
+if problems:
+    print("OVERLAY VERIFY FAILED:")
+    for p in problems:
+        print("  -", p)
+    sys.exit(1)
+print(f"overlay verify OK: FA4 from {os.path.dirname(iface.__file__)}, runtime {runtimes[0]}")
+print("CUTE_DSL_LIBS=" + os.pathsep.join(libs))
+"""
+
+
+def verify_overlay(repo_root: Path, base_env: dict[str, str], sif: str, work_dir: str, overlay: str) -> str | None:
+    """Check the provisioned overlay from a fresh session; return CUTE_DSL_LIBS to pin, or None."""
+    print("=== Verify overlay (fresh session) ===")
+    # Re-run the version floor check here too: in the provisioning session it sees the in-memory
+    # overlay state, which is always right; this session sees what the tests will see.
+    floor_check = (
+        f"python3 {shlex.quote(str(repo_root / 'tools/ci/assert_dsl_floor.py'))} "
+        f"{shlex.quote(str(repo_root / 'flash_attn/cute/pyproject.toml'))}"
+    )
+    inner = f"cd /tmp && {floor_check} && python3 -c {shlex.quote(_VERIFY_OVERLAY_PY)} {shlex.quote(str(repo_root))}"
+    cmd = ["apptainer", "exec", "--overlay", overlay, "--bind", work_dir, sif, "bash", "-c", inner]
+    proc = subprocess.run(cmd, cwd=repo_root, env=base_env, text=True, capture_output=True)
+    print(proc.stdout, end="")
+    if proc.returncode != 0:
+        print(proc.stderr[-2000:], end="")
+        return None
+    for line in proc.stdout.splitlines():
+        if line.startswith("CUTE_DSL_LIBS="):
+            return line[len("CUTE_DSL_LIBS="):]
+    return None
+
+
+def provision_verified_overlay(
+    repo_root: Path, base_env: dict[str, str], sif: str, work_dir: str, overlay: str,
+    cutlass_spec: str, quack_spec: str, dsl_variant: str,
+) -> str:
+    """Create, provision and verify the overlay; re-provision once if verification fails."""
+    for attempt in (1, 2):
+        if os.path.exists(overlay):
+            os.remove(overlay)
+        subprocess.run(["apptainer", "overlay", "create", "--size", "4096", overlay], check=True)
+        prepare_overlay(
+            repo_root=repo_root, base_env=base_env, sif=sif, work_dir=work_dir, overlay=overlay,
+            cutlass_spec=cutlass_spec, quack_spec=quack_spec, dsl_variant=dsl_variant,
+        )
+        wait_for_overlay_release(overlay)
+        dsl_libs = verify_overlay(repo_root=repo_root, base_env=base_env, sif=sif, work_dir=work_dir, overlay=overlay)
+        wait_for_overlay_release(overlay)
+        if dsl_libs:
+            return dsl_libs
+        if attempt == 1:
+            print("Overlay verification failed — recreating the overlay and provisioning again.")
+    raise SystemExit("Overlay verification failed twice; see the OVERLAY VERIFY FAILED lines above.")
+
+
+def run_step(step: Step, repo_root: Path, base_env: dict[str, str], sif: str, work_dir: str, overlay: str, dsl_libs: str | None = None) -> None:
     print(f"=== {step.name} ===")
     # Convert relative test/benchmark paths to absolute so we can run from /tmp.
     # Running from /tmp ensures Python does not insert repo_root into sys.path[0]
@@ -192,12 +338,17 @@ def run_step(step: Step, repo_root: Path, base_env: dict[str, str], sif: str, wo
         str(repo_root / arg) if (arg.startswith("tests/") or arg.startswith("benchmarks/")) else arg
         for arg in step.command
     ]
-    env_exports = " && ".join(f"export {k}={shlex.quote(v)}" for k, v in step.extra_env.items())
+    extra_env = dict(step.extra_env)
+    if dsl_libs:
+        # Pin the runtime verified by verify_overlay() so JIT linking does not depend on discovery.
+        extra_env.setdefault("CUTE_DSL_LIBS", dsl_libs)
+    env_exports = " && ".join(f"export {k}={shlex.quote(v)}" for k, v in extra_env.items())
     inner_cmd = shlex.join(command)
     shell_parts = [env_exports] if env_exports else []
     shell_parts.append(f"cd /tmp && {inner_cmd}")
     cmd = ["apptainer", "exec", "--nv", "--overlay", overlay, "--bind", work_dir, sif, "bash", "-c", " && ".join(shell_parts)]
     subprocess.run(cmd, check=True, cwd=repo_root, env=base_env)
+    wait_for_overlay_release(overlay)
 
 
 # ── CLI ───────────────────────────────────────────────────────────────────────
@@ -242,13 +393,10 @@ def main() -> None:
     dsl_variant = "cu13" if read_cuda_major() >= 13 else "cu12"
     overlay = os.path.join(work_dir, "fa4_ci_overlay.img")
     os.makedirs(work_dir, exist_ok=True)
-    if os.path.exists(overlay):
-        os.remove(overlay)
-    subprocess.run(["apptainer", "overlay", "create", "--size", "4096", overlay], check=True)
     print(f"Runtime DSL: cutlass-dsl[{dsl_variant}]{cutlass_spec} + quack-kernels{quack_spec} (into {overlay})")
 
     try:
-        prepare_overlay(
+        dsl_libs = provision_verified_overlay(
             repo_root=repo_root, base_env=base_env, sif=args.sif, work_dir=work_dir,
             overlay=overlay, cutlass_spec=cutlass_spec, quack_spec=quack_spec, dsl_variant=dsl_variant,
         )
@@ -261,7 +409,7 @@ def main() -> None:
             benchmark_visible_devices=benchmark_visible_devices,
             skip_benchmark=args.skip_benchmark,
         ):
-            run_step(step, repo_root=repo_root, base_env=base_env, sif=args.sif, work_dir=work_dir, overlay=overlay)
+            run_step(step, repo_root=repo_root, base_env=base_env, sif=args.sif, work_dir=work_dir, overlay=overlay, dsl_libs=dsl_libs)
     finally:
         # The overlay can hold gigabytes; don't leave it behind on the runner between jobs.
         if os.path.exists(overlay):