From 1254fdf3ad5e1e5f20d9e6d183ee4c04e76ac858 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 03:42:38 +0000 Subject: [PATCH] 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). --- docker/entrypoint.sh | 23 +++++++ docker/unsloth_nb_pip_magic.py | 22 ++++++- docker/unsloth_pip_shim.py | 69 +++++++++++++++++++- tests/python/test_unsloth_nb_pip_magic.py | 79 +++++++++++++++++++++++ tests/python/test_unsloth_pip_shim.py | 61 +++++++++++++++++ 5 files changed, 249 insertions(+), 5 deletions(-) create mode 100644 tests/python/test_unsloth_nb_pip_magic.py diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh index cae8e8c7fd..36138dcfd3 100755 --- a/docker/entrypoint.sh +++ b/docker/entrypoint.sh @@ -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 "$@" diff --git a/docker/unsloth_nb_pip_magic.py b/docker/unsloth_nb_pip_magic.py index 6c3d61907f..78ddb36995 100644 --- a/docker/unsloth_nb_pip_magic.py +++ b/docker/unsloth_nb_pip_magic.py @@ -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 `! -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): diff --git a/docker/unsloth_pip_shim.py b/docker/unsloth_pip_shim.py index af3b8ce74c..39f0fdebab 100644 --- a/docker/unsloth_pip_shim.py +++ b/docker/unsloth_pip_shim.py @@ -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) diff --git a/tests/python/test_unsloth_nb_pip_magic.py b/tests/python/test_unsloth_nb_pip_magic.py new file mode 100644 index 0000000000..6d4483d638 --- /dev/null +++ b/tests/python/test_unsloth_nb_pip_magic.py @@ -0,0 +1,79 @@ +"""Regression tests for docker/unsloth_nb_pip_magic.py. + +The input transformer rewrites explicit `! -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 diff --git a/tests/python/test_unsloth_pip_shim.py b/tests/python/test_unsloth_pip_shim.py index 11e3fb89c2..0620e88928 100644 --- a/tests/python/test_unsloth_pip_shim.py +++ b/tests/python/test_unsloth_pip_shim.py @@ -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