freesolo-flash-dev 0.2.26__tar.gz → 0.2.28__tar.gz

freesolo_flash_dev-0.2.28/.github/workflows/auto-rebake.yml

@@ -0,0 +1,245 @@
+name: auto-rebake per-arch kernel cache
+
+# Keeps the per-arch ghcr.io/freesolo-co/flash-worker:cu128-sm* images fresh whenever the :cu128 base
+# is rebuilt. Each cu128-smXX is physically `FROM :cu128` + a baked mega-cache, so a new base leaves
+# the per-arch tags pointing at the OLD base until rebaked. "Out of date" has two flavors, keyed off
+# the fp_cache / fp_base labels worker-image.yml stamps on :cu128 (docker/kernel_fingerprint.py):
+#
+#   * fp_cache changed -> the baked Triton/Inductor mega-cache is INVALID -> a real GPU re-warm is
+#     required (bake-kernel-cache.yml). PAID, and AUTO-fired by this detection (no manual approval) --
+#     fp_cache only moves on a deliberate toolchain bump, so spend stays tied to real need.
+#   * fp_cache same but fp_base changed -> the cache is still valid, the base just moved (FA wheel /
+#     causal-conv1d / rp_handler / a non-kernel dep) -> a FREE re-layer grafts the existing cache onto
+#     the fresh base (docker/Dockerfile.kernelcache.relayer), no GPU.
+#   * both same -> skip (this is the pure-flash-code case the old reminder over-fired on).
+#
+# The gate READS the published labels (authoritative for what was actually built, incl. a manual FA3
+# override) rather than recomputing from the checkout, and aborts unless the base's revision label ==
+# the build's head_sha (so a superseded base build or a `tag=...` dispatch that never touched :cu128
+# can't trigger a stale rebake).
+
+on:
+  workflow_run:
+    workflows: ["flash worker image"]
+    types: [completed]
+
+permissions:
+  contents: read
+  packages: write
+  actions: read
+  issues: write
+
+# Scope the group to the triggering base build's commit so a no-op completion (e.g. a tag=cu128-test
+# worker-image run, whose head_sha differs) can't cancel an in-flight real bake. The revision guard
+# already neutralizes superseded runs, so a global cancel-in-progress isn't needed.
+concurrency:
+  group: auto-rebake-${{ github.event.workflow_run.head_sha }}
+  cancel-in-progress: true
+
+env:
+  BASE: ghcr.io/freesolo-co/flash-worker:cu128
+
+jobs:
+  gate:
+    if: github.event.workflow_run.conclusion == 'success'
+    runs-on: ubuntu-latest
+    outputs:
+      proceed: ${{ steps.base.outputs.proceed }}
+      relayer_sms: ${{ steps.classify.outputs.relayer_sms }}
+      gpu_sms: ${{ steps.classify.outputs.gpu_sms }}
+      fp_cache: ${{ steps.base.outputs.fp_cache }}
+      fp_base: ${{ steps.base.outputs.fp_base }}
+      base_ref: ${{ steps.base.outputs.base_ref }}
+      base_revision: ${{ steps.base.outputs.revision }}
+    steps:
+      - uses: actions/checkout@v6
+        with:
+          ref: ${{ github.event.workflow_run.head_sha }}
+
+      - name: Install crane
+        run: |
+          curl -sSL https://github.com/google/go-containerregistry/releases/download/v0.20.2/go-containerregistry_Linux_x86_64.tar.gz \
+            | sudo tar -xz -C /usr/local/bin crane
+          crane version
+
+      - name: Log in to GHCR
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Read base labels
+        id: base
+        env:
+          HEAD_SHA: ${{ github.event.workflow_run.head_sha }}
+        run: |
+          # canonical digest ref (repo@sha256, no tag) so every tool (crane/docker/runpod) accepts it.
+          digest=$(crane digest "$BASE")
+          repo=${BASE%:*}
+          cfg=$(crane config "$repo@$digest")
+          fpc=$(echo "$cfg" | jq -r '.config.Labels["co.freesolo.flash.kernelcache.fp_cache"] // ""')
+          fpb=$(echo "$cfg" | jq -r '.config.Labels["co.freesolo.flash.kernelcache.fp_base"] // ""')
+          rev=$(echo "$cfg" | jq -r '.config.Labels["org.opencontainers.image.revision"] // ""')
+          echo "base_ref=$repo@$digest" >> "$GITHUB_OUTPUT"
+          echo "fp_cache=$fpc" >> "$GITHUB_OUTPUT"
+          echo "fp_base=$fpb" >> "$GITHUB_OUTPUT"
+          echo "revision=$rev" >> "$GITHUB_OUTPUT"
+          # Only proceed if this base IS the build that triggered us: its revision label must match the
+          # build's head_sha. Guards (a) a newer base that already superseded this run, and (b) a
+          # `tag=cu128-test` dispatch that never updated :cu128 (the workflow_run payload can't reveal
+          # the dispatch tag). An unlabeled base (pre-rollout) also no-ops here.
+          if [ -n "$fpc" ] && [ "$rev" = "$HEAD_SHA" ]; then
+            echo "proceed=true" >> "$GITHUB_OUTPUT"
+            echo "base $BASE@$digest is rev $rev (== head_sha) fp_cache=$fpc fp_base=$fpb"
+          else
+            echo "proceed=false" >> "$GITHUB_OUTPUT"
+            echo "::notice::auto-rebake no-op: base rev '$rev' != head_sha '$HEAD_SHA' (or base unlabeled)"
+          fi
+
+      - name: Classify arches
+        id: classify
+        if: steps.base.outputs.proceed == 'true'
+        env:
+          FP_CACHE: ${{ steps.base.outputs.fp_cache }}
+          FP_BASE: ${{ steps.base.outputs.fp_base }}
+          REV: ${{ steps.base.outputs.revision }}
+        run: |
+          # arch list is the single source of truth in deps.py (mirrored by the bake matrix; a test
+          # pins that mirror). textual parse -> no flash import / uv sync needed.
+          arches=$(python3 -c "import re,pathlib; t=pathlib.Path('flash/providers/runpod/train/deps.py').read_text(); m=re.search(r'BAKED_PER_SM_ARCHES\s*=\s*frozenset\(\{([^}]+)\}\)', t); print(' '.join(re.findall(r'sm\d+', m.group(1))))")
+          relayer=(); gpu=()
+          summary="| arch | published fp_cache | published fp_base | action |\n|---|---|---|---|\n"
+          for sm in $arches; do
+            img="ghcr.io/freesolo-co/flash-worker:cu128-$sm"
+            if cfg=$(crane config "$img" 2>/dev/null); then
+              pc=$(echo "$cfg" | jq -r '.config.Labels["co.freesolo.flash.kernelcache.fp_cache"] // ""')
+              pb=$(echo "$cfg" | jq -r '.config.Labels["co.freesolo.flash.kernelcache.fp_base"] // ""')
+            else
+              pc="(missing)"; pb="(missing)"
+            fi
+            if [ "$pc" != "$FP_CACHE" ]; then
+              gpu+=("$sm"); action="GPU re-warm (cache invalid)"
+            elif [ "$pb" != "$FP_BASE" ]; then
+              relayer+=("$sm"); action="re-layer (free)"
+            else
+              action="up to date"
+            fi
+            echo "$sm: pc=$pc pb=$pb -> $action"
+            summary="$summary| $sm | \`$pc\` | \`$pb\` | $action |\n"
+          done
+          if [ ${#relayer[@]} -gt 0 ]; then
+            relayer_json=$(printf '%s\n' "${relayer[@]}" | jq -R . | jq -cs .)
+          else
+            relayer_json='[]'
+          fi
+          gpu_csv=$(IFS=,; echo "${gpu[*]}")
+          echo "relayer_sms=$relayer_json" >> "$GITHUB_OUTPUT"
+          echo "gpu_sms=$gpu_csv" >> "$GITHUB_OUTPUT"
+          {
+            echo "### auto-rebake plan (base rev \`$REV\`, fp_cache \`$FP_CACHE\` fp_base \`$FP_BASE\`)"
+            echo ""
+            echo -e "$summary"
+            echo ""
+            echo "re-layer (free): \`$relayer_json\` | GPU re-warm (paid, auto-started): \`${gpu_csv:-none}\`"
+          } >> "$GITHUB_STEP_SUMMARY"
+
+      # FYI notification (not an action item -- the bake auto-runs). Keeps a heads-up + cost trail of
+      # which arches a paid GPU re-warm fired for.
+      - name: Notify GPU re-bake auto-started
+        if: steps.base.outputs.proceed == 'true' && steps.classify.outputs.gpu_sms != ''
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          GPU_SMS: ${{ steps.classify.outputs.gpu_sms }}
+          REV: ${{ steps.base.outputs.revision }}
+          RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}
+        run: |
+          title="kernel-cache: per-arch GPU re-bake auto-started"
+          body=$(printf 'The kernel toolchain changed (fp_cache), so a PAID GPU re-warm auto-started for these per-arch images (no approval needed):\n\n- arches: `%s`\n- base rev: `%s`\n\nTrack it in the run: %s' "$GPU_SMS" "$REV" "$RUN_URL")
+          # best-effort: a labeling/permission hiccup must never fail the gate. self-heal the label so
+          # no manual setup is needed, then reuse one open issue (comment) instead of spamming new ones.
+          gh label create kernel-rebake-pending --color B60205 --description "per-arch kernel cache GPU re-bake auto-started" 2>/dev/null || true
+          existing=$(gh issue list --label kernel-rebake-pending --state open --json number -q '.[0].number' 2>/dev/null || true)
+          if [ -n "${existing:-}" ]; then
+            gh issue comment "$existing" --body "$body" || true
+          else
+            gh issue create --title "$title" --body "$body" --label kernel-rebake-pending \
+              || gh issue create --title "$title" --body "$body" || true
+          fi
+
+  # FREE path: graft the existing valid cache onto the fresh base, no GPU. Runs automatically.
+  relayer:
+    needs: gate
+    if: needs.gate.outputs.proceed == 'true' && needs.gate.outputs.relayer_sms != '[]' && needs.gate.outputs.relayer_sms != ''
+    runs-on: ubuntu-24.04-8core
+    strategy:
+      fail-fast: false
+      matrix:
+        sm: ${{ fromJson(needs.gate.outputs.relayer_sms || '[]') }}
+    env:
+      BASE_REF: ${{ needs.gate.outputs.base_ref }}
+      FP_CACHE: ${{ needs.gate.outputs.fp_cache }}
+      FP_BASE: ${{ needs.gate.outputs.fp_base }}
+      BASE_REVISION: ${{ needs.gate.outputs.base_revision }}
+    steps:
+      - uses: actions/checkout@v6
+        with:
+          ref: ${{ github.event.workflow_run.head_sha }}
+
+      # The graft pulls two ~20GB images (old per-arch tag + fresh base); reclaim runner disk first.
+      - name: Free disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet /opt/ghc /usr/local/.ghcup \
+            /usr/local/lib/android /opt/hostedtoolcache /usr/share/swift \
+            /usr/local/share/boost /usr/local/lib/node_modules \
+            /usr/local/share/powershell /usr/share/miniconda || true
+          sudo apt-get clean || true
+          sudo docker image prune -af || true
+          df -h /
+
+      - uses: docker/setup-buildx-action@v3
+
+      - name: Log in to GHCR
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      # Pin OLDCACHE by digest so a concurrent overwrite of the source tag can't graft a wrong cache.
+      - name: Resolve old per-arch cache digest
+        id: old
+        run: |
+          img="ghcr.io/freesolo-co/flash-worker:cu128-${{ matrix.sm }}"
+          digest=$(docker buildx imagetools inspect "$img" --format '{{.Manifest.Digest}}')
+          # canonical digest ref (repo@sha256, no tag) for the FROM in the relayer build.
+          echo "ref=${img%:*}@$digest" >> "$GITHUB_OUTPUT"
+
+      - name: Re-layer cache onto the fresh base (no GPU)
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          file: docker/Dockerfile.kernelcache.relayer
+          push: true
+          tags: ghcr.io/freesolo-co/flash-worker:cu128-${{ matrix.sm }}
+          provenance: false
+          build-args: |
+            BASE=${{ env.BASE_REF }}
+            OLDCACHE=${{ steps.old.outputs.ref }}
+            FP_CACHE=${{ env.FP_CACHE }}
+            FP_BASE=${{ env.FP_BASE }}
+            BASE_REVISION=${{ env.BASE_REVISION }}
+
+  # PAID path: re-warm the cache on a real GPU per arch. Runs automatically -- the fingerprint gate
+  # above is the only thing that decides it fires (only when fp_cache changed), so no manual approval.
+  gpu-warm:
+    needs: gate
+    if: needs.gate.outputs.proceed == 'true' && needs.gate.outputs.gpu_sms != ''
+    uses: ./.github/workflows/bake-kernel-cache.yml
+    with:
+      sms: ${{ needs.gate.outputs.gpu_sms }}
+      fp_cache: ${{ needs.gate.outputs.fp_cache }}
+      fp_base: ${{ needs.gate.outputs.fp_base }}
+      base_ref: ${{ needs.gate.outputs.base_ref }}
+      base_revision: ${{ needs.gate.outputs.base_revision }}
+    secrets: inherit

freesolo_flash_dev-0.2.28/.github/workflows/bake-kernel-cache.yml

@@ -0,0 +1,245 @@
+name: flash worker kernel-cache (per-arch)
+
+# Builds per-SM worker images with the compiled-kernel mega-cache baked in, killing the ~10-15 min
+# first-use JIT on a cold worker (the regression #194 reintroduced). Each matrix leg:
+#   1. offloads the warmup to a RunPod GPU of that arch (GH runners have none) -> build/kernel_cache,
+#   2. docker build --build-arg BUILD_KERNEL_CACHE=true -> ghcr.io/freesolo-co/flash-worker:cu128-<sm>,
+#   3. pushes it. Activate by setting FLASH_WORKER_IMAGE_PER_SM=1 on the control plane so the
+#      allocator selects the cu128-<sm> tag per GPU class (see flash/providers/runpod/train/deps.py).
+#
+# 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: RUNPOD_API_KEY (the warmup GPU pod) + HF_TOKEN (the temp dataset that ferries code + the
+# cache artifact) are pulled from Infisical (/flash, prod) via the bootstrap INFISICAL_CLIENT_ID /
+# CLIENT_SECRET / PROJECT_ID repo secrets; 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.
+
+on:
+  workflow_dispatch:
+    inputs:
+      sms:
+        description: "comma-separated sm list to bake; default is all validated arches"
+        default: "sm80,sm86,sm89,sm90,sm120"
+        required: false
+  # auto-rebake.yml calls this for the EXPENSIVE path (the cache is invalid and must be re-warmed on a
+  # real GPU). It passes the arch subset, the fingerprints to stamp, and a digest-pinned base so the
+  # warm runs inside the exact :cu128 the per-arch tag will sit on.
+  workflow_call:
+    inputs:
+      sms:
+        type: string
+        required: true
+      fp_cache:
+        type: string
+        required: false
+        default: ""
+      fp_base:
+        type: string
+        required: false
+        default: ""
+      base_ref:
+        type: string
+        required: false
+        default: "ghcr.io/freesolo-co/flash-worker:cu128"
+      base_revision:
+        type: string
+        required: false
+        default: ""
+
+permissions:
+  contents: read
+  packages: write
+
+jobs:
+  # Ungated: the PAID RunPod bake runs as soon as it's called/dispatched. The decision of WHEN to run
+  # lives in auto-rebake.yml's fingerprint gate (it only invokes this when fp_cache changed), so a
+  # second manual approval here would be redundant. Re-add an `approve` job with `environment: gpu-bake`
+  # (and `needs: approve`) if you ever want a human checkpoint back.
+  bake:
+    runs-on: ubuntu-24.04-8core
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - {
+              sm: sm80,
+              arch: "8.0",
+              gpu_type_id: "NVIDIA A100 80GB PCIe",
+              allowed_cuda: "",
+            }
+          - {
+              sm: sm86,
+              arch: "8.6",
+              gpu_type_id: "NVIDIA RTX A6000",
+              allowed_cuda: "",
+            }
+          - {
+              sm: sm89,
+              arch: "8.9",
+              gpu_type_id: "NVIDIA GeForce RTX 4090",
+              allowed_cuda: "",
+            }
+          - {
+              sm: sm90,
+              arch: "9.0",
+              gpu_type_id: "NVIDIA H100 80GB HBM3",
+              allowed_cuda: "",
+            }
+          # Blackwell needs CUDA-13 hosts to JIT its PTX (matches min_cuda_for in the provider). Bake
+          # sm120 on the RTX Pro 6000 (Server Edition): it's the same sm120 (cache is sm-keyed, so it's
+          # valid for RTX 5090 too) but has far better secure-cloud on-demand capacity -- the RTX 5090
+          # pool repeatedly returned "machine does not have the resources" at create_pod.
+          - {
+              sm: sm120,
+              arch: "12.0",
+              gpu_type_id: "NVIDIA RTX PRO 6000 Blackwell Server Edition",
+              allowed_cuda: "13.0",
+            }
+          # Datacenter Blackwell (sm100, distinct from the sm120 RTX Pro 6000). B200 has 180 GB and
+          # good secure-cloud capacity. CUDA-13 host like the other Blackwell parts. NOT in the default
+          # `sms` list yet (B200 is unvalidated) -- bake it explicitly via `sms=sm100` until a smoke
+          # passes, then add it to the default.
+          - {
+              sm: sm100,
+              arch: "10.0",
+              gpu_type_id: "NVIDIA B200",
+              allowed_cuda: "13.0",
+            }
+    steps:
+      # Gate FIRST (no checkout needed) so skipped arches do no crane / checkout / build work.
+      - name: Should bake ${{ matrix.sm }}?
+        id: gate
+        run: |
+          req="${{ inputs.sms || 'sm80,sm86,sm89,sm90,sm120' }}"
+          req="${req// /}"   # tolerate spaces, e.g. "sm80, sm86"
+          case ",$req," in
+            *,${{ matrix.sm }},*) echo "run=true" >> "$GITHUB_OUTPUT" ;;
+            *) echo "run=false" >> "$GITHUB_OUTPUT"; echo "skipping ${{ matrix.sm }}" ;;
+          esac
+
+      - name: Install crane
+        if: steps.gate.outputs.run == 'true'
+        run: |
+          curl -sSL https://github.com/google/go-containerregistry/releases/download/v0.20.2/go-containerregistry_Linux_x86_64.tar.gz \
+            | sudo tar -xz -C /usr/local/bin crane
+          crane version
+
+      - name: Log in to GHCR
+        if: steps.gate.outputs.run == 'true'
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      # Resolve the base + fingerprints BEFORE checkout, so the checkout (hence the warmed flash code)
+      # is pinned to the EXACT commit the base was built from. auto-rebake passes everything (fast path);
+      # a manual dispatch READS the labels off the published base, so the per-arch image always matches
+      # the base it extends (same source of truth as the gate) instead of recomputing from the checkout.
+      - name: Resolve base + fingerprints
+        id: base
+        if: steps.gate.outputs.run == 'true'
+        run: |
+          if [ -n "${{ inputs.fp_cache }}" ]; then
+            echo "ref=${{ inputs.base_ref }}" >> "$GITHUB_OUTPUT"
+            echo "fp_cache=${{ inputs.fp_cache }}" >> "$GITHUB_OUTPUT"
+            echo "fp_base=${{ inputs.fp_base }}" >> "$GITHUB_OUTPUT"
+            echo "revision=${{ inputs.base_revision }}" >> "$GITHUB_OUTPUT"
+          else
+            base="${{ inputs.base_ref || 'ghcr.io/freesolo-co/flash-worker:cu128' }}"
+            digest=$(crane digest "$base"); repo=${base%:*}
+            cfg=$(crane config "$repo@$digest")
+            echo "ref=$repo@$digest" >> "$GITHUB_OUTPUT"
+            echo "fp_cache=$(echo "$cfg" | jq -r '.config.Labels["co.freesolo.flash.kernelcache.fp_cache"] // ""')" >> "$GITHUB_OUTPUT"
+            echo "fp_base=$(echo "$cfg" | jq -r '.config.Labels["co.freesolo.flash.kernelcache.fp_base"] // ""')" >> "$GITHUB_OUTPUT"
+            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@v6
+        if: steps.gate.outputs.run == 'true'
+        with:
+          ref: ${{ steps.base.outputs.revision || github.sha }}
+
+      - name: Install uv
+        if: steps.gate.outputs.run == 'true'
+        uses: astral-sh/setup-uv@v5
+        with:
+          # no cache: the post-run cache-prune calls `uv`, but the "Free disk space" step deletes
+          # /opt/hostedtoolcache (where setup-uv put uv), so the post step fails ("uv not found") and
+          # marks the whole job red even though the bake + push already succeeded.
+          enable-cache: false
+
+      - name: Sync deps (flash + runpod + hf)
+        if: steps.gate.outputs.run == 'true'
+        run: uv sync --extra server
+
+      - name: Install Infisical CLI
+        if: steps.gate.outputs.run == 'true'
+        run: |
+          curl -1sLf 'https://dl.cloudsmith.io/public/infisical/infisical-cli/setup.deb.sh' | sudo -E bash
+          sudo apt-get update && sudo apt-get install -y infisical
+
+      # RUNPOD_API_KEY + HF_TOKEN come from Infisical (/flash, prod) -- mirrors infisical-entrypoint.sh;
+      # the CLI takes the project UUID directly (the secrets-action wants a slug). `env -u
+      # FLASH_CHALK_SPEC` drops any chalk-spec override living in /flash so the bake always warms the
+      # repo-default chalk (the value the fingerprint hashed) -- otherwise the image could be warmed with
+      # one chalk and labeled with another. 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 RunPod ${{ matrix.sm }} GPU -> build/kernel_cache
+        if: steps.gate.outputs.run == 'true'
+        env:
+          INFISICAL_CLIENT_ID: ${{ secrets.INFISICAL_CLIENT_ID }}
+          INFISICAL_CLIENT_SECRET: ${{ secrets.INFISICAL_CLIENT_SECRET }}
+          INFISICAL_PROJECT_ID: ${{ secrets.INFISICAL_PROJECT_ID }}
+          BASE_IMAGE: ${{ steps.base.outputs.ref }}
+        run: |
+          INFISICAL_TOKEN="$(infisical login --method=universal-auth \
+            --client-id="$INFISICAL_CLIENT_ID" --client-secret="$INFISICAL_CLIENT_SECRET" \
+            --silent --plain)"
+          export INFISICAL_TOKEN
+          infisical run --projectId "$INFISICAL_PROJECT_ID" --env prod --path /flash -- \
+            env -u FLASH_CHALK_SPEC \
+            uv run python docker/bake_kernel_cache.py \
+              --arch "${{ matrix.arch }}" --sm "${{ matrix.sm }}" \
+              --gpu-type-id "${{ matrix.gpu_type_id }}" \
+              --allowed-cuda "${{ matrix.allowed_cuda }}" \
+              --image "$BASE_IMAGE" \
+              --out build/kernel_cache
+
+      - name: Free disk space
+        if: steps.gate.outputs.run == 'true'
+        run: |
+          sudo rm -rf /usr/share/dotnet /opt/ghc /usr/local/.ghcup \
+            /usr/local/lib/android /opt/hostedtoolcache /usr/share/swift \
+            /usr/local/share/boost /usr/local/lib/node_modules \
+            /usr/local/share/powershell /usr/share/miniconda || true
+          sudo apt-get clean || true
+          sudo docker image prune -af || true
+          df -h /
+
+      - uses: docker/setup-buildx-action@v3
+        if: steps.gate.outputs.run == 'true'
+
+      # thin layer FROM the SAME base the warmup ran inside -> the cache toolchain always matches the
+      # shipped image (rebuilding the full Dockerfile.worker from the checkout could drift from the
+      # published base and bake a cache that silently won't load). GHCR auth from the login step above.
+      - name: Build + push baked per-sm image (thin layer on the warmed base)
+        if: steps.gate.outputs.run == 'true'
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          file: docker/Dockerfile.kernelcache
+          push: true
+          tags: ghcr.io/freesolo-co/flash-worker:cu128-${{ matrix.sm }}
+          provenance: false
+          build-args: |
+            BASE=${{ steps.base.outputs.ref }}
+            FP_CACHE=${{ steps.base.outputs.fp_cache }}
+            FP_BASE=${{ steps.base.outputs.fp_base }}
+            BASE_REVISION=${{ steps.base.outputs.revision }}

