docker: close pip-shim bypasses and warn on arm64 cu13 llama.cpp mismatch

Four follow-ups to the shim/entrypoint audit fixes:

1. unsloth_pip_shim.py let a local project directory install through: `pip
   install ./transformers` / `-e ./unsloth` is not a requirement spec, so
   _canon returned None and both the arg filter and the constraints file
   (which only rejects a version MISMATCH) passed it, letting a same-version
   local build silently replace the baked wheel. _canon now resolves the
   project name from pyproject [project].name, then setup.cfg, then the
   directory basename when it is an installable project, so a local checkout
   of a protected package is dropped like every other artifact form. Names
   match exactly after normalization, so a user dir named my-torch-utils is
   untouched, and a metadata-less directory still passes through.

2. unsloth_nb_pip_magic.py only rewrote literal `!python -m pip`, so the
   `!{sys.executable} -m pip ...` form notebooks use to target the running
   kernel (and absolute interpreter paths) bypassed the PATH shim entirely.
   Input transformers see the raw cell text before IPython expands the
   braces, so the matcher now also covers {sys.executable} (quoted or bare)
   and quoted/bare interpreter paths ending in python[0-9.]*(.exe) before
   -m pip|uv.

3. unsloth_pip_shim.py did not strip uv's --exact, which performs an exact
   sync that removes every installed package outside the kept target's
   closure (vLLM, bitsandbytes, the NVIDIA libs); `uv pip install --exact
   peft` would strip the baked stack after the filter kept it. --exact now
   joins the resolver-wide destructive flags dropped in shim mode.

4. entrypoint.sh: the arm64 image bakes a CUDA 13 llama.cpp because upstream
   (unslothai/llama.cpp) publishes no CUDA 12 arm64 asset, while the torch
   stack (cu128) runs on a 570-series driver. A CUDA 13 cubin cannot load on
   a 570-579 driver, so on GH200/GB200 hosts below 580 GGUF export and Studio
   chat fail while training works. The entrypoint now warns up front on
   aarch64 + driver < 580 instead of letting llama-server fail later.

Tests: shim + nb-pip-magic suites at 81 (18 new, including local-project
name resolution, the executable/brace forms, and --exact stripping).
This commit is contained in:
Daniel Han 2026-07-13 03:42:38 +00:00
commit 1254fdf3ad
5 changed files with 249 additions and 5 deletions

View file

@ -231,5 +231,28 @@ if major < 8:
print(" Unsloth will fall back to fp16. Training works but is slightly slower.")
PY
# --- arm64 note: baked llama.cpp is a CUDA 13 build -------------------------
# Upstream publishes no CUDA 12 arm64 llama.cpp bundle (only arm64-cpu and
# arm64-cuda13), so the arm64 image bakes the cu13 build while the torch stack
# (cu128) runs fine on a 570-series driver. A CUDA 13 cubin cannot load on a
# 570-579 driver, so on GH200/GB200-class hosts below 580 GGUF export and
# Studio chat would fail even though training works -- say so up front instead
# of letting llama-server fail mysteriously later.
if [ "$(uname -m)" = "aarch64" ]; then
_drv="$(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -1)"
_drv_major="${_drv%%.*}"
case "$_drv_major" in
*[!0-9]* | "") ;; # unreadable driver version -> no claim to make
*)
if [ "$_drv_major" -lt 580 ]; then
echo "WARNING: this arm64 image bakes a CUDA 13 llama.cpp (upstream ships no CUDA 12 arm64 build)." >&2
echo " Host driver $_drv is < 580, which cannot load CUDA 13 binaries:" >&2
echo " training (torch cu128) works, but GGUF export / Studio chat will fail" >&2
echo " until the host driver is upgraded to >= 580." >&2
fi
;;
esac
fi
sync_notebooks
exec "$@"

View file

@ -21,9 +21,25 @@ 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(.*)$")
# Only the explicit `!<python> -m pip|uv ...` shell form (the `!` makes it a
# shell escape). Matched against the line with its trailing newline stripped.
# Input transformers see the RAW cell text -- IPython expands `{sys.executable}`
# later, inside the system() execution path -- so the braced form notebooks use
# to target the running kernel (`!{sys.executable} -m pip install ...`) and
# absolute interpreter paths (`!/opt/unsloth-venv/bin/python -m pip ...`),
# quoted or bare, must be matched here too or module-pip bypasses the PATH shim.
_PY_M_PIP = re.compile(
r"""^(\s*)!\s*
(?:
(?:python[0-9.]*|py) # literal python / py
| ["']?\{\s*sys\.executable\s*\}["']? # {sys.executable}, opt. quoted
| "(?:[^"]*[/\\])python[0-9.]*(?:\.exe)?" # quoted interpreter path
| '(?:[^']*[/\\])python[0-9.]*(?:\.exe)?'
| \S*[/\\]python[0-9.]*(?:\.exe)? # bare interpreter path
)
\s+-m\s+(pip|uv)\b(.*)$""",
re.VERBOSE,
)
def _rewrite_python_dash_m(lines):

