diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index b3392e7d07..a2db26e125 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -233,6 +233,7 @@ jobs: tests/sh/test_system_node_readonly.sh \ tests/sh/test_nvcc_meets_llama_minimum.sh \ tests/sh/test_resolve_cuda_archs.sh \ + tests/sh/test_select_cuda_jit_tools.sh \ tests/sh/test_tauri_install_exit_order.sh \ tests/sh/test_torch_constraint.sh \ tests/sh/test_torch_flavor.sh \ diff --git a/docker/Dockerfile b/docker/Dockerfile index 96b83aafc8..8b319b2c82 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -575,20 +575,31 @@ 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. Fix: symlink cu13 libnvrtc.so.13 over the bundled .so.12. +# 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). # # (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 # sm_120, rejects sm_103, and silently downgrades sm_121 to sm_80 per # triton-lang/triton#8335. Fix: install cu13 ptxas and point Triton at -# it with TRITON_PTXAS_PATH (ENV below). cu13.0 ptxas still spans -# sm_70..sm_90 (Volta through Hopper), so routing every JIT through it -# does not regress the older GPUs in the arch list above. +# it with TRITON_PTXAS_PATH. +# +# Both cu13 tools are activated ONLY for sm_103/sm_121, at runtime (see +# select_cuda_jit_tools in entrypoint.sh), NOT baked as a global ENV/symlink +# default: cu13 emits a cubin that a 570-579 driver cannot LOAD even when it +# targets an older arch (CUDA 13 requires a >= 580 driver), so forcing every +# host's JIT through cu13 would break the Ampere/Ada/Hopper/Turing GPUs this +# image still supports on 570+ drivers. sm_103/sm_121 launched after cu12.8 and +# only ship on >= 580 drivers, so gating cu13 to them is always safe. # # NVRTC and ptxas are CPU-side compilers; they do NOT call into libcuda, so -# cu13 installs alongside the cu128 runtime with no driver-floor bump (570+). -# Both arches carry the ~400 MB now: amd64 needs it for sm_103, arm64 for -# sm_121. +# cu13 installs alongside the cu128 runtime with no driver-floor bump at INSTALL +# time (570+). Their OUTPUT is a different story: a cu13 cubin needs a >= 580 +# driver to LOAD, so the tools are ACTIVATED per device at runtime (only for the +# sm_103/sm_121 hosts, which ship >= 580 drivers) -- see select_cuda_jit_tools +# in entrypoint.sh. Both arches carry the ~400 MB now: amd64 needs it 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,11 +621,15 @@ RUN set -eux; \ 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"; \ fi -# (2) ptxas override. Route every Triton JIT through the cu13 ptxas installed -# above (triton 3.6.0's own ptxas is cu12.8, no sm_103/sm_121). Set globally, -# not per-arch: cu13.0 ptxas spans sm_70..sm_121 so it is correct for every GPU -# this image supports, and ENV cannot be made conditional per arch. -ENV TRITON_PTXAS_PATH=/usr/local/cuda-13.0/bin/ptxas +# (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 +# global ENV: cu13 ptxas emits a cubin whose ABI a 570-579 driver cannot LOAD +# (CUDA 13 needs a >= 580 driver), even when targeting an older arch like sm_80, +# so pointing every host's Triton at it would break training on the Ampere/Ada/ +# Hopper/Turing GPUs this image still supports on 570+ drivers. TRITON_PTXAS_PATH +# is therefore selected per device at boot (only sm_103/sm_121, which ship >= 580 +# drivers, get cu13; everything else keeps Triton's bundled cu12.8 ptxas) -- see +# select_cuda_jit_tools in entrypoint.sh. # 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/entrypoint.sh b/docker/entrypoint.sh index 77523013ec..0515b50ff2 100755 --- a/docker/entrypoint.sh +++ b/docker/entrypoint.sh @@ -17,13 +17,60 @@ # docker run -e UNSLOTH_SKIP_GPU_CHECK=1 ... set -euo pipefail -# DGX Spark fix, arm64 image only: prefer the cu13 ptxas we baked into the -# image at /usr/local/cuda-13.0/bin/ptxas over Triton's bundled tools. The -# file only exists on the arm64 variant; amd64 images skip this and use -# Triton's own ptxas (cu13 in triton>=3.6.0). -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 +# --- CUDA JIT toolchain selection (device-gated) ---------------------------- +# The image bakes CUDA 13 ptxas + NVRTC ONLY so the two Blackwell datacenter +# arches the cu12.8 tools cannot target -- sm_103 (B300 / GB300) and sm_121 +# (GB10 / DGX Spark) -- can JIT Triton and torch/NVRTC kernels. Both launched +# AFTER cu12.8, so any host carrying them runs a >= 580 driver, which is exactly +# what a cu13-produced cubin needs to LOAD. +# +# Every OTHER supported arch (Turing..sm_120) works with the bundled cu12.8 +# tools and is allowed on a 570-579 driver (the documented floor). 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 routing those hosts' JIT through the cu13 tools would break ordinary +# training. ptxas/NVRTC are host-side compilers (they never link libcuda), so +# 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. +select_cuda_jit_tools() { + local cc="" nvrtc_dir orig + 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:]' )" + 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 # Make the unslothai/notebooks collection available under /workspace before the # user command runs (JupyterLab, unsloth-run, or a shell). Best-effort: it is diff --git a/docker/unsloth_pip_shim.py b/docker/unsloth_pip_shim.py index 841d172545..22c7eaf224 100644 --- a/docker/unsloth_pip_shim.py +++ b/docker/unsloth_pip_shim.py @@ -93,6 +93,13 @@ _EDITABLE_FLAGS = {"-e", "--editable"} # 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"} +# 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 +# (requirement/constraint file) or classified (-e/-P) instead of falling through +# 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"} def _canon(token): @@ -183,6 +190,33 @@ def _parse_include(stripped): 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+...` + else: + continue + return flag, (target or None), comment + return None, None, None + + def _rewrite_include(line, stripped, src_dir, depth): """Rewrite a nested `-r`/`-c` include so pip still resolves it and its protected specs are filtered too. @@ -212,6 +246,12 @@ def _rewrite_include(line, stripped, src_dir, depth): # Recursively filter the included file. Guard against cyclic / deep includes. if depth < 8: f_path, f_rec, f_drp = _filter_requirements_file(abs_target, _depth = depth + 1) + # A nested -c include is a resolver CONSTRAINT, not an install request, so + # a transformers pin inside it must NOT be recorded as a request (mirrors + # the top-level -c path in main(), which ignores _c_rec). Only a nested -r + # requirement include carries real install requests, so keep its pin. + if flag in _CONSTRAINT_FILE_FLAGS: + f_rec = None if f_path != abs_target: # The include was rewritten (protected specs dropped and/or its own # nested includes absolutised); point at the filtered copy. @@ -247,6 +287,23 @@ def _filter_requirements_file(path, _depth = 0): out.append(line) # comment / blank -> keep continue if stripped.startswith("-"): + # An editable requirement (-e/--editable ) inside the file is + # a real install target, so a protected editable such as + # `-e git+https://.../unsloth.git#egg=unsloth` would reinstall the + # 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) + if e_target is not None: + _action, _ver = _classify_flag_target(e_target) + if _action == "drop": + if _ver and not recorded: + recorded = _ver + dropped.append(e_flag + " " + e_target) + changed = True + continue + out.append(line) # kept editable -> forward the line verbatim + continue # Option or nested include. Recursively filter a nested `-r`/`-c` # include (so protected specs deep in the include tree cannot slip # past _KEEP) and repoint it so it still resolves from /tmp. @@ -394,6 +451,44 @@ def main(): else: keep_args.append(tok) # option with inline value, not a target continue + # Attached short value-flag form: pip/uv accept `-rreqs.txt`, + # `-cconstraints.txt`, `-epath` and `-Pname` as ONE token. Without this + # the token starts with "-" and falls through as an opaque option, so an + # `-r`-only cell no-ops (has_target stays False) and an attached + # `-c`/`-e`/`-P` value bypasses _KEEP. Split the 2-char flag from its + # 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: + _req_path, _req_rec, _req_drp = _filter_requirements_file(_sval) + keep_args.append(_sflag) + keep_args.append(_req_path) + has_target = True + if _req_rec and not recorded: + recorded = _req_rec + dropped.extend(_req_drp) + elif _sflag in _CONSTRAINT_FILE_FLAGS: + _c_path, _c_rec, _c_drp = _filter_requirements_file(_sval) + keep_args.append(_sflag) + keep_args.append(_c_path) + dropped.extend(_c_drp) + else: # -e / -P: the attached value is an install target / selector + _action, _ver = _classify_flag_target(_sval) + if _action == "drop": + if _ver and not recorded: + recorded = _ver + dropped.append(_sflag + " " + _sval) + else: + keep_args.append(_sflag) + keep_args.append(_sval) + if _sflag in _EDITABLE_FLAGS: + has_target = True + 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 a17692ae34..e99f26af00 100644 --- a/tests/python/test_unsloth_pip_shim.py +++ b/tests/python/test_unsloth_pip_shim.py @@ -217,3 +217,118 @@ def test_bare_transformers_recorded_and_dropped(shim): def test_index_url_value_flag_kept_verbatim(shim): execd, _ = _run(shim, "pip", ["--extra-index-url", "https://example.com/simple", "snac"]) assert execd == ["--extra-index-url", "https://example.com/simple", "snac"], execd + + +# -------------------------------------------------------------------------- +# Item 3541404842 -- filter editable entries INSIDE a requirements file. +# -------------------------------------------------------------------------- +def test_editable_protected_in_requirements_file_dropped(shim, tmp_path): + req = tmp_path / "reqs.txt" + req.write_text( + "-e git+https://github.com/unslothai/unsloth.git#egg=unsloth\n" + "snac==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 "unsloth" not in filtered # protected editable line stripped + + +def test_editable_attached_protected_in_requirements_file_dropped(shim, tmp_path): + req = tmp_path / "reqs.txt" + req.write_text( + "-egit+https://github.com/unslothai/unsloth.git#egg=unsloth\n" + "snac==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 "unsloth" not in filtered + + +def test_editable_unprotected_in_requirements_file_kept(shim, tmp_path): + # An unprotected editable survives even when the file is otherwise rewritten + # (torch dropped); only protected editables are stripped. + req = tmp_path / "reqs.txt" + req.write_text( + "-e ./localpkg\n" + "torch==2.11.0\n" + "snac==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 "./localpkg" in filtered + assert "snac==1.2.0" in filtered + assert "torch" not in filtered + + +# -------------------------------------------------------------------------- +# Item 3541404849 -- a nested -c constraint pin is not recorded as a request. +# -------------------------------------------------------------------------- +def test_nested_constraint_transformers_pin_not_recorded(shim, tmp_path): + constraints = tmp_path / "constraints.txt" + constraints.write_text("transformers==4.55.0\n", encoding = "utf-8") + req = tmp_path / "reqs.txt" + req.write_text("-c constraints.txt\nsnac==1.2.0\n", encoding = "utf-8") + execd, marker = _run(shim, "pip", ["-r", str(req)]) + assert execd is not None and execd[0] == "-r", execd + # A constraint pin is not an install request -> no sidecar marker written. + assert marker is None, marker + + +def test_nested_requirement_transformers_pin_recorded(shim, tmp_path): + # Contrast: a nested -r requirement DOES carry install requests, so its + # transformers pin is still recorded for the sidecar. + nested = tmp_path / "nested.txt" + nested.write_text("transformers==4.55.0\n", encoding = "utf-8") + req = tmp_path / "reqs.txt" + req.write_text("-r nested.txt\nsnac==1.2.0\n", encoding = "utf-8") + execd, marker = _run(shim, "pip", ["-r", str(req)]) + assert execd is not None and execd[0] == "-r", execd + assert marker == "4.55.0", marker + + +# -------------------------------------------------------------------------- +# Item 3541404845 -- handle pip's attached short options (-rfile / -cfile / etc). +# -------------------------------------------------------------------------- +def test_attached_short_requirement_file_filtered(shim, tmp_path): + # `pip install -rreqs.txt` (attached) must filter the file AND count as a + # target -- before the fix it fell through as an opaque option and no-op'd. + req = tmp_path / "reqs.txt" + req.write_text("torch==2.11.0\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 "torch" not in filtered + + +def test_attached_short_constraint_file_filtered(shim, tmp_path): + constraints = tmp_path / "constraints.txt" + constraints.write_text("torch==2.11.0\n", encoding = "utf-8") + execd, _ = _run(shim, "pip", ["-c" + str(constraints), "peft"]) + assert execd is not None and execd[0] == "-c", execd + assert "peft" in execd + filtered = Path(execd[1]).read_text(encoding = "utf-8") + 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 + assert "torch" not in execd and "-P" not in execd diff --git a/tests/run_all.sh b/tests/run_all.sh index d03f4c4d4f..95ab42a8c1 100755 --- a/tests/run_all.sh +++ b/tests/run_all.sh @@ -12,6 +12,7 @@ sh "$TESTS_DIR/sh/test_mac_intel_compat.sh" sh "$TESTS_DIR/sh/test_torch_constraint.sh" sh "$TESTS_DIR/sh/test_nvcc_meets_llama_minimum.sh" sh "$TESTS_DIR/sh/test_resolve_cuda_archs.sh" +sh "$TESTS_DIR/sh/test_select_cuda_jit_tools.sh" sh "$TESTS_DIR/sh/test_strixhalo_wsl_reroute.sh" sh "$TESTS_DIR/sh/test_uninstall_shared_icon.sh" sh "$TESTS_DIR/sh/test_torch_flavor.sh" diff --git a/tests/sh/test_select_cuda_jit_tools.sh b/tests/sh/test_select_cuda_jit_tools.sh new file mode 100755 index 0000000000..b8bd9fca0a --- /dev/null +++ b/tests/sh/test_select_cuda_jit_tools.sh @@ -0,0 +1,90 @@ +#!/bin/bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +# Unit tests for select_cuda_jit_tools() from docker/entrypoint.sh. +# +# 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. +# +# 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). +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +ENTRYPOINT_SH="$SCRIPT_DIR/../../docker/entrypoint.sh" +PASS=0 +FAIL=0 + +# Extract just the helper function (same sed range as the other function tests). +_FUNC_FILE=$(mktemp) +sed -n '/^select_cuda_jit_tools()/,/^}/p' "$ENTRYPOINT_SH" > "$_FUNC_FILE" + +assert_eq() { + _label="$1"; _expected="$2"; _actual="$3" + if [ "$_actual" = "$_expected" ]; then + echo " PASS: $_label" + PASS=$((PASS + 1)) + else + echo " FAIL: $_label (expected '$_expected', got '$_actual')" + FAIL=$((FAIL + 1)) + fi +} + +# $1 = compute_cap the mock nvidia-smi reports ("none" -> no nvidia-smi on PATH). +# 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 +# " ". +run_select() { + _cap="$1" + _tmp=$(mktemp -d) + mkdir -p "$_tmp/bin" + if [ "$_cap" != "none" ]; then + printf '#!/bin/sh\necho "%s"\n' "$_cap" > "$_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" + mkdir -p "$_nvrtc" + : > "$_nvrtc/libnvrtc.so.12.cu128.orig" + : > "$_nvrtc/libnvrtc.so.13.stub" + ln -sf libnvrtc.so.13.stub "$_nvrtc/libnvrtc.so.12" + bash -c ' + set -euo pipefail + export PATH="'"$_tmp"'/bin:/usr/bin:/bin" + export UNSLOTH_STUDIO_HOME="'"$_tmp"'/studio" + unset TRITON_PTXAS_PATH || true + . "'"$_FUNC_FILE"'" + select_cuda_jit_tools || true + printf "%s %s\n" "${TRITON_PTXAS_PATH:-UNSET}" "$(readlink "'"$_nvrtc"'/libnvrtc.so.12")" + ' + rm -rf "$_tmp" +} + +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)" + +# Blackwell datacenter: keep the build's cu13 NVRTC (NOT restored). 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)" + +rm -f "$_FUNC_FILE" + +echo "" +echo "Results: $PASS passed, $FAIL failed" +[ "$FAIL" -eq 0 ] || exit 1