{freesolo_flash_dev-0.2.26 → freesolo_flash_dev-0.2.28}/.github/workflows/worker-image.yml

@@ -36,17 +36,40 @@ on:
       # so handler changes (e.g. the weight-cache preload branch) MUST rebuild the image. flash/**
       # covers the handler source (endpoints.py); the generator lives outside flash/, so watch it too.
       - docker/make_rp_handler.py
+      # The fp_cache/fp_base labels are computed by this script at build time, so a change to the
+      # fingerprint logic must rebuild :cu128 to re-stamp them (else auto-rebake never reclassifies).
+      - docker/kernel_fingerprint.py
 
 permissions:
   contents: read
   packages: write
 
+# A newer base build supersedes an older one: cancel the in-flight run so auto-rebake (which triggers
+# on this workflow's completion) only ever reasons about the latest :cu128. The tag is in the group so
+# a non-prod `tag=cu128-test` dispatch can't cancel an in-flight prod `cu128` push on the same ref
+# (which would leave :cu128 unpublished for that commit and auto-rebake blind to it).
+concurrency:
+  group: worker-image-${{ github.ref }}-${{ github.event.inputs.tag || 'cu128' }}
+  cancel-in-progress: true
+
 jobs:
   build:
     runs-on: ubuntu-24.04-8core
