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

test: tighten apply_gpu_ids ROCm fallback assertions

Browse source 
Replace loose OR chain with exact string matches, split into three
focused tests, and add a guard check for the try/except wrapper.
This commit is contained in:
LeoBorcherding 2026-05-05 13:30:56 -05:00
commit 14fccdee48
1 changed files with 18 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -1259,28 +1259,37 @@ class TestApplyGpuIdsRocmFallback:
"""Verify apply_gpu_ids sets HIP_VISIBLE_DEVICES on ROCm hosts even when
IS_ROCM is still False (worker subprocess before detect_hardware runs)."""
def test_apply_gpu_ids_source_checks_torch_version_hip(self):
"""apply_gpu_ids should fall back to torch.version.hip when IS_ROCM is False."""
def test_apply_gpu_ids_falls_back_to_torch_version_hip(self):
"""apply_gpu_ids should probe torch.version.hip when IS_ROCM is False and no ROCm env vars are set."""
hw_path = (
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
)
source = hw_path.read_text()
func_start = source.find("def apply_gpu_ids")
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert 'getattr(_torch.version, "hip", None)' in func_body or \
"getattr(torch.version, 'hip', None)" in func_body or \
"torch.version.hip" in func_body
assert 'getattr(_torch.version, "hip", None)' in func_body
def test_apply_gpu_ids_source_sets_hip_visible_devices(self):
"""apply_gpu_ids should set HIP_VISIBLE_DEVICES and ROCR_VISIBLE_DEVICES."""
def test_apply_gpu_ids_sets_hip_and_rocr_visible_devices(self):
"""apply_gpu_ids should set both HIP_VISIBLE_DEVICES and ROCR_VISIBLE_DEVICES on ROCm."""
hw_path = (
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
)
source = hw_path.read_text()
func_start = source.find("def apply_gpu_ids")
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert "HIP_VISIBLE_DEVICES" in func_body
assert "ROCR_VISIBLE_DEVICES" in func_body
assert 'os.environ["HIP_VISIBLE_DEVICES"] = value' in func_body
assert 'os.environ["ROCR_VISIBLE_DEVICES"] = value' in func_body
def test_apply_gpu_ids_rocm_fallback_is_guarded_by_try_except(self):
"""torch import in apply_gpu_ids must be wrapped in try/except so a missing torch never crashes."""
hw_path = (
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
)
source = hw_path.read_text()
func_start = source.find("def apply_gpu_ids")
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert "import torch as _torch" in func_body
assert "except Exception" in func_body
# =============================================================================