#!/opt/unsloth-venv/bin/python # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 """pip / uv shim for the Unsloth Docker notebook environment. Sits ahead of the real tools on PATH so a notebook `!pip install` cell cannot clobber the baked cu128 stack: `transformers==X` is recorded for the unsloth_nb_compat sidecar instead of installed, other _KEEP packages are skipped WHEN THEY ARE ACTUALLY BAKED IN, everything else passes through. The real tools are invoked by absolute path, so there is no recursion. """ import os, re, sys, tempfile REAL = {"pip": "/opt/unsloth-venv/bin/pip", "uv": "/opt/unsloth-venv/bin/uv"} MARKER = os.environ.get("UNSLOTH_NB_TF_MARKER", "/tmp/unsloth_nb/requested_transformers") # Kept: packages whose replacement invalidates the tested stack (ABI/CUDA-matched wheels, # or libraries unsloth patches by version). _KEEP = { "torch", "torchvision", "torchaudio", "torchao", "torchcodec", "triton", "triton-rocm", "pytorch-triton", "xformers", "vllm", "bitsandbytes", "flashinfer", "flashinfer-python", "unsloth", "unsloth-zoo", "unsloth_zoo", "trl", "peft", "datasets", "accelerate", "huggingface-hub", "huggingface_hub", "tokenizers", "safetensors", } _KEEP_PREFIX = ("nvidia-", "nvidia_") _VALUE_FLAGS = { "-r", "--requirement", "--requirements", "-c", "--constraint", "--constraints", "-i", "--index-url", "--extra-index-url", "-f", "--find-links", "--target", "-t", "--python", "-p", "--prefix", "--index-strategy", "--upgrade-strategy", "--upgrade-package", "-P", "--reinstall-package", "--no-binary", "--only-binary", "--platform", "--python-version", "--abi", "--implementation", "-e", "--editable", "--allow-insecure-host", "--build-constraints", "-b", "--cache-dir", "--color", "--config-file", "--config-setting", "-C", "--config-settings-package", "--default-index", "--directory", "--exclude-newer", "--exclude-newer-package", "--excludes", "--extra", "--fork-strategy", "--group", "--index", "--keyring-provider", "--link-mode", "--no-build-isolation-package", "--no-sources-package", "--overrides", "--prerelease", "--project", "--python-platform", "--refresh-package", "--resolution", "--torch-backend", "--no-editable-package", "--upgrade-group", "--prerelease-package", "--build-constraint", "--cert", "--client-cert", "--config-settings", "--exists-action", "--log", "--progress-bar", "--proxy", "--report", "--resume-retries", "--retries", "--root", "--root-user-action", "--src", "--timeout", "--trusted-host", "--use-deprecated", "--use-feature", "--all-releases", "--only-final", "--requirements-from-script", "--uploaded-prior-to", "--output-format", } _REQ_FILE_FLAGS = {"-r", "--requirement", "--requirements"} _CONSTRAINT_FILE_FLAGS = {"-c", "--constraint", "--constraints"} # a protected editable must drop BOTH flag and value, or a dangling -e eats the next _EDITABLE_FLAGS = {"-e", "--editable"} _UPGRADE_PKG_FLAGS = {"-P", "--upgrade-package", "--reinstall-package"} _ATTACHED_SHORT_FLAGS = {"-r", "-c", "-e", "-P"} # these rebuild baked deps; the kept target still installs once they are dropped _REINSTALL_FLAGS = {"--force-reinstall", "--ignore-installed", "-I", "--reinstall", "--exact"} # Dropped with their value; eager would upgrade every dep of a kept target. # Hidden aliases of the flags above (absent from --help; the selfcheck reads pip's option # table). --python-preference is an undocumented uv global. _ALIAS_VALUE_FLAGS = { "--default-timeout", "--local-log", "--log-file", "--pypi-url", "--source", "--source-dir", "--source-directory", "--python-preference", } _VALUE_FLAGS |= _ALIAS_VALUE_FLAGS _DROP_VALUE_FLAGS = {"--upgrade-strategy"} # these ARE the install target, with no package on the command line: without them the # shim scans to no target and no-ops while printing "ok" _TARGET_VALUE_FLAGS = {"--group", "--upgrade-group", "--requirements-from-script"} _ARCHIVE_EXTS = (".tar.gz", ".tgz", ".tar.bz2", ".tbz2", ".tar.xz", ".txz", ".tar", ".zip") def _restore_install_privileges(): if os.environ.get("UNSLOTH_NB_ROOT_INSTALL") != "1": return try: os.setgid(0) os.setuid(0) except OSError as exc: raise SystemExit(f"unsloth notebook install could not restore root: {exc}") from exc def _norm_name(name): """PEP 503 normalised name; every _KEEP / _KEEP_PREFIX comparison goes through here.""" return re.sub(r"[-_.]+", "-", name.strip()).lower() or None def _base_site_packages(): """site-packages of the venv that owns the REAL pip/uv, or [] if it cannot be located. Every scan below MUST be scoped to this rather than to sys.path. When a notebook activates a transformers sidecar, unsloth_nb_compat and unsloth_run PREPEND that sidecar to PYTHONPATH, and the sidecars are built with `uv pip install --target`, so each one ships a real transformers-X.dist-info. A bare distributions() walks sys.path in order and therefore reports the SIDECAR's transformers first, which would pin the sidecar version into the protected constraints and let the resolver move the shared base install to it. The venv that owns the REAL pip is the authoritative answer. purelib is the fallback and is safe for the same reason: it derives from sys.prefix, which PYTHONPATH cannot influence, so it can never name a sidecar.""" import glob venv = os.path.dirname(os.path.dirname(REAL["pip"])) found = [ p for p in sorted(glob.glob(os.path.join(venv, "lib", "python*", "site-packages"))) if os.path.isdir(p) ] if found: return found import sysconfig purelib = sysconfig.get_paths().get("purelib") or "" return [purelib] if purelib and os.path.isdir(purelib) else [] _INSTALLED_NAMES = [] def _installed_names(): """Every distribution name in this venv, PEP 503 normalised, or None when the scan fails. Cached because importlib.metadata rescans sys.path on every call.""" if not _INSTALLED_NAMES: try: from importlib.metadata import distributions scope = _base_site_packages() if not scope: # unknown base: None makes _is_installed answer "assume baked", which # is the stricter of the two answers raise RuntimeError("base site-packages not found") found = set() for dist in distributions(path = scope): try: n = _norm_name(dist.metadata["Name"] or "") except Exception: n = None if n: found.add(n) _INSTALLED_NAMES.append(found) except Exception: _INSTALLED_NAMES.append(None) return _INSTALLED_NAMES[0] def _is_protected(name): """True when `name` is a baked package that is REALLY installed. Every drop decision goes through here so the three call sites cannot drift apart. An absent package is nothing to protect, and dropping it turned a recovery install into a silent success: the Dockerfile lets the torchcodec bake and the non-amd64 vLLM bake fail on purpose, so a notebook `pip install vllm` printed "kept baked versions, skipped: vllm" over an image that had no vLLM at all, and the GRPO fast_inference path stayed broken. Forwarding instead either installs it under the protected constraints or fails loudly, and both beat a false success. An unreadable metadata scan keeps the old, stricter answer.""" if name is None: return False if not (name in _KEEP or name.startswith(_KEEP_PREFIX)): return False return _is_installed(name) def _is_installed(name): """True when `name` is really in this venv, or when the scan failed and the stricter "assume baked" answer is the safe one.""" present = _installed_names() return present is None or name in present def _sdist_name(basename): """Name from a `{name}-{version}.ext` basename, split at the first hyphen before a digit so legacy hyphenated names resolve too.""" low = basename.lower() stem = None for ext in _ARCHIVE_EXTS: if low.endswith(ext): stem = basename[: -len(ext)] break if stem is None: return None m = re.match(r"^(.+?)-\d", stem) return _norm_name(m.group(1) if m else stem) def _canon(token): if token.startswith("-"): return None # IGNORECASE: URL schemes are case insensitive and pip normalises them. _dref = re.match( r"^([A-Za-z0-9][A-Za-z0-9._-]*)\s*(?:\[[^\]]*\])?\s*@(?:\s|git\+|hg\+|bzr\+|svn\+|[a-z]+://)", token, re.IGNORECASE, ) if _dref: return _norm_name(_dref.group(1)) if ( re.match(r"^[a-z]+\+", token, re.IGNORECASE) or "://" in token or token.startswith((".", "/")) ): _egg = re.search(r"[#&]egg=([A-Za-z0-9][A-Za-z0-9._-]*)", token) if _egg: return _norm_name(_egg.group(1)) _whl = re.search(r"([^/\\#?]+)\.whl(?:[#?]|$)", token) if _whl: dist = _norm_name(_whl.group(1).split("-", 1)[0]) if dist: return dist _arch = _sdist_name(token.split("#", 1)[0].split("?", 1)[0].rstrip("/").rsplit("/", 1)[-1]) if _arch: return _arch if re.match(r"^[a-z]+\+", token): _rest = token.split("#", 1)[0].split("?", 1)[0] # split path from authority first, so an SSH userinfo @ is not the ref if "://" in _rest: _authority, _slash, _path = _rest.partition("://")[2].partition("/") if "@" in _path: _path = _path.rsplit("@", 1)[0] _rest = _path if _slash else _authority _seg = _rest.rstrip("/").rsplit("/", 1)[-1] _seg = _seg.split("@", 1)[0] if _seg.endswith(".git"): _seg = _seg[:-4] _seg = _norm_name(_seg) if _seg: return _seg _local = _local_project_name(token) if _local: return _local return None if "/" in token or os.sep in token: _local = _local_project_name(token) if _local: return _local if token.lower().endswith(".whl"): dist = _norm_name(token.rsplit("/", 1)[-1][:-4].split("-", 1)[0]) if dist: return dist _barch = _sdist_name(token.rsplit("/", 1)[-1]) if _barch: return _barch name = re.split(r"[<>=!~\[\s;@]", token, maxsplit = 1)[0].strip() return _norm_name(name) def _local_project_name(token): """Name pip or uv would build for a local project dir, falling back to the basename only for an installable project; None for a metadata-less dir so ordinary paths pass through.""" # Resolve from uv's --directory, not our cwd, or the project is misclassified as unknown. path = _resolve_read_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 _norm_name(_name) except Exception: pass _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 _norm_name(_name) 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 _norm_name(_base) return None def _marker_is_false(token): """True when the PEP 508 marker fails here: pip and uv skip such a requirement, so it must not reach the sidecar marker file; without packaging, or on a parse error, nothing is vetoed.""" try: from packaging.requirements import Requirement req = Requirement(token) except Exception: return False try: return req.marker is not None and not req.marker.evaluate() except Exception: return False def _version_pin(token): if _marker_is_false(token): return None m = re.search(r"==\s*([0-9][0-9A-Za-z.\-]*)", token) return m.group(1) if m else None # pip expands ${UPPERCASE} in requirements files, so `${PKG}==` would slip _KEEP. # Expanded for CLASSIFICATION only; kept lines stay verbatim. _ENV_REF_RE = re.compile(r"\$\{([A-Z0-9_]+)\}") def _expand_env_refs(text): return _ENV_REF_RE.sub(lambda m: os.environ.get(m.group(1), m.group(0)), text) _EXTRAS_RE = re.compile( r"^(?P[A-Za-z0-9][A-Za-z0-9._-]*)\s*\[(?P[^\]]+)\]\s*(?P.*)$" ) def _extras_only_target(token): """`name[extras]` when `token` asks for EXTRAS of a baked package, else None. pip and uv treat `pkg[extra]` against an already-satisfied `pkg` as "add the optional dependencies", not "replace pkg", so dropping the whole token loses every package the extra pulls in and still prints ok: `pip install "datasets[audio]"` became a no-op and the notebook died later on the missing decoder. The version specifier is dropped along with the rest of the requirement and _protected_constraints_file() pins the baked version, so what is forwarded can only ADD. A direct reference (`pkg[extra] @ url`) or a path IS a replacement request, so those keep being dropped.""" if "://" in token or token.startswith((".", "/", "-")) or os.sep in token: return None m = _EXTRAS_RE.match(token.strip()) if not m: return None rest = m.group("rest").strip() if rest.startswith("@"): return None out = m.group("name") + "[" + m.group("extras").strip() + "]" marker = rest.split(";", 1) if len(marker) == 2 and marker[1].strip(): out += " ; " + marker[1].strip() return out def _classify_flag_target(spec): name = _canon(spec) if name == "transformers": return "drop", _version_pin(spec) if _is_protected(name): return "drop", None return "keep", None def _parse_flag_line(stripped, flags): """(flag, target_or_None, inline_comment_or_None) for a ` ` requirements-file line, in any of the forms pip accepts, else (None, None, None).""" body, sep, comment = stripped.partition(" #") body = body.rstrip() comment = ("#" + comment) if sep else None for flag in flags: if body == flag or body.startswith(flag + " "): target = body[len(flag) :].strip() elif body.startswith(flag + "="): target = body[len(flag) + 1 :].strip() elif not flag.startswith("--") and body.startswith(flag) and len(body) > len(flag): target = body[len(flag) :].strip() else: continue return flag, (target or None), comment return None, None, None def _rewrite_include(line, stripped, src_dir, depth): """Rewrite a nested `-r`/`-c` include -> (new_line, changed, recorded, dropped). pip resolves it against the directory of the file it is READING, and our filtered copy lives under /tmp, so a relative include must be absolutised.""" flag, raw_target, comment = _parse_flag_line( stripped, ("-r", "--requirement", "-c", "--constraint") ) if not raw_target: return line, False, None, [] target = _expand_env_refs(raw_target) newline_char = "\n" if line.endswith("\n") else "" def _emit(new_target): rebuilt = flag + " " + new_target if comment: rebuilt += " " + comment return rebuilt + newline_char if "://" in target: return None, True, None, [flag + " " + raw_target] abs_target = target if os.path.isabs(target) else os.path.join(src_dir, target) if depth < 8: f_path, f_rec, f_drp = _filter_requirements_file(abs_target, _depth = depth + 1) # a -c include is a constraint, not an install request: no marker if flag in _CONSTRAINT_FILE_FLAGS: f_rec = None if f_path != abs_target: return _emit(f_path), True, f_rec, f_drp if not os.path.isabs(target): return _emit(abs_target), True, None, [] return line, False, None, [] def _logical_lines(lines): """Group physical requirements-file lines into the LOGICAL lines pip parses. A trailing backslash continues onto the next line, which is exactly how `pip-compile --generate-hashes` / `uv pip compile --generate-hashes` write a pinned requirement: `torch==2.11.0 \\` then indented `--hash=...` rows. Filtering physical lines dropped only the first row and published the orphaned `--hash` rows, which uv rejects outright ("Unexpected '-', expected ... the start of a requirement"), killing the whole install.""" group = [] for line in lines: group.append(line) body = line.strip() if body.endswith("\\") and not body.startswith("#"): continue yield group group = [] if group: yield group # Where the real tool will resolve a relative `-r`/`-c` path from. None means our # own cwd, which is what pip always uses. main() sets it for uv. _WORKING_DIR = None # Set for `pip uninstall`, where `pkg[extra]` names pkg itself rather than extras to add. _UNINSTALLING = False def _uv_working_dir(tool, argv): """The directory uv will change to before resolving a relative requirements path, or None. `--directory` is a uv GLOBAL, so it is accepted in every position: before `pip`, between `pip` and `install`, and after `install`. uv rejects it more than once ("cannot be used multiple times"), so the first hit is the only hit, and it overrides UV_WORKING_DIR. pip has no equivalent.""" if tool != "uv": return None for n, tok in enumerate(argv): if tok == "--directory": return argv[n + 1] if n + 1 < len(argv) else None if tok.startswith("--directory="): return tok[len("--directory=") :] or None return os.environ.get("UV_WORKING_DIR") or None def _resolve_read_path(path): """Resolves a relative requirements path against the real tool's working dir; reading from our cwd would miss silently, forwarding an unfiltered file uv installs without recording its transformers pin.""" if not _WORKING_DIR or os.path.isabs(path): return path return os.path.join(_WORKING_DIR, path) def _filter_requirements_file(path, _depth = 0): """Strip protected packages out of a `-r` file, recursing into nested includes. Returns (path_to_use, recorded_transformers_version, dropped_specs).""" read_path = _resolve_read_path(path) try: # utf-8-sig: a BOM would hide a protected pin on line 1. with open(read_path, encoding = "utf-8-sig") as f: lines = f.readlines() except OSError: # forward what the caller passed, not the resolved path: uv resolves it # itself, and pip never had a working directory to resolve against return path, None, [] src_dir = os.path.dirname(os.path.abspath(read_path)) out, dropped, recorded, changed = [], [], None, False for group in _logical_lines(lines): line = group[0] stripped = " ".join(part.strip().rstrip("\\").strip() for part in group).strip() if not stripped or stripped.startswith("#"): out.extend(group) continue if stripped.startswith("-"): e_flag, e_target, _e_comment = _parse_flag_line(stripped, ("-e", "--editable")) if e_target is not None: _action, _ver = _classify_flag_target(_expand_env_refs(e_target)) if _action == "drop": if _ver and not recorded: recorded = _ver dropped.append(e_flag + " " + e_target) changed = True continue out.extend(group) continue if len(group) > 1: out.extend(group) continue new_line, rewrote, inc_rec, inc_drp = _rewrite_include(line, stripped, src_dir, _depth) if new_line is not None: out.append(new_line) if rewrote: changed = True if inc_rec and not recorded: recorded = inc_rec dropped.extend(inc_drp) continue spec = stripped.split(" #", 1)[0].strip() # the per-requirement options a joined `--hash` block carries are noise here report = spec.split(" --", 1)[0].strip() classified = _expand_env_refs(spec) name = _canon(classified) if name is None: out.extend(group) continue if name == "transformers": v = _version_pin(classified) if v and not recorded: recorded = v # Extras of the baked transformers are additive; only the pin is stripped. extras = ( _extras_only_target(spec) if _is_installed("transformers") and not _UNINSTALLING else None ) if extras is not None: out.append(extras + ("\n" if group[-1].endswith("\n") else "")) changed = True continue dropped.append(report) changed = True continue if _is_protected(name): extras = None if _UNINSTALLING else _extras_only_target(spec) if extras is not None: out.append(extras + ("\n" if group[-1].endswith("\n") else "")) changed = True continue dropped.append(report) changed = True continue out.extend(group) if not changed: return path, None, [] try: fd, tmp = tempfile.mkstemp(prefix = "unsloth-nb-req-", suffix = ".txt") with os.fdopen(fd, "w", encoding = "utf-8") as f: f.writelines(out) except OSError as exc: # fail CLOSED: forwarding the original hands pip the specs we must filter raise SystemExit( f"[unsloth-nb] could not write a filtered copy of {path} ({exc}); " "refusing to forward a requirements file that pins protected packages." ) return tmp, recorded, dropped def _protected_constraints_file(): """Temp constraints file pinning every INSTALLED protected package, or None when the environment has none to pin. Argument filtering does not constrain the RESOLVER: a kept package declaring `torch==99.0` would replace the baked torch, and pinning makes that fail loudly. `_extras_only_target` leans on the same pin -- it forwards `torch[opt]` with the version specifier stripped precisely because the pin holds the baked version, so an unconstrained forward can REPLACE rather than only ADD. So a failure to build the file fails CLOSED, exactly as _filter_requirements_file already does for the same mkstemp error: swallowing it returned None, which the caller cannot tell apart from "nothing to protect", and the install then ran with the whole protection silently off. Only a wholesale failure aborts; one dist whose metadata cannot be read is skipped, since a half-removed `.dist-info` left by an interrupted install is ordinary and must not block every later install.""" try: from importlib.metadata import distributions scope = _base_site_packages() if not scope: raise RuntimeError(f"could not locate the site-packages of {REAL['pip']}") # Scoped, not a bare distributions(): an active sidecar is ahead on PYTHONPATH. dists = list(distributions(path = scope)) except Exception as exc: raise SystemExit( f"[unsloth-nb] could not enumerate installed packages ({exc}); refusing to " "install without the constraints that hold the baked stack in place." ) pins = {} for dist in dists: try: raw = (dist.metadata["Name"] or "").strip() version = dist.version except Exception: continue name = _norm_name(raw) if not name or not version or name in pins: continue if name == "transformers" or name in _KEEP or name.startswith(_KEEP_PREFIX): pins[name] = f"{raw}=={version}" if not pins: return None try: fd, tmp = tempfile.mkstemp(prefix = "unsloth-nb-protected-", suffix = ".txt") with os.fdopen(fd, "w", encoding = "utf-8") as f: f.write("\n".join(pins[name] for name in sorted(pins)) + "\n") except OSError as exc: raise SystemExit( f"[unsloth-nb] could not write the protected constraints file ({exc}); " "refusing to install without the pins that hold the baked stack in place." ) return tmp def _selfcheck_value_flags(): """Assert every value-taking flag the REAL pip/uv accept is classified, else the scanner misreads a flag's VALUE as an install target and forwards the install unfiltered. Run at image build time. pip is INTROSPECTED, not scraped. Its --help prints one spelling per option, so seven hidden aliases (--default-timeout among them) were invisible to the old scrape while the check still reported OK -- a false success guarding the guard. optparse knows all of them. uv is a binary with no such table, so it is scraped at all three levels rather than just `uv pip install`, and the globals it accepts but documents nowhere are probed directly. A level we cannot inspect is a FAILURE, never a pass. """ import subprocess known = _VALUE_FLAGS | _DROP_VALUE_FLAGS missing = {} probe = ( "from pip._internal.commands import create_command\n" "v=set()\n" "for o in create_command('install').parser.option_list_all:\n" " if o.takes_value(): v.update(o._long_opts); v.update(o._short_opts)\n" "print('\\n'.join(sorted(v)))\n" ) try: out = subprocess.run( [os.path.join(os.path.dirname(REAL["pip"]), "python"), "-c", probe], capture_output = True, text = True, ) if out.returncode != 0 or not out.stdout.strip(): raise OSError(out.stderr.strip()[:200] or "no output") gap = {f for f in out.stdout.split() if f.startswith("-")} - known if gap: missing["pip"] = sorted(gap) except OSError as exc: print( f"[unsloth-nb] could not introspect pip's option table ({exc}); refusing to " "certify the value-flag list from --help alone, which cannot see aliases.", file = sys.stderr, ) sys.exit(1) seen_uv = False for cmd in ( [REAL["uv"], "--help"], [REAL["uv"], "pip", "--help"], [REAL["uv"], "pip", "install", "--help"], ): try: out = subprocess.run(cmd, capture_output = True, text = True).stdout except OSError: continue seen_uv = True flags = set() for m in re.finditer(r"^\s+(-\w)?,?\s*(--[\w-]+)[= ]<", out, re.M): if m.group(1): flags.add(m.group(1)) flags.add(m.group(2)) for m in re.finditer(r"^\s+(-\w) <", out, re.M): flags.add(m.group(1)) gap = flags - known if gap: missing.setdefault("uv", []).extend(sorted(gap)) if not seen_uv: print("[unsloth-nb] could not run uv to check its flags", file = sys.stderr) sys.exit(1) # uv globals that are accepted but documented nowhere: verify each still takes a # value, so the list rots loudly instead of silently. for flag in sorted(_ALIAS_VALUE_FLAGS): if flag == "--python-preference": r = subprocess.run( [REAL["uv"], flag, "system", "pip", "install", "--help"], capture_output = True, text = True, ) if r.returncode != 0: print( f"[unsloth-nb] uv no longer accepts `{flag} `; the entry in " "_ALIAS_VALUE_FLAGS is stale and may now swallow an install target.", file = sys.stderr, ) sys.exit(1) if missing: print(f"[unsloth-nb] value flags missing from _VALUE_FLAGS: {missing}", file = sys.stderr) sys.exit(1) print("[unsloth-nb] value-flag selfcheck OK") sys.exit(0) def _expand_short_clusters(argv): """Split `-qr reqs.txt` into `-q -r reqs.txt`, which is what both real tools see. A cluster reached no handler at all: only `tok[:2]` is tested against the attached-value flags, the exact-token comparisons never match, and the fallback keeps any unrecognised `-...` verbatim. So `pip install -qr reqs.txt` forwarded the ORIGINAL requirements file, unfiltered and with nothing recorded, and `-I` hid inside a cluster the same way and escaped the reinstall handling. A value-taking short flag consumes the REST of the cluster as its attached value, or the next argv token when it ends the cluster, exactly as getopt does. """ shorts = {f for f in _VALUE_FLAGS if len(f) == 2 and f[0] == "-" and f[1] != "-"} shorts |= _ATTACHED_SHORT_FLAGS out = [] for tok in argv: if len(tok) <= 2 or tok[0] != "-" or tok[1] == "-" or tok in _VALUE_FLAGS: out.append(tok) continue i = 1 while i < len(tok): flag = "-" + tok[i] rest = tok[i + 1 :] if flag in shorts: out.append(flag + rest) # attached value, or bare when rest is empty break out.append(flag) i += 1 else: continue return out def _install_index(tool, argv): """Index of the install SUBCOMMAND, or None when this is not a package install. pip's command path is `pip [opts] install`, but uv's is `uv [opts] pip [opts] install`, and uv has other subcommands ending in `install`. A bare argv.index("install") matched those too, which was not merely a wasted rewrite: `uv python install 3.13` and `uv tool install ruff` got the protected `--constraint` appended, which neither accepts, and `uv tool install transformers` was filtered down to nothing and reported ok, installing NO tool at all. So match the command path POSITIONALLY: `install` for pip, `pip install` for uv. Finding it needs option VALUES skipped, not just options -- `uv pip --directory /tmp install torch` is valid (uv lists --directory under Global options), and treating `/tmp` as the subcommand loses the command entirely, which is the worst outcome available here: the install runs unfiltered and unconstrained, free to replace the baked torch/CUDA stack. _VALUE_FLAGS is the same set the tail scanner uses and _selfcheck_value_flags() holds it to the real CLIs at build time.""" return _subcommand_index(tool, argv, "install") def _subcommand_index(tool, argv, sub): expect = [sub] if tool == "pip" else ["pip", sub] got = _positionals(argv) if len(got) < len(expect): return None for (idx, tok), want in zip(got, expect): if tok != want: return None return got[len(expect) - 1][0] def _uninstall_file(path, dropped): """Filtered path, or None when nothing is left: pip errors on an empty file.""" filtered, _, drp = _filter_requirements_file(path) dropped.extend(drp) if filtered == path: return filtered with open(filtered, encoding = "utf-8") as f: if any(ln.strip() and not ln.strip().startswith("#") for ln in f): return filtered return None def _uninstall(tool, argv, i): global _UNINSTALLING _UNINSTALLING = True head, tail = argv[: i + 1], _expand_short_clusters(argv[i + 1 :]) child_env = None if tool == "pip": env_tail = [x for val in os.environ.get("PIP_REQUIREMENT", "").split() for x in ("-r", val)] if env_tail: tail = env_tail + tail child_env = {k: v for k, v in os.environ.items() if k != "PIP_REQUIREMENT"} keep, dropped = [], [] has_target = False pending = None for tok in tail: if pending is not None: if pending in _REQ_FILE_FLAGS: tok = _uninstall_file(tok, dropped) if tok is not None: keep += [pending, tok] has_target = True else: keep += [pending, tok] pending = None continue if tok.startswith("--") and "=" in tok: flag, _, val = tok.partition("=") if flag in _REQ_FILE_FLAGS: filtered = _uninstall_file(val, dropped) if filtered is not None: keep.append(flag + "=" + filtered) has_target = True continue keep.append(tok) continue if len(tok) > 2 and tok[:2] in _REQ_FILE_FLAGS: filtered = _uninstall_file(tok[2:], dropped) if filtered is not None: keep += [tok[:2], filtered] has_target = True continue if tok in _VALUE_FLAGS: pending = tok continue if tok.startswith("-"): keep.append(tok) continue name = _canon(tok) if _is_protected(name) or (name == "transformers" and _is_installed(name)): dropped.append(tok) continue keep.append(tok) has_target = True if dropped: print("[unsloth-nb] kept baked versions, skipped: " + " ".join(dropped)) if not has_target and dropped: print("[unsloth-nb] nothing to uninstall after keeping the baked stack; ok.") return cmd = [REAL[tool]] + head + keep sys.stdout.flush() if child_env is not None: os.execve(REAL[tool], cmd, child_env) os.execv(REAL[tool], cmd) def _positionals(argv): """Returns (index, token) for every positional, skipping options and their values; --flag=value carries its own value while a bare value flag eats the next token; _VALUE_FLAGS is shared with the tail scanner and checked against the real CLIs at build time by _selfcheck_value_flags().""" out = [] skip = False for i, tok in enumerate(argv): if skip: skip = False continue if tok.startswith("-"): if "=" not in tok and tok in _VALUE_FLAGS: skip = True continue out.append((i, tok)) return out def _is_uv_pip_sync(argv): """uv pip sync cannot be made safe by filtering: it uninstalls every package absent from the requirements file, so keeping a protected package out deletes it, and --constraint never prevents removal.""" return [tok for _, tok in _positionals(argv)][:2] == ["pip", "sync"] # Options that install somewhere OTHER than the guarded venv. _DEST_DIR_FLAGS = {"--target", "-t", "--prefix"} # --root is a relocation prefix, not a destination: `--root /` still writes the venv. _DEST_ROOT_FLAGS = {"--root"} _DEST_PYTHON_FLAGS = {"--python", "-p"} # Per tool: pip honours PIP_TARGET/PREFIX/ROOT/PYTHON, uv only UV_PYTHON. A bypass granted # on a variable the tool ignores would install unfiltered. # --group (PEP 735) and --requirements-from-script (PEP 723) supply targets the shim # cannot read. --upgrade-group is rejected by both tools. _UNINSPECTABLE_TARGET_FLAGS = {"--group", "--requirements-from-script"} _PIP_TARGET_ENV = {"PIP_REQUIREMENT": "-r", "PIP_EDITABLE": "-e"} _DEST_DIR_ENV = {"pip": ("PIP_TARGET", "PIP_PREFIX"), "uv": ()} # PIP_ROOT is the environment spelling of --root, so it relocates rather than names a # destination and has to be resolved the same way _DEST_ROOT_ENV = {"pip": ("PIP_ROOT",), "uv": ()} _DEST_PYTHON_ENV = {"pip": ("PIP_PYTHON",), "uv": ("UV_PYTHON",)} def _flag_values(argv, flags): """Every value given for any of `flags`, in the spellings both tools accept: separated, `--flag=value`, and a short flag with its value attached.""" out, want = [], False for tok in argv: if want: out.append(tok) want = False continue if tok in flags: want = True continue for f in flags: if tok.startswith(f + "="): out.append(tok[len(f) + 1 :]) break if len(f) == 2 and f[0] == "-" and f[1] != "-" and tok.startswith(f) and len(tok) > 2: out.append(tok[2:]) break return [v for v in out if v] def _base_venv_root(): """The venv that owns the REAL pip/uv, i.e. the one this shim exists to guard.""" return os.path.realpath(os.path.dirname(os.path.dirname(REAL["pip"]))) def _relocated_under_root(root, path): """Where `path` ends up once pip re-anchors it under `--root `. pip keeps the whole absolute path and hangs it beneath the root, so `--root /` yields `path` itself and anything else lands outside the venv. Anchored on the venv ROOT rather than on site-packages so it agrees with _inside_base_venv even when the two are derived differently.""" _drive, rest = os.path.splitdrive(path) return os.path.join(root, rest.lstrip(os.sep)) def _inside_base_venv(path): root = _base_venv_root() resolved = os.path.realpath(path) return resolved == root or resolved.startswith(root + os.sep) def _installs_elsewhere(tool, argv): """True when this install is aimed at an environment that is NOT the baked venv. Every reason to filter is about protecting that venv, so filtering an install bound for somewhere else drops the packages the caller actually asked for and then reports success having installed nothing. Deliberately conservative: only a destination we can positively resolve OUTSIDE the venv counts. A bare `--python python3` or `--python 3.12` may well name the baked interpreter, so it keeps the protective behaviour. Environment variables count only for the tool that actually reads them; a flag is left unscoped because a tool that does not accept one fails loudly on it rather than silently.""" # pip uninstall has no --target/--prefix/--root and ignores their variables, so # they still leave it pointed at the baked venv dest_env = not (tool == "pip" and _subcommand_index(tool, argv, "uninstall") is not None) dirs = _flag_values(argv, _DEST_DIR_FLAGS) dirs += [os.environ[v] for v in _DEST_DIR_ENV[tool] if dest_env and os.environ.get(v)] roots = _flag_values(argv, _DEST_ROOT_FLAGS) roots += [os.environ[v] for v in _DEST_ROOT_ENV[tool] if dest_env and os.environ.get(v)] # a root re-anchors the venv's OWN paths beneath it, so resolve where each one # actually lands rather than classifying the root value itself dirs += [_relocated_under_root(r, _base_venv_root()) for r in roots] if any(not _inside_base_venv(d) for d in dirs): return True pythons = _flag_values(argv, _DEST_PYTHON_FLAGS) pythons += [os.environ[v] for v in _DEST_PYTHON_ENV[tool] if os.environ.get(v)] return any((os.path.isabs(p) or os.sep in p) and not _inside_base_venv(p) for p in pythons) def main(): tool = "uv" if os.path.basename(sys.argv[0]).startswith("uv") else "pip" argv = sys.argv[1:] if argv[:1] == ["--unsloth-selfcheck-value-flags"]: _selfcheck_value_flags() _restore_install_privileges() if os.environ.get("UNSLOTH_NB_SHIM") != "1": os.execv(REAL[tool], [REAL[tool]] + argv) return # Before the sync refusal as well: `uv pip sync --python ` only uninstalls # from that other environment, so there is no baked stack for it to strip. if _installs_elsewhere(tool, argv): os.execv(REAL[tool], [REAL[tool]] + argv) return if tool == "uv" and _is_uv_pip_sync(argv): raise SystemExit( "[unsloth-nb] refusing `uv pip sync`: it UNINSTALLS every package missing " "from the requirements file, which would strip the baked torch/CUDA stack " "this image is built on, and no --constraint can prevent a removal. Use " "`uv pip install -r `, which adds without removing." ) global _WORKING_DIR _wd = _uv_working_dir(tool, argv) _WORKING_DIR = os.path.abspath(_wd) if _wd else None i = _subcommand_index(tool, argv, "uninstall") if i is not None: _uninstall(tool, argv, i) return i = _install_index(tool, argv) if i is None: os.execv(REAL[tool], [REAL[tool]] + argv) return # after the install check, so a subcommand that merely mentions one of these is # left alone, and after the destination bypass, so a group installed into another # environment is not our problem if _flag_values(argv, _UNINSPECTABLE_TARGET_FLAGS): raise SystemExit( "[unsloth-nb] refusing an install whose targets this shim cannot read: " + " ".join(sorted(_UNINSPECTABLE_TARGET_FLAGS)) + ". A dependency group or a PEP 723 script block supplies package " "specifications the same way a requirements file does, but they live " "inside a pyproject.toml table or a script header rather than a plain " "list, so they are forwarded uninspected: a same-version local path or " "sdist naming a baked package resolves cleanly against the injected " "constraints and replaces the stack this image is built on. Pass the " "same requirements with `-r `, which IS filtered." ) # pip reads PIP_