Address reviewer findings on PR #5748: 4 release-path bugs

Round-trip with the reviewer.py 12-persona pass surfaced four real
issues. Fix all four in this PR so the new Docker release path is
self-consistent.

1. docker/smoke_test.py used `import xformers` unconditionally, which
   guarantees a failure on arm64 (built with `[huggingface]` extras to
   skip xformers since it has no aarch64 cu128 wheel). Wrap the import
   in try/except so the same smoke script validates both arches.

2. unsloth/_gpu_init.py forced `TORCHINDUCTOR_COMPILE_THREADS=1` before
   `import unsloth_zoo`, but `patch_torch_compile` in unsloth_zoo main
   pops that env var in non-debug mode. After unsloth_zoo init the
   guard was effectively undone, so cgroup-pinned `docker --gpus
   '"device=N"'` containers still spawned the Inductor subprocess pool
   that cannot enumerate the GPU. Set `torch._inductor.config.
   compile_threads = 1` directly post-import-torch and re-populate the
   env var so `determine_compile_threads()` in the zoo options dict
   also returns 1, regardless of whether the zoo-side fix from PR #694
   has shipped yet.

3. docker-publish.yml UNSLOTH_REF build-arg defaulted to `'main'` for
   tag pushes and scheduled runs, so a `v1.2.3` release image would
   contain whatever `main` happened to be at build time, not v1.2.3.
   Pick the tag's `github.ref_name` for tag events and `github.sha`
   for branch/schedule events.

4. The smoke-test job pulled `:latest` regardless of which tag the
   merge job had just published, so tag/schedule/sha publishes were
   never actually validated. Re-run docker/metadata-action with the
   same config the merge job used, then smoke-test the first tag from
   its output.

All four changes are gated and backwards-compatible.

.github/workflows/docker-publish.yml

@@ -119,7 +119,12 @@ jobs:
             CUDA_VERSION=12.8.1
             UBUNTU_VERSION=24.04
             PYTHON_VERSION=3.12
-            UNSLOTH_REF=${{ github.event.inputs.unsloth_ref || 'main' }}
+            # Workflow-dispatch: honour the explicit input. Tag pushes:
+            # bake the tag's source ref (e.g. v1.2.3) so the published
+            # tag image actually contains that release. Branch pushes and
+            # scheduled runs: bake the triggering commit SHA. Falls back
+            # to `main` for any other event class.
+            UNSLOTH_REF=${{ github.event.inputs.unsloth_ref || (startsWith(github.ref, 'refs/tags/') && github.ref_name) || github.sha || 'main' }}
             UNSLOTH_ZOO_REF=${{ github.event.inputs.unsloth_zoo_ref || 'main' }}
 
       # Stash the per-arch digest as an artifact for the merge job to pick up.
@@ -203,9 +208,30 @@ jobs:
     timeout-minutes: 20
     steps:
       - uses: actions/checkout@v4
+
+      # Re-compute the tag list deterministically from the same metadata-action
+      # config the merge job used, so tag/schedule/SHA runs pull the image
+      # they just published instead of an unrelated `:latest` from a prior run.
+      - name: Resolve published tag
+        id: meta
+        uses: docker/metadata-action@v5
+        with:
+          images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
+          tags: |
+            type=raw,value=latest,enable={{is_default_branch}}
+            type=ref,event=tag
+            type=schedule,pattern=nightly
+            type=sha,prefix=sha-,format=short
+
       - name: Pull and smoke-test
         run: |
-          docker pull ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:latest
-          docker run --rm --gpus all \
-            ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:latest \
-            python /workspace/smoke_test.py
+          # Use the first tag from the metadata output -- that is the image we
+          # just published. Falls back to :latest only when the metadata is
+          # empty (defensive; should not happen on default-branch runs).
+          TAG="$(jq -r '.tags[0] // ""' <<<"$DOCKER_METADATA_OUTPUT_JSON")"
+          if [ -z "$TAG" ]; then
+              TAG="${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:latest"
+          fi
+          echo "smoke-testing $TAG"
+          docker pull "$TAG"
+          docker run --rm --gpus all "$TAG" python /workspace/smoke_test.py

async_task_outputs/async_task_output_1rtbzl.md

@@ -0,0 +1 @@
+- Docker GPU-free build rationale done

async_task_outputs/async_task_output_1ud2z5.md

@@ -0,0 +1 @@
+- Docker ok; run `bash docker/test_locally.sh --skip-notebook`

async_task_outputs/async_task_output_1vzwo0.md

@@ -0,0 +1 @@
+- PR `#5748` pushed commit `58693c4c`

async_task_outputs/async_task_output_215bp0.md

@@ -0,0 +1 @@
+- `externally-managed-environment`

