ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix 14 stale tests under tests/studio/install/ that drifted from code (#5305)

Browse source 
* Fix 14 stale tests under tests/studio/install/ that drifted from code

All 14 failures audited locally and tracked back to test-side drift
(no production-code regressions). After these test updates the entire
tests/studio/install/ directory now passes: 346 passed, 1 skipped.

Per failure:

tests/studio/install/test_install_llama_prebuilt_logic.py (5 fails):

  * test_existing_install_matches_plan_with_fingerprint_linux
  * test_install_prebuilt_skips_download_when_existing_install_matches
  * test_install_prebuilt_skips_when_older_release_fallback_matches_existing_install
  * test_install_prebuilt_skips_same_release_fallback_attempt_when_installed
  * test_existing_install_matches_choice_fails_when_install_tree_incomplete

  All five build a fake Linux install tree via write_linux_install_shape
  and call existing_install_matches_choice. The matcher returns False
  because runtime_payload_is_healthy now requires a libllama-common.so*
  library in build/bin/ (added by PR #5135), and the fixture never wrote
  it. Add the missing library to write_linux_install_shape; matcher
  passes for all five tests.

tests/studio/install/test_rocm_support.py (8 fails after the partial
audit, one collection-tier flake):

  * TestEnsureRocmTorch::test_cpu_torch_gets_rocm_reinstall and
    TestEnsureRocmTorch::test_probe_timeout_triggers_reinstall

    _ensure_rocm_torch was refactored to call pip_install for the
    torch reinstall and pip_install_try (not pip_install) for the
    follow-up bitsandbytes install. The tests still asserted
    mock_pip.call_count == 2. Add a second @patch.object on
    pip_install_try and split the assertions across the two mocks.

  * TestInstallShStructure::test_cuda_precedence

    Asserted file-position-of-string ordering: looked for
    `if [ -z "$_smi" ]` before the first `amd-smi` literal in
    install.sh. The installer now defines top-level helpers
    `_has_amd_rocm_gpu` (uses `amd-smi`) and `_has_usable_nvidia_gpu`
    (uses `nvidia-smi`) before either is called from
    `get_torch_index_url`, so file-position ordering carries no
    semantic meaning. Rewrite the test to extract the
    `get_torch_index_url` body via a small brace-matched helper and
    assert the runtime ordering: NVIDIA call sits before the
    `if [ -z "$_smi" ]` branch and the AMD call sits inside it.

  * TestLiveRegression::test_get_torch_index_url_returns_cuda_on_nvidia

    Sed-extracted only get_torch_index_url and eval'd it -- but the
    function calls _has_amd_rocm_gpu and _has_usable_nvidia_gpu, so
    the eval'd body crashed and fell through to the CPU URL on a
    fully-loaded NVIDIA host. Extract the helpers alongside the
    function. Also pre-skip when nvidia-smi is on PATH but does not
    list a GPU (containers occasionally ship the binary without a
    driver).

  * TestWorkerRocmMambaSsm::test_probe_script_has_getattr_hip and
    TestWorkerRocmMambaSsm::test_probe_returns_hip_version_field

    The wheel-resolver probe subprocess (the only place where
    `getattr(torch.version, 'hip', None)` is emitted) was hoisted out
    of worker.py into studio/backend/utils/wheel_utils.py during the
    wheel-resolver refactor. Point the file-content assertions at
    wheel_utils.py and assert worker.py still consumes the
    `hip_version` field.

  * TestHardwareAmdBranching::test_hardware_branches_on_is_rocm_for_utilization
    TestHardwareAmdBranching::test_hardware_branches_on_is_rocm_for_visible
    TestHardwareAmdBranching::test_hardware_branches_on_is_rocm_for_physical_count

    hardware.py refactored: the IS_ROCM branch and direct
    `from . import amd` were hoisted out of get_gpu_utilization /
    get_visible_gpu_utilization into the shared `_smi_query`
    dispatcher. Update the first two tests to assert the dispatcher
    call shape (`_smi_query("get_primary_gpu_utilization", ...)` etc.)
    plus IS_ROCM + amd-import in `_smi_query` itself. Update the
    physical-count test to assert IS_ROCM + the literal `from . import
    amd` as that function still imports amd directly rather than going
    through `_smi_query`.

No production-code changes; tests-only.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-05-06 03:31:41 -07:00 • committed by GitHub
commit 94811ba75d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 120 additions and 32 deletions
Download patch file Download diff file Expand all files Collapse all files

3
tests/studio/install/test_install_llama_prebuilt_logic.py
View file

@ -754,6 +754,9 @@ def write_linux_install_shape(install_dir: Path) -> None:
(install_dir / "llama-quantize").write_text("#!/bin/sh\n", encoding = "utf-8")
(runtime_dir / "llama-server").write_text("#!/bin/sh\n", encoding = "utf-8")
(runtime_dir / "llama-quantize").write_text("#!/bin/sh\n", encoding = "utf-8")
# Mirror the runtime payload health groups in install_llama_prebuilt.py:
# libllama-common.so* was added by PR #5135 and is required.
(runtime_dir / "libllama-common.so.0").write_bytes(b"DLL")
(runtime_dir / "libllama.so.0").write_bytes(b"DLL")
(runtime_dir / "libggml.so.0").write_bytes(b"DLL")
(runtime_dir / "libggml-base.so.0").write_bytes(b"DLL")

149
tests/studio/install/test_rocm_support.py
View file

@ -53,6 +53,32 @@ _has_usable_nvidia_gpu = stack_mod._has_usable_nvidia_gpu
_ROCM_TORCH_INDEX = stack_mod._ROCM_TORCH_INDEX
def _extract_sh_function_body(source: str, name: str) -> str:
"""Return the body of a shell function from `source` by brace matching.
Used by structural tests that need to assert ordering of helper
calls inside a specific function rather than across the whole
install.sh file.
"""
needle = f"{name}() {{"
start = source.find(needle)
if start < 0:
return ""
depth = 0
i = start + len(needle) - 1 # land on the opening brace
n = len(source)
while i < n:
ch = source[i]
if ch == "{":
depth += 1
elif ch == "}":
depth -= 1
if depth == 0:
return source[start : i + 1]
i += 1
return source[start:]
# ── Helper: build HostInfo for different scenarios ──────────────────────────
@ -561,12 +587,13 @@ class TestEnsureRocmTorch:
_ensure_rocm_torch()
mock_pip.assert_not_called()
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
def test_cpu_torch_gets_rocm_reinstall(
self, mock_ver, mock_gpu, mock_nvidia, mock_pip
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
):
"""CPU-only torch on ROCm host should trigger reinstall."""
mock_probe = MagicMock()
@ -575,12 +602,11 @@ class TestEnsureRocmTorch:
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
# Should call pip_install twice: once for torch, once for bitsandbytes
assert mock_pip.call_count == 2
torch_call = mock_pip.call_args_list[0]
assert "rocm7.1" in str(torch_call)
bnb_call = mock_pip.call_args_list[1]
assert "bitsandbytes" in str(bnb_call)
# Should install torch via pip_install and bitsandbytes via pip_install_try.
assert mock_pip.call_count == 1
assert "rocm7.1" in str(mock_pip.call_args_list[0])
assert mock_pip_try.call_count >= 1
assert "bitsandbytes" in str(mock_pip_try.call_args_list[0])
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@ -642,12 +668,13 @@ class TestEnsureRocmTorch:
torch_call = mock_pip.call_args_list[0]
assert "rocm7.1" in str(torch_call)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
def test_probe_timeout_triggers_reinstall(
self, mock_ver, mock_gpu, mock_nvidia, mock_pip
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
):
"""Probe subprocess timeout should not crash; should proceed to reinstall."""
with patch("os.path.isdir", return_value = True):
@ -656,8 +683,10 @@ class TestEnsureRocmTorch:
):
_ensure_rocm_torch()
# If probe times out, the function should treat torch as unusable and reinstall
assert mock_pip.call_count == 2
# both torch (via pip_install) and bitsandbytes (via pip_install_try).
assert mock_pip.call_count == 1
assert "rocm7.1" in str(mock_pip.call_args_list[0])
assert mock_pip_try.call_count >= 1
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@ -857,15 +886,33 @@ class TestInstallShStructure:
assert "rocm" in source.lower()
def test_cuda_precedence(self):
"""ROCm detection should only run when nvidia-smi is absent."""
"""ROCm detection should only run when nvidia-smi is absent.
install.sh defines _has_amd_rocm_gpu and _has_usable_nvidia_gpu
helpers near each other (file-position order has no semantic
meaning), so check the runtime ordering inside
get_torch_index_url instead: NVIDIA branch runs first and the
AMD/ROCm branch only fires inside the `if [ -z "$_smi" ]`
block.
"""
sh_path = PACKAGE_ROOT / "install.sh"
source = sh_path.read_text()
# The ROCm block should be inside the "if [ -z "$_smi" ]" branch
smi_block_start = source.find('if [ -z "$_smi" ]')
rocm_block_start = source.find("amd-smi")
body = _extract_sh_function_body(source, "get_torch_index_url")
nvidia_call = body.find("_has_usable_nvidia_gpu")
no_nvidia_branch = body.find('if [ -z "$_smi" ]')
rocm_call = body.find("_has_amd_rocm_gpu")
assert (
smi_block_start < rocm_block_start
), "ROCm detection should be inside the 'no nvidia-smi' branch"
nvidia_call >= 0
), "get_torch_index_url should call _has_usable_nvidia_gpu"
assert (
no_nvidia_branch >= 0
), "get_torch_index_url should gate ROCm on no-nvidia-smi"
assert (
rocm_call > no_nvidia_branch
), "ROCm detection should sit inside the 'no nvidia-smi' branch"
assert (
nvidia_call < no_nvidia_branch
), "NVIDIA detection should run before the no-nvidia-smi branch"
def test_bitsandbytes_amd_install(self):
"""install.sh should install bitsandbytes for AMD when ROCm detected."""
@ -963,16 +1010,32 @@ class TestLiveRegression:
if not shutil.which("nvidia-smi"):
pytest.skip("No nvidia-smi available")
sh_path = PACKAGE_ROOT / "install.sh"
# Extract just the function (don't source the whole installer)
result = subprocess.run(
# Skip if nvidia-smi exists but does not actually list a GPU on this
# host (containers occasionally ship the binary without a driver).
check = subprocess.run(
[
"bash",
"-c",
f"eval \"$(sed -n '/^get_torch_index_url()/,/^}}/p' '{sh_path}')\"; "
"get_torch_index_url",
"nvidia-smi -L 2>/dev/null | "
"awk '/^GPU[[:space:]]+[0-9]+:/{f=1} END{exit !f}'",
],
capture_output = True,
)
if check.returncode != 0:
pytest.skip("nvidia-smi is on PATH but no GPU is listed")
sh_path = PACKAGE_ROOT / "install.sh"
# get_torch_index_url calls _has_usable_nvidia_gpu and
# _has_amd_rocm_gpu, so all three function definitions must be
# in scope when we eval the extract.
extract_cmd = (
f"sed -n '/^_has_amd_rocm_gpu()/,/^}}$/p; "
f"/^_has_usable_nvidia_gpu()/,/^}}$/p; "
f"/^get_torch_index_url()/,/^}}$/p' '{sh_path}'"
)
result = subprocess.run(
["bash", "-c", f'eval "$({extract_cmd})"; get_torch_index_url'],
capture_output = True,
text = True,
timeout = 30,
)
@ -988,19 +1051,23 @@ class TestLiveRegression:
# Load worker.py module
_WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "training" / "worker.py"
# The wheel-probe subprocess was hoisted out of worker.py into wheel_utils
# during the wheel-resolver refactor; the probe script literal lives there.
_WHEEL_UTILS_PATH = PACKAGE_ROOT / "studio" / "backend" / "utils" / "wheel_utils.py"
class TestWorkerRocmMambaSsm:
"""Verify worker.py Mamba/SSM install logic on ROCm."""
def test_probe_returns_hip_version_field(self):
"""_probe_causal_conv1d_env probe script should include hip_version."""
source = _WORKER_PATH.read_text()
assert "hip_version" in source
"""The wheel probe should include hip_version, and worker.py should
consume it."""
assert "hip_version" in _WHEEL_UTILS_PATH.read_text()
assert "hip_version" in _WORKER_PATH.read_text()
def test_probe_script_has_getattr_hip(self):
"""Probe script should use getattr for torch.version.hip (safe on CUDA)."""
source = _WORKER_PATH.read_text()
source = _WHEEL_UTILS_PATH.read_text()
assert "getattr(torch.version, 'hip', None)" in source
def test_direct_wheel_url_returns_none_without_cuda_major(self):
@ -1216,27 +1283,45 @@ class TestHardwareAmdBranching:
assert "from . import amd" in source
def test_hardware_branches_on_is_rocm_for_utilization(self):
"""get_gpu_utilization should check IS_ROCM before choosing backend."""
"""get_gpu_utilization should dispatch to amd.py via _smi_query
when IS_ROCM, and the dispatcher itself must check IS_ROCM and
import the amd backend."""
hw_path = (
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
)
source = hw_path.read_text()
# Find the get_gpu_utilization function
func_start = source.find("def get_gpu_utilization")
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert "IS_ROCM" in func_body
assert "amd.get_primary_gpu_utilization" in func_body
assert '_smi_query("get_primary_gpu_utilization"' in func_body
smi = source[
source.find("def _smi_query") : source.find(
"\ndef ", source.find("def _smi_query") + 1
)
]
assert "IS_ROCM" in smi
assert "from . import amd" in smi
def test_hardware_branches_on_is_rocm_for_visible(self):
"""get_visible_gpu_utilization should check IS_ROCM."""
"""get_visible_gpu_utilization should dispatch to amd.py via
_smi_query when IS_ROCM."""
hw_path = (
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
)
source = hw_path.read_text()
func_start = source.find("def get_visible_gpu_utilization")
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert "IS_ROCM" in func_body
assert "amd.get_visible_gpu_utilization" in func_body
# The dispatcher call may wrap onto multiple lines; allow whitespace
# between the open paren and the literal func name argument.
import re as _re
assert _re.search(r'_smi_query\(\s*"get_visible_gpu_utilization"', func_body)
smi = source[
source.find("def _smi_query") : source.find(
"\ndef ", source.find("def _smi_query") + 1
)
]
assert "IS_ROCM" in smi
assert "from . import amd" in smi
def test_hardware_branches_on_is_rocm_for_physical_count(self):
"""get_physical_gpu_count should try amd.py when IS_ROCM."""
@ -1247,7 +1332,7 @@ class TestHardwareAmdBranching:
func_start = source.find("def get_physical_gpu_count")
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert "IS_ROCM" in func_body
assert "amd.get_physical_gpu_count" in func_body
assert "from . import amd" in func_body
# =============================================================================