freesolo-flash-dev 1.4.12__tar.gz → 1.4.14__tar.gz

{freesolo_flash_dev-1.4.12 → freesolo_flash_dev-1.4.14}/.github/workflows/bake-kernel-cache.yml

@@ -2,7 +2,7 @@ name: flash worker kernel-cache (per-arch)
 
 # Builds per-SM worker images with compiled kernel caches baked in to reuse matching kernels
 # on cold workers. Each matrix leg:
-#   1. offloads the warmup to a Modal GPU of that arch (GH runners have none) -> build/kernel_cache,
+#   1. offloads the warmup to a Modal GPU of that arch (GH runners have none) -> base-source/build/kernel_cache,
 #   2. builds docker/Dockerfile.kernelcache -> ghcr.io/freesolo-co/flash-worker:cu128-<sm>,
 #   3. pushes it. After live qualification, record its immutable digest in QUALIFIED_WORKER_IMAGES
 #      (flash/providers/_lifecycle/net/worker.py); publishing a tag does not enable it for workers.
@@ -10,12 +10,9 @@ name: flash worker kernel-cache (per-arch)
 # Depends on the base ghcr.io/freesolo-co/flash-worker:cu128 image existing (worker-image.yml) — the
 # warmup runs INSIDE it so the cache matches the image's pinned torch/triton/fla/liger toolchain.
 #
-# Secrets: MODAL_TOKEN_ID + MODAL_TOKEN_SECRET (the warmup GPU sandbox) + HF_TOKEN (the temp dataset that ferries code + the
-# cache artifact) are native repo Actions secrets, synced from Infisical (prod:/flash-ci) by the
-# github-ci-flash Secret Sync; the GHCR push uses the default GITHUB_TOKEN. Heavy + paid,
-# but UNGATED: the fingerprint detection in auto-rebake.yml decides WHEN this runs (only on a real
-# kernel-toolchain change), so there is no manual approval step. Invoked by hand (workflow_dispatch)
-# OR automatically by auto-rebake.yml (workflow_call) when the kernel toolchain changed.
+# Modal credentials launch the warmup sandbox. A dedicated Azure build identity writes
+# only the build-artifacts container and issues bounded sandbox SAS credentials.
+# GHCR publication uses GITHUB_TOKEN. Automatic invocation remains fingerprint-gated.
 
 on:
   workflow_dispatch:
@@ -199,13 +196,22 @@ jobs:
             echo "revision=$(echo "$cfg" | jq -r '.config.Labels["org.opencontainers.image.revision"] // ""')" >> "$GITHUB_OUTPUT"
           fi
 
-      # Pin the checkout to the base's source commit so the warmed flash code matches the base exactly
-      # (the cache is content-keyed by the code that compiled it). Falls back to the event sha if the
-      # base predates the revision label.
-      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      # Transport and its dependencies follow this workflow, even when manually rebaking an
+      # older image whose helper still required Hugging Face artifact credentials.
+      - name: Checkout bake transport
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        if: steps.gate.outputs.run == 'true'
+        with:
+          ref: ${{ github.workflow_sha }}
+          path: bake-tools
+
+      # Only the warmup source and thin image recipe follow the base image's source revision.
+      - name: Checkout base warmup source
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
         if: steps.gate.outputs.run == 'true'
         with:
           ref: ${{ steps.base.outputs.revision || github.sha }}
+          path: base-source
 
       - name: Install uv
         if: steps.gate.outputs.run == 'true'
@@ -218,29 +224,30 @@ jobs:
           # marks the whole job red even though the bake + push already succeeded.
           enable-cache: false
 
-      - name: Sync bake dependencies
+      - name: Sync deps (flash + modal + Azure)
         if: steps.gate.outputs.run == 'true'
-        run: uv sync --extra server
+        run: uv sync --project bake-tools --extra server --locked
 
-      # MODAL_TOKEN_ID + MODAL_TOKEN_SECRET (the warmup GPU sandbox) + HF_TOKEN (upstream model/tokenizer downloads) are
-      # native repo Actions secrets, synced from Infisical (prod:/flash-ci) by the github-ci-flash
-      # Secret Sync. No Infisical CLI/login at job time -- Infisical stays the source of truth and
-      # rotations propagate via the sync. BASE_IMAGE
-      # is digest-pinned so the warm runs inside the EXACT base the per-arch tag will sit on.
-      - name: Warm kernels on a Modal ${{ matrix.sm }} GPU -> build/kernel_cache
+      # The digest-pinned base keeps the warmup toolchain identical to the published image.
+      - name: Warm kernels on a Modal ${{ matrix.sm }} GPU -> base-source/build/kernel_cache
         if: steps.gate.outputs.run == 'true'
         env:
           MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
           MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
