Refactor: deduplicate AMD detection, consolidate bitsandbytes, clean up imports
- Extract _has_amd_rocm_gpu() shell function to avoid duplicating the rocminfo/amd-smi GPU detection logic in get_torch_index_url and the Radeon auto-detect block - Consolidate bitsandbytes install into a single case block after torch install (was duplicated 4 times across Radeon success/fallback paths) - Move math and re imports to top of amd.py (were inline in functions) - Add _smi_query() helper in hardware.py to centralize IS_ROCM backend selection for get_gpu_utilization and get_visible_gpu_utilization Addresses Gemini code review suggestions.
This commit is contained in:
Daniel Han
parent
263470edca
commit
1b98f6d705
3 changed files with 78 additions and 103 deletions
57
install.sh
57
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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue