Compare commits
109 commits
main
...
bnb-rocm-f
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ec8ec8e9c7 | ||
|
|
0f3833e283 | ||
|
|
aff62c8591 | ||
|
|
1de299c673 | ||
|
|
e2a4690899 | ||
|
|
75c89eaeb5 | ||
|
|
8a3d1f7bb3 | ||
|
|
2da5678c0c | ||
|
|
2bea7862ef |
||
|
|
735fcde44c |
||
|
|
8d38165771 | ||
|
|
3455b45977 | ||
|
|
b1b499c743 | ||
|
|
b22b243e1b | ||
|
|
ef4894c559 | ||
|
|
6c3c01c0c9 | ||
|
|
5dec1284eb | ||
|
|
7beecd279e | ||
|
|
75f831616e | ||
|
|
4373fdbaee | ||
|
|
3f55e7a98c | ||
|
|
a9b3321915 | ||
|
|
a42720636b | ||
|
|
2f111da0f6 | ||
|
|
2c9587d9fc | ||
|
|
c436f3ca67 | ||
|
|
12f34025dc | ||
|
|
5e5c6c8fed | ||
|
|
287243a8fd | ||
|
|
4aa7dd5b18 | ||
|
|
c9aa094c91 | ||
|
|
1b1807ef9a | ||
|
|
a370521879 | ||
|
|
b712f9f557 | ||
|
|
667467a315 | ||
|
|
55075f691e | ||
|
|
0e929cddac | ||
|
|
ecaf3dde2a | ||
|
|
fc9d82f6bf | ||
|
|
4d2afc62a0 | ||
|
|
43c669e6ec | ||
|
|
11b69c1051 | ||
|
|
39c7a4d290 | ||
|
|
ff1088af40 | ||
|
|
6d3849b821 | ||
|
|
8e26a368e1 | ||
|
|
e31d4c6aea | ||
|
|
5688072af6 | ||
|
|
d3ac7447eb | ||
|
|
2fdfe24fb1 | ||
|
|
e2215c9d11 | ||
|
|
b11f13a710 | ||
|
|
66ae2d416c | ||
|
|
b0a35ddeab | ||
|
|
405add94a1 | ||
|
|
64dc11faa2 | ||
|
|
07e2fccf38 | ||
|
|
d482382a92 |
||
|
|
342bbc2333 | ||
|
|
e1fe3be939 | ||
|
|
d459f60458 | ||
|
|
2c4bf9ae35 | ||
|
|
a87a08610e | ||
|
|
41abe89041 | ||
|
|
d638d1bd6f | ||
|
|
2e3b2bdc27 | ||
|
|
9cc539c1b4 | ||
|
|
973c7d80c2 | ||
|
|
c27f9b99e9 | ||
|
|
0c24d61708 | ||
|
|
e9a2b5c010 | ||
|
|
e8355451ea | ||
|
|
a7b4ae19ee | ||
|
|
9066946615 | ||
|
|
b2580ae32b | ||
|
|
7ae9580ce6 | ||
|
|
9ad3b761ee | ||
|
|
8f721d28d6 | ||
|
|
396aa05ead | ||
|
|
2297f73cad | ||
|
|
a728f7c308 | ||
|
|
884152daee | ||
|
|
7437af2e44 | ||
|
|
67678d2c29 | ||
|
|
e568000a92 | ||
|
|
6700dd60f0 | ||
|
|
abf578327c | ||
|
|
a4ae80cd6a | ||
|
|
ec47b2984d | ||
|
|
d9d1a63397 | ||
|
|
0bb6379aad | ||
|
|
1c608e8ff7 | ||
|
|
f9adf6834d | ||
|
|
1fd5853741 | ||
|
|
93a70fbe4e | ||
|
|
8b77451e75 | ||
|
|
481f0618ff | ||
|
|
ddf6f6d1f9 | ||
|
|
dddc9eac92 | ||
|
|
da81c94510 | ||
|
|
4c139503c1 | ||
|
|
9d1e3c38bc | ||
|
|
8bc26f4e1d | ||
|
|
044e67f5aa | ||
|
|
ba36c12240 | ||
|
|
0a5652281f | ||
|
|
bf63f79414 | ||
|
|
2973bea3d0 | ||
|
|
ddbbfe52cf |
7 changed files with 420 additions and 698 deletions
39
install.sh
39
install.sh
|
|
@ -257,10 +257,24 @@ run_install_cmd_retry() {
|
|||
done
|
||||
}
|
||||
|
||||
# Install bitsandbytes on AMD ROCm hosts. Uses the continuous-release_main
|
||||
# wheel for the ROCm 4-bit GEMV fix (bnb PR #1887, post-0.49.2); bnb <= 0.49.2
|
||||
# NaNs at decode shape on every AMD GPU. Falls back to PyPI >=0.49.1 if the
|
||||
# pre-release URL is unreachable. Drop the pin once bnb 0.50+ ships on PyPI.
|
||||
# Install bitsandbytes on AMD ROCm hosts. bnb <= 0.49.2 NaNs at 4-bit decode
|
||||
# shape on every AMD GPU; the fix (bnb #1887) is in the continuous-release_main
|
||||
# wheel used below and first ships on PyPI in 0.50.0, hence the fallback floor.
|
||||
_BNB_ROCM_PYPI_FALLBACK="bitsandbytes>=0.50.0"
|
||||
# bitsandbytes ships no ROCm binary in its aarch64 wheel at any version: the PyPI
|
||||
# 0.50.0 and continuous-release_main aarch64 wheels both carry only
|
||||
# libbitsandbytes_cpu.so plus CUDA variants. So neither install path below gives
|
||||
# aarch64 a 4-bit backend, and the messages must not claim one. Cf. gfx906.
|
||||
_bnb_rocm_arch_has_binary() {
|
||||
case "$_ARCH" in
|
||||
aarch64|arm64) return 1 ;;
|
||||
*) return 0 ;;
|
||||
esac
|
||||
}
|
||||
_warn_bnb_no_rocm_binary() {
|
||||
_bnb_rocm_arch_has_binary && return 0
|
||||
substep "[WARN] aarch64: bitsandbytes ships no ROCm kernels on this arch; 4-bit QLoRA needs a source build -- https://docs.unsloth.ai/get-started/install-and-update/amd" "$C_WARN"
|
||||
}
|
||||
_install_bnb_rocm() {
|
||||
_label="$1"
|
||||
_venv_py="$2"
|
||||
|
|
@ -275,9 +289,8 @@ _install_bnb_rocm() {
|
|||
_bnb_whl_url=""
|
||||
;;
|
||||
esac
|
||||
# uv rejects the continuous-release_main bitsandbytes wheel because the
|
||||
# filename version (1.33.7rc0) does not match the embedded metadata version
|
||||
# (0.50.0.dev0). pip accepts the mismatch, so bootstrap pip and use it.
|
||||
# uv rejects the pre-release wheel: its filename version (1.33.7rc0) does not
|
||||
# match its metadata version (0.50.x.dev0). pip accepts it, so bootstrap pip.
|
||||
if ! "$_venv_py" -m pip --version >/dev/null 2>&1; then
|
||||
if ! run_maybe_quiet "$_venv_py" -m ensurepip --upgrade; then
|
||||
run_maybe_quiet uv pip install --python "$_venv_py" pip || \
|
||||
|
|
@ -293,6 +306,7 @@ _install_bnb_rocm() {
|
|||
--retries 8 --timeout 90 \
|
||||
"$_bnb_whl_url" >"$_bnb_log" 2>&1; then
|
||||
rm -f "$_bnb_log"
|
||||
_warn_bnb_no_rocm_binary
|
||||
return 0
|
||||
fi
|
||||
_bnb_rc=$?
|
||||
|
|
@ -301,10 +315,17 @@ _install_bnb_rocm() {
|
|||
fi
|
||||
rm -f "$_bnb_log"
|
||||
step "warning" "$_label (pre-release) failed (exit code $_bnb_rc)" "$C_WARN" >&2
|
||||
substep "[WARN] bnb pre-release install failed; falling back to PyPI (4-bit decode broken on ROCm)" "$C_WARN"
|
||||
if _bnb_rocm_arch_has_binary; then
|
||||
substep "[WARN] bnb pre-release install failed; falling back to PyPI $_BNB_ROCM_PYPI_FALLBACK, which carries the ROCm 4-bit fix" "$C_WARN"
|
||||
else
|
||||
substep "[WARN] bnb pre-release install failed; falling back to PyPI $_BNB_ROCM_PYPI_FALLBACK" "$C_WARN"
|
||||
fi
|
||||
fi
|
||||
run_install_cmd "$_label (pypi fallback)" "$_venv_py" -m pip install \
|
||||
--force-reinstall --no-cache-dir --no-deps "bitsandbytes>=0.49.1"
|
||||
--force-reinstall --no-cache-dir --no-deps "$_BNB_ROCM_PYPI_FALLBACK"
|
||||
_bnb_pypi_rc=$?
|
||||
_warn_bnb_no_rocm_binary
|
||||
return $_bnb_pypi_rc
|
||||
}
|
||||
|
||||
if [ "$_next_is_package" = true ]; then
|
||||
|
|
|
|||
790
pyproject.toml
790
pyproject.toml
File diff suppressed because it is too large
Load diff
|
|
@ -8,12 +8,19 @@ filter_sensitive_data (structlog processor for sanitization), and
|
|||
get_logger (factory for structured loggers).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import structlog
|
||||
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
||||
|
||||
# Annotations only: importing at runtime would make the ASGI stack a hard
|
||||
# dependency of every CLI command.
|
||||
if TYPE_CHECKING:
|
||||
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
||||
|
||||
from utils.native_path_leases import redact_native_paths
|
||||
|
||||
|
|
|
|||
|
|
@ -389,8 +389,7 @@ _GFX_TO_AMD_INDEX_ARCH: dict[str, str] = {
|
|||
}
|
||||
|
||||
# bitsandbytes continuous-release_main wheels with the ROCm 4-bit GEMV fix
|
||||
# (bnb PR #1887, post-0.49.2). bnb <= 0.49.2 NaNs at decode shape on every
|
||||
# AMD GPU. Drop the pin once bnb 0.50+ ships on PyPI.
|
||||
# (bnb #1887). bnb <= 0.49.2 NaNs at 4-bit decode shape on every AMD GPU.
|
||||
_BNB_ROCM_PRERELEASE_URLS: dict[str, str] = {
|
||||
"x86_64": (
|
||||
"https://github.com/bitsandbytes-foundation/bitsandbytes/releases/"
|
||||
|
|
@ -411,7 +410,9 @@ _BNB_ROCM_PRERELEASE_URLS: dict[str, str] = {
|
|||
"bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl"
|
||||
),
|
||||
}
|
||||
_BNB_ROCM_PYPI_FALLBACK = "bitsandbytes>=0.49.1"
|
||||
# First PyPI release carrying bnb #1887. Keep in step with the install.sh
|
||||
# fallback (and the amd extra, once it lands here).
|
||||
_BNB_ROCM_PYPI_FALLBACK = "bitsandbytes>=0.50.0"
|
||||
|
||||
|
||||
def _bnb_rocm_prerelease_url() -> str | None:
|
||||
|
|
@ -423,6 +424,16 @@ def _bnb_rocm_prerelease_url() -> str | None:
|
|||
return _BNB_ROCM_PRERELEASE_URLS.get(arch)
|
||||
|
||||
|
||||
def _bnb_rocm_arch_has_binary() -> bool:
|
||||
"""False on aarch64: bitsandbytes ships no ROCm kernels there at any version.
|
||||
The PyPI 0.50.0 and continuous-release_main aarch64 wheels both carry only
|
||||
libbitsandbytes_cpu.so plus CUDA variants, so neither install path gives
|
||||
aarch64 a 4-bit backend and neither message may claim one.
|
||||
"""
|
||||
arch = platform.machine().lower()
|
||||
return {"amd64": "x86_64", "arm64": "aarch64"}.get(arch, arch) != "aarch64"
|
||||
|
||||
|
||||
def _amd_smi_env() -> dict[str, str] | None:
|
||||
"""On Windows, env with __COMPAT_LAYER=RunAsInvoker; None elsewhere.
|
||||
NB: RunAsInvoker doesn't stop amd-smi's runtime elevation (its manifest is
|
||||
|
|
@ -1178,20 +1189,24 @@ _rocm_windows_torch_installed: bool = False
|
|||
|
||||
|
||||
def _install_bnb_windows_rocm() -> bool:
|
||||
"""Install the AMD Windows BNB prerelease wheel. Returns True on success.
|
||||
"""Install AMD Windows BNB, pre-release wheel first. Returns True on success.
|
||||
|
||||
The continuous-release wheel is intentionally mismatched: the filename
|
||||
encodes 1.33.7.preview (parsed as 1.33.7rc0 by PEP 440) while the wheel
|
||||
metadata reports 0.50.0.dev0. uv rejects this filename/metadata mismatch,
|
||||
metadata reports 0.50.x.dev0. uv rejects this filename/metadata mismatch,
|
||||
and bypassing it with UV_SKIP_WHEEL_FILENAME_CHECK still leaves uv mangling
|
||||
the bitsandbytes install. Per the AMD install guide
|
||||
(https://unsloth.ai/docs/get-started/install/amd/amd-hackathon) the wheel
|
||||
must be installed with plain pip, not uv, so we force pip (force_pip=True);
|
||||
plain pip performs no wheel filename/metadata check.
|
||||
|
||||
When that URL is blocked, fall back to PyPI. Its win_amd64 wheel ships
|
||||
libbitsandbytes_rocm{714,72}.dll from 0.50.0 on, so the fallback is a real
|
||||
ROCm build; before 0.50.0 it was CUDA-only, which is why there was none.
|
||||
"""
|
||||
_bnb_win_url = _BNB_ROCM_PRERELEASE_URLS.get("win_amd64")
|
||||
if _bnb_win_url is None:
|
||||
return False
|
||||
_ok = False
|
||||
if _bnb_win_url is not None:
|
||||
_ok = pip_install_try(
|
||||
"bitsandbytes (AMD Windows, pre-release main)",
|
||||
"--force-reinstall",
|
||||
|
|
@ -1201,6 +1216,22 @@ def _install_bnb_windows_rocm() -> bool:
|
|||
constrain = False,
|
||||
force_pip = True,
|
||||
)
|
||||
if not _ok:
|
||||
print(
|
||||
_red(
|
||||
" bnb pre-release install failed; falling back to PyPI "
|
||||
f"{_BNB_ROCM_PYPI_FALLBACK}, which carries the ROCm 4-bit fix"
|
||||
)
|
||||
)
|
||||
if not _ok:
|
||||
_ok = pip_install_try(
|
||||
"bitsandbytes (AMD Windows)",
|
||||
"--force-reinstall",
|
||||
"--no-cache-dir",
|
||||
"--no-deps",
|
||||
_BNB_ROCM_PYPI_FALLBACK,
|
||||
constrain = False,
|
||||
)
|
||||
if not _ok:
|
||||
return False
|
||||
# Detect the actual ROCm DLL suffix in the wheel and set BNB_ROCM_VERSION so bnb
|
||||
|
|
@ -1690,8 +1721,8 @@ def _ensure_rocm_torch() -> None:
|
|||
pass
|
||||
if _torch_ok:
|
||||
_rocm_windows_torch_installed = True
|
||||
# ROCm torch is already installed, but the AMD Windows BNB wheel is still
|
||||
# needed (the PyPI bitsandbytes ships only CUDA DLLs, fails on ROCm).
|
||||
# ROCm torch is already installed, but bnb still needs the ROCm build
|
||||
# (pre-release wheel, else PyPI >=0.50.0).
|
||||
_install_bnb_windows_rocm()
|
||||
return
|
||||
# torch was wiped between runs; fall through to the full install path
|
||||
|
|
@ -1769,12 +1800,12 @@ def _ensure_rocm_torch() -> None:
|
|||
# separate dependency -- a BNB install failure must NOT roll back the
|
||||
# torch ROCm install.
|
||||
_rocm_windows_torch_installed = True
|
||||
# Always install AMD Windows bitsandbytes -- the PyPI wheel ships only
|
||||
# CUDA DLLs and fails on ROCm. Install even when torch was already a
|
||||
# ROCm build so `studio update` repairs a broken bnb.
|
||||
# Always install AMD Windows bitsandbytes, even when torch was already a
|
||||
# ROCm build, so `studio update` repairs a broken bnb.
|
||||
if not _install_bnb_windows_rocm():
|
||||
print(
|
||||
" Warning: AMD Windows bitsandbytes install failed; "
|
||||
" Warning: AMD Windows bitsandbytes install failed "
|
||||
"(pre-release and PyPI); "
|
||||
"ROCm torch is installed but bitsandbytes may need manual install"
|
||||
)
|
||||
return
|
||||
|
|
@ -2020,10 +2051,13 @@ def _ensure_rocm_torch() -> None:
|
|||
force_pip = True,
|
||||
)
|
||||
if not _bnb_installed:
|
||||
_fallback_note = (
|
||||
", which carries the ROCm 4-bit fix" if _bnb_rocm_arch_has_binary() else ""
|
||||
)
|
||||
print(
|
||||
_red(
|
||||
" bnb pre-release install failed; falling back to PyPI "
|
||||
"(4-bit decode will be broken on ROCm)"
|
||||
f"{_BNB_ROCM_PYPI_FALLBACK}{_fallback_note}"
|
||||
)
|
||||
)
|
||||
if not _bnb_installed:
|
||||
|
|
@ -2035,6 +2069,14 @@ def _ensure_rocm_torch() -> None:
|
|||
_BNB_ROCM_PYPI_FALLBACK,
|
||||
constrain = False,
|
||||
)
|
||||
if not _bnb_rocm_arch_has_binary():
|
||||
print(
|
||||
_red(
|
||||
" aarch64: bitsandbytes ships no ROCm kernels on this arch; "
|
||||
"4-bit QLoRA needs a source build -- "
|
||||
"https://docs.unsloth.ai/get-started/install-and-update/amd"
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# _uv_safe_path is imported from backend.utils.uv_path_safety (shared with mlx_repair).
|
||||
|
|
|
|||
|
|
@ -818,3 +818,79 @@ class TestPipNoIndexScrubParity:
|
|||
text = SETUP_PS1.read_text(encoding = "utf-8")
|
||||
assert "'PIP_NO_INDEX'" in text
|
||||
assert "'PIP_INDEX_URL'" in text
|
||||
|
||||
|
||||
class TestAmdBnbFloorParity:
|
||||
"""bitsandbytes <= 0.49.2 NaNs at 4-bit decode shape on every AMD GPU; the ROCm
|
||||
4-bit GEMV fix (bnb #1887) first ships on PyPI in 0.50.0. The install.sh PyPI
|
||||
fallback and the Studio stack fallback are the two ways a supported AMD flow
|
||||
resolves bitsandbytes when the pre-release wheel URL is unreachable, so neither
|
||||
may float back into the broken range."""
|
||||
|
||||
FLOOR = "0.50.0"
|
||||
PYPROJECT = REPO_ROOT / "pyproject.toml"
|
||||
|
||||
def test_install_sh_pypi_fallback_floor(self):
|
||||
text = INSTALL_SH.read_text(encoding = "utf-8")
|
||||
assert (
|
||||
f'_BNB_ROCM_PYPI_FALLBACK="bitsandbytes>={self.FLOOR}"' in text
|
||||
), f"install.sh _install_bnb_rocm PyPI fallback must floor at {self.FLOOR}"
|
||||
|
||||
def test_stack_py_pypi_fallback_floor(self):
|
||||
text = STACK_PY.read_text(encoding = "utf-8")
|
||||
assert (
|
||||
f'_BNB_ROCM_PYPI_FALLBACK = "bitsandbytes>={self.FLOOR}"' in text
|
||||
), f"install_python_stack.py PyPI fallback must floor at {self.FLOOR}"
|
||||
|
||||
def test_amd_extra_floor_when_present(self):
|
||||
"""The amd extra is not on this branch yet (#7278). Assert it only once it is,
|
||||
so the extra and the two installer fallbacks cannot drift apart later."""
|
||||
text = self.PYPROJECT.read_text(encoding = "utf-8")
|
||||
amd = re.search(r"^amd = \[(.*?)^\]", text, re.S | re.M)
|
||||
if amd is None:
|
||||
pytest.skip("pyproject.toml has no amd extra on this branch")
|
||||
specs = re.findall(r'"(bitsandbytes[^"]*)"', amd.group(1))
|
||||
assert specs, "the amd extra must pin bitsandbytes"
|
||||
for spec in specs:
|
||||
assert spec.startswith(
|
||||
f"bitsandbytes>={self.FLOOR}"
|
||||
), f"amd extra bitsandbytes floor must be >={self.FLOOR}, got {spec!r}"
|
||||
|
||||
def test_no_installer_still_allows_the_broken_range(self):
|
||||
for path in (INSTALL_SH, INSTALL_PS1, SETUP_PS1, STACK_PY, self.PYPROJECT):
|
||||
text = path.read_text(encoding = "utf-8")
|
||||
for line in text.splitlines():
|
||||
if "bitsandbytes>=0.49" in line and not line.lstrip().startswith(("#", "//")):
|
||||
raise AssertionError(
|
||||
f"{path.name} still floors bitsandbytes in the broken ROCm range: {line.strip()!r}"
|
||||
)
|
||||
|
||||
def test_fallback_is_not_reported_as_broken(self):
|
||||
"""The fallback now installs the first fixed release, so neither installer may
|
||||
still tell the user it leaves 4-bit decode broken on ROCm."""
|
||||
for path in (INSTALL_SH, STACK_PY):
|
||||
text = path.read_text(encoding = "utf-8")
|
||||
assert (
|
||||
"4-bit decode broken on ROCm" not in text
|
||||
), f"{path.name} still reports the repaired PyPI fallback as broken"
|
||||
assert (
|
||||
"4-bit decode will be broken on ROCm" not in text
|
||||
), f"{path.name} still reports the repaired PyPI fallback as broken"
|
||||
|
||||
def test_aarch64_is_not_told_it_has_a_rocm_backend(self):
|
||||
"""bitsandbytes ships no ROCm kernels in its aarch64 wheel at any version, so
|
||||
neither installer may hand aarch64 the x86_64 "carries the ROCm 4-bit fix"
|
||||
message, and both must warn that 4-bit needs a source build there."""
|
||||
sh = INSTALL_SH.read_text(encoding = "utf-8")
|
||||
assert "_bnb_rocm_arch_has_binary()" in sh
|
||||
assert "_warn_bnb_no_rocm_binary()" in sh
|
||||
assert (
|
||||
sh.count("_warn_bnb_no_rocm_binary\n") >= 2
|
||||
), "install.sh must warn on aarch64 after both the pre-release and the fallback install"
|
||||
py = STACK_PY.read_text(encoding = "utf-8")
|
||||
assert "def _bnb_rocm_arch_has_binary(" in py
|
||||
assert "_bnb_rocm_arch_has_binary()" in py
|
||||
for text, name in ((sh, "install.sh"), (py, "install_python_stack.py")):
|
||||
assert (
|
||||
"4-bit QLoRA needs a source build" in text
|
||||
), f"{name} must tell aarch64 users 4-bit needs a source build"
|
||||
|
|
|
|||
|
|
@ -2829,12 +2829,35 @@ class TestInstallBnbWindowsRocm:
|
|||
assert result is False
|
||||
assert "BNB_ROCM_VERSION" not in os.environ
|
||||
|
||||
def test_no_op_when_win_amd64_url_missing(self):
|
||||
"""Should be silent no-op if win_amd64 key absent from _BNB_ROCM_PRERELEASE_URLS."""
|
||||
def test_falls_back_to_pypi_when_win_amd64_url_missing(self):
|
||||
"""No win_amd64 pre-release wheel must not mean no bitsandbytes: PyPI
|
||||
>=0.50.0 ships libbitsandbytes_rocm{714,72}.dll, so it is a real ROCm build."""
|
||||
with patch.object(stack_mod, "_BNB_ROCM_PRERELEASE_URLS", {}):
|
||||
with patch.object(stack_mod, "pip_install_try") as mock_pip:
|
||||
with patch.object(stack_mod, "pip_install_try", return_value = True) as mock_pip:
|
||||
stack_mod._install_bnb_windows_rocm()
|
||||
mock_pip.assert_not_called()
|
||||
assert mock_pip.call_count == 1
|
||||
assert stack_mod._BNB_ROCM_PYPI_FALLBACK in mock_pip.call_args.args
|
||||
|
||||
def test_falls_back_to_pypi_when_prerelease_install_fails(self):
|
||||
"""A blocked GitHub pre-release URL must fall through to the PyPI floor rather
|
||||
than leaving Windows ROCm with no working bitsandbytes."""
|
||||
with patch.object(stack_mod, "pip_install_try", side_effect = [False, True]) as mock_pip:
|
||||
with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "72"):
|
||||
result = stack_mod._install_bnb_windows_rocm()
|
||||
assert result is True
|
||||
assert mock_pip.call_count == 2
|
||||
assert "win_amd64" in str(mock_pip.call_args_list[0])
|
||||
assert stack_mod._BNB_ROCM_PYPI_FALLBACK in mock_pip.call_args_list[1].args
|
||||
|
||||
def test_returns_false_only_when_both_paths_fail(self):
|
||||
"""Both the pre-release wheel and the PyPI fallback must fail before the
|
||||
helper reports failure."""
|
||||
with patch.dict(os.environ, {}, clear = False):
|
||||
os.environ.pop("BNB_ROCM_VERSION", None)
|
||||
with patch.object(stack_mod, "pip_install_try", return_value = False) as mock_pip:
|
||||
result = stack_mod._install_bnb_windows_rocm()
|
||||
assert result is False
|
||||
assert mock_pip.call_count == 2
|
||||
|
||||
def test_sets_bnb_rocm_version_from_detected_dll(self):
|
||||
"""BNB_ROCM_VERSION is set from the DLL detected after install."""
|
||||
|
|
|
|||
|
|
@ -0,0 +1,85 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""The studio extra must mirror studio/backend/requirements/studio.txt.
|
||||
|
||||
Nothing else keeps them in sync, and drift reintroduces #4701 / #5260 / #7147.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pathlib
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
REPO_ROOT = pathlib.Path(__file__).resolve().parents[3]
|
||||
PYPROJECT = REPO_ROOT / "pyproject.toml"
|
||||
STUDIO_TXT = REPO_ROOT / "studio" / "backend" / "requirements" / "studio.txt"
|
||||
|
||||
# Imported at module scope by the chain every CLI command walks: structlog via
|
||||
# studio.backend, click via unsloth_cli/commands/start.py.
|
||||
CORE_RUNTIME_PACKAGES = ("structlog", "click")
|
||||
|
||||
|
||||
def _load_pyproject() -> dict:
|
||||
if sys.version_info >= (3, 11):
|
||||
import tomllib
|
||||
else:
|
||||
tomllib = pytest.importorskip("tomli")
|
||||
return tomllib.loads(PYPROJECT.read_text(encoding = "utf-8"))
|
||||
|
||||
|
||||
def _requirement_lines(path: pathlib.Path) -> list[str]:
|
||||
out = []
|
||||
for line in path.read_text(encoding = "utf-8").splitlines():
|
||||
text = line.split("#", 1)[0].strip()
|
||||
if text and not text.startswith("-"):
|
||||
out.append(text)
|
||||
return out
|
||||
|
||||
|
||||
def _normalise(name: str) -> str:
|
||||
"""PEP 503 normalisation, so PyJWT/pyjwt and nest_asyncio/nest-asyncio match."""
|
||||
head = name
|
||||
for sep in ("===", "==", ">=", "<=", "~=", "!=", ">", "<", "[", ";", " "):
|
||||
idx = head.find(sep)
|
||||
if idx > 0:
|
||||
head = head[:idx]
|
||||
return head.strip().lower().replace("_", "-").replace(".", "-")
|
||||
|
||||
|
||||
def test_studio_extra_exists():
|
||||
extras = _load_pyproject()["project"]["optional-dependencies"]
|
||||
assert "studio" in extras, (
|
||||
"pyproject.toml has no `studio` extra. The wheel ships studio/ and "
|
||||
"studio.backend*, so their dependencies need a pip-installable home."
|
||||
)
|
||||
|
||||
|
||||
def test_studio_extra_matches_requirements_file():
|
||||
extras = _load_pyproject()["project"]["optional-dependencies"]
|
||||
extra = sorted(_normalise(entry) for entry in extras["studio"])
|
||||
required = sorted(_normalise(entry) for entry in _requirement_lines(STUDIO_TXT))
|
||||
|
||||
missing = sorted(set(required) - set(extra))
|
||||
surplus = sorted(set(extra) - set(required))
|
||||
assert not missing, (
|
||||
f"studio.txt lists {missing} but the `studio` extra does not. "
|
||||
'`pip install "unsloth[studio]"` would build a venv the Studio server '
|
||||
"cannot boot in. Add them to [project.optional-dependencies] studio."
|
||||
)
|
||||
assert not surplus, (
|
||||
f"The `studio` extra lists {surplus} but studio.txt does not. "
|
||||
"Remove them, or add them to studio.txt if install.sh needs them too."
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("package", CORE_RUNTIME_PACKAGES)
|
||||
def test_cli_runtime_packages_are_core_dependencies(package):
|
||||
core = [_normalise(entry) for entry in _load_pyproject()["project"]["dependencies"]]
|
||||
assert _normalise(package) in core, (
|
||||
f"{package} is imported at module scope by the studio.backend chain "
|
||||
f"`unsloth train` / `unsloth export` walk, so a plain `pip install "
|
||||
f"unsloth` must provide it or they die with ModuleNotFoundError."
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue