docker: close more pip-shim bypasses and make cu12.8 NVRTC the default
Notebook pip/uv shim (docker/unsloth_pip_shim.py, active only under
UNSLOTH_NB_SHIM=1):
- Parse protected source archives (sdist/zip) by basename too, e.g.
`pip install https://.../unsloth-2026.7.1.tar.gz` or `./torch-2.11.0.tar.gz`,
mirroring the wheel-basename handling. A first-hyphen-before-digit split
keeps hyphenated names like flashinfer-python intact.
- Recognise uv's PLURAL long flags --requirements / --constraints, so those
files go through the same protected-package filter as the singular names.
- Drop --upgrade-strategy eager in shim mode so a kept target cannot eagerly
rebuild already-satisfied baked deps (falls back to pip's only-if-needed).
NVRTC default (docker/Dockerfile, docker/Dockerfile.studio, docker/entrypoint.sh):
- Make cu12.8 the immutable baked default (libnvrtc.so.12 -> .cu128.orig) with
a staged .cu13 alias, and have select_cuda_jit_tools retarget to cu13 ONLY
for sm_103/sm_121. Previously cu13 was baked as the default and restored to
cu12.8 at runtime, so a non-root `docker run --user` container that cannot
rewrite the symlink stayed on cu13 NVRTC and emitted cubins a 570-579 driver
cannot load. The safe default now needs no runtime write.
Adds regression tests for each case (tests/python/test_unsloth_pip_shim.py,
tests/sh/test_select_cuda_jit_tools.sh).
This commit is contained in:
Daniel Han
parent
167fdf26b9
commit
6a078b1a45
6 changed files with 233 additions and 82 deletions
|
|
@ -575,9 +575,10 @@ COPY --from=builder /opt/unsloth-venv /opt/unsloth-venv
|
|||
# (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: symlink cu13 libnvrtc.so.13 over the
|
||||
# bundled .so.12 (the .so.12 original is saved as .cu128.orig so the
|
||||
# runtime can restore it -- see below).
|
||||
# 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.
|
||||
#
|
||||
# (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
|
||||
|
|
@ -615,11 +616,21 @@ RUN set -eux; \
|
|||
cuda-nvrtc-13-0 \
|
||||
cuda-nvcc-13-0; \
|
||||
rm -rf /var/lib/apt/lists/*; \
|
||||
# (1) NVRTC swap. torch's wheel-bundled cu128 NVRTC -> cu13 NVRTC.
|
||||
# (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.
|
||||
NVRTC_DIR=/opt/unsloth-venv/lib/python${PYTHON_VERSION}/site-packages/nvidia/cuda_nvrtc/lib; \
|
||||
if [ -f "${NVRTC_DIR}/libnvrtc.so.12" ]; then \
|
||||
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 /usr/local/cuda-13.0/lib64/libnvrtc.so.13 "${NVRTC_DIR}/libnvrtc.so.12"; \
|
||||
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 own ptxas is cu12.8 (no sm_103/sm_121), so
|
||||
# those two arches need the cu13 ptxas installed above. It is NOT baked as a
|
||||
|
|
|
|||
|
|
@ -158,18 +158,20 @@ RUN set -eux \
|
|||
&& rm -rf "${UNSLOTH_STUDIO_HOME}/src/.git" \
|
||||
"${UNSLOTH_STUDIO_HOME}/src/studio/frontend/node_modules" \
|
||||
/root/.cache \
|
||||
# Swap the Studio venv's OWN bundled cu12.8 libnvrtc for cu13, mirroring the
|
||||
# base venv swap. Run on BOTH arches, not arm64 only: amd64 sm_103 (B300 /
|
||||
# GB300) needs cu13 NVRTC exactly as arm64 sm_121 (DGX Spark / GB10) does,
|
||||
# and the CUDA dedup below never touches cuda_nvrtc, so an amd64 Studio venv
|
||||
# would otherwise keep cu12.8 NVRTC and its jiterator/NVRTC JIT paths fail on
|
||||
# compute_103. The base cu13 layer installs cuda-nvrtc-13-0 on both arches,
|
||||
# so /usr/local/cuda-13.0/lib64/libnvrtc.so.13 is present here regardless of
|
||||
# TARGETARCH.
|
||||
# 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.
|
||||
&& for NVRTC_DIR in "${UNSLOTH_STUDIO_HOME}"/unsloth_studio/lib/python*/site-packages/nvidia/cuda_nvrtc/lib; do \
|
||||
if [ -f "${NVRTC_DIR}/libnvrtc.so.12" ]; then \
|
||||
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 /usr/local/cuda-13.0/lib64/libnvrtc.so.13 "${NVRTC_DIR}/libnvrtc.so.12"; \
|
||||
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; \
|
||||
done \
|
||||
&& BASE_NV=/opt/unsloth-venv/lib/python3.12/site-packages/nvidia \
|
||||
|
|
|
|||
|
|
@ -33,49 +33,50 @@ set -euo pipefail
|
|||
# they RUN under any driver -- it is only their OUTPUT the older driver rejects.
|
||||
#
|
||||
# Pick per DEVICE at boot (the compute capability is unknown at build time):
|
||||
# activate cu13 only for sm_103 / sm_121, and otherwise keep Triton on its
|
||||
# bundled cu12.8 ptxas and restore the wheel-bundled cu12.8 NVRTC the build
|
||||
# swapped for cu13. Runs before every early-exit below so the selection always
|
||||
# applies. Best-effort: a read-only / --user-dropped rootfs that cannot
|
||||
# re-point the NVRTC symlink is left unchanged.
|
||||
# cu12.8 is the immutable baked default (loadable on every supported 570+
|
||||
# driver), and only sm_103 / sm_121 -- which ship on >= 580 drivers -- switch
|
||||
# Triton to cu13 ptxas and retarget the venv NVRTC symlink to the staged cu13
|
||||
# alias. Runs before every early-exit below so the selection always applies.
|
||||
# Best-effort: because the safe default needs no write, a non-root / read-only
|
||||
# rootfs is always fine; only the rare non-root datacenter host cannot switch.
|
||||
select_cuda_jit_tools() {
|
||||
local caps="" cc nvrtc_dir orig need_cu13=0
|
||||
local caps="" cc nvrtc_dir need_cu13=0
|
||||
if command -v nvidia-smi >/dev/null 2>&1; then
|
||||
caps="$( { nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true; } )"
|
||||
fi
|
||||
# Scan EVERY visible GPU, not just the first: a Blackwell datacenter part
|
||||
# (sm_103 B300/GB300 or sm_121 GB10/DGX Spark) can sit behind an H100/B200 in
|
||||
# the nvidia-smi ordering, so keying off only the first compute_cap would
|
||||
# restore cu12.8 and leave that later device unable to JIT. If ANY visible
|
||||
# GPU needs cu13, enable it for the whole process -- those parts only ship on
|
||||
# >= 580 drivers, so the host tolerates cu13 cubins for every arch present.
|
||||
# the nvidia-smi ordering, so keying off only the first compute_cap would miss
|
||||
# it. If ANY visible GPU needs cu13, switch to it for the whole process --
|
||||
# those parts only ship on >= 580 drivers, so the host tolerates cu13 cubins
|
||||
# for every arch present.
|
||||
while IFS= read -r cc || [[ -n "${cc}" ]]; do
|
||||
cc="$(printf '%s' "${cc}" | tr -d '[:space:]')"
|
||||
case "${cc}" in
|
||||
10.3|12.1) need_cu13=1 ;;
|
||||
esac
|
||||
done <<< "${caps}"
|
||||
if [[ "${need_cu13}" -eq 1 ]]; then
|
||||
# Blackwell datacenter present: the build already points each venv's
|
||||
# libnvrtc.so.12 at cu13, so only Triton's ptxas needs redirecting.
|
||||
# -z guard leaves an explicit `docker run -e TRITON_PTXAS_PATH` win.
|
||||
if [[ -x /usr/local/cuda-13.0/bin/ptxas && -z "${TRITON_PTXAS_PATH:-}" ]]; then
|
||||
export TRITON_PTXAS_PATH=/usr/local/cuda-13.0/bin/ptxas
|
||||
fi
|
||||
else
|
||||
# No datacenter Blackwell present (or an undetectable / CPU host): leave
|
||||
# TRITON_PTXAS_PATH unset so Triton uses its bundled cu12.8 ptxas, and
|
||||
# restore the cu12.8 NVRTC in each venv that saved the original, so a
|
||||
# 570-579 driver never sees a cu13 cubin. Covers the base venv and, on
|
||||
# the Studio image, the Studio venv.
|
||||
for nvrtc_dir in \
|
||||
/opt/unsloth-venv/lib/python*/site-packages/nvidia/cuda_nvrtc/lib \
|
||||
"${UNSLOTH_STUDIO_HOME:-/opt/unsloth-studio}"/unsloth_studio/lib/python*/site-packages/nvidia/cuda_nvrtc/lib; do
|
||||
orig="${nvrtc_dir}/libnvrtc.so.12.cu128.orig"
|
||||
[[ -e "${orig}" ]] || continue
|
||||
ln -sf libnvrtc.so.12.cu128.orig "${nvrtc_dir}/libnvrtc.so.12" 2>/dev/null || true
|
||||
done
|
||||
# Non-datacenter / undetectable / CPU host: nothing to do. cu12.8 is the
|
||||
# immutable baked default (libnvrtc.so.12 -> .cu128.orig, Triton on its
|
||||
# bundled cu12.8 ptxas), loadable on every supported 570+ driver, and needs
|
||||
# NO write -- so a non-root `docker run --user` container is never left on a
|
||||
# cu13 NVRTC a 570-579 driver cannot load.
|
||||
[[ "${need_cu13}" -eq 1 ]] || return 0
|
||||
# Blackwell datacenter present: cu12.8 cannot emit compute_103/121, so point
|
||||
# Triton at cu13 ptxas and retarget each venv's libnvrtc.so.12 -> the staged
|
||||
# cu13 alias. -z guard leaves an explicit `docker run -e TRITON_PTXAS_PATH`
|
||||
# win. Best-effort: a read-only / --user rootfs that cannot rewrite the
|
||||
# symlink simply keeps cu12.8 (a rare non-root datacenter case). Covers the
|
||||
# base venv and, on the Studio image, the Studio venv.
|
||||
if [[ -x /usr/local/cuda-13.0/bin/ptxas && -z "${TRITON_PTXAS_PATH:-}" ]]; then
|
||||
export TRITON_PTXAS_PATH=/usr/local/cuda-13.0/bin/ptxas
|
||||
fi
|
||||
for nvrtc_dir in \
|
||||
/opt/unsloth-venv/lib/python*/site-packages/nvidia/cuda_nvrtc/lib \
|
||||
"${UNSLOTH_STUDIO_HOME:-/opt/unsloth-studio}"/unsloth_studio/lib/python*/site-packages/nvidia/cuda_nvrtc/lib; do
|
||||
[[ -e "${nvrtc_dir}/libnvrtc.so.12.cu13" ]] || continue
|
||||
ln -sf libnvrtc.so.12.cu13 "${nvrtc_dir}/libnvrtc.so.12" 2>/dev/null || true
|
||||
done
|
||||
}
|
||||
# Best-effort: never let JIT-tool selection block container startup.
|
||||
select_cuda_jit_tools || true
|
||||
|
|
|
|||
|
|
@ -49,8 +49,10 @@ _KEEP_PREFIX = ("nvidia-", "nvidia_")
|
|||
_VALUE_FLAGS = {
|
||||
"-r",
|
||||
"--requirement",
|
||||
"--requirements",
|
||||
"-c",
|
||||
"--constraint",
|
||||
"--constraints",
|
||||
"-i",
|
||||
"--index-url",
|
||||
"--extra-index-url",
|
||||
|
|
@ -62,6 +64,7 @@ _VALUE_FLAGS = {
|
|||
"-p",
|
||||
"--prefix",
|
||||
"--index-strategy",
|
||||
"--upgrade-strategy",
|
||||
"--upgrade-package",
|
||||
"-P",
|
||||
"--reinstall-package",
|
||||
|
|
@ -77,12 +80,15 @@ _VALUE_FLAGS = {
|
|||
# Of those value-flags, the ones whose VALUE is itself an install target: a
|
||||
# requirements file pulls real requirements. An index-url / find-links /
|
||||
# constraint / target value is an option, not something to install.
|
||||
_REQ_FILE_FLAGS = {"-r", "--requirement"}
|
||||
# uv spells the long forms in the PLURAL (`--requirements`, `--constraints`);
|
||||
# include both so a `uv pip install --requirements reqs.txt` is filtered too.
|
||||
_REQ_FILE_FLAGS = {"-r", "--requirement", "--requirements"}
|
||||
# Constraint files are not install targets, but pip applies their pins during
|
||||
# resolution, so a `-c constraints.txt` that pins torch/transformers/etc. can
|
||||
# still downgrade or reinstall a baked package when another target pulls it in.
|
||||
# Filter protected packages out of them the same way as requirement files.
|
||||
_CONSTRAINT_FILE_FLAGS = {"-c", "--constraint"}
|
||||
# (uv's long form is the plural `--constraints`.)
|
||||
_CONSTRAINT_FILE_FLAGS = {"-c", "--constraint", "--constraints"}
|
||||
# -e/--editable <path|url|vcs> takes the NEXT token as its target (pip:
|
||||
# `-e, --editable <path/url>`), and that target is a real install target. A
|
||||
# protected editable (e.g. `-e git+https://.../unsloth.git#egg=unsloth`) must
|
||||
|
|
@ -114,6 +120,36 @@ _ATTACHED_SHORT_FLAGS = {"-r", "-c", "-e", "-P"}
|
|||
# already-satisfied protected deps are left untouched. Per-package selectors
|
||||
# (--reinstall-package / -P) are handled through _UPGRADE_PKG_FLAGS instead.
|
||||
_REINSTALL_FLAGS = {"--force-reinstall", "--ignore-installed", "-I", "--reinstall"}
|
||||
# Value-flags whose flag+value pair is dropped outright in shim mode.
|
||||
# `--upgrade-strategy eager` makes pip upgrade EVERY dependency of a kept target
|
||||
# regardless of whether the installed version already satisfies it, which would
|
||||
# refresh the baked torch/transformers under the pinned CUDA stack. Dropping the
|
||||
# flag falls back to pip's default `only-if-needed`, so a kept target still
|
||||
# installs but already-satisfied protected deps stay put. (`only-if-needed` is
|
||||
# the default, so dropping a `--upgrade-strategy only-if-needed` is a no-op.)
|
||||
_DROP_VALUE_FLAGS = {"--upgrade-strategy"}
|
||||
|
||||
|
||||
# Source-distribution / archive suffixes pip accepts as an install target.
|
||||
_ARCHIVE_EXTS = (".tar.gz", ".tgz", ".tar.bz2", ".tbz2", ".tar.xz", ".txz", ".tar", ".zip")
|
||||
|
||||
|
||||
def _sdist_name(basename):
|
||||
"""Distribution name from a source-archive basename ({name}-{version}.ext),
|
||||
or None if it is not a recognised archive. Splits at the first hyphen that
|
||||
precedes a digit so legacy hyphenated names (flashinfer-python-1.0,
|
||||
pytorch-triton-2.0) resolve correctly, not just PEP 625-normalised ones."""
|
||||
low = basename.lower()
|
||||
stem = None
|
||||
for ext in _ARCHIVE_EXTS:
|
||||
if low.endswith(ext):
|
||||
stem = basename[: -len(ext)]
|
||||
break
|
||||
if stem is None:
|
||||
return None
|
||||
m = re.match(r"^(.+?)-\d", stem)
|
||||
name = (m.group(1) if m else stem).strip().lower().replace("_", "-")
|
||||
return name or None
|
||||
|
||||
|
||||
def _canon(token):
|
||||
|
|
@ -156,6 +192,14 @@ def _canon(token):
|
|||
dist = _whl.group(1).split("-", 1)[0].strip().lower().replace("_", "-")
|
||||
if dist:
|
||||
return dist
|
||||
# A source archive (sdist / zip) URL or path names its distribution the
|
||||
# same way ({name}-{version}.tar.gz etc.), so `pip install
|
||||
# https://files.pythonhosted.org/.../unsloth-2026.7.1.tar.gz` or
|
||||
# `./torch-2.11.0.tar.gz` must be matched against _KEEP too, not passed
|
||||
# through as an opaque positional that reinstalls the baked package.
|
||||
_arch = _sdist_name(token.split("#", 1)[0].split("?", 1)[0].rstrip("/").rsplit("/", 1)[-1])
|
||||
if _arch:
|
||||
return _arch
|
||||
# A VCS URL without an #egg= fragment still installs a named project:
|
||||
# pip/uv derive the distribution from the repo, and for the packages we
|
||||
# protect the repo basename equals the distribution
|
||||
|
|
@ -185,6 +229,11 @@ def _canon(token):
|
|||
dist = token.rsplit("/", 1)[-1][:-4].split("-", 1)[0].strip().lower().replace("_", "-")
|
||||
if dist:
|
||||
return dist
|
||||
# A bare source-archive filename from the CWD (`pip install torch-2.11.0.tar.gz`)
|
||||
# is a valid pip target too; parse its distribution the same way.
|
||||
_barch = _sdist_name(token.rsplit("/", 1)[-1])
|
||||
if _barch:
|
||||
return _barch
|
||||
# strip extras and any version/marker tail
|
||||
name = re.split(r"[<>=!~\[\s;@]", token, 1)[0].strip()
|
||||
return name.lower().replace("_", "-") or None
|
||||
|
|
@ -451,6 +500,13 @@ def main():
|
|||
_c_path, _c_rec, _c_drp = _filter_requirements_file(tok)
|
||||
keep_args.append(_c_path)
|
||||
dropped.extend(_c_drp)
|
||||
elif prev_flag in _DROP_VALUE_FLAGS:
|
||||
# --upgrade-strategy (eager): the flag was appended when we saw
|
||||
# it; pop it and drop the flag+value pair so pip falls back to
|
||||
# its safe only-if-needed default.
|
||||
if keep_args and keep_args[-1] == prev_flag:
|
||||
keep_args.pop()
|
||||
dropped.append(prev_flag + " " + tok)
|
||||
elif prev_flag in _EDITABLE_FLAGS or prev_flag in _UPGRADE_PKG_FLAGS:
|
||||
# The flag was held back (not appended yet): its value is an
|
||||
# install target (-e path/url/vcs) or an upgrade selector
|
||||
|
|
@ -493,6 +549,8 @@ def main():
|
|||
if _req_rec and not recorded:
|
||||
recorded = _req_rec
|
||||
dropped.extend(_req_drp)
|
||||
elif _flag in _DROP_VALUE_FLAGS:
|
||||
dropped.append(tok) # --upgrade-strategy=eager -> drop the pair
|
||||
elif _flag in _CONSTRAINT_FILE_FLAGS:
|
||||
_c_path, _c_rec, _c_drp = _filter_requirements_file(_val)
|
||||
keep_args.append(_flag + "=" + _c_path)
|
||||
|
|
|
|||
|
|
@ -455,3 +455,78 @@ 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
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 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
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Item 3542096760 -- uv's PLURAL --requirements / --constraints go through the
|
||||
# same filter as the pip-style singular names.
|
||||
# --------------------------------------------------------------------------
|
||||
def test_uv_plural_requirements_filtered(shim, tmp_path):
|
||||
req = tmp_path / "reqs.txt"
|
||||
req.write_text("torch==2.11.0\nsnac==1.2.0\n", encoding = "utf-8")
|
||||
execd, _ = _run(shim, "uv", ["--requirements", str(req)])
|
||||
assert execd is not None and execd[0] == "--requirements", execd
|
||||
filtered = Path(execd[1]).read_text(encoding = "utf-8")
|
||||
assert "snac==1.2.0" in filtered
|
||||
assert "torch" not in filtered
|
||||
|
||||
|
||||
def test_uv_plural_constraints_filtered(shim, tmp_path):
|
||||
constraints = tmp_path / "constraints.txt"
|
||||
constraints.write_text("torch==2.11.0\n", encoding = "utf-8")
|
||||
execd, _ = _run(shim, "uv", ["--constraints", str(constraints), "peft"])
|
||||
assert execd is not None and execd[0] == "--constraints", execd
|
||||
assert "peft" in execd
|
||||
filtered = Path(execd[1]).read_text(encoding = "utf-8")
|
||||
assert "torch" not in filtered
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 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
|
||||
|
|
|
|||
|
|
@ -5,15 +5,17 @@
|
|||
#
|
||||
# The image bakes CUDA 13 ptxas + NVRTC, but a cu13 cubin cannot LOAD on a
|
||||
# 570-579 driver even when it targets an old arch like sm_80 (CUDA has forward,
|
||||
# not backward, driver compatibility across major versions). So the cu13 tools
|
||||
# must be activated ONLY for the two Blackwell datacenter arches that require
|
||||
# them -- sm_103 (B300 / GB300) and sm_121 (GB10 / DGX Spark), which only ship
|
||||
# on >= 580 drivers. Every other supported arch (Turing..sm_120) keeps the
|
||||
# bundled cu12.8 tools, so a 570+ driver host is never broken.
|
||||
# not backward, driver compatibility across major versions). So cu12.8 is the
|
||||
# IMMUTABLE baked default (libnvrtc.so.12 -> .cu128.orig), and the cu13 tools are
|
||||
# switched on ONLY for the two Blackwell datacenter arches that require them --
|
||||
# sm_103 (B300 / GB300) and sm_121 (GB10 / DGX Spark), which only ship on >= 580
|
||||
# drivers. Every other supported arch (Turing..sm_120) keeps the cu12.8 default,
|
||||
# untouched, so a 570+ driver host -- including a non-root --user container that
|
||||
# cannot rewrite the symlink -- is never broken.
|
||||
#
|
||||
# The function picks per device via nvidia-smi compute_cap: DC -> keep the
|
||||
# build's cu13 NVRTC (and point Triton at cu13 ptxas); anything else -> restore
|
||||
# the wheel-bundled cu12.8 NVRTC and leave ptxas unset (bundled cu12.8).
|
||||
# The function picks per device via nvidia-smi compute_cap: DC -> retarget
|
||||
# libnvrtc.so.12 -> the staged .cu13 alias (and point Triton at cu13 ptxas);
|
||||
# anything else -> leave the cu12.8 default in place and ptxas unset.
|
||||
set -e
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
||||
|
|
@ -39,11 +41,11 @@ assert_eq() {
|
|||
# $1 = compute_cap(s) the mock nvidia-smi reports, ONE PER LINE ("none" -> no
|
||||
# nvidia-smi on PATH). A multi-line value models a mixed-GPU host so we can check
|
||||
# that every visible cap is scanned, not just the first.
|
||||
# Builds a fake Studio venv NVRTC dir (libnvrtc.so.12 symlinked to a stand-in
|
||||
# cu13 lib, with the cu128 original saved beside it exactly as the build does)
|
||||
# and runs the function against it via UNSLOTH_STUDIO_HOME. The hardcoded base
|
||||
# venv path does not exist on the test host, so its glob is skipped. Prints
|
||||
# "<PTXAS_STATE> <NVRTC_TARGET>".
|
||||
# Builds a fake Studio venv NVRTC dir exactly as the build stages it: the real
|
||||
# cu12.8 lib as .cu128.orig, libnvrtc.so.12 -> it (the immutable default), and a
|
||||
# .cu13 alias pointing at a stand-in cu13 lib. Runs the function against it via
|
||||
# UNSLOTH_STUDIO_HOME. The hardcoded base venv path does not exist on the test
|
||||
# host, so its glob is skipped. Prints "<PTXAS_STATE> <NVRTC_TARGET>".
|
||||
run_select() {
|
||||
_cap="$1"
|
||||
_tmp=$(mktemp -d)
|
||||
|
|
@ -57,9 +59,10 @@ run_select() {
|
|||
fi
|
||||
_nvrtc="$_tmp/studio/unsloth_studio/lib/python3.12/site-packages/nvidia/cuda_nvrtc/lib"
|
||||
mkdir -p "$_nvrtc"
|
||||
: > "$_nvrtc/libnvrtc.so.12.cu128.orig"
|
||||
: > "$_nvrtc/libnvrtc.so.13.stub"
|
||||
ln -sf libnvrtc.so.13.stub "$_nvrtc/libnvrtc.so.12"
|
||||
: > "$_nvrtc/libnvrtc.so.12.cu128.orig" # real cu12.8 lib
|
||||
: > "$_nvrtc/libnvrtc.so.13.stub" # stand-in cu13 lib
|
||||
ln -sf libnvrtc.so.13.stub "$_nvrtc/libnvrtc.so.12.cu13" # staged cu13 alias
|
||||
ln -sf libnvrtc.so.12.cu128.orig "$_nvrtc/libnvrtc.so.12" # immutable cu12.8 default
|
||||
bash -c '
|
||||
set -euo pipefail
|
||||
export PATH="'"$_tmp"'/bin:/usr/bin:/bin"
|
||||
|
|
@ -74,28 +77,29 @@ run_select() {
|
|||
|
||||
echo "=== test_select_cuda_jit_tools ==="
|
||||
|
||||
# Non-DC arches: restore the cu12.8 NVRTC and leave ptxas unset (Triton keeps
|
||||
# its bundled cu12.8 ptxas), so a 570-579 driver host is unaffected.
|
||||
assert_eq "sm_80 Ampere -> cu128 NVRTC restored" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select 8.0)"
|
||||
assert_eq "sm_90 Hopper -> cu128 NVRTC restored" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select 9.0)"
|
||||
assert_eq "sm_100 B200 -> cu128 NVRTC restored" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select 10.0)"
|
||||
assert_eq "sm_120 RTX50 -> cu128 NVRTC restored" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select 12.0)"
|
||||
assert_eq "no nvidia-smi -> cu128 NVRTC restored" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select none)"
|
||||
# Non-DC arches: cu12.8 default is left untouched (no write) and ptxas unset
|
||||
# (Triton keeps its bundled cu12.8 ptxas), so a 570-579 driver host -- root or
|
||||
# --user -- is unaffected.
|
||||
assert_eq "sm_80 Ampere -> cu128 default kept" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select 8.0)"
|
||||
assert_eq "sm_90 Hopper -> cu128 default kept" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select 9.0)"
|
||||
assert_eq "sm_100 B200 -> cu128 default kept" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select 10.0)"
|
||||
assert_eq "sm_120 RTX50 -> cu128 default kept" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select 12.0)"
|
||||
assert_eq "no nvidia-smi -> cu128 default kept" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select none)"
|
||||
|
||||
# Blackwell datacenter: keep the build's cu13 NVRTC (NOT restored). ptxas stays
|
||||
# Blackwell datacenter: retarget libnvrtc.so.12 -> the .cu13 alias. ptxas stays
|
||||
# UNSET here only because the test host has no /usr/local/cuda-13.0/bin/ptxas;
|
||||
# the assertion that matters is that the cu13 NVRTC is preserved for these arches.
|
||||
assert_eq "sm_103 B300 -> cu13 NVRTC kept" "UNSET libnvrtc.so.13.stub" "$(run_select 10.3)"
|
||||
assert_eq "sm_121 DGX Spark -> cu13 NVRTC kept" "UNSET libnvrtc.so.13.stub" "$(run_select 12.1)"
|
||||
# the assertion that matters is that the NVRTC switched to cu13 for these arches.
|
||||
assert_eq "sm_103 B300 -> cu13 NVRTC selected" "UNSET libnvrtc.so.12.cu13" "$(run_select 10.3)"
|
||||
assert_eq "sm_121 DGX Spark -> cu13 NVRTC selected" "UNSET libnvrtc.so.12.cu13" "$(run_select 12.1)"
|
||||
|
||||
# Mixed-GPU hosts: a datacenter Blackwell (sm_103 / sm_121) sitting BEHIND an
|
||||
# H100/B200 in the nvidia-smi ordering must still enable cu13 -- every visible
|
||||
# cap is scanned, not just the first. And a host with no datacenter Blackwell at
|
||||
# all restores cu12.8 regardless of order.
|
||||
assert_eq "H100 then B300 -> cu13 NVRTC kept" "UNSET libnvrtc.so.13.stub" "$(run_select "$(printf '9.0\n10.3')")"
|
||||
assert_eq "B200 then GB10 -> cu13 NVRTC kept" "UNSET libnvrtc.so.13.stub" "$(run_select "$(printf '10.0\n12.1')")"
|
||||
assert_eq "B300 then H100 -> cu13 NVRTC kept" "UNSET libnvrtc.so.13.stub" "$(run_select "$(printf '10.3\n9.0')")"
|
||||
assert_eq "H100 then A100 -> cu128 restored" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select "$(printf '9.0\n8.0')")"
|
||||
# H100/B200 in the nvidia-smi ordering must still switch to cu13 -- every visible
|
||||
# cap is scanned, not just the first. A host with no datacenter Blackwell at all
|
||||
# keeps the cu12.8 default regardless of order.
|
||||
assert_eq "H100 then B300 -> cu13 NVRTC selected" "UNSET libnvrtc.so.12.cu13" "$(run_select "$(printf '9.0\n10.3')")"
|
||||
assert_eq "B200 then GB10 -> cu13 NVRTC selected" "UNSET libnvrtc.so.12.cu13" "$(run_select "$(printf '10.0\n12.1')")"
|
||||
assert_eq "B300 then H100 -> cu13 NVRTC selected" "UNSET libnvrtc.so.12.cu13" "$(run_select "$(printf '10.3\n9.0')")"
|
||||
assert_eq "H100 then A100 -> cu128 default kept" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select "$(printf '9.0\n8.0')")"
|
||||
|
||||
rm -f "$_FUNC_FILE"
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue