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

@ -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