From 6a078b1a45e5cc37a50697d2331af34f4b2e8c73 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 08:06:45 +0000 Subject: [PATCH] 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). --- docker/Dockerfile | 23 +++++--- docker/Dockerfile.studio | 22 ++++---- docker/entrypoint.sh | 61 ++++++++++----------- docker/unsloth_pip_shim.py | 62 ++++++++++++++++++++- tests/python/test_unsloth_pip_shim.py | 75 ++++++++++++++++++++++++++ tests/sh/test_select_cuda_jit_tools.sh | 72 +++++++++++++------------ 6 files changed, 233 insertions(+), 82 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index 8b319b2c82..3e250b6c9d 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -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 diff --git a/docker/Dockerfile.studio b/docker/Dockerfile.studio index bf41bc9d24..f392cb36cf 100644 --- a/docker/Dockerfile.studio +++ b/docker/Dockerfile.studio @@ -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 \ diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh index 23362888bf..5639103147 100755 --- a/docker/entrypoint.sh +++ b/docker/entrypoint.sh @@ -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 diff --git a/docker/unsloth_pip_shim.py b/docker/unsloth_pip_shim.py index 1a834e9054..afe51e6c8b 100644 --- a/docker/unsloth_pip_shim.py +++ b/docker/unsloth_pip_shim.py @@ -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 takes the NEXT token as its target (pip: # `-e, --editable `), 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) diff --git a/tests/python/test_unsloth_pip_shim.py b/tests/python/test_unsloth_pip_shim.py index cdade99991..dcb617255d 100644 --- a/tests/python/test_unsloth_pip_shim.py +++ b/tests/python/test_unsloth_pip_shim.py @@ -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 diff --git a/tests/sh/test_select_cuda_jit_tools.sh b/tests/sh/test_select_cuda_jit_tools.sh index f02f31bc03..578e0c2cb5 100755 --- a/tests/sh/test_select_cuda_jit_tools.sh +++ b/tests/sh/test_select_cuda_jit_tools.sh @@ -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 -# " ". +# 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 " ". 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"