Add AMD ROCm gaps: Mamba/SSM source builds, GPU monitoring, Windows messaging, RDNA expansion
- worker.py: Add HIP detection to causal-conv1d/mamba-ssm probe, check for hipcc before ROCm source builds, improve status messages and error reporting, add timeout and uv support for the source build fallback - amd.py: New AMD GPU monitoring module via amd-smi metric --json, mirroring nvidia.py structure (utilization, temperature, power, VRAM) - hardware.py: Branch to amd.py when IS_ROCM is True for GPU utilization, visible GPU queries, and physical GPU count - install_python_stack.py: Detect AMD GPUs on Windows and warn that ROCm-enabled PyTorch must be installed manually - kernels/utils.py: Expand is_rdna() to cover RDNA2 (gfx1030-1032), RDNA3 (gfx1102-1103), RDNA3.5 (gfx1150-1152) alongside existing entries - tests: Add 32 new tests covering all changes (95/95 pass)
This commit is contained in:
Daniel Han
parent
726fab1f37
commit
f17e007caf
6 changed files with 766 additions and 50 deletions
|
|
@ -86,6 +86,7 @@ def _probe_causal_conv1d_env() -> dict[str, str] | None:
|
|||
"'python_tag': f'cp{sys.version_info.major}{sys.version_info.minor}', "
|
||||
"'torch_mm': torch_mm, "
|
||||
"'cuda_major': str(int(str(torch.version.cuda).split('.', 1)[0])) if torch.version.cuda else '', "
|
||||
"'hip_version': str(torch.version.hip) if getattr(torch.version, 'hip', None) else '', "
|
||||
"'cxx11abi': str(torch._C._GLIBCXX_USE_CXX11_ABI).upper()"
|
||||
"}))"
|
||||
),
|
||||
|
|
@ -237,25 +238,88 @@ def _install_package_wheel_first(
|
|||
else:
|
||||
logger.info("No published %s wheel found: %s", display_name, wheel_url)
|
||||
|
||||
_send_status(event_queue, f"Installing {display_name} from PyPI...")
|
||||
pypi_cmd = [
|
||||
sys.executable,
|
||||
"-m",
|
||||
"pip",
|
||||
"install",
|
||||
"--no-build-isolation",
|
||||
"--no-deps",
|
||||
"--no-cache-dir",
|
||||
f"{pypi_name}=={pypi_version}",
|
||||
]
|
||||
result = _sp.run(
|
||||
pypi_cmd,
|
||||
stdout = _sp.PIPE,
|
||||
stderr = _sp.STDOUT,
|
||||
text = True,
|
||||
)
|
||||
is_hip = env and env.get("hip_version")
|
||||
if is_hip and not shutil.which("hipcc"):
|
||||
logger.error(
|
||||
"%s requires hipcc for source compilation on ROCm. "
|
||||
"Install the ROCm HIP SDK: https://rocm.docs.amd.com",
|
||||
display_name,
|
||||
)
|
||||
_send_status(
|
||||
event_queue,
|
||||
f"{display_name}: hipcc not found (ROCm HIP SDK required)",
|
||||
)
|
||||
return
|
||||
|
||||
if is_hip:
|
||||
_send_status(
|
||||
event_queue,
|
||||
f"Compiling {display_name} from source for ROCm "
|
||||
"(this may take several minutes)...",
|
||||
)
|
||||
else:
|
||||
_send_status(event_queue, f"Installing {display_name} from PyPI...")
|
||||
|
||||
# Prefer uv for faster dependency resolution when available
|
||||
if shutil.which("uv"):
|
||||
pypi_cmd = [
|
||||
"uv", "pip", "install",
|
||||
"--python", sys.executable,
|
||||
"--no-build-isolation",
|
||||
"--no-deps",
|
||||
f"{pypi_name}=={pypi_version}",
|
||||
]
|
||||
else:
|
||||
pypi_cmd = [
|
||||
sys.executable,
|
||||
"-m",
|
||||
"pip",
|
||||
"install",
|
||||
"--no-build-isolation",
|
||||
"--no-deps",
|
||||
"--no-cache-dir",
|
||||
f"{pypi_name}=={pypi_version}",
|
||||
]
|
||||
|
||||
# Source compilation on ROCm can take 5-10 minutes; use a generous timeout
|
||||
timeout = 600 if is_hip else 300
|
||||
|
||||
try:
|
||||
result = _sp.run(
|
||||
pypi_cmd,
|
||||
stdout = _sp.PIPE,
|
||||
stderr = _sp.STDOUT,
|
||||
text = True,
|
||||
timeout = timeout,
|
||||
)
|
||||
except _sp.TimeoutExpired:
|
||||
logger.error(
|
||||
"%s installation timed out after %ds", display_name, timeout,
|
||||
)
|
||||
_send_status(
|
||||
event_queue,
|
||||
f"{display_name} installation timed out after {timeout}s",
|
||||
)
|
||||
return
|
||||
|
||||
if result.returncode != 0:
|
||||
logger.error("Failed to install %s from PyPI:\n%s", display_name, result.stdout)
|
||||
if is_hip:
|
||||
# Surface a clear error for ROCm source build failures
|
||||
error_lines = (result.stdout or "").strip().splitlines()
|
||||
snippet = "\n".join(error_lines[-5:]) if error_lines else "(no output)"
|
||||
logger.error(
|
||||
"Failed to compile %s for ROCm:\n%s", display_name, result.stdout,
|
||||
)
|
||||
_send_status(
|
||||
event_queue,
|
||||
f"Failed to compile {display_name} for ROCm. "
|
||||
"Check that hipcc and ROCm development headers are installed.\n"
|
||||
f"{snippet}",
|
||||
)
|
||||
else:
|
||||
logger.error(
|
||||
"Failed to install %s from PyPI:\n%s", display_name, result.stdout,
|
||||
)
|
||||
return
|
||||
|
||||
logger.info("Installed %s from PyPI", display_name)
|
||||
|
|
|
|||
224
studio/backend/utils/hardware/amd.py
Normal file
224
studio/backend/utils/hardware/amd.py
Normal file
|
|
@ -0,0 +1,224 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""AMD GPU monitoring via amd-smi.
|
||||
|
||||
Mirrors the nvidia.py module structure so hardware.py can swap backends
|
||||
based on IS_ROCM. All functions return the same dict shapes as their
|
||||
nvidia.py counterparts.
|
||||
"""
|
||||
|
||||
import json
|
||||
import subprocess
|
||||
from typing import Any, Optional
|
||||
|
||||
from loggers import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[dict]:
|
||||
"""Run amd-smi with the given arguments and return parsed JSON, or None."""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["amd-smi", *args, "--json"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=timeout,
|
||||
)
|
||||
except (OSError, subprocess.TimeoutExpired) as e:
|
||||
logger.warning("amd-smi query failed: %s", e)
|
||||
return None
|
||||
if result.returncode != 0 or not result.stdout.strip():
|
||||
logger.warning("amd-smi returned code %d", result.returncode)
|
||||
return None
|
||||
try:
|
||||
return json.loads(result.stdout)
|
||||
except json.JSONDecodeError:
|
||||
logger.warning("Failed to parse amd-smi JSON output")
|
||||
return None
|
||||
|
||||
|
||||
def _parse_numeric(value: Any) -> Optional[float]:
|
||||
"""Extract a numeric value from amd-smi output (may be str, int, float, or dict)."""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, (int, float)):
|
||||
return float(value)
|
||||
if isinstance(value, str):
|
||||
# Strip units like "W", "C", "%", "MB" etc.
|
||||
cleaned = value.strip().rstrip("WCMBGb% ").strip()
|
||||
if not cleaned or cleaned.lower() in ("n/a", "none", "unknown"):
|
||||
return None
|
||||
try:
|
||||
return float(cleaned)
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]:
|
||||
"""Extract standardized metrics from a single GPU's amd-smi data."""
|
||||
# amd-smi metric output structure varies by version; try common paths
|
||||
usage = gpu_data.get("usage", gpu_data.get("gpu_activity", {}))
|
||||
if isinstance(usage, dict):
|
||||
gpu_util = _parse_numeric(usage.get("gfx_activity", usage.get("gpu_use_percent")))
|
||||
else:
|
||||
gpu_util = _parse_numeric(usage)
|
||||
|
||||
# Temperature
|
||||
temp_data = gpu_data.get("temperature", {})
|
||||
if isinstance(temp_data, dict):
|
||||
temp = _parse_numeric(
|
||||
temp_data.get("edge", temp_data.get("temperature_edge",
|
||||
temp_data.get("hotspot", temp_data.get("temperature_hotspot"))))
|
||||
)
|
||||
else:
|
||||
temp = _parse_numeric(temp_data)
|
||||
|
||||
# Power
|
||||
power_data = gpu_data.get("power", {})
|
||||
if isinstance(power_data, dict):
|
||||
power_draw = _parse_numeric(
|
||||
power_data.get("current_socket_power",
|
||||
power_data.get("average_socket_power",
|
||||
power_data.get("socket_power")))
|
||||
)
|
||||
power_limit = _parse_numeric(
|
||||
power_data.get("power_cap", power_data.get("max_power_limit"))
|
||||
)
|
||||
else:
|
||||
power_draw = None
|
||||
power_limit = None
|
||||
|
||||
# VRAM
|
||||
vram_data = gpu_data.get("vram", gpu_data.get("fb_memory_usage", {}))
|
||||
if isinstance(vram_data, dict):
|
||||
vram_used_bytes = _parse_numeric(
|
||||
vram_data.get("vram_used", vram_data.get("used"))
|
||||
)
|
||||
vram_total_bytes = _parse_numeric(
|
||||
vram_data.get("vram_total", vram_data.get("total"))
|
||||
)
|
||||
else:
|
||||
vram_used_bytes = None
|
||||
vram_total_bytes = None
|
||||
|
||||
# Convert VRAM from bytes to MB if values are large (>10000 = likely bytes)
|
||||
vram_used_mb = None
|
||||
vram_total_mb = None
|
||||
if vram_used_bytes is not None:
|
||||
if vram_used_bytes > 100000: # Likely bytes
|
||||
vram_used_mb = vram_used_bytes / (1024 * 1024)
|
||||
else: # Likely already MB
|
||||
vram_used_mb = vram_used_bytes
|
||||
if vram_total_bytes is not None:
|
||||
if vram_total_bytes > 100000: # Likely bytes
|
||||
vram_total_mb = vram_total_bytes / (1024 * 1024)
|
||||
else: # Likely already MB
|
||||
vram_total_mb = vram_total_bytes
|
||||
|
||||
# Build the standardized dict (same shape as nvidia._build_gpu_metrics)
|
||||
vram_used_gb = round(vram_used_mb / 1024, 2) if vram_used_mb is not None else None
|
||||
vram_total_gb = round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None
|
||||
vram_util = (
|
||||
round((vram_used_mb / vram_total_mb) * 100, 1)
|
||||
if vram_used_mb is not None and vram_total_mb and vram_total_mb > 0
|
||||
else None
|
||||
)
|
||||
power_util = (
|
||||
round((power_draw / power_limit) * 100, 1)
|
||||
if power_draw is not None and power_limit and power_limit > 0
|
||||
else None
|
||||
)
|
||||
|
||||
return {
|
||||
"gpu_utilization_pct": gpu_util,
|
||||
"temperature_c": temp,
|
||||
"vram_used_gb": vram_used_gb,
|
||||
"vram_total_gb": vram_total_gb,
|
||||
"vram_utilization_pct": vram_util,
|
||||
"power_draw_w": power_draw,
|
||||
"power_limit_w": power_limit,
|
||||
"power_utilization_pct": power_util,
|
||||
}
|
||||
|
||||
|
||||
def get_physical_gpu_count() -> Optional[int]:
|
||||
"""Return physical AMD GPU count via amd-smi, or None on failure."""
|
||||
data = _run_amd_smi("list")
|
||||
if data is None:
|
||||
return None
|
||||
if isinstance(data, list):
|
||||
return len(data)
|
||||
# Some versions return a dict with a "gpu" key
|
||||
gpus = data.get("gpu", data.get("gpus", []))
|
||||
if isinstance(gpus, list):
|
||||
return len(gpus)
|
||||
return None
|
||||
|
||||
|
||||
def get_primary_gpu_utilization() -> dict[str, Any]:
|
||||
"""Return utilization metrics for the primary AMD GPU."""
|
||||
data = _run_amd_smi("metric", "-g", "0")
|
||||
if data is None:
|
||||
return {"available": False}
|
||||
|
||||
# amd-smi may return a list with one entry or a dict
|
||||
if isinstance(data, list):
|
||||
if len(data) == 0:
|
||||
return {"available": False}
|
||||
gpu_data = data[0]
|
||||
else:
|
||||
gpu_data = data
|
||||
|
||||
metrics = _extract_gpu_metrics(gpu_data)
|
||||
metrics["available"] = True
|
||||
return metrics
|
||||
|
||||
|
||||
def get_visible_gpu_utilization(
|
||||
parent_visible_ids: Optional[list[int]],
|
||||
parent_cuda_visible_devices: Optional[str] = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Return utilization metrics for visible AMD GPUs."""
|
||||
if parent_visible_ids is None:
|
||||
return {
|
||||
"available": False,
|
||||
"backend_cuda_visible_devices": parent_cuda_visible_devices,
|
||||
"parent_visible_gpu_ids": [],
|
||||
"devices": [],
|
||||
"index_kind": "unresolved",
|
||||
}
|
||||
|
||||
data = _run_amd_smi("metric")
|
||||
if data is None:
|
||||
return {
|
||||
"available": False,
|
||||
"backend_cuda_visible_devices": parent_cuda_visible_devices,
|
||||
"parent_visible_gpu_ids": parent_visible_ids or [],
|
||||
"devices": [],
|
||||
"index_kind": "physical",
|
||||
}
|
||||
|
||||
gpu_list = data if isinstance(data, list) else data.get("gpus", [data])
|
||||
visible_set = set(parent_visible_ids)
|
||||
ordinal_map = {gpu_id: ordinal for ordinal, gpu_id in enumerate(parent_visible_ids)}
|
||||
|
||||
devices = []
|
||||
for idx, gpu_data in enumerate(gpu_list):
|
||||
if idx not in visible_set:
|
||||
continue
|
||||
metrics = _extract_gpu_metrics(gpu_data)
|
||||
metrics["index"] = idx
|
||||
metrics["index_kind"] = "physical"
|
||||
metrics["visible_ordinal"] = ordinal_map.get(idx, len(devices))
|
||||
devices.append(metrics)
|
||||
|
||||
return {
|
||||
"available": len(devices) > 0,
|
||||
"backend_cuda_visible_devices": parent_cuda_visible_devices,
|
||||
"parent_visible_gpu_ids": parent_visible_ids or [],
|
||||
"devices": devices,
|
||||
"index_kind": "physical",
|
||||
}
|
||||
|
|
@ -405,15 +405,26 @@ def get_gpu_utilization() -> Dict[str, Any]:
|
|||
device = get_device()
|
||||
|
||||
if device == DeviceType.CUDA:
|
||||
try:
|
||||
from . import nvidia
|
||||
if IS_ROCM:
|
||||
try:
|
||||
from . import amd
|
||||
|
||||
result = nvidia.get_primary_gpu_utilization()
|
||||
if result.get("available"):
|
||||
result["backend"] = device.value
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.warning("nvidia-smi utilization query failed: %s", e)
|
||||
result = amd.get_primary_gpu_utilization()
|
||||
if result.get("available"):
|
||||
result["backend"] = device.value
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.warning("amd-smi utilization query failed: %s", e)
|
||||
else:
|
||||
try:
|
||||
from . import nvidia
|
||||
|
||||
result = nvidia.get_primary_gpu_utilization()
|
||||
if result.get("available"):
|
||||
result["backend"] = device.value
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.warning("nvidia-smi utilization query failed: %s", e)
|
||||
|
||||
mem = get_gpu_memory_info()
|
||||
if device != DeviceType.CPU and mem.get("available"):
|
||||
|
|
@ -438,18 +449,32 @@ def get_visible_gpu_utilization() -> Dict[str, Any]:
|
|||
|
||||
if device == DeviceType.CUDA:
|
||||
parent_visible_spec = _get_parent_visible_gpu_spec()
|
||||
try:
|
||||
from . import nvidia
|
||||
if IS_ROCM:
|
||||
try:
|
||||
from . import amd
|
||||
|
||||
result = nvidia.get_visible_gpu_utilization(
|
||||
parent_visible_spec["numeric_ids"],
|
||||
parent_cuda_visible_devices = parent_visible_spec["raw"],
|
||||
)
|
||||
if result.get("available"):
|
||||
result["backend"] = device.value
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.warning("nvidia-smi visible GPU utilization query failed: %s", e)
|
||||
result = amd.get_visible_gpu_utilization(
|
||||
parent_visible_spec["numeric_ids"],
|
||||
parent_cuda_visible_devices = parent_visible_spec["raw"],
|
||||
)
|
||||
if result.get("available"):
|
||||
result["backend"] = device.value
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.warning("amd-smi visible GPU utilization query failed: %s", e)
|
||||
else:
|
||||
try:
|
||||
from . import nvidia
|
||||
|
||||
result = nvidia.get_visible_gpu_utilization(
|
||||
parent_visible_spec["numeric_ids"],
|
||||
parent_cuda_visible_devices = parent_visible_spec["raw"],
|
||||
)
|
||||
if result.get("available"):
|
||||
result["backend"] = device.value
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.warning("nvidia-smi visible GPU utilization query failed: %s", e)
|
||||
|
||||
# Torch-based fallback for CUDA (nvidia-smi unavailable, AMD ROCm) and XPU (Intel)
|
||||
if device in (DeviceType.CUDA, DeviceType.XPU):
|
||||
|
|
@ -1121,16 +1146,27 @@ def get_physical_gpu_count() -> int:
|
|||
device = get_device()
|
||||
|
||||
if device == DeviceType.CUDA:
|
||||
try:
|
||||
from . import nvidia
|
||||
if IS_ROCM:
|
||||
try:
|
||||
from . import amd
|
||||
|
||||
count = nvidia.get_physical_gpu_count()
|
||||
if count is not None:
|
||||
_physical_gpu_count = count
|
||||
return _physical_gpu_count
|
||||
except Exception:
|
||||
pass
|
||||
# nvidia-smi unavailable or failed — fall back to torch
|
||||
count = amd.get_physical_gpu_count()
|
||||
if count is not None:
|
||||
_physical_gpu_count = count
|
||||
return _physical_gpu_count
|
||||
except Exception:
|
||||
pass
|
||||
else:
|
||||
try:
|
||||
from . import nvidia
|
||||
|
||||
count = nvidia.get_physical_gpu_count()
|
||||
if count is not None:
|
||||
_physical_gpu_count = count
|
||||
return _physical_gpu_count
|
||||
except Exception:
|
||||
pass
|
||||
# SMI tool unavailable or failed -- fall back to torch
|
||||
count = _torch_get_physical_gpu_count()
|
||||
_physical_gpu_count = count if count is not None else 1
|
||||
return _physical_gpu_count
|
||||
|
|
@ -1153,8 +1189,8 @@ def get_backend_visible_gpu_info() -> Dict[str, Any]:
|
|||
device = get_device()
|
||||
if device in (DeviceType.CUDA, DeviceType.XPU):
|
||||
parent_visible_ids = get_parent_visible_gpu_ids()
|
||||
# Try nvidia-smi first (NVIDIA only)
|
||||
if device == DeviceType.CUDA:
|
||||
# Try native SMI tool first (nvidia-smi for NVIDIA, skipped for ROCm)
|
||||
if device == DeviceType.CUDA and not IS_ROCM:
|
||||
try:
|
||||
from . import nvidia
|
||||
|
||||
|
|
|
|||
|
|
@ -705,6 +705,19 @@ def install_python_stack() -> int:
|
|||
_progress("ROCm torch check")
|
||||
_ensure_rocm_torch()
|
||||
|
||||
# Windows + AMD GPU: PyTorch does not publish ROCm wheels for Windows.
|
||||
# Detect and warn so users know manual steps are needed for GPU training.
|
||||
if IS_WINDOWS and not NO_TORCH:
|
||||
if shutil.which("hipinfo") or shutil.which("amd-smi"):
|
||||
_safe_print(
|
||||
_dim(" Note:"),
|
||||
"AMD GPU detected on Windows. ROCm-enabled PyTorch must be",
|
||||
)
|
||||
_safe_print(
|
||||
" " * 8,
|
||||
"installed manually. See: https://docs.unsloth.ai/get-started/install-and-update/amd",
|
||||
)
|
||||
|
||||
# 3. Extra dependencies
|
||||
_progress("unsloth extras")
|
||||
pip_install(
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ All tests use mocks -- no AMD hardware required.
|
|||
"""
|
||||
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
|
|
@ -958,5 +959,371 @@ class TestLiveRegression:
|
|||
assert "cu1" in url or "cuda" in url.lower(), f"Expected CUDA URL, got: {url}"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TEST: worker.py -- ROCm Mamba/SSM source build path
|
||||
# =============================================================================
|
||||
|
||||
# Load worker.py module
|
||||
_WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "training" / "worker.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
|
||||
|
||||
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()
|
||||
assert "getattr(torch.version, 'hip', None)" in source
|
||||
|
||||
def test_direct_wheel_url_returns_none_without_cuda_major(self):
|
||||
"""_direct_wheel_url should return None when cuda_major is empty (ROCm)."""
|
||||
# Load module for function access
|
||||
_worker_spec = importlib.util.spec_from_file_location(
|
||||
"test_worker", _WORKER_PATH
|
||||
)
|
||||
assert _worker_spec is not None and _worker_spec.loader is not None
|
||||
worker_mod = importlib.util.module_from_spec(_worker_spec)
|
||||
|
||||
# Mock all the imports worker.py needs
|
||||
sys.modules["structlog"] = MagicMock()
|
||||
sys.modules["loggers"] = MagicMock()
|
||||
sys.modules["loggers"].get_logger = MagicMock(return_value=MagicMock())
|
||||
sys.modules["utils"] = MagicMock()
|
||||
sys.modules["utils.hardware"] = MagicMock()
|
||||
|
||||
try:
|
||||
_worker_spec.loader.exec_module(worker_mod)
|
||||
except Exception:
|
||||
pytest.skip("Could not load worker module in test environment")
|
||||
|
||||
env_rocm = {
|
||||
"python_tag": "cp312",
|
||||
"torch_mm": "2.6",
|
||||
"cuda_major": "",
|
||||
"hip_version": "7.1.12345",
|
||||
"cxx11abi": "TRUE",
|
||||
}
|
||||
result = worker_mod._direct_wheel_url(
|
||||
filename_prefix="causal_conv1d",
|
||||
package_version="1.6.1",
|
||||
release_tag="v1.6.1.post4",
|
||||
release_base_url="https://github.com/Dao-AILab/causal-conv1d/releases/download",
|
||||
env=env_rocm,
|
||||
)
|
||||
assert result is None
|
||||
|
||||
def test_hipcc_check_exists_in_source(self):
|
||||
"""worker.py should check for hipcc before ROCm source builds."""
|
||||
source = _WORKER_PATH.read_text()
|
||||
assert "hipcc" in source
|
||||
|
||||
def test_rocm_source_build_status_message(self):
|
||||
"""worker.py should send a specific status for ROCm source compilation."""
|
||||
source = _WORKER_PATH.read_text()
|
||||
assert "Compiling" in source and "from source for ROCm" in source
|
||||
|
||||
def test_rocm_build_failure_message(self):
|
||||
"""worker.py should send a clear error on ROCm build failure."""
|
||||
source = _WORKER_PATH.read_text()
|
||||
assert "Failed to compile" in source and "for ROCm" in source
|
||||
|
||||
def test_timeout_on_install(self):
|
||||
"""worker.py should have a timeout on pip install subprocess."""
|
||||
source = _WORKER_PATH.read_text()
|
||||
assert "TimeoutExpired" in source
|
||||
assert "timeout" in source
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TEST: amd.py -- AMD GPU monitoring
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestAmdGpuMonitoring:
|
||||
"""Verify amd.py module structure and mock behavior."""
|
||||
|
||||
def test_amd_py_exists(self):
|
||||
"""amd.py should exist in the hardware directory."""
|
||||
amd_path = (
|
||||
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
|
||||
)
|
||||
assert amd_path.exists()
|
||||
|
||||
def test_amd_py_has_required_functions(self):
|
||||
"""amd.py should export the same function signatures as nvidia.py."""
|
||||
amd_path = (
|
||||
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
|
||||
)
|
||||
source = amd_path.read_text()
|
||||
assert "def get_physical_gpu_count" in source
|
||||
assert "def get_primary_gpu_utilization" in source
|
||||
assert "def get_visible_gpu_utilization" in source
|
||||
|
||||
def test_amd_smi_json_parsing(self):
|
||||
"""Verify _extract_gpu_metrics parses amd-smi JSON correctly."""
|
||||
amd_path = (
|
||||
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
|
||||
)
|
||||
_amd_spec = importlib.util.spec_from_file_location("test_amd", amd_path)
|
||||
assert _amd_spec is not None and _amd_spec.loader is not None
|
||||
amd_mod = importlib.util.module_from_spec(_amd_spec)
|
||||
|
||||
sys.modules["loggers"] = MagicMock()
|
||||
sys.modules["loggers"].get_logger = MagicMock(return_value=MagicMock())
|
||||
|
||||
try:
|
||||
_amd_spec.loader.exec_module(amd_mod)
|
||||
except Exception:
|
||||
pytest.skip("Could not load amd module in test environment")
|
||||
|
||||
# Simulate amd-smi metric JSON output
|
||||
gpu_data = {
|
||||
"usage": {"gfx_activity": "85"},
|
||||
"temperature": {"edge": "72"},
|
||||
"power": {
|
||||
"current_socket_power": "200.5",
|
||||
"power_cap": "300",
|
||||
},
|
||||
"vram": {
|
||||
"vram_used": 8192, # MB
|
||||
"vram_total": 16384, # MB
|
||||
},
|
||||
}
|
||||
metrics = amd_mod._extract_gpu_metrics(gpu_data)
|
||||
assert metrics["gpu_utilization_pct"] == 85.0
|
||||
assert metrics["temperature_c"] == 72.0
|
||||
assert metrics["power_draw_w"] == 200.5
|
||||
assert metrics["power_limit_w"] == 300.0
|
||||
assert metrics["vram_used_gb"] == round(8192 / 1024, 2)
|
||||
assert metrics["vram_total_gb"] == round(16384 / 1024, 2)
|
||||
assert metrics["vram_utilization_pct"] is not None
|
||||
assert metrics["power_utilization_pct"] is not None
|
||||
|
||||
def test_amd_primary_gpu_with_mock(self):
|
||||
"""get_primary_gpu_utilization returns correct dict with mocked amd-smi."""
|
||||
amd_path = (
|
||||
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
|
||||
)
|
||||
_amd_spec = importlib.util.spec_from_file_location("test_amd2", amd_path)
|
||||
assert _amd_spec is not None and _amd_spec.loader is not None
|
||||
amd_mod = importlib.util.module_from_spec(_amd_spec)
|
||||
|
||||
sys.modules["loggers"] = MagicMock()
|
||||
sys.modules["loggers"].get_logger = MagicMock(return_value=MagicMock())
|
||||
|
||||
try:
|
||||
_amd_spec.loader.exec_module(amd_mod)
|
||||
except Exception:
|
||||
pytest.skip("Could not load amd module")
|
||||
|
||||
mock_json = json.dumps([{
|
||||
"usage": {"gfx_activity": "50"},
|
||||
"temperature": {"edge": "65"},
|
||||
"power": {"current_socket_power": "150", "power_cap": "250"},
|
||||
"vram": {"vram_used": 4096, "vram_total": 16384},
|
||||
}])
|
||||
mock_result = MagicMock()
|
||||
mock_result.returncode = 0
|
||||
mock_result.stdout = mock_json
|
||||
|
||||
with patch.object(subprocess, "run", return_value=mock_result):
|
||||
result = amd_mod.get_primary_gpu_utilization()
|
||||
assert result["available"] is True
|
||||
assert result["gpu_utilization_pct"] == 50.0
|
||||
assert result["temperature_c"] == 65.0
|
||||
|
||||
def test_amd_smi_not_found_returns_unavailable(self):
|
||||
"""get_primary_gpu_utilization returns available=False when amd-smi is missing."""
|
||||
amd_path = (
|
||||
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
|
||||
)
|
||||
_amd_spec = importlib.util.spec_from_file_location("test_amd3", amd_path)
|
||||
assert _amd_spec is not None and _amd_spec.loader is not None
|
||||
amd_mod = importlib.util.module_from_spec(_amd_spec)
|
||||
|
||||
sys.modules["loggers"] = MagicMock()
|
||||
sys.modules["loggers"].get_logger = MagicMock(return_value=MagicMock())
|
||||
|
||||
try:
|
||||
_amd_spec.loader.exec_module(amd_mod)
|
||||
except Exception:
|
||||
pytest.skip("Could not load amd module")
|
||||
|
||||
with patch.object(
|
||||
subprocess, "run", side_effect=OSError("amd-smi not found")
|
||||
):
|
||||
result = amd_mod.get_primary_gpu_utilization()
|
||||
assert result["available"] is False
|
||||
|
||||
def test_amd_timeout_returns_unavailable(self):
|
||||
"""get_primary_gpu_utilization handles timeout gracefully."""
|
||||
amd_path = (
|
||||
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
|
||||
)
|
||||
_amd_spec = importlib.util.spec_from_file_location("test_amd4", amd_path)
|
||||
assert _amd_spec is not None and _amd_spec.loader is not None
|
||||
amd_mod = importlib.util.module_from_spec(_amd_spec)
|
||||
|
||||
sys.modules["loggers"] = MagicMock()
|
||||
sys.modules["loggers"].get_logger = MagicMock(return_value=MagicMock())
|
||||
|
||||
try:
|
||||
_amd_spec.loader.exec_module(amd_mod)
|
||||
except Exception:
|
||||
pytest.skip("Could not load amd module")
|
||||
|
||||
with patch.object(
|
||||
subprocess,
|
||||
"run",
|
||||
side_effect=subprocess.TimeoutExpired("amd-smi", 5),
|
||||
):
|
||||
result = amd_mod.get_primary_gpu_utilization()
|
||||
assert result["available"] is False
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TEST: hardware.py -- IS_ROCM branching to amd.py
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestHardwareAmdBranching:
|
||||
"""Verify hardware.py branches to amd.py when IS_ROCM is True."""
|
||||
|
||||
def test_hardware_imports_amd_module(self):
|
||||
"""hardware.py should import from amd module when IS_ROCM."""
|
||||
hw_path = (
|
||||
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
|
||||
)
|
||||
source = hw_path.read_text()
|
||||
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."""
|
||||
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
|
||||
|
||||
def test_hardware_branches_on_is_rocm_for_visible(self):
|
||||
"""get_visible_gpu_utilization should check 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
|
||||
|
||||
def test_hardware_branches_on_is_rocm_for_physical_count(self):
|
||||
"""get_physical_gpu_count should try amd.py when IS_ROCM."""
|
||||
hw_path = (
|
||||
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
|
||||
)
|
||||
source = hw_path.read_text()
|
||||
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
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TEST: install_python_stack.py -- Windows AMD warning
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestWindowsRocmWarning:
|
||||
"""Verify Windows AMD GPU detection and warning message."""
|
||||
|
||||
def test_windows_amd_warning_in_source(self):
|
||||
"""install_python_stack.py should warn Windows AMD users."""
|
||||
source = _STACK_PATH.read_text()
|
||||
assert "AMD GPU detected on Windows" in source
|
||||
|
||||
def test_windows_amd_warning_checks_hipinfo_or_amdsmi(self):
|
||||
"""Warning should check for hipinfo or amd-smi."""
|
||||
source = _STACK_PATH.read_text()
|
||||
assert "hipinfo" in source
|
||||
assert "amd-smi" in source
|
||||
|
||||
def test_windows_amd_warning_has_docs_link(self):
|
||||
"""Warning should include AMD docs link."""
|
||||
source = _STACK_PATH.read_text()
|
||||
assert "docs.unsloth.ai/get-started/install-and-update/amd" in source
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TEST: unsloth/kernels/utils.py -- is_rdna() expansion
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestIsRdnaExpansion:
|
||||
"""Verify is_rdna() covers RDNA2, RDNA3, RDNA3.5, RDNA4 architectures."""
|
||||
|
||||
def test_is_rdna_source_has_rdna2(self):
|
||||
"""is_rdna() should include RDNA2 architectures."""
|
||||
utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py"
|
||||
source = utils_path.read_text()
|
||||
func_start = source.find("def is_rdna()")
|
||||
func_body = source[func_start:source.find("\ndef ", func_start + 1)]
|
||||
assert "gfx1030" in func_body
|
||||
assert "gfx1031" in func_body
|
||||
assert "gfx1032" in func_body
|
||||
|
||||
def test_is_rdna_source_has_rdna3(self):
|
||||
"""is_rdna() should include RDNA3 architectures."""
|
||||
utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py"
|
||||
source = utils_path.read_text()
|
||||
func_start = source.find("def is_rdna()")
|
||||
func_body = source[func_start:source.find("\ndef ", func_start + 1)]
|
||||
assert "gfx1100" in func_body
|
||||
assert "gfx1101" in func_body
|
||||
assert "gfx1102" in func_body
|
||||
assert "gfx1103" in func_body
|
||||
|
||||
def test_is_rdna_source_has_rdna35(self):
|
||||
"""is_rdna() should include RDNA3.5 architectures."""
|
||||
utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py"
|
||||
source = utils_path.read_text()
|
||||
func_start = source.find("def is_rdna()")
|
||||
func_body = source[func_start:source.find("\ndef ", func_start + 1)]
|
||||
assert "gfx1150" in func_body
|
||||
assert "gfx1151" in func_body
|
||||
assert "gfx1152" in func_body
|
||||
|
||||
def test_is_rdna_source_has_rdna4(self):
|
||||
"""is_rdna() should include RDNA4 architectures."""
|
||||
utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py"
|
||||
source = utils_path.read_text()
|
||||
func_start = source.find("def is_rdna()")
|
||||
func_body = source[func_start:source.find("\ndef ", func_start + 1)]
|
||||
assert "gfx1200" in func_body
|
||||
assert "gfx1201" in func_body
|
||||
|
||||
def test_is_cdna_not_changed(self):
|
||||
"""is_cdna() should remain unchanged (no RDNA architectures added)."""
|
||||
utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py"
|
||||
source = utils_path.read_text()
|
||||
func_start = source.find("def is_cdna()")
|
||||
func_body = source[func_start:source.find("\ndef ", func_start + 1)]
|
||||
assert "gfx940" in func_body
|
||||
assert "gfx941" in func_body
|
||||
assert "gfx942" in func_body
|
||||
assert "gfx950" in func_body
|
||||
# RDNA architectures should NOT be in is_cdna
|
||||
assert "gfx1030" not in func_body
|
||||
assert "gfx1100" not in func_body
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
|
|
|
|||
|
|
@ -88,10 +88,22 @@ def is_cdna():
|
|||
|
||||
@functools.lru_cache(1)
|
||||
def is_rdna():
|
||||
"""Detect ROCm-supported RDNA consumer/workstation GPUs (RDNA3, RDNA4)."""
|
||||
"""Detect ROCm-supported RDNA consumer/workstation GPUs (RDNA2, RDNA3, RDNA3.5, RDNA4)."""
|
||||
return is_hip() and triton.runtime.driver.active.get_current_target().arch in (
|
||||
# RDNA2 (Navi 21-24)
|
||||
"gfx1030",
|
||||
"gfx1031",
|
||||
"gfx1032",
|
||||
# RDNA3 (Navi 31-33)
|
||||
"gfx1100",
|
||||
"gfx1101",
|
||||
"gfx1102",
|
||||
"gfx1103",
|
||||
# RDNA3.5 (Strix Point / Strix Halo)
|
||||
"gfx1150",
|
||||
"gfx1151",
|
||||
"gfx1152",
|
||||
# RDNA4 (Navi 48-44)
|
||||
"gfx1200",
|
||||
"gfx1201",
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue