From 10ec0cdefabc484dbeed2178d63ed3f233397d4a Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 31 Mar 2026 09:51:58 +0000
Subject: [PATCH] Strengthen AMD GPU detection and add NVIDIA precedence guard

- Change amd-smi list detection from any-non-empty-output to requiring
  "gpu" marker in output, matching the shell-side NR>1 check. Prevents
  false positives from header-only amd-smi list output.
- Add nvidia-smi check at the top of _ensure_rocm_torch() so mixed
  AMD+NVIDIA hosts preserve NVIDIA precedence (matching install.sh and
  install_llama_prebuilt.py behavior).
- Apply the same amd-smi marker fix to install_llama_prebuilt.py
  detect_host() for consistency.
---
 studio/install_llama_prebuilt.py |  4 ++--
 studio/install_python_stack.py   | 10 +++++++---
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py
index 3b46373eb0..fc3c65bc4b 100755
--- a/studio/install_llama_prebuilt.py
+++ b/studio/install_llama_prebuilt.py
@@ -1436,7 +1436,7 @@ def detect_host() -> HostInfo:
     if not is_macos:
         for _cmd, _marker in (
             (["rocminfo"], "gfx"),
-            (["amd-smi", "list"], None),
+            (["amd-smi", "list"], "gpu"),
         ):
             _exe = shutil.which(_cmd[0])
             if not _exe:
@@ -1446,7 +1446,7 @@ def detect_host() -> HostInfo:
             except Exception:
                 continue
             if _result.returncode == 0 and _result.stdout.strip():
-                if _marker is None or _marker in _result.stdout.lower():
+                if _marker in _result.stdout.lower():
                     has_rocm = True
                     break
 
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index 8196b03ddb..461f83d9b9 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -88,7 +88,7 @@ def _has_rocm_gpu() -> bool:
     """Return True only if an actual AMD GPU is visible (not just ROCm tools installed)."""
     for cmd, marker in (
         (["rocminfo"], "gfx"),
-        (["amd-smi", "list"], None),
+        (["amd-smi", "list"], "gpu"),
     ):
         exe = shutil.which(cmd[0])
         if not exe:
@@ -104,7 +104,7 @@ def _has_rocm_gpu() -> bool:
         except Exception:
             continue
         if result.returncode == 0 and result.stdout.strip():
-            if marker is None or marker in result.stdout.lower():
+            if marker in result.stdout.lower():
                 return True
     return False
 
@@ -114,9 +114,13 @@ def _ensure_rocm_torch() -> None:
 
     Runs only on Linux hosts where ROCm is installed and an AMD GPU is
     present.  No-op when torch already links against HIP (ROCm) or CUDA
-    (NVIDIA).  Skips on Windows/macOS.
+    (NVIDIA).  Skips on Windows/macOS and on mixed AMD+NVIDIA hosts
+    (NVIDIA takes precedence).
     Uses pip_install() to respect uv, constraints, and --python targeting.
     """
+    # NVIDIA takes precedence on mixed hosts
+    if shutil.which("nvidia-smi"):
+        return
     rocm_root = os.environ.get("ROCM_PATH") or "/opt/rocm"
     if not os.path.isdir(rocm_root) and not shutil.which("hipcc"):
         return  # no ROCm toolchain