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 <target>) 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.
This commit is contained in:
Daniel Han 2026-07-08 06:20:31 +00:00
commit 251e3edf93
7 changed files with 383 additions and 19 deletions

View file

@ -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