From 251e3edf93e43ed6e5cf2dfc74d6e481d509f0c6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 06:20:31 +0000 Subject: [PATCH] docker: address review round 3 (requirement-file shim edges + device-gate cu13 JIT tools) unsloth_pip_shim.py: close three more ways a protected package slipped past _KEEP. An editable line (-e/--editable ) inside a -r requirements file is a real install target, so a protected editable there is now classified and dropped like the command-line case (new _parse_editable). pip/uv accept the attached short forms -rreqs.txt / -cconstraints.txt / -epath / -Pname as one token; these were falling through as opaque options (so an attached -r-only cell no-op'd and an attached -c/-e/-P value bypassed _KEEP), so the 2-char flag is now split from its value and routed through the separated-form handling. And a nested -c constraint inside a -r file no longer records its transformers pin as an install request (a constraint is not a request; mirrors the top-level -c path). entrypoint.sh / Dockerfile: gate the CUDA 13 ptxas + NVRTC to sm_103 / sm_121 at runtime instead of a global build-time default. A cu13 cubin needs a >= 580 driver to LOAD even when it targets an older arch (CUDA has forward, not backward, cross-major driver compatibility), but the image supports Turing.. sm_120 on a 570+ driver, so the previous global TRITON_PTXAS_PATH ENV + cu13 NVRTC symlink would break ordinary Triton/NVRTC JIT on 570-579 driver hosts. The build still bakes cu13 (saving the cu12.8 NVRTC as .cu128.orig); a new select_cuda_jit_tools() in the entrypoint reads the device compute_cap and only activates cu13 for sm_103/sm_121 (which ship >= 580 drivers), otherwise leaving Triton on its bundled cu12.8 ptxas and restoring the cu12.8 NVRTC in both the base and Studio venvs. The base ENTRYPOINT runs for the Studio image too. Adds 9 pip-shim regression tests and tests/sh/test_select_cuda_jit_tools.sh (7 device-gating cases); registers the latter in CI and tests/run_all.sh. --- .github/workflows/studio-backend-ci.yml | 1 + docker/Dockerfile | 39 +++++--- docker/entrypoint.sh | 61 +++++++++++-- docker/unsloth_pip_shim.py | 95 ++++++++++++++++++++ tests/python/test_unsloth_pip_shim.py | 115 ++++++++++++++++++++++++ tests/run_all.sh | 1 + tests/sh/test_select_cuda_jit_tools.sh | 90 +++++++++++++++++++ 7 files changed, 383 insertions(+), 19 deletions(-) create mode 100755 tests/sh/test_select_cuda_jit_tools.sh 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