+    # Single source for the FlashAttention wheels: used BOTH as build-args and as fingerprint inputs
+    # (so a wheel bump moves fp_base). Keep FA3_SPEC in sync with Dockerfile.worker's ARG default.
+    env:
+      FA2_SPEC: https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.0/flash_attn-2.8.3%2Bcu128torch2.10-cp312-cp312-linux_x86_64.whl
+      FA3_SPEC: ${{ github.event.inputs.flash_attn_3_spec || 'https://github.com/windreamer/flash-attention3-wheels/releases/download/2026.03.19-850211f/flash_attn_3-3.0.0%2B20260318.cu128torch2100cxx11abitrue.8afc61-cp39-abi3-linux_x86_64.whl' }}
     steps:
       - uses: actions/checkout@v6
 
+      # Compute the kernel-cache fingerprints (stdlib-only, no uv sync) and stamp them as labels on
+      # :cu128 below, so auto-rebake.yml can read what this base was built from. Fails the build if a
+      # pin can't be parsed -- a silent miss would let a stale per-arch cache ship as a cold-JIT.
+      - name: Compute kernel fingerprints
+        id: fp
+        run: python3 docker/kernel_fingerprint.py --format github --fa2-spec "$FA2_SPEC" --fa3-spec "$FA3_SPEC" >> "$GITHUB_OUTPUT"
+
       - name: Free disk space
         run: |
           sudo rm -rf /usr/share/dotnet /opt/ghc /usr/local/.ghcup \
@@ -77,21 +100,20 @@ jobs:
           # so a fork run must not push to its own owner namespace and leave the real image stale.
           tags: ghcr.io/freesolo-co/flash-worker:${{ github.event.inputs.tag || 'cu128' }}
           provenance: false
+          # Stamp the kernel-cache fingerprints + the source commit so auto-rebake.yml can read what
+          # this base was built from (docker/kernel_fingerprint.py is the source of truth).
+          labels: |
+            co.freesolo.flash.kernelcache.fp_cache=${{ steps.fp.outputs.fp_cache }}
+            co.freesolo.flash.kernelcache.fp_base=${{ steps.fp.outputs.fp_base }}
+            org.opencontainers.image.revision=${{ github.sha }}
           # Install a PREBUILT flash-attn wheel instead of compiling from source (a multi-arch
           # source compile is huge and OOM-kills the runner). The wheel matches the base exactly
-          # (cu128 / torch2.10 / cp312).
+          # (cu128 / torch2.10 / cp312). FA specs come from the job-level env (single source).
           build-args: |
-            FLASH_ATTN_SPEC=https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.0/flash_attn-2.8.3%2Bcu128torch2.10-cp312-cp312-linux_x86_64.whl
-            FLASH_ATTN_3_SPEC=${{ github.event.inputs.flash_attn_3_spec || 'https://github.com/windreamer/flash-attention3-wheels/releases/download/2026.03.19-850211f/flash_attn_3-3.0.0%2B20260318.cu128torch2100cxx11abitrue.8afc61-cp39-abi3-linux_x86_64.whl' }}
+            FLASH_ATTN_SPEC=${{ env.FA2_SPEC }}
+            FLASH_ATTN_3_SPEC=${{ env.FA3_SPEC }}
 
