docker: dedupe repeated rationale comments and parametrize the pip-shim tests
Comment-only consolidation: the sm_103/sm_121 + cu13 JIT story and the xformers-aarch64 note were each told four times across docker/Dockerfile; keep the header telling canonical and cross-reference it elsewhere (same for the workflow's six retellings of the resolve-refs-once rationale and Dockerfile.studio's NVRTC block). Comments that pointed at the removed dev scripts now name the underlying command or artifact instead. Non-comment lines of both Dockerfiles and the workflow are byte-identical. unsloth_sync_notebooks.sh folds the three copies of the override -> PATH -> sibling helper resolution into one resolve_helper(), behavior verified for all four modes including graceful absence under set -u. unsloth_pip_shim.py collapses an if/else whose branches were identical and merges the structurally duplicate _parse_include/_parse_editable into one _parse_flag_line. The test suite folds 35 near-duplicate tests into 8 parametrized groups with exact case-count parity (69 collected before and after, 81 passing including the nb-pip-magic suite). Cuts another 144 lines with zero behavior change outside the two refactors.
This commit is contained in:
Daniel Han
parent
da0e908d55
commit
cd982a121d
6 changed files with 261 additions and 405 deletions
80
.github/workflows/docker-publish.yml
vendored
80
.github/workflows/docker-publish.yml
vendored
|
|
@ -80,12 +80,15 @@ permissions:
|
|||
|
||||
jobs:
|
||||
# ---------------------------------------------------------------------------
|
||||
# Resolve the llama.cpp prebuilt release ONCE, up front, so both arch legs of
|
||||
# the base build bake the identical GGUF binaries. Resolving "latest" inside
|
||||
# each leg would let upstream publish a new release between the amd64 and
|
||||
# arm64 builds, putting different binaries under one published image tag.
|
||||
# An explicit dispatch input pins a frozen release; otherwise we follow the
|
||||
# /releases/latest redirect to a concrete tag (mirrors docker/build.sh).
|
||||
# Resolve every upstream ref ONCE, up front -- the llama.cpp prebuilt tag plus
|
||||
# one unsloth sha, one zoo sha and one notebooks commit -- so both arch legs
|
||||
# of the base build AND the Studio build bake identical bits. Resolving
|
||||
# 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. An explicit dispatch input pins a frozen value; otherwise a
|
||||
# branch/tag is frozen to a sha via ls-remote (falling back to the bare ref
|
||||
# on a lookup miss so the Dockerfile can still fetch it by name), and the
|
||||
# llama "latest" follows the /releases/latest redirect (mirrors build.sh).
|
||||
# ---------------------------------------------------------------------------
|
||||
prepare:
|
||||
runs-on: ubuntu-latest
|
||||
|
|
@ -94,11 +97,7 @@ jobs:
|
|||
contents: read
|
||||
outputs:
|
||||
llama_tag: ${{ steps.llama.outputs.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.
|
||||
# Resolved once, shared by every consumer -- see the job header.
|
||||
unsloth_ref: ${{ steps.unsloth_ref.outputs.ref }}
|
||||
zoo_ref: ${{ steps.zoo_ref.outputs.ref }}
|
||||
notebooks_commit: ${{ steps.notebooks.outputs.commit }}
|
||||
|
|
@ -117,16 +116,10 @@ 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 that
|
||||
# advances during the ~4h base + Studio run. Same requested-ref precedence
|
||||
# the inline build-arg used: the dispatch input wins (blank by default,
|
||||
# so stable tags stay enabled), 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.
|
||||
# Requested-ref precedence (same as the old inline build-arg): the
|
||||
# dispatch input wins (blank by default, so stable tags stay enabled),
|
||||
# else the pushed tag, else the triggering commit sha, else main --
|
||||
# then frozen to one sha per the job header.
|
||||
- name: Resolve unsloth ref
|
||||
id: unsloth_ref
|
||||
env:
|
||||
|
|
@ -150,9 +143,7 @@ jobs:
|
|||
# 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
|
||||
# made every tag publish fail inside the Dockerfile's zoo install. Resolved
|
||||
# once here and forwarded to the base build AND the Studio build, so the
|
||||
# full image's Studio venv runs the same zoo as the base image.
|
||||
# made every tag publish fail inside the Dockerfile's zoo install.
|
||||
- name: Resolve unsloth-zoo ref
|
||||
id: zoo_ref
|
||||
run: |
|
||||
|
|
@ -164,11 +155,7 @@ jobs:
|
|||
fi
|
||||
fi
|
||||
REF="${REF:-main}"
|
||||
# Freeze a branch/tag ref to ONE concrete sha before the matrix fans
|
||||
# out, so both arch legs (and the base vs Studio builds) bake the
|
||||
# identical unsloth-zoo even if main advances mid-build. A 40-char sha
|
||||
# is already frozen; resolve anything else via ls-remote, as the
|
||||
# notebooks step does, falling back to the bare ref on a lookup miss.
|
||||
# Freeze to one sha per the job header; a 40-char sha already is one.
|
||||
if printf '%s' "$REF" | grep -Eq '^[0-9a-f]{40}$'; then
|
||||
SHA="$REF"
|
||||
else
|
||||
|
|
@ -178,12 +165,9 @@ jobs:
|
|||
echo "ref=${SHA}" >> "$GITHUB_OUTPUT"
|
||||
echo "unsloth-zoo ref: ${SHA}"
|
||||
|
||||
# Freeze unslothai/notebooks to ONE concrete commit so both arch legs (and
|
||||
# release reruns) bake the identical baked-notebook templates and
|
||||
# .unsloth_template_commit, even if upstream advances mid-build. A 40-char
|
||||
# sha input is already frozen; a branch/tag (default main) is resolved to
|
||||
# its current sha via ls-remote, falling back to the bare ref on a lookup
|
||||
# miss so the Dockerfile can still fetch it by name.
|
||||
# Freeze unslothai/notebooks to ONE commit per the job header, so the
|
||||
# baked templates + .unsloth_template_commit are identical across legs
|
||||
# and release reruns.
|
||||
- name: Resolve unsloth/notebooks commit
|
||||
id: notebooks
|
||||
env:
|
||||
|
|
@ -276,20 +260,9 @@ jobs:
|
|||
outputs: type=image,name=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }},push-by-digest=true,name-canonical=true,push=true
|
||||
# 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 (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 (an explicit
|
||||
# dispatch unsloth_ref) 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
|
||||
# shared with the Studio build so both venvs run the same zoo.
|
||||
# LLAMA_PREBUILT_TAG / UNSLOTH_NOTEBOOKS_REF (from the prepare job):
|
||||
# one concrete tag / commit shared by both arch legs so the
|
||||
# published manifest is byte-reproducible across platforms.
|
||||
# passed as a bogus --build-arg. All four values come from the prepare
|
||||
# job: resolved once so both arch legs and the Studio build bake
|
||||
# identical bits (precedence rules live on prepare's steps).
|
||||
build-args: |
|
||||
CUDA_VERSION=12.8.1
|
||||
UBUNTU_VERSION=24.04
|
||||
|
|
@ -465,12 +438,9 @@ 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 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 --
|
||||
# Both refs are the SAME resolved shas the base build baked (prepare
|
||||
# job), so the Studio tree + its zoo overlay match the base venv even
|
||||
# if the branch moved mid-run. (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 }}
|
||||
|
|
|
|||
|
|
@ -30,7 +30,7 @@
|
|||
#
|
||||
# Cross-arch build (DGX Spark / GB10 / sm_121):
|
||||
# The arm64 image is built via QEMU binfmt emulation on an x86_64 host:
|
||||
# bash docker/setup_qemu.sh # one-time host setup
|
||||
# docker run --privileged --rm tonistiigi/binfmt --install all # one-time host setup
|
||||
# docker buildx build --platform linux/arm64 -t unsloth-blackwell:arm64 .
|
||||
# The resulting arm64 image runs NATIVELY on aarch64 hosts (DGX Spark, Grace).
|
||||
# QEMU is only used at build time -- runtime emulation does NOT work for CUDA.
|
||||
|
|
@ -41,7 +41,7 @@
|
|||
# * Docker with buildkit (default since 23.x)
|
||||
# * docker buildx (mandatory for multi-platform; install: apt install docker-buildx)
|
||||
# * nvidia-container-toolkit (only needed for `docker run --gpus all` at test time)
|
||||
# * For arm64 builds on x86_64 hosts: QEMU binfmt (see docker/setup_qemu.sh)
|
||||
# * For arm64 builds on x86_64 hosts: QEMU binfmt (the one-time setup above)
|
||||
# * A GPU is NOT required at build time.
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
|
@ -56,7 +56,7 @@ FROM nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu${UBUNTU_VERSION} AS builder
|
|||
|
||||
# TARGETARCH is auto-populated by buildx ("amd64" or "arm64"). We use it to
|
||||
# select an unsloth extras set that matches the wheels actually available for
|
||||
# the target platform (xformers has no cu128 aarch64 wheel as of 0.0.34).
|
||||
# the target platform (the xformers aarch64 gap -- see header).
|
||||
ARG TARGETARCH
|
||||
ARG PYTHON_VERSION
|
||||
ENV DEBIAN_FRONTEND=noninteractive \
|
||||
|
|
@ -147,10 +147,8 @@ RUN python -m venv ${VENV} && ${VENV}/bin/pip install -U pip wheel setuptools
|
|||
#
|
||||
# Why arm64 uses a different extra:
|
||||
# `cu128-ampere-torch2110` transitively pulls `cu128onlytorch2110` whose
|
||||
# xformers wheel URL is hardcoded to manylinux_2_28_x86_64. There is no
|
||||
# cu128 aarch64 wheel for xformers as of 0.0.35. We use the plain
|
||||
# `huggingface` extra on arm64 -- Unsloth falls back to its native SDPA
|
||||
# kernels (a ~5-10% slowdown vs xformers; functionally complete).
|
||||
# xformers wheel URL is hardcoded to manylinux_2_28_x86_64 (the aarch64
|
||||
# wheel gap -- see header), so arm64 takes the plain `huggingface` extra.
|
||||
#
|
||||
# Why no `flash-attn` here:
|
||||
# - FA3 is hard-refused on Blackwell (Dao-AILab/flash-attention#1810).
|
||||
|
|
@ -194,7 +192,7 @@ RUN set -eux \
|
|||
# since 0.17, so the SAME pass now runs on the arm64 leg (DGX Spark /
|
||||
# GB10 class). amd64 failures abort the build; arm64 is fail-soft
|
||||
# because aarch64 wheels are newer and the GPU-side kernels there are
|
||||
# validated on Spark hardware via docker_confirm.sh, not in CI.
|
||||
# validated manually on Spark hardware, not in CI.
|
||||
#
|
||||
# https://docs.vllm.ai/en/latest/getting_started/installation/gpu/
|
||||
# https://wheels.vllm.ai/nightly
|
||||
|
|
@ -279,11 +277,9 @@ RUN set -eux \
|
|||
# JupyterLab so the published image runs unslothai/notebooks out of the box:
|
||||
# docker run --gpus all -p 8888:8888 unsloth/unsloth \
|
||||
# jupyter lab --ip 0.0.0.0 --port 8888 --allow-root --no-browser
|
||||
# Installed as a separate pass AFTER the torch-pinned resolves on purpose:
|
||||
# jupyterlab's dependency closure is pure-Python (tornado, jinja2, nbconvert,
|
||||
# nbclient, ipykernel, ...) and never names torch, so uv cannot disturb the
|
||||
# cu128 pin set here. Naming torch in this pass would be actively dangerous:
|
||||
# without the cu128 extra index uv could swap in the PyPI CPU wheel.
|
||||
# Separate pass AFTER the torch-pinned resolves: this closure is pure-Python
|
||||
# and never names torch, so uv cannot disturb the cu128 pin set (naming torch
|
||||
# without the cu128 index could swap in the PyPI CPU wheel).
|
||||
# matplotlib rides along for the notebook crowd: plotting is table stakes in
|
||||
# a Jupyter image, and several model repos' trust_remote_code modeling files
|
||||
# (e.g. DeepSeek-OCR) import it unconditionally.
|
||||
|
|
@ -301,15 +297,14 @@ RUN set -eux \
|
|||
# einx TTS codec tensor-rearrange (Llasa / Oute / Spark TTS)
|
||||
# librosa Whisper audio feature extraction (pairs with soundfile + torchcodec)
|
||||
# ftfy Oute TTS text normalisation
|
||||
# decord (ERNIE-VL video decode) is installed separately below: it ships no
|
||||
# aarch64 wheel, so a hard install here would break the arm64 build.
|
||||
# decord is installed separately below (no aarch64 wheel; see that block).
|
||||
# librosa pulls numba/soxr/audioread; numba is already pinned >=0.65 (numpy 2.4
|
||||
# compatible) by the vLLM pass, so the resolve must NOT move torch/numpy/numba --
|
||||
# the assertion below fails the build loudly if it did.
|
||||
# Pinned (==) to the resolved, tested versions for reproducible rebuilds -- the
|
||||
# same convention as the cu128 core (torch/torchvision/torchaudio). Bump these
|
||||
# deliberately, not silently on the next build. Transitive deps of these are
|
||||
# captured by the full venv lockfile (docker/freeze.sh -> requirements.lock.txt).
|
||||
# captured by the in-image pin record (/opt/unsloth-venv/requirements.lock.txt).
|
||||
RUN ${VENV}/bin/uv pip install \
|
||||
--python ${VENV}/bin/python \
|
||||
"jupyterlab==4.6.0" "notebook==7.6.0" "ipywidgets==8.1.8" "matplotlib==3.11.0" \
|
||||
|
|
@ -462,8 +457,7 @@ assert "sm_120" in arches, f"sm_120 missing: {arches}"
|
|||
print(f"OK: torch 2.11.0+cu128 with sm_100 + sm_120 native SASS intact ({target})")
|
||||
|
||||
from importlib.metadata import version, PackageNotFoundError
|
||||
# xformers has no cu128 aarch64 wheel as of 0.0.34, so we only require it
|
||||
# on amd64. Everything else is platform-agnostic.
|
||||
# xformers is amd64-only (the aarch64 wheel gap -- see Dockerfile header).
|
||||
REQUIRED = ["torch", "triton", "bitsandbytes", "unsloth",
|
||||
"unsloth_zoo", "transformers", "trl", "peft", "accelerate"]
|
||||
if target == "amd64":
|
||||
|
|
@ -515,14 +509,10 @@ ENV DEBIAN_FRONTEND=noninteractive \
|
|||
PATH=/opt/unsloth-venv/bin:${PATH} \
|
||||
HF_HOME=/workspace/.cache/huggingface \
|
||||
TRITON_CACHE_DIR=/workspace/.cache/triton \
|
||||
# Keep the arch list visible at runtime in case the user source-builds anything
|
||||
# extra inside the container (e.g. a custom CUDA op). Same list as the builder
|
||||
# stage so a `pip install some-cuda-ext` inside the container gets a SASS blob
|
||||
# that covers every supported arch. 10.3 (B300) is intentionally omitted: it
|
||||
# runs sm_100 SASS, and the bundled CUDA 12.8 nvcc cannot compile compute_103
|
||||
# (added in CUDA 12.9), so listing it would fail any such in-container build.
|
||||
# The same cu12.8 limit affects runtime Triton/NVRTC JIT on amd64 sm_103 (see
|
||||
# the header note); precompiled SASS still runs there via sm_100 forward-compat.
|
||||
# Keep the arch list visible at runtime so an in-container source build of a
|
||||
# custom CUDA op gets the same SASS coverage as the builder stage. 10.3 is
|
||||
# omitted for the same cu12.8-cannot-emit-compute_103 reason as the builder
|
||||
# list + header (sm_103 runs sm_100 SASS via forward-compat).
|
||||
TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0;10.0;12.0+PTX"
|
||||
|
||||
# zstd: the official Ollama notebooks run `curl ollama.com/install.sh | sh`
|
||||
|
|
@ -565,20 +555,14 @@ RUN CUDA_PKG="$(echo "${CUDA_VERSION}" | awk -F. '{print $1"-"$2}')" \
|
|||
|
||||
COPY --from=builder /opt/unsloth-venv /opt/unsloth-venv
|
||||
|
||||
# Blackwell JIT fix for sm_103 (B300/GB300, amd64) and sm_121 (DGX Spark /
|
||||
# GB10, arm64). Precompiled SASS already runs on both via forward-compat
|
||||
# (sm_100 SASS -> sm_103, sm_120 SASS -> sm_121); this covers the JIT gap.
|
||||
#
|
||||
# Two cu12.8 compilers baked into the stack cannot emit compute_103 /
|
||||
# compute_121, so JIT-heavy paths error out or silently downgrade:
|
||||
# Blackwell JIT fix for sm_103 (amd64) and sm_121 (arm64) -- the cu12.8 JIT
|
||||
# gap described in the header. Two JIT paths need the cu13 override:
|
||||
#
|
||||
# (1) torch's bundled libnvrtc.so.12 is CUDA 12.8. The jiterator C++ side
|
||||
# queries the device cap directly, so any NVRTC JIT path (e.g.
|
||||
# torch.fft.rfft(complex).abs(), used inside mel-spectrogram code)
|
||||
# errors out on sm_103/sm_121. Fix: keep cu12.8 as the immutable default
|
||||
# (real lib saved as .cu128.orig, libnvrtc.so.12 -> it) and stage a cu13
|
||||
# alias (.cu13); the runtime retargets libnvrtc.so.12 -> .cu13 for those
|
||||
# two arches only -- see below.
|
||||
# errors out on sm_103/sm_121. Fix: stage a cu13 NVRTC alias beside the
|
||||
# immutable cu12.8 default (mechanics at the staging step below).
|
||||
#
|
||||
# (2) Triton's nvidia backend invokes its OWN bundled ptxas, which in the
|
||||
# triton 3.6.0 we pin is still CUDA 12.8 (V12.8.93): it tops out at
|
||||
|
|
@ -586,15 +570,13 @@ COPY --from=builder /opt/unsloth-venv /opt/unsloth-venv
|
|||
# triton-lang/triton#8335. Fix: install cu13 ptxas and point Triton at
|
||||
# it with TRITON_PTXAS_PATH.
|
||||
#
|
||||
# Both cu13 tools are CPU-side compilers (no libcuda call), so they install
|
||||
# alongside the cu128 runtime with no driver-floor bump at INSTALL time (570+).
|
||||
# But their OUTPUT cubin needs a >= 580 driver to LOAD, so they are NOT baked as a
|
||||
# global ENV/symlink default -- forcing every host's JIT through cu13 would break
|
||||
# the Ampere/Ada/Hopper/Turing GPUs this image still supports on 570-579 drivers.
|
||||
# They are activated per device at runtime only for sm_103/sm_121 (which launched
|
||||
# after cu12.8 and only ship on >= 580 drivers, so gating cu13 to them is always
|
||||
# safe) -- see select_cuda_jit_tools in entrypoint.sh. Both arches carry the
|
||||
# ~400 MB: amd64 needs it for sm_103, arm64 for sm_121.
|
||||
# Both cu13 tools are CPU-side compilers (no driver-floor bump at INSTALL
|
||||
# time), but their OUTPUT cubin needs a >= 580 driver to LOAD, so neither is
|
||||
# baked as a global ENV/symlink default -- that would break the Ampere/Ada/
|
||||
# Hopper/Turing GPUs this image still supports on 570-579 drivers. Instead
|
||||
# select_cuda_jit_tools in entrypoint.sh activates them per device, only for
|
||||
# sm_103/sm_121 (which only ship on >= 580 drivers, so the gate is always
|
||||
# safe). Both arches carry the ~400 MB: amd64 for sm_103, arm64 for sm_121.
|
||||
RUN set -eux; \
|
||||
# The nvidia/cuda base already configures the CUDA apt repo (x86_64 or
|
||||
# sbsa) with its own Signed-By keyring at
|
||||
|
|
@ -610,26 +592,21 @@ RUN set -eux; \
|
|||
cuda-nvrtc-13-0 \
|
||||
cuda-nvcc-13-0; \
|
||||
rm -rf /var/lib/apt/lists/*; \
|
||||
# (1) NVRTC staging. cu12.8 stays the IMMUTABLE default; a cu13 alias is
|
||||
# staged beside it for the runtime switch. Keep the wheel's real
|
||||
# cu12.8 lib as .cu128.orig, point libnvrtc.so.12 at it (relative
|
||||
# symlink), and add .cu13 -> the cu13 lib. select_cuda_jit_tools in
|
||||
# entrypoint.sh retargets libnvrtc.so.12 -> .cu13 ONLY for sm_103/
|
||||
# sm_121 hosts. Because the default needs no runtime write, a non-root
|
||||
# `docker run --user` container -- which cannot rewrite the symlink --
|
||||
# keeps cu12.8, which every supported 570+ driver can load; a baked
|
||||
# cu13 default would instead leave those hosts on a cubin a 570-579
|
||||
# driver cannot load.
|
||||
# (1) NVRTC staging: keep the wheel's real cu12.8 lib as .cu128.orig,
|
||||
# point libnvrtc.so.12 at it (relative symlink), and stage
|
||||
# .cu13 -> the cu13 lib; select_cuda_jit_tools retargets the
|
||||
# symlink ONLY on sm_103/sm_121 hosts. The default needs no
|
||||
# runtime write, so a non-root `docker run --user` container
|
||||
# (which cannot rewrite the symlink) keeps cu12.8, loadable on
|
||||
# every supported 570+ driver.
|
||||
NVRTC_DIR=/opt/unsloth-venv/lib/python${PYTHON_VERSION}/site-packages/nvidia/cuda_nvrtc/lib; \
|
||||
if [ -f "${NVRTC_DIR}/libnvrtc.so.12" ] && [ ! -L "${NVRTC_DIR}/libnvrtc.so.12" ]; then \
|
||||
mv "${NVRTC_DIR}/libnvrtc.so.12" "${NVRTC_DIR}/libnvrtc.so.12.cu128.orig"; \
|
||||
ln -s libnvrtc.so.12.cu128.orig "${NVRTC_DIR}/libnvrtc.so.12"; \
|
||||
ln -s /usr/local/cuda-13.0/lib64/libnvrtc.so.13 "${NVRTC_DIR}/libnvrtc.so.12.cu13"; \
|
||||
fi
|
||||
# (2) ptxas override. triton 3.6.0's ptxas is cu12.8 (no sm_103/sm_121), so those
|
||||
# two arches need the cu13 ptxas installed above. Not baked as a global ENV for the
|
||||
# same driver-floor reason as NVRTC (a cu13 cubin needs a >= 580 driver to load);
|
||||
# TRITON_PTXAS_PATH is selected per device at boot -- see select_cuda_jit_tools.
|
||||
# (2) ptxas: the cu13 nvcc package above provides it; TRITON_PTXAS_PATH is set
|
||||
# per device at boot (select_cuda_jit_tools) for the same driver-floor reason.
|
||||
|
||||
# Register the venv's torch + NVIDIA lib dirs with the loader so torchcodec
|
||||
# (installed in the builder, see the bake comment there) can dlopen them.
|
||||
|
|
|
|||
|
|
@ -157,15 +157,12 @@ RUN set -eux \
|
|||
&& rm -rf "${UNSLOTH_STUDIO_HOME}/src/.git" \
|
||||
"${UNSLOTH_STUDIO_HOME}/src/studio/frontend/node_modules" \
|
||||
/root/.cache \
|
||||
# Stage the Studio venv's NVRTC the same way as the base venv: cu12.8 stays
|
||||
# the immutable default (real lib as .cu128.orig, libnvrtc.so.12 -> it) with
|
||||
# a cu13 alias (.cu13) beside it; select_cuda_jit_tools retargets it to cu13
|
||||
# only for sm_103/sm_121. 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 have no cu13 alias to switch to 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 regardless of TARGETARCH.
|
||||
# Stage the Studio venv's NVRTC exactly like the base venv (see
|
||||
# docker/Dockerfile: immutable .cu128.orig default + staged .cu13 alias,
|
||||
# retargeted per device by select_cuda_jit_tools). Run on BOTH arches:
|
||||
# amd64 sm_103 needs cu13 NVRTC exactly as arm64 sm_121 does, the CUDA
|
||||
# dedup below never touches cuda_nvrtc, and the base cu13 layer installs
|
||||
# cuda-nvrtc-13-0 on both arches so libnvrtc.so.13 always exists.
|
||||
&& 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" ] && [ ! -L "${NVRTC_DIR}/libnvrtc.so.12" ]; then \
|
||||
mv "${NVRTC_DIR}/libnvrtc.so.12" "${NVRTC_DIR}/libnvrtc.so.12.cu128.orig"; \
|
||||
|
|
|
|||
|
|
@ -328,47 +328,27 @@ def _classify_flag_target(spec):
|
|||
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)."""
|
||||
def _parse_flag_line(stripped, flags):
|
||||
"""If `stripped` is a `<flag> <target>` requirements-file line for one of
|
||||
`flags`, return (flag, target_or_None, inline_comment_or_None); else
|
||||
(None, None, None).
|
||||
|
||||
Shared by the `-r`/`--requirement`/`-c`/`--constraint` include parse and
|
||||
the `-e`/`--editable` install-line parse. Handles the separated
|
||||
(`-r <t>` / `--editable <t>`), inline (`--editable=<t>` / `-e=<t>`) and
|
||||
attached short (`-rextras.txt`, `-egit+...`) forms pip accepts from a
|
||||
requirement file, so a protected include or editable there is handled
|
||||
exactly like the command-line case."""
|
||||
body, sep, comment = stripped.partition(" #")
|
||||
body = body.rstrip()
|
||||
comment = ("#" + comment) if sep else None
|
||||
for flag in ("-r", "--requirement", "-c", "--constraint"):
|
||||
target = 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() # attached short form, e.g. `-rextras.txt`
|
||||
else:
|
||||
continue
|
||||
return flag, (target or None), comment
|
||||
return None, None, None
|
||||
|
||||
|
||||
def _parse_editable(stripped):
|
||||
"""If `stripped` is an `-e`/`--editable` install line, return
|
||||
(flag, target, inline_comment_or_None); else (None, None, None).
|
||||
|
||||
Handles the separated (`-e <t>` / `--editable <t>`), attached (`-e<t>`),
|
||||
long inline (`--editable=<t>`) and short inline (`-e=<t>`) forms pip accepts
|
||||
from a requirement file, so a protected editable there is dropped exactly
|
||||
like the command-line -e case."""
|
||||
body, sep, comment = stripped.partition(" #")
|
||||
body = body.rstrip()
|
||||
comment = ("#" + comment) if sep else None
|
||||
for flag in ("-e", "--editable"):
|
||||
target = None
|
||||
if body == flag:
|
||||
target = None
|
||||
elif 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() # attached short form, e.g. `-egit+...`
|
||||
target = body[len(flag) :].strip() # attached short form
|
||||
else:
|
||||
continue
|
||||
return flag, (target or None), comment
|
||||
|
|
@ -386,7 +366,9 @@ def _rewrite_include(line, stripped, src_dir, depth):
|
|||
parent at that filtered copy. URLs and unreadable/absolute-unfiltered files
|
||||
fall back to an absolutised path so they still resolve. Returns
|
||||
(new_line, changed, recorded, dropped)."""
|
||||
flag, raw_target, comment = _parse_include(stripped)
|
||||
flag, raw_target, comment = _parse_flag_line(
|
||||
stripped, ("-r", "--requirement", "-c", "--constraint")
|
||||
)
|
||||
if not raw_target:
|
||||
return line, False, None, []
|
||||
# Resolve pip's ${VAR} references so the include we read/filter is the file
|
||||
|
|
@ -458,7 +440,7 @@ def _filter_requirements_file(path, _depth = 0):
|
|||
# baked stack. Classify it through _KEEP exactly like the
|
||||
# command-line -e case and drop the whole line (flag + target) when
|
||||
# the target is protected; a transformers pin is still recorded.
|
||||
e_flag, e_target, _e_comment = _parse_editable(stripped)
|
||||
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":
|
||||
|
|
@ -561,13 +543,10 @@ def main():
|
|||
os.execv(REAL[tool], [REAL[tool]] + argv)
|
||||
return
|
||||
|
||||
# Locate the `install` verb (uv: `uv pip install ...`; pip: `pip install ...`).
|
||||
# Locate the `install` verb (pip: `pip install ...`; uv: `uv pip install ...`
|
||||
# -- index() already skips uv's leading `pip` subcommand).
|
||||
try:
|
||||
if tool == "uv":
|
||||
# skip a leading `pip` subcommand
|
||||
i = argv.index("install")
|
||||
else:
|
||||
i = argv.index("install")
|
||||
i = argv.index("install")
|
||||
except ValueError:
|
||||
os.execv(REAL[tool], [REAL[tool]] + argv) # not an install -> passthrough
|
||||
return
|
||||
|
|
|
|||
|
|
@ -37,45 +37,28 @@ STATE="$DEST/.unsloth_sync_state" # "sha256 relpath" of what we last wrote
|
|||
SYNCED="$DEST/.unsloth_sync_commit" # upstream commit we last synced to
|
||||
TIMEOUT="${UNSLOTH_NOTEBOOK_FETCH_TIMEOUT:-60}"
|
||||
|
||||
# Helper that compares the *content* (the middle, ignoring the auto-generated
|
||||
# install header / announcements / footer) of two notebooks. Used so a refresh
|
||||
# doesn't rewrite an untouched notebook when only that boilerplate moved
|
||||
# upstream. Resolved from an explicit override, then PATH, then a sibling file.
|
||||
# Resolve a helper script ($1 explicit override, $2 PATH command name, $3
|
||||
# sibling filename next to this script), echoing the resolved path or nothing.
|
||||
# An empty result leaves the caller's guard to degrade gracefully. Used for the
|
||||
# content-sig comparator (SIG), categorized-view builder (VIEW) and Docker-only
|
||||
# Colab-intro stripper (STRIP).
|
||||
PYBIN="$(command -v python3 2>/dev/null || command -v python 2>/dev/null || true)"
|
||||
SIG_HELPER="${UNSLOTH_NB_SIG_HELPER:-}"
|
||||
if [ -z "$SIG_HELPER" ]; then
|
||||
if command -v unsloth-nb-content-sig >/dev/null 2>&1; then
|
||||
SIG_HELPER="$(command -v unsloth-nb-content-sig)"
|
||||
else
|
||||
_self_dir="$(cd "$(dirname "$0")" 2>/dev/null && pwd)"
|
||||
[ -n "$_self_dir" ] && [ -f "$_self_dir/unsloth_nb_content_sig.py" ] \
|
||||
&& SIG_HELPER="$_self_dir/unsloth_nb_content_sig.py"
|
||||
fi
|
||||
fi
|
||||
_self_dir="$(cd "$(dirname "$0")" 2>/dev/null && pwd)"
|
||||
resolve_helper() {
|
||||
if [ -n "$1" ]; then printf '%s' "$1"; return 0; fi
|
||||
if command -v "$2" >/dev/null 2>&1; then command -v "$2"; return 0; fi
|
||||
[ -n "$_self_dir" ] && [ -f "$_self_dir/$3" ] && printf '%s' "$_self_dir/$3"
|
||||
return 0
|
||||
}
|
||||
SIG_HELPER="$(resolve_helper "${UNSLOTH_NB_SIG_HELPER:-}" unsloth-nb-content-sig unsloth_nb_content_sig.py)"
|
||||
VIEW_HELPER="$(resolve_helper "${UNSLOTH_NB_VIEW_HELPER:-}" unsloth-nb-view unsloth_nb_view.py)"
|
||||
STRIP_HELPER="$(resolve_helper "${UNSLOTH_NB_STRIP_HELPER:-}" unsloth-nb-strip-colab unsloth_nb_strip_colab.py)"
|
||||
|
||||
# Same resolution (override -> PATH -> sibling file) for the categorized-view
|
||||
# builder and the Docker-only Colab-intro stripper.
|
||||
_self_dir="${_self_dir:-$(cd "$(dirname "$0")" 2>/dev/null && pwd)}"
|
||||
VIEW_HELPER="${UNSLOTH_NB_VIEW_HELPER:-}"
|
||||
if [ -z "$VIEW_HELPER" ]; then
|
||||
if command -v unsloth-nb-view >/dev/null 2>&1; then
|
||||
VIEW_HELPER="$(command -v unsloth-nb-view)"
|
||||
elif [ -n "$_self_dir" ] && [ -f "$_self_dir/unsloth_nb_view.py" ]; then
|
||||
VIEW_HELPER="$_self_dir/unsloth_nb_view.py"
|
||||
fi
|
||||
fi
|
||||
STRIP_HELPER="${UNSLOTH_NB_STRIP_HELPER:-}"
|
||||
if [ -z "$STRIP_HELPER" ]; then
|
||||
if command -v unsloth-nb-strip-colab >/dev/null 2>&1; then
|
||||
STRIP_HELPER="$(command -v unsloth-nb-strip-colab)"
|
||||
elif [ -n "$_self_dir" ] && [ -f "$_self_dir/unsloth_nb_strip_colab.py" ]; then
|
||||
STRIP_HELPER="$_self_dir/unsloth_nb_strip_colab.py"
|
||||
fi
|
||||
fi
|
||||
|
||||
# True only when BOTH are .ipynb, the helper is usable, and it reports the
|
||||
# non-boilerplate middle is identical (so only the header/footer changed).
|
||||
# Any failure returns false, so the caller falls back to a normal refresh.
|
||||
# True only when BOTH are .ipynb, the SIG helper is usable, and it reports the
|
||||
# non-boilerplate middle (ignoring the auto-generated install header /
|
||||
# announcements / footer) is identical -- so a refresh doesn't rewrite an
|
||||
# untouched notebook when only that boilerplate moved upstream. Any failure
|
||||
# returns false, so the caller falls back to a normal refresh.
|
||||
middle_unchanged() {
|
||||
case "$1" in *.ipynb) : ;; *) return 1 ;; esac
|
||||
[ -n "$PYBIN" ] && [ -n "$SIG_HELPER" ] || return 1
|
||||
|
|
@ -176,14 +159,12 @@ if [ ! -f "$STATE" ]; then
|
|||
echo "[unsloth-nb] notebooks ready at $DEST"
|
||||
fi
|
||||
|
||||
# 1b) Every-boot OFFLINE restore of deleted notebooks. A file we previously wrote
|
||||
# that the user has since DELETED is restored from the baked template -- works
|
||||
# with no network and even when upstream has not advanced. Files that still exist
|
||||
# (edited or not) are never touched, so this cannot resurrect or clobber an edit;
|
||||
# the GitHub refresh below then bumps any restored file to the latest upstream.
|
||||
# The restored file's recorded hash is reset to the template's so the refresh
|
||||
# treats it as pristine (not as a user edit). Opt out with
|
||||
# UNSLOTH_KEEP_DELETED_NOTEBOOKS=1 (for users who prune notebooks on purpose).
|
||||
# 1b) Every-boot OFFLINE restore of deleted notebooks: a file we previously
|
||||
# wrote that the user has since DELETED comes back from the baked template (no
|
||||
# network needed). Files that still exist (edited or not) are never touched, so
|
||||
# this cannot clobber an edit; the restored file's recorded hash is reset to
|
||||
# the template's so the GitHub refresh below treats it as pristine and bumps it
|
||||
# to latest. Opt out with UNSLOTH_KEEP_DELETED_NOTEBOOKS=1.
|
||||
if [ -f "$STATE" ] && [ "${UNSLOTH_KEEP_DELETED_NOTEBOOKS:-0}" != "1" ]; then
|
||||
restored=0
|
||||
RS_TMP="$(mktemp)"
|
||||
|
|
|
|||
|
|
@ -103,100 +103,81 @@ def _run(shim, tool, args):
|
|||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Item 3541142907 -- pair -e/--editable with its target.
|
||||
# Item 3541142907 -- pair -e/--editable with its target (the attached short
|
||||
# `-e<target>` form from item 3541404845 is folded in here). A protected
|
||||
# editable such as `pip install -e git+...unsloth...#egg=unsloth peft` must
|
||||
# NOT become `pip install -e peft` (which pip rejects): the flag drops WITH
|
||||
# its value, and an unprotected editable is forwarded verbatim.
|
||||
# --------------------------------------------------------------------------
|
||||
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
|
||||
UNSLOTH_VCS = "git+https://github.com/unslothai/unsloth.git#egg=unsloth"
|
||||
|
||||
# Sentinel expectation: the whole command line is forwarded verbatim (execd == args).
|
||||
KEPT = object()
|
||||
|
||||
|
||||
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
|
||||
@pytest.mark.parametrize(
|
||||
"args, expected",
|
||||
[
|
||||
pytest.param(["-e", UNSLOTH_VCS, "peft"], ["peft"], id = "sep-protected"),
|
||||
# nothing left to install -> no-op, no dangling -e
|
||||
pytest.param(["-e", UNSLOTH_VCS], None, id = "sep-only-protected-noop"),
|
||||
pytest.param(["-e", "./localpkg"], KEPT, id = "sep-unprotected-kept"),
|
||||
pytest.param(["--editable=" + UNSLOTH_VCS, "peft"], ["peft"], id = "inline-protected"),
|
||||
pytest.param(["--editable=./localpkg"], KEPT, id = "inline-unprotected-kept"),
|
||||
pytest.param(["-e" + UNSLOTH_VCS, "peft"], ["peft"], id = "attached-protected"),
|
||||
],
|
||||
)
|
||||
def test_editable_forms(shim, args, expected):
|
||||
execd, _ = _run(shim, "pip", args)
|
||||
assert execd == (args if expected is KEPT else expected), execd
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Item 3541142906 -- filter uv -P/--upgrade-package values.
|
||||
# Item 3541142906 -- filter uv -P/--upgrade-package values. `uv pip install
|
||||
# -P torch peft` must not let uv refresh baked torch; a pinned transformers
|
||||
# upgrade selector still feeds the sidecar marker.
|
||||
# --------------------------------------------------------------------------
|
||||
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
|
||||
@pytest.mark.parametrize(
|
||||
"args, expected, expected_marker",
|
||||
[
|
||||
pytest.param(["-P", "torch", "peft"], ["peft"], None, id = "protected-dropped"),
|
||||
pytest.param(["--upgrade-package=transformers", "peft"], ["peft"], None, id = "inline"),
|
||||
pytest.param(["-P", "transformers==4.55.0", "peft"], ["peft"], "4.55.0", id = "tf-pin"),
|
||||
pytest.param(["-P", "requests", "requests"], KEPT, None, id = "unprotected-kept"),
|
||||
# -P is not itself a target
|
||||
pytest.param(["-P", "torch"], None, None, id = "only-protected-noop"),
|
||||
],
|
||||
)
|
||||
def test_upgrade_package_forms(shim, args, expected, expected_marker):
|
||||
execd, marker = _run(shim, "uv", args)
|
||||
assert execd == (args if expected is KEPT else expected), execd
|
||||
assert marker == expected_marker, marker
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Item 3541142908 -- parse protected wheel basenames before URL passthrough.
|
||||
# Item 3541142908 -- parse protected wheel basenames before URL passthrough
|
||||
# (a recognised protected wheel URL/path is dropped -> no-op).
|
||||
# --------------------------------------------------------------------------
|
||||
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
|
||||
NUMPY_WHEEL_URL = "https://example.com/wheels/numpy-2.1.0-cp312-cp312-linux_x86_64.whl"
|
||||
|
||||
|
||||
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
|
||||
@pytest.mark.parametrize(
|
||||
"args, expected",
|
||||
[
|
||||
pytest.param([TORCH_WHEEL_URL], None, id = "direct-url"),
|
||||
pytest.param(
|
||||
["/tmp/torch-2.11.0+cu128-cp312-cp312-linux_x86_64.whl"], None, id = "local-path"
|
||||
),
|
||||
# unsloth_zoo-*.whl normalises to unsloth-zoo, which is protected.
|
||||
pytest.param(
|
||||
["https://example.com/unsloth_zoo-1.0-py3-none-any.whl"], None, id = "normalised"
|
||||
),
|
||||
pytest.param([NUMPY_WHEEL_URL], KEPT, id = "unprotected-kept"),
|
||||
],
|
||||
)
|
||||
def test_wheel_url_and_path_forms(shim, args, expected):
|
||||
execd, _ = _run(shim, "pip", args)
|
||||
assert execd == (args if expected is KEPT else expected), execd
|
||||
|
||||
|
||||
def test_protected_wheel_in_requirements_file_dropped(shim, tmp_path):
|
||||
|
|
@ -305,7 +286,8 @@ def test_nested_requirement_transformers_pin_recorded(shim, tmp_path):
|
|||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Item 3541404845 -- handle pip's attached short options (-rfile / -cfile / etc).
|
||||
# Item 3541404845 -- handle pip's attached short options (-rfile / -cfile /
|
||||
# etc). The attached `-e<target>` case lives in test_editable_forms above.
|
||||
# --------------------------------------------------------------------------
|
||||
def test_attached_short_requirement_file_filtered(shim, tmp_path):
|
||||
# `pip install -rreqs.txt` (attached) must filter the file AND count as a
|
||||
|
|
@ -329,15 +311,6 @@ def test_attached_short_constraint_file_filtered(shim, tmp_path):
|
|||
assert "torch" not in filtered
|
||||
|
||||
|
||||
def test_attached_short_editable_protected_dropped(shim):
|
||||
execd, _ = _run(
|
||||
shim,
|
||||
"pip",
|
||||
["-egit+https://github.com/unslothai/unsloth.git#egg=unsloth", "peft"],
|
||||
)
|
||||
assert execd == ["peft"], execd
|
||||
|
||||
|
||||
def test_attached_short_upgrade_package_protected_dropped(shim):
|
||||
execd, _ = _run(shim, "uv", ["-Ptorch", "peft"])
|
||||
assert execd == ["peft"], execd
|
||||
|
|
@ -346,22 +319,20 @@ def test_attached_short_upgrade_package_protected_dropped(shim):
|
|||
|
||||
# --------------------------------------------------------------------------
|
||||
# Item 3541773143 -- a bare wheel filename (no ./ or / prefix) is still a pip
|
||||
# target from the CWD, so its protected distribution must be parsed too.
|
||||
# target from the CWD, so its protected distribution must be parsed too
|
||||
# (`pip install torch-2.11.0-...whl` must not reinstall torch).
|
||||
# --------------------------------------------------------------------------
|
||||
def test_bare_torch_wheel_filename_dropped(shim):
|
||||
# `pip install torch-2.11.0-...whl` from the CWD must not reinstall torch.
|
||||
execd, _ = _run(shim, "pip", ["torch-2.11.0+cu128-cp312-cp312-linux_x86_64.whl"])
|
||||
assert execd is None, execd
|
||||
|
||||
|
||||
def test_bare_wheel_in_subdir_dropped(shim):
|
||||
execd, _ = _run(shim, "pip", ["dist/torch-2.11.0-cp312-cp312-linux_x86_64.whl"])
|
||||
assert execd is None, execd
|
||||
|
||||
|
||||
def test_bare_unprotected_wheel_filename_kept(shim):
|
||||
execd, _ = _run(shim, "pip", ["numpy-2.1.0-cp312-cp312-linux_x86_64.whl"])
|
||||
assert execd == ["numpy-2.1.0-cp312-cp312-linux_x86_64.whl"], execd
|
||||
@pytest.mark.parametrize(
|
||||
"args, expected",
|
||||
[
|
||||
pytest.param(["torch-2.11.0+cu128-cp312-cp312-linux_x86_64.whl"], None, id = "bare-torch"),
|
||||
pytest.param(["dist/torch-2.11.0-cp312-cp312-linux_x86_64.whl"], None, id = "subdir-torch"),
|
||||
pytest.param(["numpy-2.1.0-cp312-cp312-linux_x86_64.whl"], KEPT, id = "unprotected-kept"),
|
||||
],
|
||||
)
|
||||
def test_bare_wheel_filename_forms(shim, args, expected):
|
||||
execd, _ = _run(shim, "pip", args)
|
||||
assert execd == (args if expected is KEPT else expected), execd
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
|
|
@ -389,29 +360,23 @@ def test_vcs_url_without_egg_unprotected_kept(shim):
|
|||
# Item 3541773153 -- refuse remote (URL) requirement / constraint files in shim
|
||||
# mode; their protected pins cannot be inspected before the real tool installs.
|
||||
# --------------------------------------------------------------------------
|
||||
def test_remote_requirement_url_only_noops(shim):
|
||||
execd, _ = _run(shim, "pip", ["-r", "https://example.com/reqs.txt"])
|
||||
assert execd is None, execd # dropped, and no dangling -r left behind
|
||||
R_URL = "https://example.com/reqs.txt"
|
||||
|
||||
|
||||
def test_remote_requirement_url_with_other_target_kept(shim):
|
||||
execd, _ = _run(shim, "pip", ["-r", "https://example.com/reqs.txt", "peft"])
|
||||
assert execd == ["peft"], execd
|
||||
|
||||
|
||||
def test_remote_requirement_inline_form_dropped(shim):
|
||||
execd, _ = _run(shim, "pip", ["--requirement=https://example.com/reqs.txt", "peft"])
|
||||
assert execd == ["peft"], execd
|
||||
|
||||
|
||||
def test_remote_requirement_attached_form_dropped(shim):
|
||||
execd, _ = _run(shim, "pip", ["-rhttps://example.com/reqs.txt", "peft"])
|
||||
assert execd == ["peft"], execd
|
||||
|
||||
|
||||
def test_remote_constraint_url_dropped_target_kept(shim):
|
||||
execd, _ = _run(shim, "pip", ["-c", "https://example.com/constraints.txt", "peft"])
|
||||
assert execd == ["peft"], execd
|
||||
@pytest.mark.parametrize(
|
||||
"args, expected",
|
||||
[
|
||||
# dropped, and no dangling -r left behind
|
||||
pytest.param(["-r", R_URL], None, id = "sep-r-only-noop"),
|
||||
pytest.param(["-r", R_URL, "peft"], ["peft"], id = "sep-r-target-kept"),
|
||||
pytest.param(["--requirement=" + R_URL, "peft"], ["peft"], id = "inline-r"),
|
||||
pytest.param(["-r" + R_URL, "peft"], ["peft"], id = "attached-r"),
|
||||
pytest.param(["-c", "https://example.com/constraints.txt", "peft"], ["peft"], id = "sep-c"),
|
||||
],
|
||||
)
|
||||
def test_remote_requirement_and_constraint_urls_refused(shim, args, expected):
|
||||
execd, _ = _run(shim, "pip", args)
|
||||
assert execd == expected, execd
|
||||
|
||||
|
||||
def test_nested_remote_include_dropped(shim, tmp_path):
|
||||
|
|
@ -449,56 +414,43 @@ def test_uv_reinstall_flag_stripped(shim):
|
|||
# Item 3541773168 -- uv's --reinstall-package selector is filtered through _KEEP
|
||||
# exactly like -P/--upgrade-package (both forms, no dangling flag).
|
||||
# --------------------------------------------------------------------------
|
||||
def test_reinstall_package_protected_separated_dropped(shim):
|
||||
execd, _ = _run(shim, "uv", ["--reinstall-package", "torch", "peft"])
|
||||
assert execd == ["peft"], execd
|
||||
assert "torch" not in execd and "--reinstall-package" not in execd
|
||||
|
||||
|
||||
def test_reinstall_package_protected_inline_dropped(shim):
|
||||
execd, _ = _run(shim, "uv", ["--reinstall-package=torch", "peft"])
|
||||
assert execd == ["peft"], execd
|
||||
|
||||
|
||||
def test_reinstall_package_unprotected_kept(shim):
|
||||
execd, _ = _run(shim, "uv", ["--reinstall-package", "requests", "requests"])
|
||||
assert execd == ["--reinstall-package", "requests", "requests"], execd
|
||||
|
||||
|
||||
def test_reinstall_package_transformers_pin_recorded(shim):
|
||||
execd, marker = _run(shim, "uv", ["--reinstall-package", "transformers==4.55.0", "peft"])
|
||||
assert execd == ["peft"], execd
|
||||
assert marker == "4.55.0", marker
|
||||
@pytest.mark.parametrize(
|
||||
"args, expected, expected_marker",
|
||||
[
|
||||
pytest.param(["--reinstall-package", "torch", "peft"], ["peft"], None, id = "sep-protected"),
|
||||
pytest.param(["--reinstall-package=torch", "peft"], ["peft"], None, id = "inline-protected"),
|
||||
pytest.param(["--reinstall-package", "requests", "requests"], KEPT, None, id = "unprotected"),
|
||||
pytest.param(
|
||||
["--reinstall-package", "transformers==4.55.0", "peft"], ["peft"], "4.55.0", id = "tf-pin"
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_reinstall_package_forms(shim, args, expected, expected_marker):
|
||||
execd, marker = _run(shim, "uv", args)
|
||||
assert execd == (args if expected is KEPT else expected), execd
|
||||
assert marker == expected_marker, marker
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Item 3542096750 -- parse protected source archives (sdist / zip) too.
|
||||
# --------------------------------------------------------------------------
|
||||
def test_sdist_url_protected_dropped(shim):
|
||||
url = "https://files.pythonhosted.org/packages/aa/unsloth-2026.7.1.tar.gz"
|
||||
execd, _ = _run(shim, "pip", [url, "peft"])
|
||||
assert execd == ["peft"], execd
|
||||
SDIST_URL = "https://files.pythonhosted.org/packages/aa/unsloth-2026.7.1.tar.gz"
|
||||
|
||||
|
||||
def test_sdist_bare_protected_dropped(shim):
|
||||
execd, _ = _run(shim, "pip", ["torch-2.11.0.tar.gz"])
|
||||
assert execd is None, execd
|
||||
|
||||
|
||||
def test_sdist_zip_protected_dropped(shim):
|
||||
execd, _ = _run(shim, "pip", ["./transformers-4.55.0.zip", "peft"])
|
||||
assert execd == ["peft"], execd
|
||||
|
||||
|
||||
def test_sdist_hyphenated_name_protected_dropped(shim):
|
||||
# flashinfer-python is protected; the name must survive the hyphen split.
|
||||
execd, _ = _run(shim, "pip", ["flashinfer-python-0.5.0.tar.gz"])
|
||||
assert execd is None, execd
|
||||
|
||||
|
||||
def test_sdist_unprotected_kept(shim):
|
||||
execd, _ = _run(shim, "pip", ["numpy-2.1.0.tar.gz"])
|
||||
assert execd == ["numpy-2.1.0.tar.gz"], execd
|
||||
@pytest.mark.parametrize(
|
||||
"args, expected",
|
||||
[
|
||||
pytest.param([SDIST_URL, "peft"], ["peft"], id = "url-protected"),
|
||||
pytest.param(["torch-2.11.0.tar.gz"], None, id = "bare-protected"),
|
||||
pytest.param(["./transformers-4.55.0.zip", "peft"], ["peft"], id = "zip-protected"),
|
||||
# flashinfer-python is protected; the name must survive the hyphen split.
|
||||
pytest.param(["flashinfer-python-0.5.0.tar.gz"], None, id = "hyphenated-name"),
|
||||
pytest.param(["numpy-2.1.0.tar.gz"], KEPT, id = "unprotected-kept"),
|
||||
],
|
||||
)
|
||||
def test_source_archive_forms(shim, args, expected):
|
||||
execd, _ = _run(shim, "pip", args)
|
||||
assert execd == (args if expected is KEPT else expected), execd
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
|
|
@ -529,21 +481,21 @@ def test_uv_plural_constraints_filtered(shim, tmp_path):
|
|||
# Item 3542096764 -- neutralise --upgrade-strategy eager so a kept target cannot
|
||||
# eagerly rebuild already-satisfied baked deps.
|
||||
# --------------------------------------------------------------------------
|
||||
def test_upgrade_strategy_eager_dropped(shim):
|
||||
execd, _ = _run(shim, "pip", ["-U", "--upgrade-strategy", "eager", "peft"])
|
||||
assert execd == ["-U", "peft"], execd
|
||||
|
||||
|
||||
def test_upgrade_strategy_eager_inline_dropped(shim):
|
||||
execd, _ = _run(shim, "pip", ["--upgrade-strategy=eager", "peft"])
|
||||
assert execd == ["peft"], execd
|
||||
|
||||
|
||||
def test_upgrade_strategy_only_if_needed_also_dropped(shim):
|
||||
# only-if-needed is pip's default, so dropping it is a harmless no-op that
|
||||
# keeps the kept target installing normally.
|
||||
execd, _ = _run(shim, "pip", ["--upgrade-strategy", "only-if-needed", "peft"])
|
||||
assert execd == ["peft"], execd
|
||||
@pytest.mark.parametrize(
|
||||
"args, expected",
|
||||
[
|
||||
pytest.param(["-U", "--upgrade-strategy", "eager", "peft"], ["-U", "peft"], id = "eager"),
|
||||
pytest.param(["--upgrade-strategy=eager", "peft"], ["peft"], id = "inline-eager"),
|
||||
# only-if-needed is pip's default, so dropping it is a harmless no-op that
|
||||
# keeps the kept target installing normally.
|
||||
pytest.param(
|
||||
["--upgrade-strategy", "only-if-needed", "peft"], ["peft"], id = "only-if-needed"
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_upgrade_strategy_forms(shim, args, expected):
|
||||
execd, _ = _run(shim, "pip", args)
|
||||
assert execd == expected, execd
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
|
|
@ -572,9 +524,9 @@ def test_forwarded_install_carries_protected_constraints(shim):
|
|||
assert all("==" in pin for pin in pins), pins
|
||||
names = {pin.split("==", 1)[0].lower().replace("_", "-") for pin in pins}
|
||||
protected = {"transformers"} | shim._KEEP | {"nvidia-"}
|
||||
assert all(
|
||||
n in shim._KEEP or n == "transformers" or n.startswith("nvidia-") for n in names
|
||||
), names
|
||||
assert all(n in shim._KEEP or n == "transformers" or n.startswith("nvidia-") for n in names), (
|
||||
names
|
||||
)
|
||||
|
||||
|
||||
def test_noop_install_gets_no_constraints(shim):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue