From b3649d40ccae9f9db3a2e1776c3edd8e0ab2d30b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 07:16:30 +0000 Subject: [PATCH] docker: close notebook pip-shim bypasses and scan all GPUs for cu13 Notebook pip/uv shim (docker/unsloth_pip_shim.py), all active only under UNSLOTH_NB_SHIM=1: - Parse a bare wheel filename (torch-*.whl in the CWD, no ./ or / prefix) so it is matched against _KEEP instead of passing through as an opaque positional and reinstalling the baked torch. - Infer the distribution from an egg-less VCS URL by repo basename (git+https://github.com/huggingface/transformers.git -> transformers) so the egg-less form the repo itself recommends cannot clobber the baked stack. - Refuse remote (URL) -r/-c requirement/constraint files -- top-level and nested includes -- since their pins cannot be inspected before the real tool would fetch and install them. - Strip resolver-wide reinstall/ignore-installed switches (--force-reinstall, --ignore-installed, -I, uv --reinstall) so they cannot rebuild already-satisfied baked deps pulled in by a kept target. - Route uv --reinstall-package through the same _KEEP handling as -P/--upgrade-package (both attached and separated forms; no dangling flag). Entrypoint (docker/entrypoint.sh): select_cuda_jit_tools() now scans every visible GPU's compute_cap instead of only the first, so a datacenter Blackwell (sm_103/sm_121) behind an H100/B200 still enables the cu13 JIT tools it needs. Adds regression tests for each case (tests/python/test_unsloth_pip_shim.py, tests/sh/test_select_cuda_jit_tools.sh). --- docker/entrypoint.sh | 62 +++++++----- docker/unsloth_pip_shim.py | 131 +++++++++++++++++++------ tests/python/test_unsloth_pip_shim.py | 129 ++++++++++++++++++++++++ tests/sh/test_select_cuda_jit_tools.sh | 18 +++- 4 files changed, 282 insertions(+), 58 deletions(-) diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh index 0515b50ff2..23362888bf 100755 --- a/docker/entrypoint.sh +++ b/docker/entrypoint.sh @@ -39,35 +39,43 @@ set -euo pipefail # applies. Best-effort: a read-only / --user-dropped rootfs that cannot # re-point the NVRTC symlink is left unchanged. select_cuda_jit_tools() { - local cc="" nvrtc_dir orig + local caps="" cc nvrtc_dir orig need_cu13=0 if command -v nvidia-smi >/dev/null 2>&1; then - cc="$( { nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true; } \ - | head -n1 | tr -d '[:space:]' )" + 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. + 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 fi - case "${cc}" in - 10.3|12.1) - # Blackwell datacenter: 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 - ;; - *) - # Every other arch (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 - ;; - esac } # 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 1a6ce4ab97..1a834e9054 100644 --- a/docker/unsloth_pip_shim.py +++ b/docker/unsloth_pip_shim.py @@ -64,6 +64,7 @@ _VALUE_FLAGS = { "--index-strategy", "--upgrade-package", "-P", + "--reinstall-package", "--no-binary", "--only-binary", "--platform", @@ -88,11 +89,15 @@ _CONSTRAINT_FILE_FLAGS = {"-c", "--constraint"} # drop BOTH the flag and its value; dropping the value alone leaves pip a # dangling `-e` that swallows the next kept package and fails the whole cell. _EDITABLE_FLAGS = {"-e", "--editable"} -# -P/--upgrade-package is uv's selective-upgrade flag: naming a baked -# package (e.g. `uv pip install -P torch peft`) lets an ordinary install target -# refresh that package and clobber the pinned stack. Filter its value through -# _KEEP too. Unlike -e it is not itself an install target (no has_target). -_UPGRADE_PKG_FLAGS = {"-P", "--upgrade-package"} +# -P/--upgrade-package is uv's selective-upgrade flag and +# --reinstall-package is uv's selective-reinstall flag: naming a baked +# package (e.g. `uv pip install -P torch peft` or +# `uv pip install --reinstall-package torch peft`) lets an ordinary install +# target refresh/reinstall that package and clobber the pinned stack. Filter the +# value through _KEEP too, dropping the flag+value pair for a protected name so +# no dangling selector is left to swallow the next kept target. Unlike -e none of +# these is itself an install target (no has_target). +_UPGRADE_PKG_FLAGS = {"-P", "--upgrade-package", "--reinstall-package"} # Short value-flags pip/uv accept in the ATTACHED form, i.e. the 2-char flag # glued to its value in one token: `-rreqs.txt`, `-cconstraints.txt`, `-epath`, # `-Pname`. The scanner splits the flag from the value so the value is filtered @@ -100,6 +105,15 @@ _UPGRADE_PKG_FLAGS = {"-P", "--upgrade-package"} # as an opaque option -- otherwise an attached `-r`-only cell no-ops and an # attached `-c`/`-e`/`-P` value bypasses _KEEP. _ATTACHED_SHORT_FLAGS = {"-r", "-c", "-e", "-P"} +# Resolver-wide reinstall / ignore-installed switches (pip --force-reinstall, +# --ignore-installed, -I; uv --reinstall) force the tool to REINSTALL packages +# that are already satisfied -- including the baked torch/transformers pulled in +# as dependencies of a kept target. Drop them in shim mode so a +# `pip install --force-reinstall peft` cannot rebuild the pinned stack under the +# guise of installing an unprotected package. The kept target still installs; its +# 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"} def _canon(token): @@ -142,7 +156,35 @@ def _canon(token): dist = _whl.group(1).split("-", 1)[0].strip().lower().replace("_", "-") if dist: return dist - return None # vcs / url / local path -> let it pass through + # 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 + # (huggingface/transformers.git -> transformers, + # unslothai/unsloth-zoo.git -> unsloth-zoo). Infer it from the last path + # segment so a bare `pip install git+https://github.com/huggingface/ + # transformers.git` -- an egg-less form this repo itself recommends in + # unsloth/models/loader.py -- cannot reinstall the baked package past + # _KEEP. A non-protected repo returns its basename and the caller keeps + # the token as a normal target either way. + if re.match(r"^[a-z]+\+", token): + _seg = token.split("#", 1)[0].split("?", 1)[0].rstrip("/").rsplit("/", 1)[-1] + _seg = _seg.split("@", 1)[0] # drop a @branch / @tag / @commit ref + if _seg.endswith(".git"): + _seg = _seg[:-4] + _seg = _seg.strip().lower().replace("_", "-") + if _seg: + return _seg + return None # plain url / local path -> let it pass through + # A bare wheel filename (no ./ or / prefix and no scheme) is still a valid + # pip target from the CWD: `pip install torch-2.11.0-cp312-...-linux.whl`. + # It reaches here because it starts with neither `.`/`/` nor a scheme, so + # without this it would fall through as the whole filename and miss _KEEP, + # reinstalling the baked torch. Parse its PEP 427 distribution the same way + # as the URL/path wheel case above. + if token.lower().endswith(".whl"): + dist = token.rsplit("/", 1)[-1][:-4].split("-", 1)[0].strip().lower().replace("_", "-") + if dist: + return dist # strip extras and any version/marker tail name = re.split(r"[<>=!~\[\s;@]", token, 1)[0].strip() return name.lower().replace("_", "-") or None @@ -239,9 +281,13 @@ def _rewrite_include(line, stripped, src_dir, depth): rebuilt += " " + comment return rebuilt + newline_char - # A URL include cannot be filtered locally; leave it verbatim. + # A remote (URL) nested include cannot be fetched/filtered here, so its + # protected pins would reach the real tool untouched. Drop the include line + # instead of letting pip pull an unfiltered requirements file off the network + # (mirrors the top-level remote `-r`/`-c` refusal in main). new_line=None + # tells the caller to remove the line entirely. if "://" in target: - return line, False, None, [] + return None, True, None, [flag + " " + target] abs_target = target if os.path.isabs(target) else os.path.join(src_dir, target) # Recursively filter the included file. Guard against cyclic / deep includes. if depth < 8: @@ -308,7 +354,8 @@ def _filter_requirements_file(path, _depth = 0): # include (so protected specs deep in the include tree cannot slip # past _KEEP) and repoint it so it still resolves from /tmp. new_line, rewrote, inc_rec, inc_drp = _rewrite_include(line, stripped, src_dir, _depth) - out.append(new_line) + if new_line is not None: + out.append(new_line) # None -> a remote include was dropped if rewrote: changed = True if inc_rec and not recorded: @@ -376,24 +423,34 @@ def main(): # The value of -r/--requirement pulls real requirements (a target); the # value of an index-url / find-links / constraint / etc. flag is an # option, not something to install. - if prev_flag in _REQ_FILE_FLAGS: - # Filter baked/protected packages out of the requirements file so a - # notebook `pip install -r reqs.txt` cannot clobber the cu128 stack - # or push transformers into the base venv. - _req_path, _req_rec, _req_drp = _filter_requirements_file(tok) - keep_args.append(_req_path) - has_target = True - if _req_rec and not recorded: - recorded = _req_rec - dropped.extend(_req_drp) - elif prev_flag in _CONSTRAINT_FILE_FLAGS: - # Strip protected pins from the constraint file so it cannot - # downgrade the baked stack, but a constraint is not an install - # target and its transformers pin is not an install request, so - # do not set has_target / recorded here. - _c_path, _c_rec, _c_drp = _filter_requirements_file(tok) - keep_args.append(_c_path) - dropped.extend(_c_drp) + if prev_flag in _REQ_FILE_FLAGS or prev_flag in _CONSTRAINT_FILE_FLAGS: + if "://" in tok: + # Remote requirement/constraint file: it cannot be inspected + # or filtered, so refuse it in shim mode rather than let the + # real tool fetch and install protected pins off the network. + # The flag was appended when we first saw it; pop it so pip/uv + # is not left a dangling -r/-c. + if keep_args and keep_args[-1] == prev_flag: + keep_args.pop() + dropped.append(prev_flag + " " + tok) + elif prev_flag in _REQ_FILE_FLAGS: + # Filter baked/protected packages out of the requirements file + # so a notebook `pip install -r reqs.txt` cannot clobber the + # cu128 stack or push transformers into the base venv. + _req_path, _req_rec, _req_drp = _filter_requirements_file(tok) + keep_args.append(_req_path) + has_target = True + if _req_rec and not recorded: + recorded = _req_rec + dropped.extend(_req_drp) + else: + # Strip protected pins from the constraint file so it cannot + # downgrade the baked stack, but a constraint is not an install + # target and its transformers pin is not an install request, so + # do not set has_target / recorded here. + _c_path, _c_rec, _c_drp = _filter_requirements_file(tok) + keep_args.append(_c_path) + dropped.extend(_c_drp) elif prev_flag in _EDITABLE_FLAGS or prev_flag in _UPGRADE_PKG_FLAGS: # The flag was held back (not appended yet): its value is an # install target (-e path/url/vcs) or an upgrade selector @@ -424,7 +481,12 @@ def main(): if tok.startswith("--") and "=" in tok: _flag, _, _val = tok.partition("=") if _flag in _VALUE_FLAGS: - if _flag in _REQ_FILE_FLAGS: + if (_flag in _REQ_FILE_FLAGS or _flag in _CONSTRAINT_FILE_FLAGS) and "://" in _val: + # Remote requirement/constraint file in `--flag=URL` form: + # refuse it in shim mode (the flag rides in the same token, so + # dropping the token leaves nothing dangling). + dropped.append(tok) + elif _flag in _REQ_FILE_FLAGS: _req_path, _req_rec, _req_drp = _filter_requirements_file(_val) keep_args.append(_flag + "=" + _req_path) has_target = True @@ -459,7 +521,12 @@ def main(): # value and reuse the separated-form handling. if len(tok) > 2 and tok[0] == "-" and tok[1] != "-" and tok[:2] in _ATTACHED_SHORT_FLAGS: _sflag, _sval = tok[:2], tok[2:] - if _sflag in _REQ_FILE_FLAGS: + if (_sflag in _REQ_FILE_FLAGS or _sflag in _CONSTRAINT_FILE_FLAGS) and "://" in _sval: + # Remote requirement/constraint file in attached `-rURL`/`-cURL` + # form: refuse it in shim mode (nothing was appended yet, so just + # drop the whole token). + dropped.append(_sflag + " " + _sval) + elif _sflag in _REQ_FILE_FLAGS: _req_path, _req_rec, _req_drp = _filter_requirements_file(_sval) keep_args.append(_sflag) keep_args.append(_req_path) @@ -484,6 +551,12 @@ def main(): if _sflag in _EDITABLE_FLAGS: has_target = True continue + if tok in _REINSTALL_FLAGS: + # Resolver-wide reinstall / ignore-installed switch: drop it so pip/uv + # cannot rebuild already-satisfied baked deps (torch/transformers + # pulled in by a kept target). The kept target still installs. + dropped.append(tok) + continue if tok in _VALUE_FLAGS: # -e/--editable and -P/--upgrade-package carry a value that is a # potential install target, so hold the flag back and let the diff --git a/tests/python/test_unsloth_pip_shim.py b/tests/python/test_unsloth_pip_shim.py index 0ea57373af..bd6291de69 100644 --- a/tests/python/test_unsloth_pip_shim.py +++ b/tests/python/test_unsloth_pip_shim.py @@ -328,3 +328,132 @@ def test_attached_short_upgrade_package_protected_dropped(shim): execd, _ = _run(shim, "uv", ["-Ptorch", "peft"]) assert execd == ["peft"], execd assert "torch" not in execd and "-P" not in execd + + +# -------------------------------------------------------------------------- +# 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. +# -------------------------------------------------------------------------- +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 + + +# -------------------------------------------------------------------------- +# Item 3541773157 -- a protected VCS URL WITHOUT an #egg= fragment (the egg-less +# form this repo recommends) must be dropped via its repo basename. +# -------------------------------------------------------------------------- +def test_vcs_url_without_egg_protected_dropped(shim): + # git+https://github.com/huggingface/transformers.git -> transformers. + execd, _ = _run(shim, "pip", ["git+https://github.com/huggingface/transformers.git", "peft"]) + assert execd == ["peft"], execd + + +def test_vcs_url_without_egg_with_ref_dropped(shim): + execd, _ = _run( + shim, "pip", ["git+https://github.com/unslothai/unsloth-zoo.git@main", "peft"] + ) + assert execd == ["peft"], execd + + +def test_vcs_url_without_egg_unprotected_kept(shim): + url = "git+https://github.com/someone/coolpkg.git" + execd, _ = _run(shim, "pip", [url]) + assert execd == [url], execd + + +# -------------------------------------------------------------------------- +# 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 + + +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 + + +def test_nested_remote_include_dropped(shim, tmp_path): + # A local reqs file that pulls a remote include must have that include + # stripped, not passed through for the real pip to fetch unfiltered. + req = tmp_path / "reqs.txt" + req.write_text("-r https://example.com/evil.txt\nsnac==1.2.0\n", encoding = "utf-8") + execd, _ = _run(shim, "pip", ["-r", str(req)]) + assert execd is not None and execd[0] == "-r", execd + filtered = Path(execd[1]).read_text(encoding = "utf-8") + assert "snac==1.2.0" in filtered + assert "example.com" not in filtered and "://" not in filtered + + +# -------------------------------------------------------------------------- +# Item 3541773164 -- resolver-wide reinstall / ignore-installed flags are +# stripped so they cannot rebuild already-satisfied baked deps. +# -------------------------------------------------------------------------- +def test_force_reinstall_flag_stripped(shim): + execd, _ = _run(shim, "pip", ["--force-reinstall", "peft"]) + assert execd == ["peft"], execd + + +def test_ignore_installed_short_flag_stripped(shim): + execd, _ = _run(shim, "pip", ["-I", "peft"]) + assert execd == ["peft"], execd + + +def test_uv_reinstall_flag_stripped(shim): + execd, _ = _run(shim, "uv", ["--reinstall", "peft"]) + assert execd == ["peft"], execd + + +# -------------------------------------------------------------------------- +# 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 diff --git a/tests/sh/test_select_cuda_jit_tools.sh b/tests/sh/test_select_cuda_jit_tools.sh index b8bd9fca0a..f02f31bc03 100755 --- a/tests/sh/test_select_cuda_jit_tools.sh +++ b/tests/sh/test_select_cuda_jit_tools.sh @@ -36,7 +36,9 @@ assert_eq() { fi } -# $1 = compute_cap the mock nvidia-smi reports ("none" -> no nvidia-smi on PATH). +# $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 @@ -47,7 +49,10 @@ run_select() { _tmp=$(mktemp -d) mkdir -p "$_tmp/bin" if [ "$_cap" != "none" ]; then - printf '#!/bin/sh\necho "%s"\n' "$_cap" > "$_tmp/bin/nvidia-smi" + # nvidia-smi --query-gpu=compute_cap prints one cap per line; cat a file + # so an embedded newline in $_cap survives into the mock's output. + printf '%s\n' "$_cap" > "$_tmp/caps.txt" + printf '#!/bin/sh\ncat "%s"\n' "$_tmp/caps.txt" > "$_tmp/bin/nvidia-smi" chmod +x "$_tmp/bin/nvidia-smi" fi _nvrtc="$_tmp/studio/unsloth_studio/lib/python3.12/site-packages/nvidia/cuda_nvrtc/lib" @@ -83,6 +88,15 @@ assert_eq "no nvidia-smi -> cu128 NVRTC restored" "UNSET libnvrtc.so.12.cu128.or 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)" +# 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')")" + rm -f "$_FUNC_FILE" echo ""