-          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+          FLASH_ARTIFACT_BLOB_ACCOUNT_URL: https://freesoloenvs.blob.core.windows.net
+          FLASH_ARTIFACT_BLOB_CONTAINER: build-artifacts
+          AZURE_TENANT_ID: ${{ secrets.AZURE_TENANT_ID }}
+          AZURE_CLIENT_ID: ${{ secrets.AZURE_BUILD_CLIENT_ID }}
+          AZURE_CLIENT_SECRET: ${{ secrets.AZURE_BUILD_CLIENT_SECRET }}
           BASE_IMAGE: ${{ steps.base.outputs.ref }}
           BAKE_DEADLINE_MINUTES: ${{ inputs.deadline_minutes }}
         run: |
-          uv run python docker/bake_kernel_cache.py \
+          uv run --project bake-tools --locked python bake-tools/docker/bake_kernel_cache.py \
             --arch "${{ matrix.arch }}" --sm "${{ matrix.sm }}" \
             --image "$BASE_IMAGE" \
+            --source-root base-source \
             --deadline-min "$BAKE_DEADLINE_MINUTES" \
-            --out build/kernel_cache
+            --out base-source/build/kernel_cache
 
       - name: Free disk space
         if: steps.gate.outputs.run == 'true'
@@ -263,8 +270,8 @@ jobs:
         if: steps.gate.outputs.run == 'true'
         uses: docker/build-push-action@c3c9e263c25d99ce0380d002d59b67737d91b0dc # v7.4.0
         with:
-          context: .
-          file: docker/Dockerfile.kernelcache
+          context: base-source
+          file: base-source/docker/Dockerfile.kernelcache
           push: true
           tags: ghcr.io/freesolo-co/flash-worker:${{ inputs.target_tag_prefix || 'cu128' }}-${{ matrix.sm }}
           provenance: false

{freesolo_flash_dev-1.4.12 → freesolo_flash_dev-1.4.14}/.github/workflows/dev-kernel-cache.yml

@@ -2,6 +2,7 @@ name: dev kernel-cache candidate
 
 # dev pushes only run the metadata detector. the scheduled or manual run builds the newest
 # missing candidate, so several daytime changes coalesce into one off-peak gpu bake.
+# a dev -> main promotion builds missing candidates immediately and retries its readiness gate.
 
 on:
   push:
@@ -11,11 +12,12 @@ on:
       - .github/workflows/worker-image.yml
       - .github/workflows/bake-kernel-cache.yml
       - .github/workflows/dev-kernel-cache.yml
-      - docker/kernel_fingerprint.py
-      - docker/make_rp_handler.py
-      - docker/extract_cudart_fix.py
+      - docker/**
       - flash/**
       - pyproject.toml
+  pull_request:
+    branches: [main]
+    types: [opened, synchronize, reopened, ready_for_review]
   schedule:
     # 09:00 utc is 01:00 pst / 02:00 pdt.
     - cron: "0 9 * * *"
@@ -33,7 +35,15 @@ permissions:
   packages: write
 
 concurrency:
-  group: dev-kernel-cache
+  # metadata pushes and rejected PRs must not replace a queued bake in the pending slot.
+  group: >-
+    dev-kernel-cache-${{
+      (github.event_name == 'workflow_dispatch' ||
+       (github.event_name == 'pull_request' &&
+        github.event.pull_request.head.repo.full_name == github.repository &&
+        github.event.pull_request.head.ref == 'dev' && !github.event.pull_request.draft))
+      && 'build' || 'detect'
+    }}
   cancel-in-progress: false
 
 jobs:
@@ -48,7 +58,11 @@ jobs:
         run: gh workflow run dev-kernel-cache.yml --repo "$GITHUB_REPOSITORY" --ref dev
 
   detect:
-    if: github.repository == 'freesolo-co/flash' && github.event_name != 'schedule'
+    if: >-
+      github.repository == 'freesolo-co/flash' && github.event_name != 'schedule' &&
+      (github.event_name != 'pull_request' ||
+       (github.event.pull_request.head.repo.full_name == github.repository &&
+        github.event.pull_request.head.ref == 'dev' && !github.event.pull_request.draft))
     runs-on: ubuntu-latest
     timeout-minutes: 15
     outputs:
@@ -63,7 +77,7 @@ jobs:
     steps:
       - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
         with:
-          ref: ${{ github.sha }}
+          ref: ${{ github.event.pull_request.head.sha || github.sha }}
 
       - name: Resolve source metadata
         id: source
@@ -143,7 +157,7 @@ jobs:
 
   build-base:
     needs: detect
-    if: github.event_name == 'workflow_dispatch' && needs.detect.outputs.needs_bake == 'true'
+    if: (github.event_name == 'workflow_dispatch' || github.event_name == 'pull_request') && needs.detect.outputs.needs_bake == 'true'
     uses: ./.github/workflows/worker-image.yml
     with:
       tag: cu128-dev-${{ needs.detect.outputs.revision }}
@@ -152,11 +166,11 @@ jobs:
 
   bake:
     needs: [detect, build-base]
-    if: github.event_name == 'workflow_dispatch' && needs.detect.outputs.needs_bake == 'true'
+    if: (github.event_name == 'workflow_dispatch' || github.event_name == 'pull_request') && needs.detect.outputs.needs_bake == 'true'
     uses: ./.github/workflows/bake-kernel-cache.yml
     with:
       sms: ${{ needs.detect.outputs.sms }}
-      deadline_minutes: ${{ fromJSON(format('{0}', inputs.deadline_minutes)) }}
+      deadline_minutes: ${{ fromJSON(github.event_name == 'pull_request' && '90' || format('{0}', inputs.deadline_minutes)) }}
       fp_cache: ${{ needs.build-base.outputs.fp_cache }}
       fp_base: ${{ needs.build-base.outputs.fp_base }}
       base_ref: ${{ needs.build-base.outputs.base_ref }}
@@ -166,7 +180,7 @@ jobs:
 
   publish-ready:
     needs: [detect, bake]
-    if: github.event_name == 'workflow_dispatch' && needs.detect.outputs.needs_bake == 'true'
+    if: (github.event_name == 'workflow_dispatch' || github.event_name == 'pull_request') && needs.detect.outputs.needs_bake == 'true'
     runs-on: ubuntu-latest
     timeout-minutes: 15
     steps:
@@ -213,7 +227,26 @@ jobs:
               "$repo:cu128-dev-ready-$sm"
           done
 
-      - name: Re-run production readiness for dev
+  refresh-readiness:
+    needs: [detect, publish-ready]
+    # a queued revision may find compatible caches published by the preceding bake.
+    # always() permits the skipped publication dependency; failed/cancelled work stays red.
+    if: >-
+      always() && !cancelled() && needs.detect.result == 'success' &&
+      (github.event_name == 'workflow_dispatch' || github.event_name == 'pull_request') &&
+      (needs.detect.outputs.needs_bake == 'false' || needs['publish-ready'].result == 'success')
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    steps:
+      - name: Re-run failed production readiness for this revision
         env:
           GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-        run: gh workflow run production-kernel-cache-ready.yml --repo "$GITHUB_REPOSITORY" --ref dev
+          REVISION: ${{ needs.detect.outputs.revision }}
+        run: |
+          set -euo pipefail
+          runs=$(gh run list --repo "$GITHUB_REPOSITORY" \
+            --workflow production-kernel-cache-ready.yml --commit "$REVISION" \
+            --event pull_request --status failure --json databaseId --jq '.[].databaseId')
+          for run_id in $runs; do
+            gh run rerun "$run_id" --repo "$GITHUB_REPOSITORY" --failed
+          done

{freesolo_flash_dev-1.4.12 → freesolo_flash_dev-1.4.14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: freesolo-flash-dev
-Version: 1.4.12
+Version: 1.4.14
 Summary: Flash — managed LoRA post-training (SFT/GRPO) for Freesolo environments, driven by the `flash` CLI
 Project-URL: Homepage, https://github.com/freesolo-co/flash
 Project-URL: Repository, https://github.com/freesolo-co/flash

freesolo_flash_dev-1.4.14/docker/bake_kernel_cache.py

@@ -0,0 +1,270 @@
+"""Bake per-architecture worker kernel caches on Modal and verify the returned artifact."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import shutil
+import tarfile
+import time
+import uuid
+
+ARTIFACT_NAMESPACE = "Freesolo-Co"
+
+# One Modal GPU per compiler architecture; the worker verifies the requested SM.
+GPU_BY_SM = {
+    "sm80": "A100-80GB",
+    "sm86": "A10",
+    "sm89": "L40S",
+    "sm90": "H100!",
+    "sm100": "B200",
+    "sm120": "RTX-PRO-6000",
+}
+
+
+def log(msg: str) -> None:
+    print(f"[bake-ci] {msg}", flush=True)
+
+
+def _upload_flash_code(api, repo: str, source_root) -> None:
+    """Upload the base-pinned warmup package without importing it into the transport process."""
+    from pathlib import Path
+
+    pkg_dir = Path(source_root).resolve() / "flash"
+    if not (pkg_dir / "engine/worker/runtime/kernel_warmup.py").is_file():
+        raise ValueError("bake source must contain the base image's Flash warmup package")
+    api.create_repo(repo_id=repo, repo_type="dataset", private=True, exist_ok=True)
+    api.upload_folder(
+        folder_path=str(pkg_dir),
+        path_in_repo="code/flash",
+        repo_id=repo,
+        repo_type="dataset",
+        ignore_patterns=["__pycache__/*", "*.pyc"],
+        delete_patterns=["**"],
+    )
+    log(f"uploaded flash code -> {repo}:code/flash (from {pkg_dir})")
+
+
+def _bake_image(modal, image_ref, entry):
+    """Seed the same standalone Azure client used by managed worker bootstrap."""
+    from pathlib import Path
+
+    import flash.artifacts
+
+    image = (
+        modal.Image.from_registry(image_ref)
+        .pip_install("azure-storage-blob==12.31.1", "azure-identity==1.27.0")
+        .add_local_file(str(entry), "/opt/flash/bake.py", copy=True)
+    )
+    package = Path(flash.artifacts.__file__).parent
+    for module in ("api", "config", "manifest", "storage"):
+        image = image.add_local_file(
+            str(package / f"{module}.py"), f"/opt/flash/artifact_{module}.py", copy=True
+        )
+    return image
+
+
+def _delete_artifacts(repo):
+    from flash.artifacts.config import open_container
+    from flash.artifacts.manifest import repo_prefix
+
+    with open_container() as container:
+        for blob in container.list_blobs(name_starts_with=repo_prefix(repo) + "/"):
+            container.delete_blob(blob.name)
+
+
+def main() -> int:
+    from pathlib import Path
+
+    from flash.artifacts.api import ArtifactStore, download_snapshot
+    from flash.artifacts.config import runtime_environment
+    from flash.providers.modal import api as modal_api
+
+    parser = argparse.ArgumentParser(description="Bake worker kernels on Modal")
+    parser.add_argument("--arch", required=True)
+    parser.add_argument("--sm", required=True, choices=tuple(GPU_BY_SM))
+    parser.add_argument("--image", required=True)
+    parser.add_argument("--out", default="build/kernel_cache")
+    parser.add_argument(
+        "--source-root",
+        type=Path,
+        default=Path(__file__).resolve().parents[1],
+        help="Checkout matching the base image; transport runs from this helper's checkout",
+    )
+    parser.add_argument("--deadline-min", type=int, default=90)
+    parser.add_argument("--run-id", default="")
+    args = parser.parse_args()
+    if not 1 <= args.deadline_min <= 1440:
+        parser.error("deadline must be between 1 and 1440 minutes")
+    artifacts = ArtifactStore()
+    worker_secrets = runtime_environment()
+    if not worker_secrets:
+        raise RuntimeError("Azure build artifact storage is not configured")
+    suffix = args.run_id or f"{int(time.time())}-{uuid.uuid4().hex[:8]}"
+    repo = f"{os.environ.get('FLASH_ARTIFACT_NAMESPACE', ARTIFACT_NAMESPACE)}/kernel-bake-{args.sm}-{suffix}"
+    entry = Path(__file__).with_name("bake_sandbox_entry.py")
+
+    async def bake(modal, client):
+        app = await modal.App.lookup.aio(
+            "flash-kernel-bake",
+            environment_name=modal_api.environment(),
+            create_if_missing=True,
+            client=client,
+        )
+        image = _bake_image(modal, args.image, entry)
+        sandbox = await modal.Sandbox.create.aio(
+            "python",
+            "-u",
+            "/opt/flash/bake.py",
+            app=app,
+            image=image,
+            name=f"bake-{args.sm}-{suffix}",
+            gpu=GPU_BY_SM[args.sm],
+            timeout=args.deadline_min * 60,
+            cpu=(4, 8),
+            # Keep the reservation at 32 GiB, but allow compiler memory to burst. The
+            # H100 warmup was evicted for thrashing at the previous 64 GiB limit.
+            memory=(32768, 131072),
+            secrets=[
+                modal.Secret.from_dict(
+                    {"BAKE_ARTIFACT_REPO": repo, "BAKE_ARCH": args.arch, **worker_secrets}
+                )
+            ],
+            client=client,
+        )
+        try:
+            await sandbox.wait.aio()
+            returncode = await sandbox.poll.aio()
+            if returncode != 0:
+                # The provider wrapper sanitizes ordinary exceptions. Preserve this
+                # credential-free diagnostic so CI names the failed sandbox and exit.
+                raise modal_api.ModalApiError(
+                    f"Kernel warmup exited with code {returncode}; "
+                    f"inspect Modal sandbox {sandbox.object_id}"
+                )
+        finally:
+            await sandbox.terminate.aio(wait=True)
+
+    try:
+        _upload_flash_code(artifacts, repo, args.source_root)
+        modal_api.call(bake, timeout=args.deadline_min * 60 + 120)
+        if not _status_present(artifacts, repo, retries=3):
+            raise RuntimeError("kernel warmup finished without its completion marker")
+        import tempfile
+
+        with tempfile.TemporaryDirectory(prefix="flash-kernel-") as temp:
+            download_snapshot(
+                repo_id=repo,
+                repo_type="dataset",
+                allow_patterns=["out/**"],
+                local_dir=temp,
+            )
+            staging = Path(temp) / "out"
+            status = (staging / "STATUS").read_text()
+            if status != f"rc=0\narch={args.arch}\n":
+                raise RuntimeError("kernel warmup returned an invalid completion marker")
+            archive = staging / "kernel_cache.tar.gz"
+            current_log = (staging / "bake-warmup.log").read_bytes()
+            _extract_cache_archive(archive, staging)
+            archive.unlink()
+            # A cache may contain an older log. Keep this run's subprocess log authoritative.
+            (staging / "bake-warmup.log").write_bytes(current_log)
+            (staging / "warmup.log").write_bytes(current_log)
+            result = _verify(str(staging), args.sm)
+            if result == 0:
+                shutil.copytree(staging, args.out, dirs_exist_ok=True)
+            return result
+    finally:
+        try:
+            _delete_artifacts(repo)
+        except Exception:
+            log(
+                f"WARNING: could not delete temporary Azure artifacts for {repo}; remove it manually"
+            )
+
+
+def _extract_cache_archive(archive_path, destination) -> None:
+    """Restore cache layout without allowing archive links or paths outside the output tree."""
+    from pathlib import Path
+
+    root = Path(destination).resolve()
+    with tarfile.open(archive_path, "r:gz") as archive:
+        members = archive.getmembers()
+        total = 0
+        for member in members:
+            target = (root / member.name).resolve()
+            if not target.is_relative_to(root) or not (member.isfile() or member.isdir()):
+                raise ValueError("invalid kernel cache archive member")
+            total += member.size
+            if total > 32 * 1024**3:
+                raise ValueError("kernel cache archive exceeds 32 GiB")
+        archive.extractall(root, members=members, filter="data")
+
+
+def _status_present(api, repo: str, retries: int = 1) -> bool:
+    """Authoritative out/STATUS check, retried to ride out a transient Azure read error."""
+    for i in range(retries):
+        try:
+            if "out/STATUS" in api.list_repo_files(repo, repo_type="dataset"):
+                return True
+        except Exception as e:
+            log(f"artifact status recheck {i + 1}/{retries}: {type(e).__name__}")
+        if i < retries - 1:
+            time.sleep(5)
+    return False
+
+
+def _verify(out: str, sm: str) -> int:
+    """Confirm the artifact actually landed and its metadata matches the requested arch."""
+    from flash.engine.worker.runtime.kernel_warmup import CACHE_META_FILENAME, has_raw_kernel_cache
+
+    meta = os.path.join(out, CACHE_META_FILENAME)
+    if not has_raw_kernel_cache(out):
+        log(f"FAIL: no compiled kernel artifacts in {out}; what the warmup actually produced:")
+        for root, _, files in os.walk(out):
+            for f in sorted(files):
+                p = os.path.join(root, f)
+                log(f"   present: {os.path.relpath(p, out)} ({os.path.getsize(p)} b)")
+        wl = os.path.join(out, "warmup.log")
+        if os.path.isfile(wl):
+            log("   --- warmup.log tail ---")
+            with open(wl, errors="replace") as wlf:
+                tail = wlf.read().splitlines()[-40:]
+            for line in tail:
+                log(f"   | {line}")
+        return 1
+    try:
+        with open(meta) as f:
+            got = json.load(f).get("sm")
+    except Exception as e:
+        log(f"FAIL: metadata unreadable ({e})")
+        return 1
+    if got != sm:
+        log(f"FAIL: baked arch {got!r} != requested {sm!r}")
+        return 1
+    try:
+        with open(os.path.join(out, "cache_validation.json")) as file:
+            validation = json.load(file)
+        with open(os.path.join(out, "cache_reload.json")) as file:
+            reload = json.load(file)
+        if (
+            validation.get("first_exit_code") != 0
+            or validation.get("reload_exit_code") != 0
+            or validation.get("exported_cache_replayed") is not True
+            or reload.get("cached_groups_compilation_forbidden") is not True
+            or reload.get("cached_groups") != ["flash_attn", "fla_gdn"]
+            or reload.get("startup_jit_groups") != (["flashinfer_gdn"] if sm == "sm90" else [])
+            or reload.get("compile_attempts") != 0
+            or reload.get("passed") is not True
+        ):
+            raise ValueError("fresh-process cache reuse was not proven")
+    except (OSError, ValueError, TypeError, AttributeError) as error:
+        log(f"FAIL: cache reload validation failed ({error})")
+        return 1
+    log(f"OK: compiled kernel artifacts for {sm}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

freesolo_flash_dev-1.4.14/docker/bake_sandbox_entry.py

@@ -0,0 +1,151 @@
+"""Warm kernels using private Azure source and output artifacts.
+
+The CI helper seeds the standalone Azure client into the sandbox image. The entry
+fetches code/**, compiles and replays the exported cache, then publishes out/STATUS
+last. BAKE_ARTIFACT_REPO names the disposable namespace; BAKE_ARCH selects CUDA.
+Only a bounded Azure container SAS reaches the sandbox.
+"""
+
+import io
+import json
+import os
+import shutil
+import subprocess
+import tarfile
+import time
+
+CACHE_DIR = "/opt/flash/kernelcache"
+OUT_DIR = "/out"
+CACHE_ARCHIVE = "/tmp/kernel_cache.tar.gz"
+
+
+def _replace_cache_with_exported_copy(cache_dir: str, archive_path: str) -> None:
+    """Archive the cache, delete the originals, and restore files at the same absolute paths."""
+    with tarfile.open(archive_path, "w:gz", compresslevel=1) as archive:
+        archive.add(cache_dir, arcname=".")
+    shutil.rmtree(cache_dir)
+    with tarfile.open(archive_path, "r:gz") as archive:
+        archive.extractall(cache_dir, filter="data")
+
+
+def main() -> int:
+    from artifact_api import ArtifactStore, download_snapshot
+
+    repo = os.environ["BAKE_ARTIFACT_REPO"]
+    arch = os.environ.get("BAKE_ARCH", "")
+    api = ArtifactStore()
+
+    # Upload a STARTED marker FIRST -- before the code download / warmup -- so a CI-side
+    # timeout can tell whether this entrypoint even ran. Present on a timeout = the warmup started (so it
+    # hung or was slow); absent = the pod ran the wrong entrypoint, OR this upload itself failed while
+    # the warmup is still running. Retry a few times so a transient Azure error doesn't drop the marker and
+    # make the helper misread a slow run as a wrong entrypoint.
+    started_uploaded = False
+    for attempt in range(3):
+        try:
+            api.upload_file(
+                path_or_fileobj=io.BytesIO(f"started arch={arch}\n".encode()),
+                path_in_repo="out/STARTED",
+                repo_id=repo,
+                repo_type="dataset",
+            )
+            started_uploaded = True
+            print("[bake] uploaded out/STARTED", flush=True)
+            break
+        except Exception as e:
+            # keep the "out/STARTED" token in the line so it greps the same as the success/WARNING lines
+            print(
+                f"[bake] out/STARTED upload attempt {attempt + 1}/3 failed: {type(e).__name__}",
+                flush=True,
+            )
+            if attempt < 2:  # no point sleeping after the last attempt
+                time.sleep(3)
+    if not started_uploaded:
+        # The marker never landed; warn loudly so a later "no out/STARTED" is read as upload-failed,
+        # not as a wrong-entrypoint run.
+        print("[bake] WARNING: out/STARTED never uploaded; warmup still proceeds", flush=True)
+
+    # the flash code the helper uploaded (upload_code -> code/flash); run the warmup from it.
+    download_snapshot(
+        repo_id=repo,
+        repo_type="dataset",
+        allow_patterns=["code/**"],
+        local_dir="/runcode",
+    )
+
+    env = {**os.environ, "PYTHONPATH": "/runcode/code"}
+    python_bin = os.environ.get("FLASH_VERL_PYTHON", "/opt/verl-venv/bin/python")
+    if not os.path.isfile(python_bin):
+        raise RuntimeError(f"training interpreter missing: {python_bin}")
+    cache_dir = CACHE_DIR
+    # An already-baked input image must not supply stale artifacts or validation receipts.
+    shutil.rmtree(cache_dir, ignore_errors=True)
+    os.makedirs(cache_dir)
+    cmd = [
+        python_bin,
+        "-m",
+        "flash.engine.worker.runtime.kernel_warmup",
+        "--out",
+        cache_dir,
+    ]
+    if arch:
+        cmd += ["--arch", arch]
+    print(f"[bake] running: {' '.join(cmd)}", flush=True)
+    # capture the warmup output into /out so it ships back with the cache -- lets the CI helper show
+    # which warm steps ran and why a bake failed to produce compiled kernel files.
+    os.makedirs(OUT_DIR, exist_ok=True)
+    started = time.monotonic()
+    with open(os.path.join(OUT_DIR, "warmup.log"), "wb") as lf:
+        rc = subprocess.call(cmd, env=env, stdout=lf, stderr=subprocess.STDOUT)
+    first_seconds = time.monotonic() - started
+    validation = {"first_process_seconds": first_seconds, "first_exit_code": rc}
+    if rc == 0:
+        # Remove the compiler outputs and restore only the exported archive. A fresh process
+        # alone could otherwise reuse files that never make it into the shipped artifact.
+        _replace_cache_with_exported_copy(cache_dir, CACHE_ARCHIVE)
+        started = time.monotonic()
+        with open(os.path.join(OUT_DIR, "warmup.log"), "ab") as lf:
+            lf.write(b"\n[bake] fresh-process cache reload validation\n")
+            rc = subprocess.call(
+                [*cmd, "--verify-cache"], env=env, stdout=lf, stderr=subprocess.STDOUT
+            )
+        validation.update(
+            reload_process_seconds=time.monotonic() - started,
+            reload_exit_code=rc,
+            exported_cache_replayed=True,
+        )
+    with open(os.path.join(cache_dir, "cache_validation.json"), "w") as output:
+        json.dump(validation, output)
+    print(f"[bake] kernel_warmup rc={rc}", flush=True)
+    # Preserve nested cache paths and their absolute compilation locations in one archive.
+    try:
+        with tarfile.open(CACHE_ARCHIVE, "w:gz", compresslevel=1) as archive:
+            archive.add(cache_dir, arcname=".")
+        api.upload_file(
+            path_or_fileobj=CACHE_ARCHIVE,
+            path_in_repo="out/kernel_cache.tar.gz",
+            repo_id=repo,
+            repo_type="dataset",
+        )
+        api.upload_file(
+            path_or_fileobj=os.path.join(OUT_DIR, "warmup.log"),
+            path_in_repo="out/bake-warmup.log",
+            repo_id=repo,
+            repo_type="dataset",
+        )
+    except Exception as e:
+        rc = 1
+        print(f"[bake] artifact upload failed: {type(e).__name__}", flush=True)
+    # STATUS marker LAST: its presence is the helper's completion signal.
+    api.upload_file(
+        path_or_fileobj=io.BytesIO(f"rc={rc}\narch={arch}\n".encode()),
+        path_in_repo="out/STATUS",
+        repo_id=repo,
+        repo_type="dataset",
+    )
+    print(f"[bake] ENTRY_DONE rc={rc}", flush=True)
+    return rc
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())