ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

docker: protect the tested training stack from notebook install cells

Browse source 
The pip shim fronts pip/uv inside the notebook kernel so an install cell cannot
replace the baked cu128 stack, but _KEEP only covered torch/vLLM/unsloth. Across
the 433 shipped notebooks that left the training half wide open:

  trl         382 pin an older release, 378 of them ending the install cell with
              `pip install --no-deps trl==0.22.2`, against a baked trl 0.24.0
  torchao     273 reinstall it and 2 pin 0.15.0, replacing 0.17.0+cu128
  torchcodec   92 reinstall it and 26 pin 0.5 or 0.7.0, replacing the
              0.11.0+cu128 wheel the Dockerfile pairs with torch 2.11
  datasets    254 reinstall it, observed falling from 4.3.0 to 3.0.0
  peft        225 reinstall it, observed falling from 0.19.1 to 0.14.0
  accelerate  225 reinstall it
  hf hub      240 reinstall it and tokenizers 64, both version-locked to
              transformers and shipped in matched copies inside every sidecar

So every notebook run mutated the stack the image was validated with, while the
shim printed that it was keeping the baked versions.

The membership criterion is "replacing this invalidates the tested stack or
breaks unsloth", not "a notebook mentions it": snac, causal-conv1d, mamba-ssm,
omegaconf, protobuf, sentencepiece and the rest still install normally.

Verified in the rebuilt image by running the Gemma3 (270M) install cell verbatim:
trl 0.24.0, peft 0.19.1, datasets 4.3.0, accelerate 1.14.0, torchao 0.17.0+cu128,
transformers 5.14.1 and huggingface_hub 1.24.0 are all unchanged afterwards, the
requested transformers pin is still recorded for the sidecar, and a package the
image does not bake still installs.

The existing shim tests used peft as their "unprotected package" sentinel, so
they move to snac.
This commit is contained in:
Daniel Han 2026-07-26 17:28:15 +00:00
commit 6162d4d87d
3 changed files with 289 additions and 57 deletions
Download patch file Download diff file Expand all files Collapse all files

196
tests/python/test_docker_pip_shim_training_stack.py Normal file
View file

@ -0,0 +1,196 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0
"""Regression guard for what the Docker pip shim protects.
The shim fronts pip/uv inside the notebook kernel so a `!pip install` cell cannot
replace the baked, ABI-matched stack. It protected torch/vLLM/unsloth and stopped
there, which left the training stack wide open. Measured over the 433 shipped
notebooks (probe_notebook_pins.py against the baked image):
trl 382 notebooks pin an older release -- 378 of them end their
install cell with `!pip install --no-deps trl==0.22.2`, against
a baked and tested trl 0.24.0
torchao 273 reinstall it, 2 pin 0.15.0, replacing 0.17.0+cu128 with a
generic PyPI build
torchcodec 92 reinstall it, 26 pin 0.5 / 0.7.0, replacing the 0.11.0+cu128
wheel the Dockerfile deliberately paired with torch 2.11
datasets 254 reinstall it; a trl 0.22.2 resolve was observed pulling it
back from 4.3.0 to 3.0.0
peft 225 reinstall it; observed dropping 0.19.1 -> 0.14.0
accelerate 225 reinstall it
hf_hub 240 reinstall it, tokenizers 64 -- both version-locked to
transformers, and the sidecars ship their own matched copies
So EVERY notebook run silently mutated the stack the image was validated with,
and printed "Successfully installed trl-0.22.2 peft-0.14.0 datasets-3.0.0" while
the shim reported it was keeping the baked versions.
The criterion for _KEEP is "replacing this invalidates the tested stack or breaks
unsloth", not "any package a notebook mentions": a package the notebook genuinely
needs and the image does not bake still has to install normally.
Static: drives the shim's main() with os.execv captured. No docker, no GPU, no
network.
"""
from __future__ import annotations
import importlib.util
import os
from pathlib import Path
import pytest
REPO_ROOT = Path(__file__).resolve().parents[2]
SHIM_PATH = REPO_ROOT / "docker" / "unsloth_pip_shim.py"
# The install cell 378 of the 433 shipped notebooks actually end on.
SHIPPED_TRL_CELL = ["--no-deps", "trl==0.22.2"]
# A package the image does NOT bake: must keep installing normally.
UNBAKED = "snac"
class _Exec(Exception):
def __init__(self, path, argv):
self.path = path
self.argv = list(argv)
@pytest.fixture()
def shim(tmp_path, monkeypatch):
monkeypatch.setenv("UNSLOTH_NB_TF_MARKER", str(tmp_path / "requested_transformers"))
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_stack_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)
return mod
def _run(shim, args, tool = "pip"):
"""Return the args that reached the real tool after `install`, or None when
the shim no-op'd. The always-injected protected-constraints pair is dropped."""
argv = ["uv", "pip", "install", *args] if tool == "uv" else ["pip", "install", *args]
with pytest.MonkeyPatch.context() as mp:
mp.setattr(shim.sys, "argv", argv)
try:
shim.main()
return None
except _Exec as exc:
i = exc.argv.index("install")
execd = exc.argv[i + 1 :]
if (
len(execd) >= 2
and execd[-2] == "--constraint"
and os.path.basename(execd[-1]).startswith("unsloth-nb-protected-")
):
execd = execd[:-2]
return execd
# --------------------------------------------------------------------------
# Membership
# --------------------------------------------------------------------------
@pytest.mark.parametrize(
"pkg",
["trl", "peft", "datasets", "accelerate", "torchao", "torchcodec",
"huggingface-hub", "tokenizers", "safetensors"],
)
def test_training_stack_is_protected(shim, pkg):
assert pkg in shim._KEEP, (
f"{pkg} is baked and tested; a notebook pin replacing it invalidates the image"
)
def test_the_original_gpu_stack_is_still_protected(shim):
for pkg in ["torch", "torchvision", "torchaudio", "triton", "xformers",
"vllm", "bitsandbytes", "unsloth", "unsloth-zoo"]:
assert pkg in shim._KEEP
def test_unrelated_packages_are_not_swept_in(shim):
# The criterion is "invalidates the tested stack", not "a notebook mentions
# it". These are all installed by shipped notebooks and must stay installable.
for pkg in ["snac", "causal-conv1d", "mamba-ssm", "omegaconf", "timm",
"librosa", "trackio", "open-spiel", "protobuf", "sentencepiece"]:
assert pkg not in shim._KEEP, f"{pkg} must still install for the notebooks that need it"
# --------------------------------------------------------------------------
# Behaviour
# --------------------------------------------------------------------------
def test_the_shipped_trl_cell_installs_nothing(shim):
# `!pip install --no-deps trl==0.22.2` is the last line of 378 notebooks.
assert _run(shim, SHIPPED_TRL_CELL) is None
def test_a_mixed_cell_keeps_only_the_unbaked_package(shim):
execd = _run(
shim,
["--no-deps", "trl==0.22.2", "peft==0.14.0", "datasets==3.0.0",
"accelerate==1.0.0", UNBAKED],
)
assert execd == ["--no-deps", UNBAKED], execd
def test_cuda_matched_wheels_are_not_replaced_by_pypi_builds(shim):
# torchao 0.17.0+cu128 and torchcodec 0.11.0+cu128 are resolved from the
# cu128 index; a PyPI pin swaps in a generic (or cu13) build.
assert _run(shim, ["torchao==0.15.0", "torchcodec==0.5"]) is None
def test_transformers_companions_cannot_desynchronise_the_sidecars(shim):
# Each sidecar ships its own matched huggingface_hub/tokenizers/safetensors;
# replacing the base-venv copies desynchronises every sidecar at once.
assert _run(shim, ["huggingface_hub==0.30.0", "tokenizers==0.20.0",
"safetensors==0.4.0"]) is None
def test_an_unbaked_package_still_installs(shim):
assert _run(shim, [UNBAKED]) == [UNBAKED]
assert _run(shim, [UNBAKED], tool = "uv") == [UNBAKED]
def test_protection_survives_a_requirements_file(shim, tmp_path):
req = tmp_path / "requirements.txt"
req.write_text(f"trl==0.22.2\npeft==0.14.0\ndatasets==3.0.0\n{UNBAKED}\n")
execd = _run(shim, ["-r", str(req)])
assert execd is not None and execd[0] == "-r"
filtered = Path(execd[1]).read_text()
assert UNBAKED in filtered
for dropped in ("trl", "peft", "datasets"):
assert dropped not in filtered, f"{dropped} slipped through the requirements file"
def test_protection_survives_a_direct_wheel_url(shim):
url = "https://files.pythonhosted.org/x/trl-0.22.2-py3-none-any.whl"
assert _run(shim, [url, UNBAKED]) == [UNBAKED]
def test_protection_survives_an_editable_vcs_install(shim):
assert _run(shim, ["-e", "git+https://github.com/huggingface/trl.git", UNBAKED]) == [UNBAKED]
def test_forwarded_installs_pin_the_protected_set_for_the_resolver(shim):
# Argument filtering alone does not stop a dependency of the kept target from
# dragging peft/datasets back down -- which is how peft 0.19.1 became 0.14.0
# with no notebook ever naming peft. Every forwarded install carries pins.
with pytest.MonkeyPatch.context() as mp:
mp.setattr(shim.sys, "argv", ["pip", "install", UNBAKED])
with pytest.raises(_Exec) as exc:
shim.main()
argv = exc.value.argv
assert "--constraint" in argv
pins = Path(argv[argv.index("--constraint") + 1]).read_text()
names = {line.split("==")[0].lower().replace("_", "-") for line in pins.splitlines() if line}
# only the installed subset is pinned, but nothing outside the protected set
assert names, "the constraints file must not be empty"
assert all(
n in shim._KEEP or n == "transformers" or n.startswith("nvidia-") for n in names
), sorted(names)