docker: address review round 2 (CI ref freeze, Studio NVRTC amd64, pip-shim edges)
docker-publish.yml: freeze the requested unsloth ref to one sha in the prepare job before the matrix fans out. UNSLOTH_REF / UNSLOTH_STUDIO_REF were raw expressions re-evaluated per base arch leg and in the Studio build, so a mutable branch (the workflow_dispatch default unsloth_ref=main) advancing during the run could bake different unsloth commits under one manifest. Resolve once (same precedence: dispatch input, else pushed tag, else triggering sha, else main; ls-remote a branch/tag to a sha, mirroring the zoo/notebooks steps) and read needs.prepare.outputs.unsloth_ref everywhere. Dockerfile.studio: run the Studio venv NVRTC cu13 swap on both arches, not arm64 only. amd64 sm_103 (B300/GB300) needs cu13 NVRTC just as arm64 sm_121 does, and the CUDA dedup never touches cuda_nvrtc, so an amd64 Studio venv would otherwise keep its bundled cu12.8 libnvrtc and fail NVRTC/jiterator JIT on compute_103. The base cu13 layer installs cuda-nvrtc-13-0 on both arches, so the target .so.13 exists here regardless of TARGETARCH. unsloth_pip_shim.py: close three ways a protected package slipped past _KEEP. Treat -e/--editable as a value-taking flag paired with its target and drop both when the target is protected (was leaving a dangling -e that failed the cell); filter -P/--upgrade-package values through _KEEP (a named baked package could be refreshed while installing another target); and parse the PEP 427 distribution name out of a wheel URL/path so a bare `pip install https://.../torch-...whl` drops instead of reinstalling the baked torch. Non-protected editables, upgrade selectors, and wheels are unchanged. Adds tests/python/test_unsloth_pip_shim.py (18 regression tests, exec captured via a patched os.execv).
This commit is contained in:
Daniel Han
parent
386d3a7c74
commit
d4dc8b6391
4 changed files with 369 additions and 25 deletions
59
.github/workflows/docker-publish.yml
vendored
59
.github/workflows/docker-publish.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 \
|
||||
|
|
|
|||
|
|
@ -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 <path|url|vcs> takes the NEXT token as its target (pip:
|
||||
# `-e, --editable <path/url>`), 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 <name> 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=<target> / --upgrade-package=<name>: 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
|
||||
|
|
|
|||
224
tests/python/test_unsloth_pip_shim.py
Normal file
224
tests/python/test_unsloth_pip_shim.py
Normal file
|
|
@ -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 <args>` 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 <filtered.txt>`.
|
||||
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
|
||||
Loading…
Add table
Add a link
Reference in a new issue