diff --git a/install.sh b/install.sh index bd3029b2e6..7f393165e8 100755 --- a/install.sh +++ b/install.sh @@ -978,6 +978,21 @@ _find_no_torch_runtime() { fi } +# ── AMD ROCm GPU detection helper ── +# Returns 0 (true) if an actual AMD GPU is present, 1 (false) otherwise. +# Checks rocminfo for gfx[1-9]* (excludes gfx000 CPU agent) and +# amd-smi list for GPU data rows (excludes header-only output). +_has_amd_rocm_gpu() { + if command -v rocminfo >/dev/null 2>&1 && \ + rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9]/{found=1} END{exit !found}'; then + return 0 + elif command -v amd-smi >/dev/null 2>&1 && \ + amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[]/{ found=1 } END{ exit !found }'; then + return 0 + fi + return 1 +} + # ── Detect GPU and choose PyTorch index URL ── # Mirrors Get-TorchIndexUrl in install.ps1. # On CPU-only machines this returns the cpu index, avoiding the solver @@ -995,16 +1010,7 @@ get_torch_index_url() { fi if [ -z "$_smi" ]; then # No NVIDIA GPU -- check for AMD ROCm GPU - # First confirm an actual AMD GPU is present (not just ROCm tools installed) - _has_rocm_gpu=false - if command -v rocminfo >/dev/null 2>&1 && \ - rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9]/{found=1} END{exit !found}'; then - _has_rocm_gpu=true - elif command -v amd-smi >/dev/null 2>&1 && \ - amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[]/{ found=1 } END{ exit !found }'; then - _has_rocm_gpu=true - fi - if [ "$_has_rocm_gpu" != true ]; then + if ! _has_amd_rocm_gpu; then echo "$_base/cpu"; return fi # AMD GPU confirmed -- detect ROCm version @@ -1155,16 +1161,8 @@ TORCH_INDEX_URL=$(get_torch_index_url) # (gfx in rocminfo or amd-smi list). Then require rocminfo "Marketing Name:.*Radeon". case "$TORCH_INDEX_URL" in */rocm*) - _amd_gpu_here=false _amd_gpu_radeon=false - if command -v rocminfo >/dev/null 2>&1 && \ - rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9]/{found=1} END{exit !found}'; then - _amd_gpu_here=true - elif command -v amd-smi >/dev/null 2>&1 && \ - amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[]/{ found=1 } END{ exit !found }'; then - _amd_gpu_here=true - fi - if [ "$_amd_gpu_here" = true ] && command -v rocminfo >/dev/null 2>&1 && \ + if _has_amd_rocm_gpu && command -v rocminfo >/dev/null 2>&1 && \ rocminfo 2>/dev/null | grep -q 'Marketing Name:.*Radeon'; then _amd_gpu_radeon=true fi @@ -1255,37 +1253,30 @@ elif [ -n "$TORCH_INDEX_URL" ]; then run_install_cmd "install triton + PyTorch" uv pip install --python "$_VENV_PY" \ --find-links "$_RADEON_BASE_URL" \ $_radeon_pkgs - substep "installing bitsandbytes for AMD Radeon..." - run_install_cmd "install bitsandbytes (AMD)" uv pip install --python "$_VENV_PY" \ - "bitsandbytes>=0.49.1" else substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" - substep "installing bitsandbytes for AMD ROCm..." - run_install_cmd "install bitsandbytes (AMD)" uv pip install --python "$_VENV_PY" "bitsandbytes>=0.49.1" fi else substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN" run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" - substep "installing bitsandbytes for AMD ROCm..." - run_install_cmd "install bitsandbytes (AMD)" uv pip install --python "$_VENV_PY" "bitsandbytes>=0.49.1" fi else substep "installing PyTorch ($TORCH_INDEX_URL)..." run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" - # AMD ROCm: install bitsandbytes with AMD support - case "$TORCH_INDEX_URL" in - */rocm*) - substep "installing bitsandbytes for AMD ROCm..." - run_install_cmd "install bitsandbytes (AMD)" uv pip install --python "$_VENV_PY" "bitsandbytes>=0.49.1" - ;; - esac fi + # AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths) + case "$TORCH_INDEX_URL" in + */rocm*) + substep "installing bitsandbytes for AMD ROCm..." + run_install_cmd "install bitsandbytes (AMD)" uv pip install --python "$_VENV_PY" "bitsandbytes>=0.49.1" + ;; + esac # Fresh: Step 2 - install unsloth, preserving pre-installed torch substep "installing unsloth (this may take a few minutes)..." if [ "$SKIP_TORCH" = true ]; then diff --git a/studio/backend/utils/hardware/amd.py b/studio/backend/utils/hardware/amd.py index c3a0536a56..5cfd6bde97 100644 --- a/studio/backend/utils/hardware/amd.py +++ b/studio/backend/utils/hardware/amd.py @@ -9,6 +9,8 @@ nvidia.py counterparts. """ import json +import math +import re import subprocess from typing import Any, Optional @@ -47,14 +49,10 @@ def _parse_numeric(value: Any) -> Optional[float]: if isinstance(value, dict): return _parse_numeric(value.get("value")) if isinstance(value, (int, float)): - import math - f = float(value) return f if math.isfinite(f) else None if isinstance(value, str): # Strip units like "W", "C", "%", "MB", "MiB", "GB", "GiB" etc. - import re - cleaned = re.sub(r"\s*[A-Za-z/%]+$", "", value.strip()) if not cleaned or cleaned.lower() in ("n/a", "none", "unknown"): return None diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index edeada0190..6712513aa8 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -402,31 +402,44 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any] # ========== Live GPU Utilization ========== +def _smi_query(func_name: str, *args, **kwargs) -> Optional[Dict[str, Any]]: + """Run a query against the appropriate SMI backend (amd-smi or nvidia-smi). + + Returns the result dict if available, or None on failure/unavailability. + """ + if IS_ROCM: + backend_name = "amd-smi" + try: + from . import amd as _backend + except Exception as e: + logger.warning("%s import failed: %s", backend_name, e) + return None + else: + backend_name = "nvidia-smi" + try: + from . import nvidia as _backend + except Exception as e: + logger.warning("%s import failed: %s", backend_name, e) + return None + try: + func = getattr(_backend, func_name) + result = func(*args, **kwargs) + if result.get("available"): + return result + except Exception as e: + logger.warning("%s %s query failed: %s", backend_name, func_name, e) + return None + + def get_gpu_utilization() -> Dict[str, Any]: """Return a live snapshot of device utilization information.""" device = get_device() if device == DeviceType.CUDA: - if IS_ROCM: - try: - from . import amd - - 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) + result = _smi_query("get_primary_gpu_utilization") + if result is not None: + result["backend"] = device.value + return result mem = get_gpu_memory_info() if device != DeviceType.CPU and mem.get("available"): @@ -451,32 +464,14 @@ def get_visible_gpu_utilization() -> Dict[str, Any]: if device == DeviceType.CUDA: parent_visible_spec = _get_parent_visible_gpu_spec() - if IS_ROCM: - try: - from . import amd - - 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) + result = _smi_query( + "get_visible_gpu_utilization", + parent_visible_spec["numeric_ids"], + parent_cuda_visible_devices = parent_visible_spec["raw"], + ) + if result is not None: + result["backend"] = device.value + return result # Torch-based fallback for CUDA (nvidia-smi unavailable, AMD ROCm) and XPU (Intel) if device in (DeviceType.CUDA, DeviceType.XPU): @@ -1162,26 +1157,17 @@ def get_physical_gpu_count() -> int: device = get_device() if device == DeviceType.CUDA: - if IS_ROCM: - try: - from . import amd - - 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 + try: + if IS_ROCM: + from . import amd as _smi_mod + else: + from . import nvidia as _smi_mod + count = _smi_mod.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