async_task_outputs/async_task_output_31h69u.md

@@ -0,0 +1 @@
+- PR `#5748`: added/pushed script.

async_task_outputs/async_task_output_3aetsv.md

@@ -0,0 +1 @@
+- Buildx required; edited/pushed `56d2701a`

async_task_outputs/async_task_output_3k07zq.md

@@ -0,0 +1 @@
+- PR `unslothai/unsloth`: `sm_120`

async_task_outputs/async_task_output_3xv17f.md

@@ -0,0 +1 @@
+- Build OK; smoke `ImportError`

async_task_outputs/async_task_output_5etrdc.md

@@ -0,0 +1 @@
+- Pushed `f7b34793`; pending retry.

async_task_outputs/async_task_output_5ih1fr.md

@@ -0,0 +1 @@
+- Pushed `23a5b431`; pending retry.

async_task_outputs/async_task_output_778ozt.md

@@ -0,0 +1 @@
+- `Failed to find C compiler`; exit 1

async_task_outputs/async_task_output_8cu1bq.md

@@ -0,0 +1 @@
+- Asked: `test on other machines`

async_task_outputs/async_task_output_8rdvq1.md

@@ -0,0 +1 @@
+- Asked: `gzip` multiprocessing? slow

async_task_outputs/async_task_output_924viw.md

@@ -0,0 +1 @@
+- Asked: `So what did you do to make the docker work?`

async_task_outputs/async_task_output_98ywdm.md

@@ -0,0 +1 @@
+- Wants executable `sh`-like script

async_task_outputs/async_task_output_9gyz0z.md

@@ -0,0 +1 @@
+- Chose apt `docker-buildx`; run `bash docker/test_locally.sh --skip-notebook`

async_task_outputs/async_task_output_9stc7n.md

@@ -0,0 +1 @@
+- Recommended `pigz`; gzip-compatible.

async_task_outputs/async_task_output_a6uutw.md

@@ -0,0 +1 @@
+- Ran `docker info`; Docker works

async_task_outputs/async_task_output_a72mjm.md

@@ -0,0 +1,2 @@
+- `dde5170e` pushed
+- GPU tests pending

async_task_outputs/async_task_output_af1yr3.md

@@ -0,0 +1 @@
+- Asked why RTX 6000/5090 needed

async_task_outputs/async_task_output_b6xdx8.md

@@ -0,0 +1 @@
+- `docker build --progress` unsupported; build failed exit `125`

async_task_outputs/async_task_output_c76tz2.md

@@ -0,0 +1 @@
+- Fixed Triton JIT via `1cdc5f17`

async_task_outputs/async_task_output_d9vccx.md

@@ -0,0 +1 @@
+- Pending: upload `/tmp/unsloth-blackwell.tar.gz`

async_task_outputs/async_task_output_dmy1yd.md

@@ -0,0 +1 @@
+- Gave docker build/run commands; pending output

async_task_outputs/async_task_output_e7m19a.md

@@ -0,0 +1 @@
+- Asked `docker` GPU for `unsloth`

async_task_outputs/async_task_output_eu96wu.md

@@ -0,0 +1 @@
+- `56d2701a`: buildx fix OK; done

async_task_outputs/async_task_output_f3n8gk.md

@@ -0,0 +1 @@
+- Asked `So what should we try next`

async_task_outputs/async_task_output_hfha2k.md

@@ -0,0 +1 @@
+- Build fails: `docker buildx` missing

async_task_outputs/async_task_output_hni3eq.md

@@ -0,0 +1 @@
+- fixed Dockerfile PEP668; pushed `fd55ed0a`

async_task_outputs/async_task_output_j4howx.md

@@ -0,0 +1 @@
+- `latest docker`; vllm web search

async_task_outputs/async_task_output_j5k6f1.md

@@ -0,0 +1 @@
+- `Ok you upload to danielhanchen`

async_task_outputs/async_task_output_jogf0i.md

@@ -0,0 +1 @@
+- pushed `4bfb4b89` add `docker/hf_*.sh`

async_task_outputs/async_task_output_k4wy4g.md

@@ -0,0 +1 @@
+- HF repo made; upload tar pending.

async_task_outputs/async_task_output_lls2b3.md

@@ -0,0 +1 @@
+Continue the user’s latest directive: make the Docker setup work for DGX Spark via emulation. The WebSearch tool for Docker buildx multi-platform/QEMU CUDA build context is already done, so do not re-run it. Next, inspect the search result context and continue researching current linux/arm64 CUDA/PyTorch/xformers/bitsandbytes/flash-attn wheel availability, then update the Docker/build/test/publish files as needed for a build-time QEMU arm64 path.

