docker: address review round 4 (jupyter probe, CPU messaging, llama EXDEV, %pip shim)

- docker-publish smoke + docker_confirm.sh probe Jupyter /login, not /api: the
  launcher always configures a password hash so /api returns 403 and curl -f
  would never flip the health flag (false build failure).
- entrypoint.sh CPU messaging: CPU mode covers Jupyter, GGUF tooling and
  llama.cpp (GGUF) Studio chat; training AND loading an Unsloth model
  (FastLanguageModel) still need a GPU, since from_pretrained runs CUDA probes.
- install_llama_prebuilt.py: rollback/activation moves used bare os.replace,
  which fails with EXDEV across overlayfs in a Docker build and fell back to a
  broken source build (no nvcc). Add is_cross_device_error + move_install_dir_aside
  (os.replace fast path, copy+remove on EXDEV; busy errors still re-raise).
- notebooks: %pip / %uv line magics and the `!python -m pip` form bypassed the
  PATH pip/uv shim and could overwrite the baked cu128 torch/vLLM stack. Add
  unsloth_nb_pip_magic.py to re-point them at the shim, wired via the IPython
  startup hook and installed into the venv site-packages.

.github/workflows/docker-publish.yml

@@ -576,12 +576,15 @@ jobs:
           ok_studio=0; ok_jupyter=0
           for i in $(seq 1 60); do
               if curl -fsS http://localhost:18000/api/health >/dev/null 2>&1; then ok_studio=1;  fi
-              if curl -fsS http://localhost:18888/api        >/dev/null 2>&1; then ok_jupyter=1; fi
+              # Probe /login, not /api: the launcher always sets a Jupyter password
+              # hash, so /api returns 403 (curl -f would never flip ok_jupyter).
+              # /login is the unauthenticated page and 200s once the server is up.
+              if curl -fsS http://localhost:18888/login      >/dev/null 2>&1; then ok_jupyter=1; fi
               [ "$ok_studio" = 1 ] && [ "$ok_jupyter" = 1 ] && break
               sleep 5
           done
           [ "$ok_studio"  = 1 ] || { echo "Studio /api/health never went healthy"; exit 1; }
-          [ "$ok_jupyter" = 1 ] || { echo "Jupyter /api never responded"; exit 1; }
+          [ "$ok_jupyter" = 1 ] || { echo "Jupyter /login never responded"; exit 1; }
           echo "Studio + Jupyter healthy"
         env:
           STEPS_META_STUDIO_JSON: ${{ steps.meta_studio.outputs.json }}

docker/.dockerignore

@@ -10,6 +10,7 @@
 !unsloth_jupyter_tunnel.sh
 !unsloth_nb_compat.py
 !unsloth_pip_shim.py
+!unsloth_nb_pip_magic.py
 !unsloth_ipython_startup.py
 !unsloth_run.py
 !unsloth_sync_notebooks.sh

docker/Dockerfile

@@ -664,15 +664,19 @@ RUN mkdir -p ${HF_HOME} ${TRITON_CACHE_DIR}
 #     `!pip install ...` / `!uv pip install ...` cell becomes SAFE + idempotent
 #     (keeps the baked torch/vLLM stack; records the requested transformers so
 #     its sidecar is activated for the model cells).
+#   * unsloth_nb_pip_magic.py -> site-packages: re-points the IPython `%pip` /
+#     `%uv` line magics and the `!python -m pip` form at the same shim, so the
+#     in-process / module install paths cannot bypass PATH and clobber the stack.
 #   * IPython startup hook: activates the right sidecar before the first model
 #     cell in manual JupyterLab.
 #   * unsloth-run: headless `unsloth-run <notebook|url>` that auto-picks the
 #     sidecar and executes every cell -- the robust driven path.
 # ---------------------------------------------------------------------------
-COPY unsloth_nb_compat.py unsloth_pip_shim.py unsloth_ipython_startup.py unsloth_run.py unsloth_sync_notebooks.sh unsloth_nb_content_sig.py /opt/unsloth-nb/
+COPY unsloth_nb_compat.py unsloth_pip_shim.py unsloth_nb_pip_magic.py unsloth_ipython_startup.py unsloth_run.py unsloth_sync_notebooks.sh unsloth_nb_content_sig.py /opt/unsloth-nb/
 RUN set -eux \
  && SP=/opt/unsloth-venv/lib/python${PYTHON_VERSION}/site-packages \
  && cp /opt/unsloth-nb/unsloth_nb_compat.py "$SP/unsloth_nb_compat.py" \
+ && cp /opt/unsloth-nb/unsloth_nb_pip_magic.py "$SP/unsloth_nb_pip_magic.py" \
  && chmod +x /opt/unsloth-nb/unsloth_pip_shim.py /opt/unsloth-nb/unsloth_run.py /opt/unsloth-nb/unsloth_sync_notebooks.sh /opt/unsloth-nb/unsloth_nb_content_sig.py \
  && mkdir -p /opt/unsloth-nb/bin \
  && for t in pip pip3 uv; do ln -sf /opt/unsloth-nb/unsloth_pip_shim.py /opt/unsloth-nb/bin/$t; done \

docker/docker_confirm.sh

@@ -236,12 +236,13 @@ else
   ok_studio=0; ok_jupyter=0
   for _ in $(seq 1 60); do
     if [ "$ok_studio" = 0 ] && curl -fsS "http://localhost:$PORT_STUDIO/api/health" >/dev/null 2>&1; then ok_studio=1; fi
-    if [ "$ok_jupyter" = 0 ] && curl -fsS "http://localhost:$PORT_JUPYTER/api" >/dev/null 2>&1; then ok_jupyter=1; fi
+    # /login, not /api: a password hash is always configured so /api returns 403.
+    if [ "$ok_jupyter" = 0 ] && curl -fsS "http://localhost:$PORT_JUPYTER/login" >/dev/null 2>&1; then ok_jupyter=1; fi
     [ "$ok_studio" = 1 ] && [ "$ok_jupyter" = 1 ] && break
     sleep 5
   done
   [ "$ok_studio" = 1 ]  && ok "Studio /api/health healthy" || { bad "Studio /api/health never went healthy (docker logs ${STUDIO_CID:0:12})"; docker logs --tail 15 "$STUDIO_CID" 2>&1 | sed 's/^/         /'; }
-  [ "$ok_jupyter" = 1 ] && ok "JupyterLab /api responding"  || bad "JupyterLab /api never responded"
+  [ "$ok_jupyter" = 1 ] && ok "JupyterLab /login responding"  || bad "JupyterLab /login never responded"
 fi
 hr
 

docker/entrypoint.sh

@@ -45,16 +45,19 @@ warn() { printf "\033[1;33mWARN:\033[0m %s\n"  "$*" >&2; }
 
 # CPU mode for hosts that cannot pass a GPU into a Linux container at all:
 # Docker Desktop on macOS (no Metal passthrough), Docker Desktop on Windows
-# without WSL2 GPU support, plain CPU Linux boxes, and CI runners. Training
-# needs an NVIDIA GPU, but Jupyter, GGUF tooling (the baked llama.cpp), and
-# Studio chat / Data Recipes all work on CPU. With UNSLOTH_ALLOW_CPU=1 a
-# missing GPU degrades to a warning instead of the hard pre-flight failure;
-# when a GPU IS visible the normal checks below still run so a broken GPU
-# setup is not silently ignored.
+# without WSL2 GPU support, plain CPU Linux boxes, and CI runners. CPU mode
+# covers Jupyter, the GGUF tooling and llama.cpp-backed Studio chat (llama.cpp
+# runs on CPU), and Data Recipes. It does NOT cover training or loading an
+# Unsloth model for chat (FastLanguageModel.from_pretrained runs CUDA probes
+# like torch.cuda.get_device_properties and raises without a GPU). With
+# UNSLOTH_ALLOW_CPU=1 a missing GPU degrades to a warning instead of the hard
+# pre-flight failure; when a GPU IS visible the normal checks below still run so
+# a broken GPU setup is not silently ignored.
 if [[ "${UNSLOTH_ALLOW_CPU:-0}" == "1" ]]; then
     if ! command -v nvidia-smi >/dev/null 2>&1 || ! nvidia-smi -L 2>/dev/null | grep -q '^GPU'; then
         warn "UNSLOTH_ALLOW_CPU=1 and no GPU visible -- continuing on CPU."
-        warn "Training requires an NVIDIA GPU. CPU mode covers Jupyter, GGUF tooling and Studio chat."
+        warn "CPU mode covers Jupyter, GGUF tooling and llama.cpp (GGUF) Studio chat."
+        warn "Training and loading Unsloth models (FastLanguageModel) still require an NVIDIA GPU."
         sync_notebooks
         exec "$@"
     fi
@@ -91,8 +94,9 @@ Likely causes (in order of frequency):
        k8s:        nvidia.com/gpu resource request + GPU operator
 
   5. This host has no NVIDIA GPU at all (Docker Desktop on macOS, Windows
-     without WSL2 GPU support, CPU-only Linux). Training needs a GPU, but
-     Jupyter, GGUF tooling and Studio chat work on CPU:
+     without WSL2 GPU support, CPU-only Linux). Training and loading Unsloth
+     models need a GPU, but Jupyter, GGUF tooling and llama.cpp (GGUF) Studio
+     chat work on CPU:
        docker run -e UNSLOTH_ALLOW_CPU=1 ...
 
 To bypass this check entirely (e.g. offline tooling), set UNSLOTH_SKIP_GPU_CHECK=1.

docker/unsloth_ipython_startup.py

@@ -16,6 +16,13 @@ try:
     import unsloth_nb_compat
 
     unsloth_nb_compat.register_ipython()
+
+    # Re-point the %pip / %uv line magics and `!python -m pip` at the same shim,
+    # so the in-process / module install paths cannot bypass the PATH shim and
+    # overwrite the baked torch/vLLM stack. Independent of the sidecar hook.
+    import unsloth_nb_pip_magic
+
+    unsloth_nb_pip_magic.register_ipython()
 except Exception as _e:  # never break a kernel because of the helper
     import sys
     print(f"[unsloth-nb] startup hook skipped: {_e!r}", file = sys.stderr)

docker/unsloth_nb_pip_magic.py

@@ -0,0 +1,68 @@
+"""Route notebook `%pip` / `%uv` / `python -m pip` installs through the shim.
+
+The PATH shim (/opt/unsloth-nb/bin/{pip,pip3,uv} -> unsloth_pip_shim.py) only
+intercepts `!pip` / `!uv` shell cells. IPython's `%pip` / `%uv` LINE MAGICS run
+pip in-process, and `python -m pip` runs pip as a module -- both bypass PATH, so
+a notebook could still reinstall torch / transformers / vLLM and clobber the
+baked cu128 stack the shim is meant to protect.
+
+This closes that gap two ways, with no clobbering of the shell-escape path:
+  * `%pip` / `%pip3` / `%uv` are re-registered as line magics that delegate to
+    the shell (`get_ipython().system("pip ...")`); since /opt/unsloth-nb/bin is
+    first on PATH, that resolves to the shim. Overriding the real magic (rather
+    than rewriting cell text) means we only act when IPython actually dispatches
+    the magic -- a `%pip` inside a string is left untouched.
+  * a narrow input transformer rewrites an explicit `!python -m pip` /
+    `!python -m uv` shell line to `!pip` / `!uv`, so that form hits the shim too.
+
+UNSLOTH_NB_SHIM=1 is already exported by the startup hook and inherited by the
+subprocess, so the shim applies. Safe no-op outside IPython.
+"""
+
+import re
+
+# Only the explicit `!python -m pip|uv ...` shell form (the `!` makes it a shell
+# escape). Matched against the line with its trailing newline stripped.
+_PY_M_PIP = re.compile(r"^(\s*)!\s*(?:python[0-9.]*|py)\s+-m\s+(pip|uv)\b(.*)$")
+
+
+def _rewrite_python_dash_m(lines):
+    """`!python -m pip install X` -> `!pip install X` (so it hits the PATH shim)."""
+    try:
+        out = []
+        for line in lines:
+            body = line.rstrip("\n")
+            tail = line[len(body):]  # preserve the trailing newline(s), if any
+            m = _PY_M_PIP.match(body)
+            if m:
+                out.append(m.group(1) + "!" + m.group(2) + m.group(3) + tail)
+            else:
+                out.append(line)
+        return out
+    except Exception:
+        return lines
+
+
+def register_ipython():
+    try:
+        ip = get_ipython()  # noqa: F821 (provided by IPython)
+    except Exception:
+        ip = None
+    if ip is None or getattr(ip, "_unsloth_pip_magic", False):
+        return
+
+    def _make(tool):
+        def _magic(line):
+            # /opt/unsloth-nb/bin is first on PATH, so `pip`/`uv` here is the shim.
+            return ip.system(tool + " " + line)
+        return _magic
+
+    # Override the built-in %pip / %uv so they route through the shim too.
+    ip.register_magic_function(_make("pip"), "line", "pip")
+    ip.register_magic_function(_make("pip"), "line", "pip3")
+    ip.register_magic_function(_make("uv"), "line", "uv")
+
+    if _rewrite_python_dash_m not in ip.input_transformers_cleanup:
+        ip.input_transformers_cleanup.append(_rewrite_python_dash_m)
+
+    ip._unsloth_pip_magic = True

studio/install_llama_prebuilt.py

@@ -431,6 +431,17 @@ def _os_error_messages(exc: BaseException) -> list[str]:
     return [message.lower() for message in messages if message]
 
 
+def is_cross_device_error(exc: BaseException) -> bool:
+    """True for an EXDEV "cross-device link" rename failure.
+
+    os.replace / os.rename cannot move across filesystems -- e.g. inside a Docker
+    build where the staging tree and the install dir land on different overlayfs
+    layers (Errno 18). Unlike a busy/in-use error, a cross-device move is safely
+    completed by a copy + remove of the (idle) source.
+    """
+    return isinstance(exc, OSError) and exc.errno == errno.EXDEV
+
+
 def is_busy_lock_error(exc: BaseException) -> bool:
     if isinstance(exc, BusyInstallConflict):
         return True
@@ -4767,13 +4778,36 @@ def activate_staged_dir(staging_dir: Path, dst: Path) -> None:
     try:
         os.replace(staging_dir, dst)
     except OSError as exc:
-        if not is_busy_lock_error(exc):
+        # Busy/in-use (Windows AV holding a DLL) OR cross-device (overlayfs in a
+        # Docker build): both are safe to complete by copying the freshly
+        # extracted staging tree and removing it. Anything else (disk full,
+        # missing path) re-raises so we never leave a partial install behind.
+        if not (is_busy_lock_error(exc) or is_cross_device_error(exc)):
             raise
         log(f"os.replace failed ({exc!r}); falling back to file-by-file copy of staging tree")
         shutil.copytree(staging_dir, dst, dirs_exist_ok = True)
         remove_tree(staging_dir)
 
 
+def move_install_dir_aside(src: Path, dst: Path) -> None:
+    """Move an existing install dir to ``dst`` (a unique, non-existent sibling).
+
+    os.replace is the fast path. On a cross-device link (EXDEV -- e.g. moving the
+    base-image llama.cpp aside during a Docker studio build, where the rollback
+    path is on a different overlay) fall back to copy + remove. A busy/in-use
+    failure is deliberately NOT copy-faked here: the source is a live install and
+    a partial copy + rmtree would be worse than failing, so it re-raises.
+    """
+    try:
+        os.replace(src, dst)
+    except OSError as exc:
+        if not is_cross_device_error(exc):
+            raise
+        log(f"os.replace cross-device ({exc!r}); copy+remove {src} -> {dst}")
+        shutil.copytree(src, dst, dirs_exist_ok = True)
+        remove_tree(src)
+
+
 def activate_install_tree(staging_dir: Path, install_dir: Path, host: HostInfo) -> None:
     rollback_dir: Path | None = None
     failed_dir: Path | None = None
@@ -4781,7 +4815,7 @@ def activate_install_tree(staging_dir: Path, install_dir: Path, host: HostInfo)
         if install_dir.exists():
             rollback_dir = unique_install_side_path(install_dir, "rollback")
             log(f"moving existing install to rollback path {rollback_dir}")
-            os.replace(install_dir, rollback_dir)
+            move_install_dir_aside(install_dir, rollback_dir)
             log(f"moved existing install to rollback path {rollback_dir.name}")
 
         log(f"activating staged install {staging_dir} -> {install_dir}")
@@ -4796,7 +4830,7 @@ def activate_install_tree(staging_dir: Path, install_dir: Path, host: HostInfo)
             if install_dir.exists():
                 failed_dir = unique_install_side_path(install_dir, "failed")
                 log(f"moving failed active install to {failed_dir}")
-                os.replace(install_dir, failed_dir)
+                move_install_dir_aside(install_dir, failed_dir)
             elif staging_dir.exists():
                 failed_dir = staging_dir
                 staging_dir = None
@@ -4804,7 +4838,7 @@ def activate_install_tree(staging_dir: Path, install_dir: Path, host: HostInfo)
 
             if rollback_dir and rollback_dir.exists():
                 log(f"restoring rollback path {rollback_dir} -> {install_dir}")
-                os.replace(rollback_dir, install_dir)
+                move_install_dir_aside(rollback_dir, install_dir)
                 log(f"restored previous install from rollback path {rollback_dir.name}")
                 if is_busy_lock_error(exc):
                     raise BusyInstallConflict(