View file

@ -119,7 +119,11 @@ _ATTACHED_SHORT_FLAGS = {"-r", "-c", "-e", "-P"}
# guise of installing an unprotected package. The kept target still installs; its
# already-satisfied protected deps are left untouched. Per-package selectors
# (--reinstall-package / -P) are handled through _UPGRADE_PKG_FLAGS instead.
_REINSTALL_FLAGS = {"--force-reinstall", "--ignore-installed", "-I", "--reinstall"}
# uv's --exact is destructive the other way around: it performs an exact SYNC,
# REMOVING every installed package outside the kept target's closure (vLLM,
# bitsandbytes, the NVIDIA libs, ...), so `uv pip install --exact peft` would
# strip the baked stack after the argument filter kept it off the command line.
_REINSTALL_FLAGS = {"--force-reinstall", "--ignore-installed", "-I", "--reinstall", "--exact"}
# Value-flags whose flag+value pair is dropped outright in shim mode.
# `--upgrade-strategy eager` makes pip upgrade EVERY dependency of a kept target
# regardless of whether the installed version already satisfies it, which would
@ -218,7 +222,25 @@ def _canon(token):
_seg = _seg.strip().lower().replace("_", "-")
if _seg:
return _seg
return None # plain url / local path -> let it pass through
# A local project DIRECTORY (`pip install ./transformers`,
# `pip install -e ./unsloth`) installs the project it contains, and a
# same-version dev build slips past even the protected constraints file
# (constraints only reject a version MISMATCH), silently swapping the
# baked, tested wheel for a local build. Resolve the project name from
# its metadata so _KEEP applies to this form like every other artifact
# form (wheel/sdist/VCS/egg). Non-directories and metadata-less dirs
# pass through as before.
_local = _local_project_name(token)
if _local:
return _local
return None # plain url / metadata-less local path -> let it pass through
# A local project dir referenced without ./ or / (`pip install subdir/proj`)
# is still a path target to pip when it exists on disk; classify it the same
# way before the spec parse below mangles the separator.
if "/" in token or os.sep in token:
_local = _local_project_name(token)
if _local:
return _local
# A bare wheel filename (no ./ or / prefix and no scheme) is still a valid
# pip target from the CWD: `pip install torch-2.11.0-cp312-...-linux.whl`.
# It reaches here because it starts with neither `.`/`/` nor a scheme, so
@ -239,6 +261,49 @@ def _canon(token):
return name.lower().replace("_", "-") or None
def _local_project_name(token):
"""Distribution name of a local project directory install target, else None.
Reads the name pip/uv would build: pyproject.toml [project].name, falling
back to setup.cfg [metadata] name, falling back to the directory basename
when a setup.py exists (a bare basename guess is used ONLY when the dir is
an installable project at all). A directory without any project metadata is
not a pip target and returns None so ordinary paths pass through untouched.
Names are exact after normalization: a user's own `my-torch-utils` dir never
matches the protected `torch`.
"""
path = token.split("#", 1)[0]
if not os.path.isdir(path):
return None
_pyproject = os.path.join(path, "pyproject.toml")
if os.path.isfile(_pyproject):
try:
import tomllib
with open(_pyproject, "rb") as f:
_name = (tomllib.load(f).get("project") or {}).get("name")
if _name:
return _name.strip().lower().replace("_", "-") or None
except Exception:
pass # unparseable metadata -> fall through to the other signals
_setup_cfg = os.path.join(path, "setup.cfg")
if os.path.isfile(_setup_cfg):
try:
import configparser
_cp = configparser.ConfigParser()
_cp.read(_setup_cfg)
_name = _cp.get("metadata", "name", fallback = None)
if _name:
return _name.strip().lower().replace("_", "-") or None
except Exception:
pass
if os.path.isfile(os.path.join(path, "setup.py")) or os.path.isfile(_pyproject):
_base = os.path.basename(os.path.normpath(path))
return _base.strip().lower().replace("_", "-") or None
return None
def _version_pin(token):
"""Return the pinned version for a `pkg==X` token, else None."""
m = re.search(r"==\s*([0-9][0-9A-Za-z.\-]*)", token)

View file

@ -0,0 +1,79 @@
"""Regression tests for docker/unsloth_nb_pip_magic.py.
The input transformer rewrites explicit `!<python> -m pip|uv ...` shell lines
to `!pip|uv ...` so they resolve to the PATH shim. IPython input transformers
see the RAW cell text (brace expansion like `{sys.executable}` happens later,
in the system() execution path), so the braced and absolute-interpreter forms
notebooks use to target the running kernel must be rewritten too (item
3567875025); only matching literal `python`/`py` let module-pip bypass the
shim entirely.
"""
import importlib.util
import pathlib
_MOD_PATH = pathlib.Path(__file__).resolve().parents[2] / "docker" / "unsloth_nb_pip_magic.py"
_spec = importlib.util.spec_from_file_location("unsloth_nb_pip_magic", _MOD_PATH)
magic = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(magic)
def _rewrite(line):
return magic._rewrite_python_dash_m([line])[0]
def test_literal_python_rewritten():
assert _rewrite("!python -m pip install peft\n") == "!pip install peft\n"
def test_literal_python_version_rewritten():
assert _rewrite("!python3.12 -m pip install peft") == "!pip install peft"
def test_sys_executable_braces_rewritten():
assert _rewrite("!{sys.executable} -m pip install peft\n") == "!pip install peft\n"
def test_sys_executable_braces_quoted_rewritten():
assert _rewrite('!"{sys.executable}" -m pip install peft') == "!pip install peft"
def test_sys_executable_braces_spaced_rewritten():
assert _rewrite("!{ sys.executable } -m pip install peft") == "!pip install peft"
def test_absolute_interpreter_path_rewritten():
assert (
_rewrite("!/opt/unsloth-venv/bin/python -m pip install peft\n")
== "!pip install peft\n"
)
def test_absolute_interpreter_versioned_path_rewritten():
assert _rewrite("!/usr/bin/python3.11 -m uv pip install peft") == "!uv pip install peft"
def test_quoted_interpreter_path_rewritten():
assert (
_rewrite('!"/opt/unsloth venv/bin/python" -m pip install peft')
== "!pip install peft"
)
def test_indent_preserved():
assert _rewrite(" !{sys.executable} -m pip install peft") == " !pip install peft"
def test_python_script_not_rewritten():
line = "!python train.py --epochs 3"
assert _rewrite(line) == line
def test_module_other_than_pip_not_rewritten():
line = "!python -m venv .venv"
assert _rewrite(line) == line
def test_non_shell_line_not_rewritten():
line = "x = '{sys.executable} -m pip install peft'"
assert _rewrite(line) == line

View file

@ -642,3 +642,64 @@ def test_filter_write_failure_clean_file_passes_through(shim, tmp_path, monkeypa
monkeypatch.setattr(shim.tempfile, "mkstemp", denied)
path, recorded, dropped = shim._filter_requirements_file(str(req))
assert path == str(req) and recorded is None and dropped == []
# --------------------------------------------------------------------------
# Item 3567875029 -- uv's --exact performs an exact SYNC (removes packages
# outside the kept target's closure), so it is stripped like the other
# resolver-wide destructive switches.
# --------------------------------------------------------------------------
def test_uv_exact_flag_stripped(shim):
execd, _ = _run(shim, "uv", ["--exact", "peft"])
assert execd == ["peft"], execd
# --------------------------------------------------------------------------
# Item 3567875023 -- a local project directory naming a protected package
# (pip install ./transformers, pip install -e ./unsloth) is filtered like the
# wheel/sdist/VCS forms: a same-version dev build slips past the constraints
# file, so the name must come from the project metadata.
# --------------------------------------------------------------------------
def _make_local_project(tmp_path, dirname, project_name):
proj = tmp_path / dirname
proj.mkdir()
(proj / "pyproject.toml").write_text(
f'[project]\nname = "{project_name}"\nversion = "1.0"\n'
)
return str(proj)
def test_local_dir_protected_by_metadata_dropped(shim, tmp_path):
# Directory name is innocuous; pyproject names a protected package.
path = _make_local_project(tmp_path, "my-checkout", "transformers")
execd, _ = _run(shim, "pip", [path, "peft"])
assert execd == ["peft"], execd
def test_local_dir_protected_editable_dropped(shim, tmp_path):
path = _make_local_project(tmp_path, "unsloth", "unsloth")
execd, _ = _run(shim, "pip", ["-e", path, "peft"])
assert execd == ["peft"], execd
assert "-e" not in execd
def test_local_dir_basename_fallback_setup_py(shim, tmp_path):
# No parseable name in metadata: setup.py + protected basename still drops.
proj = tmp_path / "torch"
proj.mkdir()
(proj / "setup.py").write_text("from setuptools import setup\nsetup()\n")
execd, _ = _run(shim, "pip", [str(proj), "peft"])
assert execd == ["peft"], execd
def test_local_dir_unprotected_kept(shim, tmp_path):
path = _make_local_project(tmp_path, "my-torch-utils", "my-torch-utils")
execd, _ = _run(shim, "pip", [path])
assert execd == [path], execd
def test_local_dir_without_metadata_passes_through(shim, tmp_path):
plain = tmp_path / "datadir"
plain.mkdir()
execd, _ = _run(shim, "pip", [str(plain)])
assert execd == [str(plain)], execd