async_task_outputs/async_task_output_ltr06m.md

@@ -0,0 +1 @@
+- Need support SM 7.5..12.1

async_task_outputs/async_task_output_nx17q0.md

@@ -0,0 +1 @@
+- LoRA train OK; all blocks completed

async_task_outputs/async_task_output_oo1b5k.md

@@ -0,0 +1 @@
+- User said `Yes try it out`

async_task_outputs/async_task_output_ptu313.md

@@ -0,0 +1 @@
+- PR `#5748` opened; push used scoped token

async_task_outputs/async_task_output_q9q3pp.md

@@ -0,0 +1 @@
+- Docker push via `huggingface_hub`

async_task_outputs/async_task_output_qe4w44.md

@@ -0,0 +1 @@
+- Ran `docker save`; completed

async_task_outputs/async_task_output_qp3v3e.md

@@ -0,0 +1 @@
+- Docker build needn't GPU; sm_120 only test

async_task_outputs/async_task_output_r3h3hd.md

@@ -0,0 +1 @@
+Continue the latest directive: make the Docker image work for DGX Spark via build-time emulation, targeting a proper linux/arm64 image rather than CUDA runtime emulation. The assistant had started checking current multi-platform/QEMU and arm64 CUDA/PyTorch wheel availability; one WebSearch result for Docker buildx multi-platform builds has returned, so do not re-run that completed search. Next action is to continue the remaining live research from the in-flight WebSearch context, especially whether PyTorch/CUDA, xformers, bitsandbytes, flash-attn, and Unsloth dependencies have usable arm64 wheels or need conditional Dockerfile handling.

async_task_outputs/async_task_output_rccvqp.md

@@ -0,0 +1 @@
+- PR #5748 smoke passed; comment posted.

async_task_outputs/async_task_output_re5vqk.md

@@ -0,0 +1 @@
+- Validated; fixed `Dockerfile`; pending none

async_task_outputs/async_task_output_s3bkm5.md

@@ -0,0 +1 @@
+- PENDING: test Docker notebook

async_task_outputs/async_task_output_tqe3ko.md

@@ -0,0 +1 @@
+- Pushed `00cbc825`; retry `bash docker/test_locally.sh --skip-notebook`

async_task_outputs/async_task_output_uudy1y.md

@@ -0,0 +1 @@
+- `--skip-notebook`; Docker smoke OK

async_task_outputs/async_task_output_wcgi8q.md

@@ -0,0 +1 @@
+- Asked if Blackwell image supports GPUs

async_task_outputs/async_task_output_xgyhvf.md

@@ -0,0 +1 @@
+- Installed `docker-buildx`; build failed Docker socket perm

async_task_outputs/async_task_output_yt3d7j.md

@@ -0,0 +1 @@
+- Answered GPU support matrix; DGX pending

async_task_outputs/async_task_output_zsqsui.md

@@ -0,0 +1 @@
+- B200 OK; pending HF upload cmds

docker/smoke_test.py

@@ -63,9 +63,15 @@ def check_imports() -> None:
     import unsloth_zoo
 
     print(f"unsloth_zoo {unsloth_zoo.__version__}")
-    import xformers
-
-    print(f"xformers    {xformers.__version__}")
+    # xformers is not built for aarch64 cu128 as of this writing; the arm64
+    # variant of this image installs unsloth with `[huggingface]` extras
+    # which omits it. Treat the import as best-effort so the same script
+    # smoke-tests both arches.
+    try:
+        import xformers
+        print(f"xformers    {xformers.__version__}")
+    except ImportError:
+        print("xformers    (missing -- expected on arm64 [huggingface] extras)")
     import bitsandbytes as bnb
 
     print(f"bnb         {bnb.__version__}")

unsloth/_gpu_init.py

@@ -138,6 +138,21 @@ except ModuleNotFoundError:
 except:
     raise
 
+# Re-assert the single-compile-worker policy after unsloth_zoo has had a
+# chance to run its patch_torch_compile (which historically popped
+# TORCHINDUCTOR_COMPILE_THREADS in non-debug mode). Force the Inductor
+# config value directly so the Docker --gpus '"device=N"' subprocess-pool
+# bug is fixed even when the installed unsloth_zoo predates the
+# corresponding zoo-side patch. No-op when the user opted out.
+if os.environ.get("UNSLOTH_FORCE_SINGLE_COMPILE_WORKER", "0") == "1":
+    try:
+        torch._inductor.config.compile_threads = 1
+    except Exception:
+        pass
+    # Re-populate the env var so determine_compile_threads in the zoo
+    # options dict also sees it; cheap and forward-compatible.
+    os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"
+
 from unsloth_zoo.device_type import (
     is_hip,
     get_device_type,