-      # The per-arch baked images (cu128-<sm>) are rebaked MANUALLY (bake-kernel-cache.yml). When this
-      # rebuilds :cu128, those tags go stale -> prompt a rebake. Only when the cu128 tag was (re)built.
-      - name: Remind to rebake per-SM kernel-cache images
-        if: ${{ success() && (github.event.inputs.tag || 'cu128') == 'cu128' }}
-        run: |
-          echo "::warning title=Rebake per-SM images::A new :cu128 base was published. If FLASH_WORKER_IMAGE_PER_SM=1 is enabled, run bake-kernel-cache.yml to refresh cu128-sm{80,86,89,90,120}, else workers run stale baked deps."
-          {
-            echo "### ⚠️ Per-SM kernel-cache images may now be stale"
-            echo ""
-            echo "A fresh \`:cu128\` base was just published. If \`FLASH_WORKER_IMAGE_PER_SM=1\` is enabled, re-run **bake-kernel-cache.yml** to refresh the per-arch \`cu128-sm*\` tags, or workers will run the previous baked deps."
-          } >> "$GITHUB_STEP_SUMMARY"
+      # When :cu128 is republished the per-arch cu128-sm* tags go stale. auto-rebake.yml now handles
+      # that automatically on this workflow's completion: a free re-layer when only the base moved, or
+      # a one-click GPU re-warm when the kernel toolchain changed (it reads the fp_cache/fp_base labels
+      # stamped above vs each published per-arch tag). The old manual reminder is gone.

{freesolo_flash_dev-0.2.26 → freesolo_flash_dev-0.2.28}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 0.2.26
+Version: 0.2.28
 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-0.2.26 → freesolo_flash_dev-0.2.28}/docker/Dockerfile.kernelcache

@@ -16,3 +16,13 @@ ENV TRITON_CACHE_DIR=/opt/flash/kernelcache/triton \
     TORCHINDUCTOR_CACHE_DIR=/opt/flash/kernelcache/inductor
 
 COPY build/kernel_cache/ /opt/flash/kernelcache/
+
+# Fingerprint + base-revision stamps so auto-rebake.yml can read what this image was baked from and
+# decide skip / re-layer / re-warm (docker/kernel_fingerprint.py is the source of truth for the
+# values). Default to "unknown" so a hand build still works; CI always passes the real ones.
+ARG FP_CACHE=unknown
+ARG FP_BASE=unknown
+ARG BASE_REVISION=unknown
+LABEL co.freesolo.flash.kernelcache.fp_cache=$FP_CACHE \
+      co.freesolo.flash.kernelcache.fp_base=$FP_BASE \
+      org.opencontainers.image.revision=$BASE_REVISION

freesolo_flash_dev-0.2.28/docker/Dockerfile.kernelcache.relayer

@@ -0,0 +1,27 @@
+# Cheap per-arch RE-LAYER: graft the EXISTING baked mega-cache onto a freshly rebuilt :cu128 base,
+# with NO GPU. Used by auto-rebake.yml when fp_cache is unchanged (the cache is still valid) but the
+# base moved (fp_base changed) -- e.g. an FA2/FA3/causal-conv1d/rp_handler bump. The valid cache
+# already lives inside the currently-published per-arch tag, so OLDCACHE IS the cache store; we just
+# re-stack it on the new base. This is what makes a base-only change free (vs a paid GPU re-warm).
+#
+# BASE and OLDCACHE are passed DIGEST-PINNED by the gate so a concurrent base/tag overwrite can't
+# graft the cache onto the wrong base.
+ARG BASE=ghcr.io/freesolo-co/flash-worker:cu128
+ARG OLDCACHE=ghcr.io/freesolo-co/flash-worker:cu128-sm89
+FROM ${OLDCACHE} AS oldcache
+FROM ${BASE}
+
+# Match Dockerfile.worker / Dockerfile.kernelcache so the layer is self-contained even on an older base.
+ENV TRITON_CACHE_DIR=/opt/flash/kernelcache/triton \
+    TORCHINDUCTOR_CACHE_DIR=/opt/flash/kernelcache/inductor
+
+# Pull the warmed cache straight out of the old per-arch image -- no GPU, no external cache store.
+COPY --from=oldcache /opt/flash/kernelcache /opt/flash/kernelcache
+
+# Re-stamp with the CURRENT fingerprints (fp_cache unchanged, fp_base + revision are the new base's).
+ARG FP_CACHE=unknown
+ARG FP_BASE=unknown
+ARG BASE_REVISION=unknown
+LABEL co.freesolo.flash.kernelcache.fp_cache=$FP_CACHE \
+      co.freesolo.flash.kernelcache.fp_base=$FP_BASE \
+      org.opencontainers.image.revision=$BASE_REVISION