From 06f28e4d03468c58e3d46f6b02cdead90c538932 Mon Sep 17 00:00:00 2001 From: LeoBorcherding Date: Tue, 19 May 2026 14:35:40 -0500 Subject: [PATCH] fix(studio/rocm): code review hardening pass - main.py: numeric DLL sort (string sort picked rocm72 over rocm713); add basename() to regex; log warning on detection failure; log info when BNB_ROCM_VERSION is set (mirrors worker.py) - worker.py: explicit len-guard in _hip_ver_at_least() with warning logs instead of silent IndexError/ValueError swallow - hardware.py: isinstance(result, dict) guard before result.get() in _smi_query() to prevent AttributeError on non-dict backend returns - amd.py: round() before int() on parsed GPU IDs; log warning when truncation occurs (defensive against malformed amd-smi output) - setup.sh: quote --gcc-install-dir value in CMAKE_HIP_FLAGS so paths with spaces do not break the CMake argument - install.ps1, setup.ps1: apply colon-split + ToLower() to hipinfo gcnArchName match (consistent with each other and with setup.sh) - install.sh: tighten ROCm tag case patterns to explicit rocmX.Y|rocmX.Y.* to avoid unintended prefix matches --- install.ps1 | 2 +- install.sh | 16 +++++++------- studio/backend/core/training/worker.py | 18 +++++++++++++++- studio/backend/main.py | 26 ++++++++++++++++++----- studio/backend/utils/hardware/amd.py | 12 +++++++++-- studio/backend/utils/hardware/hardware.py | 2 +- studio/setup.ps1 | 2 +- studio/setup.sh | 7 +++--- 8 files changed, 63 insertions(+), 22 deletions(-) diff --git a/install.ps1 b/install.ps1 index a3309485f3..a0922a14c9 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1241,7 +1241,7 @@ shell.Run cmd, 0, False if ($LASTEXITCODE -eq 0 -and $hipOut -match "(?i)gcnArchName") { $HasROCm = $true if ($hipOut -match "(?im)^\s*gcnArchName\s*:\s*(\S+)") { - $ROCmGfxArch = $Matches[1].Trim() + $ROCmGfxArch = ($Matches[1] -split ':')[0].Trim().ToLower() $ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)" } else { $ROCmGpuLabel = "AMD ROCm" diff --git a/install.sh b/install.sh index c52f2669a9..3ea19df56c 100755 --- a/install.sh +++ b/install.sh @@ -1647,14 +1647,14 @@ get_torch_index_url() { # PyTorch publishes major.minor URLs only (no patch level), so # rocm7.2.1 / rocm6.0.2 / etc. must normalise to rocm7.2 / rocm6.0. case "$_rocm_tag" in - rocm6.0*) echo "$_base/rocm6.0" ;; - rocm6.1*) echo "$_base/rocm6.1" ;; - rocm6.2*) echo "$_base/rocm6.2" ;; - rocm6.3*) echo "$_base/rocm6.3" ;; - rocm6.4*) echo "$_base/rocm6.4" ;; - rocm7.0*) echo "$_base/rocm7.0" ;; - rocm7.1*) echo "$_base/rocm7.1" ;; - rocm7.2*) echo "$_base/rocm7.2" ;; + rocm6.0|rocm6.0.*) echo "$_base/rocm6.0" ;; + rocm6.1|rocm6.1.*) echo "$_base/rocm6.1" ;; + rocm6.2|rocm6.2.*) echo "$_base/rocm6.2" ;; + rocm6.3|rocm6.3.*) echo "$_base/rocm6.3" ;; + rocm6.4|rocm6.4.*) echo "$_base/rocm6.4" ;; + rocm7.0|rocm7.0.*) echo "$_base/rocm7.0" ;; + rocm7.1|rocm7.1.*) echo "$_base/rocm7.1" ;; + rocm7.2|rocm7.2.*) echo "$_base/rocm7.2" ;; rocm6.*) # ROCm 6.5+ (no published PyTorch wheels): clip down # to the last supported 6.x wheel set. diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index ac811d32e9..ddd1fff990 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -2149,8 +2149,24 @@ def run_training_process( return False try: _parts = [int(x) for x in str(_hip_str).split(".")[:2]] + if len(_parts) < 2: + logger.warning( + "Windows ROCm: torch.version.hip %r has fewer than " + "two components; cannot compare against %d.%d", + _hip_str, + major, + minor, + ) + return False return (_parts[0], _parts[1]) >= (major, minor) - except (ValueError, IndexError): + except ValueError: + logger.warning( + "Windows ROCm: could not parse torch.version.hip %r as " + "a version number; assuming HIP < %d.%d", + _hip_str, + major, + minor, + ) return False # _grouped_mm HIP kernel was null on gfx1200 in ROCm ≤ 7.12, diff --git a/studio/backend/main.py b/studio/backend/main.py index fac528f1e1..482b31d689 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -88,6 +88,7 @@ if sys.platform == "win32": pass if _is_rocm_host and "BNB_ROCM_VERSION" not in os.environ: import glob as _glob + import logging as _logging _bnb_rocm_ver = None try: @@ -99,14 +100,29 @@ if sys.platform == "win32": _dlls = _glob.glob(os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll")) import re as _re_bnb - for _dll in sorted(_dlls): - _m = _re_bnb.search(r"libbitsandbytes_rocm(\d+)\.dll", _dll) + def _bnb_ver_key(p: str) -> int: + _km = _re_bnb.search(r"rocm(\d+)", os.path.basename(p)) + return int(_km.group(1)) if _km else -1 + + for _dll in sorted(_dlls, key=_bnb_ver_key, reverse=True): + _m = _re_bnb.search( + r"libbitsandbytes_rocm(\d+)\.dll", os.path.basename(_dll) + ) if _m: _bnb_rocm_ver = _m.group(1) break - except Exception: - pass - os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver or "72" + except Exception as _e: + _logging.getLogger(__name__).warning( + "Windows ROCm: BNB DLL detection failed (%s); falling back to version '72'", + _e, + ) + _bnb_rocm_ver_final = _bnb_rocm_ver or "72" + os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver_final + _logging.getLogger(__name__).info( + "Windows ROCm: set BNB_ROCM_VERSION=%s " + "(detected from installed BNB wheel; overrides torch.version.hip auto-detection)", + _bnb_rocm_ver_final, + ) # Ensure backend dir is on sys.path so _platform_compat is importable when # main.py is launched directly (e.g. `uvicorn main:app`). diff --git a/studio/backend/utils/hardware/amd.py b/studio/backend/utils/hardware/amd.py index bcacdb57ea..7804836fd2 100644 --- a/studio/backend/utils/hardware/amd.py +++ b/studio/backend/utils/hardware/amd.py @@ -383,7 +383,7 @@ def get_visible_gpu_utilization( ) parsed_id = _parse_numeric(raw_id) if parsed_id is None: - logger.debug( + logger.warning( "amd-smi GPU id %r could not be parsed; falling back to " "enumeration index %d", raw_id, @@ -391,7 +391,15 @@ def get_visible_gpu_utilization( ) idx = fallback_idx else: - idx = int(parsed_id) + rounded = round(parsed_id) + if rounded != parsed_id: + logger.warning( + "amd-smi GPU id %r parsed as non-integer %r; truncating to %d", + raw_id, + parsed_id, + rounded, + ) + idx = int(rounded) if idx not in visible_set: continue metrics = _extract_gpu_metrics(gpu_data) diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 09c097a687..a81159324c 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -467,7 +467,7 @@ def _smi_query(func_name: str, *args, **kwargs) -> Optional[Dict[str, Any]]: try: func = getattr(_backend, func_name) result = func(*args, **kwargs) - if result.get("available"): + if isinstance(result, dict) and result.get("available"): return result except Exception as e: logger.warning("%s %s query failed: %s", backend_name, func_name, e) diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 6ee20143ec..cddc252b86 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -715,7 +715,7 @@ if (-not $HasNvidiaSmi) { if ($LASTEXITCODE -eq 0 -and $hipOut -match "(?i)gcnArchName") { $HasROCm = $true if ($hipOut -match "(?im)^\s*gcnArchName\s*:\s*(\S+)") { - $script:ROCmGfxArch = $Matches[1].Trim() + $script:ROCmGfxArch = ($Matches[1] -split ':')[0].Trim().ToLower() $ROCmGpuLabel = "AMD ROCm ($script:ROCmGfxArch)" } else { $ROCmGpuLabel = "AMD ROCm" diff --git a/studio/setup.sh b/studio/setup.sh index 0227bfcaf7..60af1e26b4 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -1043,15 +1043,16 @@ else # runtime dir AND /usr/include/c++/ headers, then pass it # to clang via --gcc-install-dir so HIP builds succeed. _GCC_INSTALL_DIR="" + _GCC_MULTIARCH="$(gcc -print-multiarch 2>/dev/null || uname -m)-linux-gnu" for _gcc_ver in 14 13 12 11; do - if [ -d "/usr/lib/gcc/x86_64-linux-gnu/$_gcc_ver/include" ] && \ + if [ -d "/usr/lib/gcc/$_GCC_MULTIARCH/$_gcc_ver/include" ] && \ [ -d "/usr/include/c++/$_gcc_ver" ]; then - _GCC_INSTALL_DIR="/usr/lib/gcc/x86_64-linux-gnu/$_gcc_ver" + _GCC_INSTALL_DIR="/usr/lib/gcc/$_GCC_MULTIARCH/$_gcc_ver" break fi done if [ -n "$_GCC_INSTALL_DIR" ]; then - CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_HIP_FLAGS=--gcc-install-dir=$_GCC_INSTALL_DIR" + CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_HIP_FLAGS=--gcc-install-dir=\"$_GCC_INSTALL_DIR\"" substep "ROCm HIP gcc install dir: $_GCC_INSTALL_DIR" fi