From cd982a121dffd6728208f5a93f21e91a129678ff Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 16 Jul 2026 05:27:22 +0000 Subject: [PATCH] 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. --- .github/workflows/docker-publish.yml | 80 ++----- docker/Dockerfile | 95 +++----- docker/Dockerfile.studio | 15 +- docker/unsloth_pip_shim.py | 61 ++--- docker/unsloth_sync_notebooks.sh | 71 ++---- tests/python/test_unsloth_pip_shim.py | 328 +++++++++++--------------- 6 files changed, 253 insertions(+), 397 deletions(-) diff --git a/.github/workflows/docker-publish.yml b/.github/workflows/docker-publish.yml index f568689c1e..6288d59a7d 100644 --- a/.github/workflows/docker-publish.yml +++ b/.github/workflows/docker-publish.yml @@ -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 }} diff --git a/docker/Dockerfile b/docker/Dockerfile index 13bdb884e3..746332159d 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -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. diff --git a/docker/Dockerfile.studio b/docker/Dockerfile.studio index 9c85d8d745..122742037b 100644 --- a/docker/Dockerfile.studio +++ b/docker/Dockerfile.studio @@ -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"; \ diff --git a/docker/unsloth_pip_shim.py b/docker/unsloth_pip_shim.py index 454bf28303..46888d81f1 100644 --- a/docker/unsloth_pip_shim.py +++ b/docker/unsloth_pip_shim.py @@ -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 ` ` 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 ` / `--editable `), inline (`--editable=` / `-e=`) 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 ` / `--editable `), attached (`-e`), - long inline (`--editable=`) and short inline (`-e=`) 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 diff --git a/docker/unsloth_sync_notebooks.sh b/docker/unsloth_sync_notebooks.sh index 1c66427783..fab8430027 100644 --- a/docker/unsloth_sync_notebooks.sh +++ b/docker/unsloth_sync_notebooks.sh @@ -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)" diff --git a/tests/python/test_unsloth_pip_shim.py b/tests/python/test_unsloth_pip_shim.py index 3dd002e707..818edf784e 100644 --- a/tests/python/test_unsloth_pip_shim.py +++ b/tests/python/test_unsloth_pip_shim.py @@ -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` 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` 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):