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