diff --git a/.github/workflows/docker-publish.yml b/.github/workflows/docker-publish.yml index 6682e09113..e2cce67c52 100644 --- a/.github/workflows/docker-publish.yml +++ b/.github/workflows/docker-publish.yml @@ -89,10 +89,12 @@ jobs: contents: read outputs: llama_tag: ${{ steps.llama.outputs.tag }} - # One zoo ref + one notebooks commit, resolved here so BOTH arch legs of - # the base build (and the Studio build) bake the identical bits. Resolving - # them per-leg would let upstream advance between the amd64 and arm64 - # builds, putting different content under one published tag. + # One unsloth ref + one zoo ref + one notebooks commit, resolved here so + # BOTH arch legs of the base build AND the Studio build bake the identical + # bits. Resolving them per-leg would let upstream advance between the amd64 + # and arm64 builds (or between the base and Studio builds), putting + # different content under one published tag. + unsloth_ref: ${{ steps.unsloth_ref.outputs.ref }} zoo_ref: ${{ steps.zoo_ref.outputs.ref }} notebooks_commit: ${{ steps.notebooks.outputs.commit }} steps: @@ -110,6 +112,36 @@ jobs: echo "tag=${TAG:-latest}" >> "$GITHUB_OUTPUT" echo "llama.cpp prebuilt tag: ${TAG:-latest}" + # Freeze the requested unsloth ref to ONE concrete sha before the matrix + # fans out, so both base arch legs AND the Studio build bake the identical + # unsloth commit even when the requested ref is a mutable branch (the + # workflow_dispatch default is unsloth_ref=main) that advances during the + # ~4h base + Studio run. Same requested-ref precedence the inline build-arg + # used: the dispatch input wins (default main), else the pushed tag, else + # the triggering commit sha, else main. A 40-char sha (branch/schedule + # push) is already frozen; a branch/tag is resolved via ls-remote, exactly + # like the zoo and notebooks steps, falling back to the bare ref on a + # lookup miss so the Dockerfile can still fetch it by name. + - name: Resolve unsloth ref + id: unsloth_ref + env: + INPUT_REF: ${{ github.event.inputs.unsloth_ref }} + TAG_REF: ${{ startsWith(github.ref, 'refs/tags/') && github.ref_name || '' }} + PUSH_SHA: ${{ github.sha }} + run: | + REF="$INPUT_REF" + [ -n "$REF" ] || REF="$TAG_REF" + [ -n "$REF" ] || REF="$PUSH_SHA" + REF="${REF:-main}" + if printf '%s' "$REF" | grep -Eq '^[0-9a-f]{40}$'; then + SHA="$REF" + else + SHA="$(git ls-remote https://github.com/unslothai/unsloth "$REF" | awk 'NR==1{print $1}')" + [ -n "$SHA" ] || SHA="$REF" + fi + echo "ref=${SHA}" >> "$GITHUB_OUTPUT" + echo "unsloth ref: ${SHA}" + # Mirror the unsloth tag into the zoo ONLY when that tag actually exists # there. unsloth's v* tags are Studio releases the zoo never cuts (the zoo # repo currently has no tags at all), so blindly mirroring github.ref_name @@ -242,10 +274,12 @@ jobs: # NOTE: keep prose OUT of build-args -- docker/build-push-action # forwards every non-empty line verbatim, so a leading-# line would be # passed as a bogus --build-arg. Explanations live here instead: - # UNSLOTH_REF: workflow-dispatch honours the explicit input; tag - # pushes bake the tag's source ref (e.g. v1.2.3) so the published - # image actually contains that release; branch + scheduled runs bake - # the triggering commit SHA; any other event falls back to main. + # UNSLOTH_REF (from the prepare job): resolved to ONE sha before the + # matrix fans out, so both arch legs and the Studio build bake the + # identical unsloth commit even if a mutable branch (dispatch's + # unsloth_ref=main default) advances mid-run. Same requested-ref + # precedence as before: dispatch input, else the pushed tag, else + # the triggering commit sha, else main. # UNSLOTH_ZOO_REF (from the prepare job): explicit dispatch input, # else the pushed tag IF the zoo repo has it, else main -- a branch # SHA does not exist in the zoo repo. Resolved once in `prepare` and @@ -257,7 +291,7 @@ jobs: CUDA_VERSION=12.8.1 UBUNTU_VERSION=24.04 PYTHON_VERSION=3.12 - UNSLOTH_REF=${{ github.event.inputs.unsloth_ref || (startsWith(github.ref, 'refs/tags/') && github.ref_name) || github.sha || 'main' }} + UNSLOTH_REF=${{ needs.prepare.outputs.unsloth_ref }} UNSLOTH_ZOO_REF=${{ needs.prepare.outputs.zoo_ref }} LLAMA_PREBUILT_TAG=${{ needs.prepare.outputs.llama_tag }} UNSLOTH_NOTEBOOKS_REF=${{ needs.prepare.outputs.notebooks_commit }} @@ -430,15 +464,16 @@ jobs: cache-from: type=gha,scope=studio-${{ matrix.platform }} cache-to: type=gha,scope=studio-${{ matrix.platform }},mode=min outputs: type=image,name=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }},push-by-digest=true,name-canonical=true,push=true - # UNSLOTH_STUDIO_REF mirrors the base job's UNSLOTH_REF resolution so the - # Studio tree matches the unsloth baked into the base venv. + # UNSLOTH_STUDIO_REF is the SAME resolved unsloth sha the base build + # baked (needs.prepare.outputs.unsloth_ref), so the Studio tree matches + # the unsloth in the base venv even if the branch moved mid-run. # UNSLOTH_STUDIO_ZOO_REF is the SAME resolved zoo ref the base build # baked, so install.sh --local overlays the Studio venv with that zoo # instead of always tracking main. (Prose stays out of build-args -- # forwarded lines must be KEY=VALUE only.) build-args: | BASE_IMAGE=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}@${{ needs.merge.outputs.digest }} - UNSLOTH_STUDIO_REF=${{ github.event.inputs.unsloth_ref || (startsWith(github.ref, 'refs/tags/') && github.ref_name) || github.sha || 'main' }} + UNSLOTH_STUDIO_REF=${{ needs.prepare.outputs.unsloth_ref }} UNSLOTH_STUDIO_ZOO_REF=${{ needs.prepare.outputs.zoo_ref }} - name: Export digest diff --git a/docker/Dockerfile.studio b/docker/Dockerfile.studio index 2f51481830..bf41bc9d24 100644 --- a/docker/Dockerfile.studio +++ b/docker/Dockerfile.studio @@ -104,10 +104,11 @@ RUN apt-get update \ # probing would land on cpu or cu126 wheels depending on which host built # the image. cu128 on BOTH arches, mirroring the base venv: cu130 wheels # would silently lift the arm64 driver floor to 580+ while the base venv -# keeps the documented 570+ floor. DGX Spark / GB10 (sm_121) support comes -# from the same NVRTC cu13 swap the base image applies to its venv -- -# repeated below for the Studio venv's own bundled libnvrtc (the base's -# arm64 layer already installed cuda-nvrtc-13-0, so the cu13 .so exists). +# keeps the documented 570+ floor. Blackwell JIT (amd64 sm_103 B300/GB300 and +# arm64 sm_121 DGX Spark / GB10) support comes from the same NVRTC cu13 swap +# the base image applies to its venv -- repeated below for the Studio venv's +# own bundled libnvrtc, on BOTH arches (the base cu13 layer installed +# cuda-nvrtc-13-0 on both, so the cu13 .so exists here regardless of arch). # # UNSLOTH_PYTHON=3.12 pins the Studio venv to the SAME Python minor as the base # venv (install.sh defaults Linux to 3.13). Matching minors makes the two venvs' @@ -157,14 +158,20 @@ RUN set -eux \ && rm -rf "${UNSLOTH_STUDIO_HOME}/src/.git" \ "${UNSLOTH_STUDIO_HOME}/src/studio/frontend/node_modules" \ /root/.cache \ - && if [ "${TARGETARCH:-amd64}" = "arm64" ]; then \ - for NVRTC_DIR in "${UNSLOTH_STUDIO_HOME}"/unsloth_studio/lib/python*/site-packages/nvidia/cuda_nvrtc/lib; do \ - if [ -f "${NVRTC_DIR}/libnvrtc.so.12" ]; then \ - mv "${NVRTC_DIR}/libnvrtc.so.12" "${NVRTC_DIR}/libnvrtc.so.12.cu128.orig"; \ - ln -s /usr/local/cuda-13.0/lib64/libnvrtc.so.13 "${NVRTC_DIR}/libnvrtc.so.12"; \ - fi; \ - done; \ - fi \ + # Swap the Studio venv's OWN bundled cu12.8 libnvrtc for cu13, mirroring the + # base venv swap. Run on BOTH arches, not arm64 only: amd64 sm_103 (B300 / + # GB300) needs cu13 NVRTC exactly as arm64 sm_121 (DGX Spark / GB10) does, + # and the CUDA dedup below never touches cuda_nvrtc, so an amd64 Studio venv + # would otherwise keep cu12.8 NVRTC and its jiterator/NVRTC JIT paths fail on + # compute_103. The base cu13 layer installs cuda-nvrtc-13-0 on both arches, + # so /usr/local/cuda-13.0/lib64/libnvrtc.so.13 is present here regardless of + # TARGETARCH. + && for NVRTC_DIR in "${UNSLOTH_STUDIO_HOME}"/unsloth_studio/lib/python*/site-packages/nvidia/cuda_nvrtc/lib; do \ + if [ -f "${NVRTC_DIR}/libnvrtc.so.12" ]; then \ + mv "${NVRTC_DIR}/libnvrtc.so.12" "${NVRTC_DIR}/libnvrtc.so.12.cu128.orig"; \ + ln -s /usr/local/cuda-13.0/lib64/libnvrtc.so.13 "${NVRTC_DIR}/libnvrtc.so.12"; \ + fi; \ + done \ && BASE_NV=/opt/unsloth-venv/lib/python3.12/site-packages/nvidia \ && STU_NV="${UNSLOTH_STUDIO_HOME}/unsloth_studio/lib/python3.12/site-packages/nvidia" \ && if [ ! -d "${STU_NV}" ] || [ ! -d "${BASE_NV}" ]; then \ diff --git a/docker/unsloth_pip_shim.py b/docker/unsloth_pip_shim.py index b78aac7c62..841d172545 100644 --- a/docker/unsloth_pip_shim.py +++ b/docker/unsloth_pip_shim.py @@ -70,6 +70,8 @@ _VALUE_FLAGS = { "--python-version", "--abi", "--implementation", + "-e", + "--editable", } # Of those value-flags, the ones whose VALUE is itself an install target: a # requirements file pulls real requirements. An index-url / find-links / @@ -80,6 +82,17 @@ _REQ_FILE_FLAGS = {"-r", "--requirement"} # still downgrade or reinstall a baked package when another target pulls it in. # Filter protected packages out of them the same way as requirement files. _CONSTRAINT_FILE_FLAGS = {"-c", "--constraint"} +# -e/--editable takes the NEXT token as its target (pip: +# `-e, --editable `), and that target is a real install target. A +# protected editable (e.g. `-e git+https://.../unsloth.git#egg=unsloth`) must +# drop BOTH the flag and its value; dropping the value alone leaves pip a +# dangling `-e` that swallows the next kept package and fails the whole cell. +_EDITABLE_FLAGS = {"-e", "--editable"} +# -P/--upgrade-package is uv's selective-upgrade flag: naming a baked +# package (e.g. `uv pip install -P torch peft`) lets an ordinary install target +# refresh that package and clobber the pinned stack. Filter its value through +# _KEEP too. Unlike -e it is not itself an install target (no has_target). +_UPGRADE_PKG_FLAGS = {"-P", "--upgrade-package"} def _canon(token): @@ -109,6 +122,19 @@ def _canon(token): _egg = re.search(r"[#&]egg=([A-Za-z0-9][A-Za-z0-9._-]*)", token) if _egg: return _egg.group(1).lower().replace("_", "-") or None + # A direct wheel URL or local wheel path still names its distribution in + # the PEP 427 filename ({distribution}-{version}-...-...-....whl), so a + # bare `pip install https://.../torch-2.11.0+cu128-...whl` would slip a + # protected package past _KEEP as an opaque positional and reinstall the + # baked torch. Dashes cannot appear inside the distribution component (a + # run of -_. normalises to a single -), so the leading dash-split of the + # basename is the distribution name; pull it so _KEEP can drop it. A + # non-protected wheel returns its name and the caller keeps the token. + _whl = re.search(r"([^/\\#?]+)\.whl(?:[#?]|$)", token) + if _whl: + dist = _whl.group(1).split("-", 1)[0].strip().lower().replace("_", "-") + if dist: + return dist return None # vcs / url / local path -> let it pass through # strip extras and any version/marker tail name = re.split(r"[<>=!~\[\s;@]", token, 1)[0].strip() @@ -121,6 +147,22 @@ def _version_pin(token): return m.group(1) if m else None +def _classify_flag_target(spec): + """Classify the value that rides on -e/--editable or -P/--upgrade-package. + + Returns ("drop", version_or_None) when the value names a protected package + (so the flag+value pair must be dropped, closing the same bypass the bare + positional spec closes) or ("keep", None) when it is safe to forward. + transformers is reported as "drop" with any pinned version so its sidecar + marker is still recorded, mirroring the bare-spec handling in main().""" + name = _canon(spec) + if name == "transformers": + return "drop", _version_pin(spec) + if name is not None and (name in _KEEP or name.startswith(_KEEP_PREFIX)): + return "drop", None + return "keep", None + + def _parse_include(stripped): """If `stripped` is an `-r`/`--requirement`/`-c`/`--constraint` include, return (flag, target_path, inline_comment_or_None); else (None, None, None).""" @@ -295,6 +337,23 @@ def main(): _c_path, _c_rec, _c_drp = _filter_requirements_file(tok) keep_args.append(_c_path) dropped.extend(_c_drp) + elif prev_flag in _EDITABLE_FLAGS or prev_flag in _UPGRADE_PKG_FLAGS: + # The flag was held back (not appended yet): its value is an + # install target (-e path/url/vcs) or an upgrade selector + # (-P name), both filtered through _KEEP. Dropping a protected + # value drops the flag with it, so pip/uv is never left a + # dangling `-e`/`-P` that fails the cell or refreshes a baked + # package. A kept editable target sets has_target; -P does not. + _action, _ver = _classify_flag_target(tok) + if _action == "drop": + if _ver and not recorded: + recorded = _ver + dropped.append(prev_flag + " " + tok) + else: + keep_args.append(prev_flag) + keep_args.append(tok) + if prev_flag in _EDITABLE_FLAGS: + has_target = True else: keep_args.append(tok) skip_next = False @@ -319,11 +378,30 @@ def main(): _c_path, _c_rec, _c_drp = _filter_requirements_file(_val) keep_args.append(_flag + "=" + _c_path) dropped.extend(_c_drp) + elif _flag in _EDITABLE_FLAGS or _flag in _UPGRADE_PKG_FLAGS: + # --editable= / --upgrade-package=: filter the + # inline value through _KEEP just like the space-separated + # form, dropping the whole token for a protected package. + _action, _ver = _classify_flag_target(_val) + if _action == "drop": + if _ver and not recorded: + recorded = _ver + dropped.append(tok) + else: + keep_args.append(tok) + if _flag in _EDITABLE_FLAGS: + has_target = True else: keep_args.append(tok) # option with inline value, not a target continue if tok in _VALUE_FLAGS: - keep_args.append(tok) + # -e/--editable and -P/--upgrade-package carry a value that is a + # potential install target, so hold the flag back and let the + # skip_next handler emit or drop the flag+value pair together. Every + # other value-flag keeps its flag verbatim; only its value (an + # index-url / find-links / target dir / etc.) is an opaque option. + if tok not in _EDITABLE_FLAGS and tok not in _UPGRADE_PKG_FLAGS: + keep_args.append(tok) skip_next = True prev_flag = tok continue diff --git a/tests/python/test_unsloth_pip_shim.py b/tests/python/test_unsloth_pip_shim.py new file mode 100644 index 0000000000..c466c8c7be --- /dev/null +++ b/tests/python/test_unsloth_pip_shim.py @@ -0,0 +1,224 @@ +"""Regression tests for docker/unsloth_pip_shim.py. + +The shim sits ahead of the real pip/uv on PATH inside the Unsloth Docker +notebook environment so a notebook `!pip install ...` / `!uv pip install ...` +cell cannot clobber the baked, ABI-matched cu128 torch/vLLM/transformers stack. +These tests drive main() with UNSLOTH_NB_SHIM=1 and capture the command it would +os.execv, so we can assert what actually reaches the real tool. They cover: + + * -e/--editable paired with its target (a protected editable drops the flag + too, so pip is never left a dangling `-e`); + * -P/--upgrade-package values filtered through the protected set (uv cannot be + told to refresh a baked package); + * direct wheel URL / local wheel path basenames parsed for protected + distribution names before URL passthrough. + +No GPU or network is required. +""" + +from __future__ import annotations + +import importlib.util +import os +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +SHIM_PATH = REPO_ROOT / "docker" / "unsloth_pip_shim.py" + +TORCH_WHEEL_URL = ( + "https://download.pytorch.org/whl/cu128/" + "torch-2.11.0%2Bcu128-cp312-cp312-linux_x86_64.whl" +) + + +class _Exec(Exception): + """Raised by the patched os.execv so main() stops at the exec point and the + intended command is captured instead of replacing the test process.""" + + def __init__(self, path, argv): + self.path = path + self.argv = list(argv) + + +@pytest.fixture() +def shim(tmp_path, monkeypatch): + """Load a fresh copy of the shim with the transformers marker pointed at a + temp file and os.execv patched to capture (not perform) the exec.""" + marker = tmp_path / "requested_transformers" + monkeypatch.setenv("UNSLOTH_NB_TF_MARKER", str(marker)) + monkeypatch.setenv("UNSLOTH_NB_SHIM", "1") + + assert SHIM_PATH.is_file(), f"missing shim: {SHIM_PATH}" + spec = importlib.util.spec_from_file_location("unsloth_pip_shim_under_test", SHIM_PATH) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + + def _fake_execv(path, argv): + raise _Exec(path, argv) + + monkeypatch.setattr(mod.os, "execv", _fake_execv) + mod._marker_path = marker # convenience for assertions + return mod + + +def _run(shim, tool, args): + """Invoke the shim as `tool install ` and return (execd_tail, marker). + + execd_tail is the argument list after the `install` verb that reached the + real tool, or None when the shim no-op'd (nothing left to install). marker is + the recorded transformers version, or None. + """ + if tool == "uv": + argv = ["uv", "pip", "install", *args] + else: + argv = ["pip", "install", *args] + with pytest.MonkeyPatch.context() as mp: + mp.setattr(shim.sys, "argv", argv) + try: + shim.main() + execd = None + except _Exec as exc: + # main() builds [REAL[tool]] + head + keep_args; head ends with the + # `install` verb, so everything after it is what we asserted on. + i = exc.argv.index("install") + execd = exc.argv[i + 1 :] + marker = shim._marker_path.read_text() if shim._marker_path.exists() else None + return execd, marker + + +# -------------------------------------------------------------------------- +# Item 3541142907 -- pair -e/--editable with its target. +# -------------------------------------------------------------------------- +def test_editable_protected_target_drops_flag_and_value(shim): + # `pip install -e git+...unsloth...#egg=unsloth peft` must NOT become + # `pip install -e peft` (which pip rejects); it must install just peft. + execd, _ = _run( + shim, + "pip", + ["-e", "git+https://github.com/unslothai/unsloth.git#egg=unsloth", "peft"], + ) + assert execd == ["peft"], execd + assert "-e" not in execd + + +def test_editable_only_protected_target_noops(shim): + execd, _ = _run( + shim, "pip", ["-e", "git+https://github.com/unslothai/unsloth.git#egg=unsloth"] + ) + assert execd is None # nothing left to install -> no-op, no dangling -e + + +def test_editable_unprotected_target_is_kept(shim): + execd, _ = _run(shim, "pip", ["-e", "./localpkg"]) + assert execd == ["-e", "./localpkg"], execd + + +def test_editable_long_form_inline_protected(shim): + execd, _ = _run( + shim, + "pip", + ["--editable=git+https://github.com/unslothai/unsloth.git#egg=unsloth", "peft"], + ) + assert execd == ["peft"], execd + + +def test_editable_long_form_inline_unprotected_kept(shim): + execd, _ = _run(shim, "pip", ["--editable=./localpkg"]) + assert execd == ["--editable=./localpkg"], execd + + +# -------------------------------------------------------------------------- +# Item 3541142906 -- filter uv -P/--upgrade-package values. +# -------------------------------------------------------------------------- +def test_upgrade_package_protected_short_flag_dropped(shim): + # `uv pip install -P torch peft` must not let uv refresh baked torch. + execd, _ = _run(shim, "uv", ["-P", "torch", "peft"]) + assert execd == ["peft"], execd + assert "torch" not in execd and "-P" not in execd + + +def test_upgrade_package_protected_long_inline_dropped(shim): + execd, marker = _run(shim, "uv", ["--upgrade-package=transformers", "peft"]) + assert execd == ["peft"], execd + assert "--upgrade-package=transformers" not in execd + + +def test_upgrade_package_transformers_pin_recorded(shim): + # A pinned transformers upgrade selector still feeds the sidecar marker. + execd, marker = _run(shim, "uv", ["-P", "transformers==4.55.0", "peft"]) + assert execd == ["peft"], execd + assert marker == "4.55.0" + + +def test_upgrade_package_unprotected_kept(shim): + execd, _ = _run(shim, "uv", ["-P", "requests", "requests"]) + assert execd == ["-P", "requests", "requests"], execd + + +def test_upgrade_package_only_protected_noops(shim): + execd, _ = _run(shim, "uv", ["-P", "torch"]) + assert execd is None # -P is not itself a target + + +# -------------------------------------------------------------------------- +# Item 3541142908 -- parse protected wheel basenames before URL passthrough. +# -------------------------------------------------------------------------- +def test_direct_torch_wheel_url_dropped(shim): + execd, _ = _run(shim, "pip", [TORCH_WHEEL_URL]) + assert execd is None # torch wheel URL recognised + dropped -> no-op + + +def test_local_torch_wheel_path_dropped(shim): + execd, _ = _run( + shim, "pip", ["/tmp/wheels/torch-2.11.0+cu128-cp312-cp312-linux_x86_64.whl"] + ) + assert execd is None + + +def test_normalised_wheel_name_dropped(shim): + # unsloth_zoo-*.whl normalises to unsloth-zoo, which is protected. + execd, _ = _run(shim, "pip", ["https://example.com/unsloth_zoo-1.0-py3-none-any.whl"]) + assert execd is None + + +def test_unprotected_wheel_url_kept(shim): + url = "https://example.com/wheels/numpy-2.1.0-cp312-cp312-linux_x86_64.whl" + execd, _ = _run(shim, "pip", [url]) + assert execd == [url], execd + + +def test_protected_wheel_in_requirements_file_dropped(shim, tmp_path): + req = tmp_path / "reqs.txt" + req.write_text( + TORCH_WHEEL_URL + "\n" + "snac==1.2.0\n", + encoding = "utf-8", + ) + execd, _ = _run(shim, "pip", ["-r", str(req)]) + # The filtered requirements copy still installs snac; torch's wheel line is + # stripped. execd is `-r `. + assert execd is not None and execd[0] == "-r" + filtered = Path(execd[1]).read_text(encoding = "utf-8") + assert "snac==1.2.0" in filtered + assert "torch" not in filtered + + +# -------------------------------------------------------------------------- +# Guardrails: the ordinary happy paths still work unchanged. +# -------------------------------------------------------------------------- +def test_plain_package_passes_through(shim): + execd, _ = _run(shim, "pip", ["omegaconf==2.3.1"]) + assert execd == ["omegaconf==2.3.1"], execd + + +def test_bare_transformers_recorded_and_dropped(shim): + execd, marker = _run(shim, "pip", ["transformers==4.55.0"]) + assert execd is None + assert marker == "4.55.0" + + +def test_index_url_value_flag_kept_verbatim(shim): + execd, _ = _run(shim, "pip", ["--extra-index-url", "https://example.com/simple", "snac"]) + assert execd == ["--extra-index-url", "https://example.com/simple", "snac"], execd