fix(studio/rocm): respect HIP_VISIBLE_DEVICES when picking lemonade gfx target

The rocminfo / hipinfo regex took the first gfx match in the agent listing.
On mixed APU + dGPU hosts (e.g. Strix Halo gfx1151 + discrete RX 7900 gfx1100)
this picked whichever GPU appeared first in the tool's stdout, not the one
HIP actually runs on. The downloaded lemonade asset could then be a binary
for a different arch than the active device.

Extracted a module-level _pick_rocm_gfx_target() helper that:
- collects every gfx token in order via re.findall (skips gfx000 / generic ISAs)
- if HIP_VISIBLE_DEVICES or ROCR_VISIBLE_DEVICES is set, parses the first
  comma-separated entry as an integer index into that list
- falls back to the first GPU for non-integer (UUID-style) or out-of-range
  values, matching the previous default behaviour

Both Linux (rocminfo) and Windows (hipinfo) branches use the helper.
Existing 18 lemonade tests pass; no behavioural change for single-GPU hosts.
This commit is contained in:
Daniel Han 2026-05-19 08:02:23 +00:00
commit 8793aef026

View file

@ -2626,6 +2626,29 @@ def run_capture(
return result
def _pick_rocm_gfx_target(out: str) -> str | None:
"""Choose the gfx target rocminfo / hipinfo report for the active GPU.
A bare first-match picked the wrong device on mixed APU + dGPU hosts
(e.g. Strix Halo gfx1151 + discrete RX 7900 gfx1100). Respect
HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES so the asset matches what HIP
actually runs on. Falls back to the first GPU when no env var is set.
"""
_tokens = re.findall(r"gfx[1-9][0-9a-z]{2,3}", out.lower())
if not _tokens:
return None
_vis = os.environ.get("HIP_VISIBLE_DEVICES") or os.environ.get("ROCR_VISIBLE_DEVICES") or ""
if _vis:
_first = _vis.split(",")[0].strip()
try:
_idx = int(_first)
if 0 <= _idx < len(_tokens):
return _tokens[_idx]
except ValueError:
pass
return _tokens[0]
def detect_host() -> HostInfo:
system = platform.system()
machine = platform.machine().lower()
@ -2746,9 +2769,7 @@ def detect_host() -> HostInfo:
if _result.returncode == 0 and _result.stdout.strip():
if _check(_result.stdout):
has_rocm = True
_gfx_m = re.search(r"gfx[1-9][0-9a-z]{2,3}", _result.stdout.lower())
if _gfx_m:
rocm_gfx_target = _gfx_m.group(0)
rocm_gfx_target = _pick_rocm_gfx_target(_result.stdout)
break
elif is_windows:
# Windows: prefer active probes that validate GPU presence
@ -2767,9 +2788,7 @@ def detect_host() -> HostInfo:
if _check(_result.stdout):
has_rocm = True
# hipinfo reports "gcnArchName: gfx1100" -- extract if present
_gfx_m = re.search(r"gfx[1-9][0-9a-z]{2,3}", _result.stdout.lower())
if _gfx_m:
rocm_gfx_target = _gfx_m.group(0)
rocm_gfx_target = _pick_rocm_gfx_target(_result.stdout)
break
# Note: amdhip64.dll presence alone is NOT treated as GPU evidence
# since the HIP SDK can be installed without an AMD GPU.