diff --git a/iris/drivers/fabric/amd.py b/iris/drivers/fabric/amd.py index 489adf924..07c7865f6 100644 --- a/iris/drivers/fabric/amd.py +++ b/iris/drivers/fabric/amd.py @@ -26,6 +26,7 @@ PeerMapping, ) from iris.host.distributed.topology import InterconnectLevel +from iris.host.platform.libs import load_vendor_library logger = logging.getLogger("iris.drivers.fabric.amd") @@ -38,29 +39,21 @@ ] -def _load_cdll(*names: Optional[str]) -> Any: - for name in names: - if not name: - continue - try: - return ctypes.CDLL(name) - except OSError: - continue - return None - - -_hip = _load_cdll( - # Let LD_LIBRARY_PATH select the active ROCm SDK before consulting ldconfig. - # find_library can return a stale system SONAME (e.g. HIP 5 beside HIP 7), - # and loading both runtimes can abort inside hipSetDevice. +_hip = load_vendor_library( + "libamdhip64", + # When nothing is loaded yet, let LD_LIBRARY_PATH select the active ROCm SDK + # before consulting ldconfig. find_library can return a stale system SONAME + # (e.g. HIP 5 beside HIP 7), and loading both runtimes can abort inside hipSetDevice. "libamdhip64.so", ctypes.util.find_library("amdhip64"), "/opt/rocm/lib/libamdhip64.so", ) -_amdsmi = _load_cdll( +_amdsmi = load_vendor_library( + "libamd_smi", "libamd_smi.so", ctypes.util.find_library("amd_smi"), "/opt/rocm/lib/libamd_smi.so", + beside="libamdhip64", ) HIP_SUCCESS = 0 diff --git a/iris/drivers/local/amd.py b/iris/drivers/local/amd.py index ff346f942..718ea1ca5 100644 --- a/iris/drivers/local/amd.py +++ b/iris/drivers/local/amd.py @@ -22,6 +22,7 @@ PeerMapping, ) from iris.host.distributed.topology import InterconnectLevel +from iris.host.platform.libs import load_vendor_library logger = logging.getLogger("iris.drivers.local.amd") @@ -31,11 +32,7 @@ "LocalHipDriver", ] -_hip = None -try: - _hip = ctypes.cdll.LoadLibrary("libamdhip64.so") -except OSError: - pass +_hip = load_vendor_library("libamdhip64", "libamdhip64.so") HIP_SUCCESS = 0 HIP_ERROR_NOT_SUPPORTED = 801 diff --git a/iris/host/distributed/fd_passing.py b/iris/host/distributed/fd_passing.py index fe5ad3b77..f408fce61 100644 --- a/iris/host/distributed/fd_passing.py +++ b/iris/host/distributed/fd_passing.py @@ -18,6 +18,17 @@ from typing import Dict, Tuple from contextlib import contextmanager +# Higher ranks retry connect() for 10 s after the path exchange, so a peer that +# has not connected well past that will not. +_ACCEPT_TIMEOUT_S = 60.0 +# The connecting side sends its rank immediately after connect(). +_HANDSHAKE_TIMEOUT_S = 10.0 + + +def _close_all(socks) -> None: + for sock in socks: + sock.close() + @contextmanager def managed_fd(fd: int): @@ -132,19 +143,38 @@ def setup_fd_mesh(rank: int, world_size: int, all_paths: Dict[int, str]) -> Dict conns[peer] = s # Accept connections from higher ranks - for _ in range(rank + 1, world_size): - client, _ = listener.accept() - peer_rank_bytes = recv_exact(client, 4) - peer_rank = int.from_bytes(peer_rank_bytes, "little", signed=False) - conns[peer_rank] = client - - # Close listener and clean up socket path - listener.close() + listener.settimeout(_ACCEPT_TIMEOUT_S) try: - os.unlink(path) - except OSError: - # Best effort cleanup - pass + for _ in range(rank + 1, world_size): + try: + client, _ = listener.accept() + except socket.timeout as e: + missing = sorted(set(range(rank + 1, world_size)) - set(conns)) + _close_all(conns.values()) + raise TimeoutError( + f"Timed out after {_ACCEPT_TIMEOUT_S:.0f}s waiting for ranks {missing} to connect to rank " + f"{rank} at {path}. FD passing only works between ranks on the same host." + ) from e + client.settimeout(_HANDSHAKE_TIMEOUT_S) + try: + peer_rank_bytes = recv_exact(client, 4) + except socket.timeout as e: + _close_all([client, *conns.values()]) + raise TimeoutError( + f"A peer connected to rank {rank} at {path} but did not send its rank within " + f"{_HANDSHAKE_TIMEOUT_S:.0f}s." + ) from e + client.settimeout(None) + peer_rank = int.from_bytes(peer_rank_bytes, "little", signed=False) + conns[peer_rank] = client + finally: + # Close listener and clean up socket path + listener.close() + try: + os.unlink(path) + except OSError: + # Best effort cleanup + pass return conns diff --git a/iris/host/distributed/topology.py b/iris/host/distributed/topology.py index af1406d5a..625f465d9 100644 --- a/iris/host/distributed/topology.py +++ b/iris/host/distributed/topology.py @@ -17,6 +17,8 @@ import torch import torch.distributed as dist +from iris.host.platform.libs import load_vendor_library + logger = logging.getLogger("iris.topology") @@ -374,21 +376,15 @@ def _amd_get_gpu_fabric_info(gpu_id: int, pci_bus_id: str = "") -> FabricInfo: return FabricInfo() try: - amdsmi = None - for amdsmi_path in ( - ctypes.util.find_library("amd_smi"), + amdsmi = load_vendor_library( + "libamd_smi", "libamd_smi.so", + ctypes.util.find_library("amd_smi"), "/opt/rocm/lib/libamd_smi.so", - ): - if not amdsmi_path: - continue - try: - amdsmi = ctypes.CDLL(amdsmi_path) - break - except OSError: - continue + beside="libamdhip64", + ) if amdsmi is None: - logger.debug("libamd_smi.so not available, skipping AMD fabric info") + _fabric_unavailable("libamd_smi.so could not be loaded") return FabricInfo() amdsmi.amdsmi_init.argtypes = [ctypes.c_uint64] @@ -406,7 +402,7 @@ def _amd_get_gpu_fabric_info(gpu_id: int, pci_bus_id: str = "") -> FabricInfo: ret = int(amdsmi.amdsmi_init(_AMDSMI_INIT_AMD_GPUS)) if ret != _AMDSMI_STATUS_SUCCESS: - logger.debug("amdsmi_init() failed for AMD fabric info query with status %d", ret) + _fabric_unavailable("amdsmi_init() failed with status %d", ret) return FabricInfo() processor = ctypes.c_void_p() @@ -417,24 +413,24 @@ def _amd_get_gpu_fabric_info(gpu_id: int, pci_bus_id: str = "") -> FabricInfo: ) ) if ret != _AMDSMI_STATUS_SUCCESS or processor.value is None: - logger.debug( - "amdsmi_get_processor_handle_from_bdf(%s) failed for GPU %d with status %d", - pci_bus_id, - gpu_id, - ret, + _fabric_unavailable( + "amdsmi_get_processor_handle_from_bdf(%s) failed for GPU %d with status %d", pci_bus_id, gpu_id, ret ) return FabricInfo() raw = _AmdSmiFabricInfo() ret = int(amdsmi.amdsmi_get_gpu_fabric_info(processor, ctypes.byref(raw))) if ret != _AMDSMI_STATUS_SUCCESS: - logger.debug("amdsmi_get_gpu_fabric_info failed for GPU %d with status %d", gpu_id, ret) + _fabric_unavailable("amdsmi_get_gpu_fabric_info failed for GPU %d with status %d", gpu_id, ret) return FabricInfo() v1 = raw.fabric_info.fabric_version.v1 - return _amd_fabric_info_from_raw({"ppod_id": bytes(v1.ppod_id), "vpod_id": int(v1.vpod_id)}) + fabric = _amd_fabric_info_from_raw({"ppod_id": bytes(v1.ppod_id), "vpod_id": int(v1.vpod_id)}) + if not fabric.domain_key: + _fabric_unavailable("amdsmi reports no pod membership for GPU %d (PPOD_ID/VPOD_ID unset)", gpu_id) + return fabric except Exception as e: - logger.debug("AMDSMI fabric info query failed for GPU %d: %s", gpu_id, e) + _fabric_unavailable("AMD fabric info query failed for GPU %d: %s", gpu_id, e) return FabricInfo() @@ -885,6 +881,38 @@ def _get_total_memory_mb(gpu_id: int) -> int: return total_bytes // (1024 * 1024) +_fabric_failure: Optional[str] = None + + +def _fabric_unavailable(msg: str, *args) -> None: + # Most AMD systems have no fabric, so this is not worth a warning on its own; + # a multi-host job that needs the fabric puts the reason in its error. + global _fabric_failure + _fabric_failure = msg % args + logger.debug(msg, *args) + + +def _clear_fabric_failure() -> None: + global _fabric_failure + _fabric_failure = None + + +def fabric_failure_reason() -> Optional[str]: + """Why this process's last topology discovery found no AMD fabric domain, or None.""" + return _fabric_failure + + +def _torch_pci_bus_id(device_idx: int) -> Optional[str]: + try: + props = torch.cuda.get_device_properties(device_idx) + domain, bus, device = props.pci_domain_id, props.pci_bus_id, props.pci_device_id + except Exception as e: + logger.debug("torch PCI query failed for device %d: %s", device_idx, e) + return None + # torch does not report the PCI function; GPUs are function 0. + return f"{domain:04x}:{bus:02x}:{device:02x}.0" + + def _get_pci_bus_id(device_idx: int, vendor: str) -> str: """ Get the PCI bus ID for a GPU device. @@ -893,6 +921,9 @@ def _get_pci_bus_id(device_idx: int, vendor: str) -> str: logical device index to the physical index expected by NVML/AMDSMI, then queries NVML/AMDSMI by that physical index to obtain and normalize the busId/BDF string. + + On AMD, falls back to torch's device properties when the amdsmi Python + bindings are unavailable; the ROCm SDK wheels do not put them on sys.path. """ if device_idx < 0: logger.debug("Invalid device index: %d", device_idx) @@ -945,6 +976,13 @@ def _get_pci_bus_id(device_idx: int, vendor: str) -> str: e, ) + bus_id = _torch_pci_bus_id(device_idx) + if bus_id is not None: + return bus_id + _fabric_unavailable( + "no PCI bus ID for GPU %d (amdsmi Python bindings unavailable and torch reports none)", device_idx + ) + else: logger.debug("Unknown vendor: %s", vendor) @@ -1394,6 +1432,7 @@ def discover(self) -> TopologyMap: # Probe local GPU info device_name = torch.cuda.get_device_name(self.gpu_id) total_memory_mb = _get_total_memory_mb(self.gpu_id) + _clear_fabric_failure() pci_bus_id = _get_pci_bus_id(self.gpu_id, vendor) # Pass PCI bus ID to UUID query for PCI-based handle resolution gpu_uuid = _get_gpu_uuid(self.gpu_id, vendor, pci_bus_id=pci_bus_id) @@ -1402,9 +1441,13 @@ def discover(self) -> TopologyMap: # Query fabric info — pass PCI bus ID for PCI-based handle resolution fabric_info = _get_gpu_fabric_info(self.gpu_id, vendor, pci_bus_id=pci_bus_id) - logger.debug( - f"[Rank {self.rank}] Fabric info: cluster_uuid={fabric_info.cluster_uuid}, " - f"clique_id={fabric_info.clique_id}, domain_key={fabric_info.domain_key}" + logger.info( + "[Rank %d] %s GPU %d: pci_bus_id=%s fabric_domain=%s", + self.rank, + hostname, + self.gpu_id, + pci_bus_id, + fabric_info.domain_key or "none", ) local_gpu_info = GPUInfo( diff --git a/iris/host/memory/allocators/vmem_allocator.py b/iris/host/memory/allocators/vmem_allocator.py index c6a07459c..a8ac08c18 100644 --- a/iris/host/memory/allocators/vmem_allocator.py +++ b/iris/host/memory/allocators/vmem_allocator.py @@ -61,6 +61,15 @@ def __init__( va_multiplier: float = 1.0, ): super().__init__(heap_size, device_id, rank, world_size) + # Peer access is granted to local device ids 0..world_size-1. + local_devices = torch.cuda.device_count() + if world_size > local_devices: + self._closed = True + raise RuntimeError( + f"VMemAllocator needs every rank's GPU visible in each process, but world_size={world_size} and " + f"this process sees {local_devices} GPUs. It is single-host; use IRIS_ALLOCATOR=vmem_chunked " + "for multi-host jobs." + ) _log_rank( logging.INFO, "VMemAllocator: init heap_size=%.1fGB device=%d", diff --git a/iris/host/memory/symmetric_heap.py b/iris/host/memory/symmetric_heap.py index 8c6a73b16..5fab77f51 100644 --- a/iris/host/memory/symmetric_heap.py +++ b/iris/host/memory/symmetric_heap.py @@ -8,9 +8,12 @@ hiding the details of allocators and inter-process memory sharing. """ +import itertools import logging import os +import socket import struct +from typing import List import numpy as np import torch @@ -75,6 +78,55 @@ def _validate_local_handle(handle_bytes: bytes) -> None: ) +_hostname_exchanges = itertools.count(1) + + +def _hostnames_by_rank(num_ranks: int) -> List[str]: + """Every rank's hostname, exchanged through the rendezvous store so no NCCL collective is issued.""" + import torch.distributed as dist + + store = dist.distributed_c10d._get_default_store() + prefix = f"iris_hostname_v{next(_hostname_exchanges)}/" + store.set(f"{prefix}{dist.get_rank()}", socket.gethostname()) + return [store.get(f"{prefix}{r}").decode("utf-8") for r in range(num_ranks)] + + +def _require_single_host(allocator_type: str, hostnames: List[str]) -> None: + hosts = sorted(set(hostnames)) + if len(hosts) > 1: + raise RuntimeError( + f"IRIS_ALLOCATOR={allocator_type!r} only maps peers on the same host, but this job spans " + f"{len(hosts)} hosts ({', '.join(hosts)}). Use IRIS_ALLOCATOR=vmem_chunked, which maps peers " + "on other hosts through the GPU fabric driver." + ) + + +def _require_shared_fabric_domain(topology, cur_rank: int) -> None: + from iris.host.distributed.topology import fabric_failure_reason + + domains = {info.fabric_info.domain_key for info in topology.gpu_info.values()} + if len(domains) == 1 and "" not in domains: + return + + ranks_by_host = {} + for rank, info in sorted(topology.gpu_info.items()): + ranks_by_host.setdefault(info.hostname, []).append(rank) + lines = [] + for host, ranks in sorted(ranks_by_host.items()): + host_domains = sorted({topology.gpu_info[r].fabric_info.domain_key or "" for r in ranks}) + lines.append(f" {host}: ranks {ranks}, fabric domain {', '.join(host_domains)}") + reason = fabric_failure_reason() + if reason: + lines.append(f"Rank {cur_rank} found no fabric domain: {reason}.") + raise RuntimeError( + f"This job spans {len(ranks_by_host)} hosts, which needs every GPU in one fabric domain, but found:\n" + + "\n".join(lines) + + "\n means fabric discovery failed for that GPU; each rank's own error gives its reason. " + "Different domains mean the hosts are not in the same pod: on AMD, `amd-smi fabric -i` must show " + "the same PPOD_ID and VPOD_ID and ACCEL_STATE ACTIVE on every host." + ) + + class SymmetricHeap: """ High-level symmetric heap abstraction. @@ -134,6 +186,13 @@ def __init__( if is_simulation_env(): allocator_type = "torch" + multi_host = False + if num_ranks > 1 and not is_simulation_env(): + hostnames = _hostnames_by_rank(num_ranks) + multi_host = len(set(hostnames)) > 1 + if allocator_type in ("torch", "vmem"): + _require_single_host(allocator_type, hostnames) + if allocator_type == "torch": self.allocator = TorchAllocator(heap_size, device_id, cur_rank, num_ranks) elif allocator_type == "vmem": @@ -144,11 +203,18 @@ def __init__( try: topology = TopologyDiscovery().discover() except Exception as exc: + if multi_host: + raise RuntimeError( + f"TopologyDiscovery.discover() failed ({exc}). This job spans hosts, and " + "VMemChunkedAllocator needs topology discovery to find the GPUs' fabric domain." + ) from exc logger.warning( "TopologyDiscovery.discover() failed (%s); VMemChunkedAllocator will default to INTRA_NODE driver.", exc, ) topology = None + if multi_host: + _require_shared_fabric_domain(topology, cur_rank) self.allocator = VMemChunkedAllocator( heap_size, device_id, diff --git a/iris/host/platform/hip.py b/iris/host/platform/hip.py index a23ea8f9e..7b7257a48 100644 --- a/iris/host/platform/hip.py +++ b/iris/host/platform/hip.py @@ -9,20 +9,17 @@ import os from iris.host.logging.logging import _log_rank +from iris.host.platform.libs import load_vendor_library # Auto-detect backend _is_amd_backend = True -try: - rt_path = "libamdhip64.so" - gpu_runtime = ctypes.cdll.LoadLibrary(rt_path) -except OSError: +gpu_runtime = load_vendor_library("libamdhip64", "libamdhip64.so") +if gpu_runtime is None: try: - rt_path = "libcudart.so" - gpu_runtime = ctypes.cdll.LoadLibrary(rt_path) + gpu_runtime = ctypes.cdll.LoadLibrary("libcudart.so") _is_amd_backend = False except OSError: - rt_path = "libamdhip64.so" - gpu_runtime = ctypes.cdll.LoadLibrary(rt_path) + gpu_runtime = ctypes.cdll.LoadLibrary("libamdhip64.so") def gpu_try(err): diff --git a/iris/host/platform/libs.py b/iris/host/platform/libs.py new file mode 100644 index 000000000..9d794f80e --- /dev/null +++ b/iris/host/platform/libs.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +"""Load vendor runtime libraries, reusing the copies already mapped into the process.""" + +import ctypes +import glob +import logging +import os +from typing import Any, List, Optional, Set, Tuple + +logger = logging.getLogger("iris.platform") + + +def _mapped_library_paths(stem: str) -> List[str]: + """Distinct files named ``.so*`` mapped into this process, in /proc/self/maps order.""" + try: + with open("/proc/self/maps") as maps: + lines = maps.readlines() + except OSError: + return [] + + paths: List[str] = [] + seen = set() + for line in lines: + fields = line.split(maxsplit=5) + if len(fields) < 6: + continue + path = fields[5].strip() + if not os.path.basename(path).startswith(f"{stem}.so"): + continue + try: + st = os.stat(path) + key = (st.st_dev, st.st_ino) + except OSError: + key = path + if key not in seen: + seen.add(key) + paths.append(path) + return paths + + +_reported: Set[Tuple[str, Tuple[str, ...]]] = set() + + +def _sibling_library(stem: str, beside: str) -> Optional[str]: + """``.so*`` in the directory of the first mapped copy of ``beside``.""" + anchors = _mapped_library_paths(beside) + if not anchors: + return None + candidates = sorted(glob.glob(os.path.join(os.path.dirname(anchors[0]), f"{stem}.so*"))) + return candidates[0] if candidates else None + + +def _ld_library_path_copy(name: str) -> Optional[str]: + for directory in os.environ.get("LD_LIBRARY_PATH", "").split(":"): + path = os.path.join(directory, name) + if directory and os.path.exists(path): + return path + return None + + +def _same_file(a: str, b: str) -> bool: + try: + return os.path.samefile(a, b) + except OSError: + return False + + +def load_vendor_library(stem: str, *fallbacks: Optional[str], beside: Optional[str] = None) -> Any: + """ + Load ``.so``, preferring a copy that is already mapped, then one in the + directory of the mapped ``beside`` library, then the first of ``fallbacks`` + that loads. Returns None if none does. + + torch's ROCm wheels bring their own HIP runtime and amd_smi. Resolving the + library by name instead (LD_LIBRARY_PATH, ldconfig, /opt/rocm) can map a + second copy from a different ROCm into the same process, which can abort + during device initialization. + """ + mapped = _mapped_library_paths(stem) + preferred = mapped[:1] or ([_sibling_library(stem, beside)] if beside else []) + for name in preferred + list(fallbacks): + if not name: + continue + try: + lib = ctypes.CDLL(name) + except OSError: + continue + loaded = tuple(_mapped_library_paths(stem)) or (name,) + if (stem, loaded) not in _reported: + _reported.add((stem, loaded)) + if len(loaded) > 1: + logger.warning( + "Multiple copies of %s are loaded in this process: %s. Point LD_LIBRARY_PATH at the ROCm " + "that torch uses.", + stem, + list(loaded), + ) + else: + shadowed = _ld_library_path_copy(f"{stem}.so") + if shadowed and not _same_file(shadowed, loaded[0]): + logger.info("Using %s from %s, not %s from LD_LIBRARY_PATH", stem, loaded[0], shadowed) + else: + logger.debug("Using %s from %s", stem, loaded[0]) + return lib + return None diff --git a/tests/unittests/test_drivers.py b/tests/unittests/test_drivers.py index 1d237da05..eaac998b1 100644 --- a/tests/unittests/test_drivers.py +++ b/tests/unittests/test_drivers.py @@ -43,6 +43,7 @@ _round_up, ) from iris.host.distributed.topology import InterconnectLevel +from iris.host.platform import libs as platform_libs @pytest.mark.parametrize("unversioned_available", [True, False]) @@ -57,6 +58,7 @@ def load(name): monkeypatch.setattr(amd_driver_module.ctypes, "CDLL", load) monkeypatch.setattr(amd_driver_module.ctypes.util, "find_library", lambda name: f"lib{name}.so.5") + monkeypatch.setattr(platform_libs, "_mapped_library_paths", lambda stem: []) spec = importlib.util.spec_from_file_location("_iris_amd_library_selection_test", amd_driver_module.__file__) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) @@ -67,6 +69,25 @@ def load(name): assert "libamdhip64.so.5" not in loaded +def test_amd_fabric_library_selection_prefers_loaded_copy(monkeypatch): + """The HIP runtime and amd_smi torch already loaded win over any name lookup.""" + loaded = {} + mapped = { + "libamdhip64": ["/sdk/core/lib/libamdhip64.so.7"], + "libamd_smi": ["/sdk/core/lib/libamd_smi.so.27"], + } + + monkeypatch.setattr(amd_driver_module.ctypes, "CDLL", lambda name: loaded.setdefault(name, object())) + monkeypatch.setattr(amd_driver_module.ctypes.util, "find_library", lambda name: f"lib{name}.so.5") + monkeypatch.setattr(platform_libs, "_mapped_library_paths", lambda stem: mapped[stem]) + spec = importlib.util.spec_from_file_location("_iris_amd_loaded_library_test", amd_driver_module.__file__) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + assert module._hip is loaded["/sdk/core/lib/libamdhip64.so.7"] + assert module._amdsmi is loaded["/sdk/core/lib/libamd_smi.so.27"] + assert set(loaded) == {"/sdk/core/lib/libamdhip64.so.7", "/sdk/core/lib/libamd_smi.so.27"} + + class TestExceptions: def test_cuda_fabric_error_uses_driver_error_hierarchy(self): assert issubclass(CudaFabricError, DriverError) diff --git a/tests/unittests/test_multihost_guard.py b/tests/unittests/test_multihost_guard.py new file mode 100644 index 000000000..fd12c810a --- /dev/null +++ b/tests/unittests/test_multihost_guard.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +"""Multi-host jobs must fail fast instead of falling back to the single-host FD-passing mesh.""" + +import itertools +import socket +import threading +import time +import types + +import pytest +import torch.distributed as dist + +import iris.host.distributed.topology as topology_module +from iris.host.distributed import fd_passing +from iris.host.distributed.topology import FabricInfo +from iris.host.memory import symmetric_heap +from iris.host.memory.symmetric_heap import SymmetricHeap + +POD = "3e568d58317a4bff84fcb661a751878f:2" + + +class _Constructed(Exception): + """Raised by a faked allocator: the guard let the heap reach allocator construction.""" + + +def _constructed(*args, **kwargs): + raise _Constructed() + + +def _fake_topology(hosts_and_domains): + gpu_info = { + rank: types.SimpleNamespace( + hostname=host, fabric_info=FabricInfo(*domain.split(":")) if domain else FabricInfo() + ) + for rank, (host, domain) in enumerate(hosts_and_domains) + } + return types.SimpleNamespace(gpu_info=gpu_info) + + +@pytest.fixture +def heap_env(monkeypatch): + """Fake everything SymmetricHeap.__init__ touches before the allocator exists.""" + monkeypatch.delenv("IRIS_ALLOCATOR", raising=False) + monkeypatch.setattr(symmetric_heap, "is_simulation_env", lambda: False) + monkeypatch.setattr(topology_module, "_fabric_failure", None) + for name in ("TorchAllocator", "VMemAllocator", "VMemChunkedAllocator"): + monkeypatch.setattr(symmetric_heap, name, _constructed) + + def use(hostnames, topology=None, discover_error=None): + monkeypatch.setattr(symmetric_heap, "_hostnames_by_rank", lambda num_ranks: list(hostnames)) + + class FakeDiscovery: + def discover(self): + if discover_error is not None: + raise discover_error + return topology + + monkeypatch.setattr(topology_module, "TopologyDiscovery", FakeDiscovery) + + return use + + +def _heap(allocator_type, num_ranks=4): + return SymmetricHeap(1 << 20, 0, 0, num_ranks, allocator_type) + + +class TestSingleHostAllocators: + @pytest.mark.parametrize("allocator_type", ["torch", "vmem"]) + def test_multi_host_rejected_before_allocator(self, heap_env, allocator_type): + heap_env(["tray-0", "tray-0", "tray-1", "tray-1"]) + + with pytest.raises(RuntimeError, match="IRIS_ALLOCATOR=vmem_chunked") as err: + _heap(allocator_type) + assert f"IRIS_ALLOCATOR='{allocator_type}'" in str(err.value) + assert "tray-0, tray-1" in str(err.value) + + @pytest.mark.parametrize("allocator_type", ["torch", "vmem"]) + def test_single_host_allowed(self, heap_env, allocator_type): + heap_env(["tray-0"] * 4) + + with pytest.raises(_Constructed): + _heap(allocator_type) + + def test_single_rank_skips_hostname_exchange(self, monkeypatch, heap_env): + def no_exchange(num_ranks): + raise AssertionError("hostname exchange should not run for one rank") + + monkeypatch.setattr(symmetric_heap, "_hostnames_by_rank", no_exchange) + + with pytest.raises(_Constructed): + _heap("vmem", num_ranks=1) + + +class TestVMemChunkedMultiHost: + def test_shared_fabric_domain_allowed(self, heap_env): + hosts = [("tray-0", POD), ("tray-0", POD), ("tray-1", POD), ("tray-1", POD)] + heap_env([h for h, _ in hosts], _fake_topology(hosts)) + + with pytest.raises(_Constructed): + _heap("vmem_chunked") + + def test_missing_fabric_domain_rejected(self, monkeypatch, heap_env): + hosts = [("tray-0", ""), ("tray-0", ""), ("tray-1", ""), ("tray-1", "")] + heap_env([h for h, _ in hosts], _fake_topology(hosts)) + monkeypatch.setattr(topology_module, "_fabric_failure", "amdsmi_init() failed with status 8") + + with pytest.raises(RuntimeError, match="one fabric domain") as err: + _heap("vmem_chunked") + message = str(err.value) + assert "tray-0: ranks [0, 1], fabric domain " in message + assert "tray-1: ranks [2, 3], fabric domain " in message + assert "Rank 0 found no fabric domain: amdsmi_init() failed with status 8." in message + assert "amd-smi fabric -i" in message + + def test_different_pods_rejected(self, heap_env): + other = "aa" * 16 + ":5" + hosts = [("tray-0", POD), ("tray-1", other)] + heap_env([h for h, _ in hosts], _fake_topology(hosts)) + + with pytest.raises(RuntimeError, match="one fabric domain") as err: + _heap("vmem_chunked", num_ranks=2) + assert f"tray-1: ranks [1], fabric domain {other}" in str(err.value) + + def test_discovery_failure_is_fatal_on_multi_host(self, heap_env): + heap_env(["tray-0", "tray-1"], discover_error=RuntimeError("boom")) + + with pytest.raises(RuntimeError, match="needs topology discovery") as err: + _heap("vmem_chunked", num_ranks=2) + assert isinstance(err.value.__cause__, RuntimeError) + + def test_discovery_failure_still_warns_on_single_host(self, heap_env): + heap_env(["tray-0", "tray-0"], discover_error=RuntimeError("boom")) + + with pytest.raises(_Constructed): + _heap("vmem_chunked", num_ranks=2) + + +def test_hostnames_exchanged_through_store(monkeypatch): + class FakeStore: + def __init__(self): + self.data = {} + + def set(self, key, value): + self.data[key] = value.encode("utf-8") + + def get(self, key): + return self.data[key] + + store = FakeStore() + monkeypatch.setattr(dist.distributed_c10d, "_get_default_store", lambda: store) + monkeypatch.setattr(dist, "get_rank", lambda: 1) + monkeypatch.setattr(symmetric_heap.socket, "gethostname", lambda: "tray-1") + monkeypatch.setattr(symmetric_heap, "_hostname_exchanges", itertools.count(1)) + store.data["iris_hostname_v1/0"] = b"tray-0" + + assert symmetric_heap._hostnames_by_rank(2) == ["tray-0", "tray-1"] + + +def test_fd_mesh_accept_times_out_on_missing_peer(monkeypatch, tmp_path): + monkeypatch.setattr(fd_passing, "_ACCEPT_TIMEOUT_S", 0.2) + paths = {0: str(tmp_path / "r0.sock"), 1: str(tmp_path / "r1.sock")} + + with pytest.raises(TimeoutError, match=r"ranks \[1\] to connect to rank 0"): + fd_passing.setup_fd_mesh(0, 2, paths) + assert not (tmp_path / "r0.sock").exists() + + +def test_fd_mesh_handshake_times_out_on_silent_peer(monkeypatch, tmp_path): + monkeypatch.setattr(fd_passing, "_HANDSHAKE_TIMEOUT_S", 0.2) + paths = {0: str(tmp_path / "r0.sock"), 1: str(tmp_path / "r1.sock")} + silent = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + + def connect_without_sending_rank(): + deadline = time.time() + 5 + while time.time() < deadline: + try: + silent.connect(paths[0]) + return + except (FileNotFoundError, ConnectionRefusedError): + time.sleep(0.01) + + peer = threading.Thread(target=connect_without_sending_rank) + peer.start() + try: + with pytest.raises(TimeoutError, match="did not send its rank"): + fd_passing.setup_fd_mesh(0, 2, paths) + finally: + peer.join() + silent.close() + assert not (tmp_path / "r0.sock").exists() + + +def test_vmem_allocator_rejects_more_ranks_than_local_gpus(monkeypatch): + from iris.host.memory.allocators import vmem_allocator + + monkeypatch.setattr(vmem_allocator.torch.cuda, "device_count", lambda: 4) + + with pytest.raises(RuntimeError, match="world_size=8 and this process sees 4 GPUs"): + vmem_allocator.VMemAllocator(1 << 21, 0, 0, 8) diff --git a/tests/unittests/test_platform_libs.py b/tests/unittests/test_platform_libs.py new file mode 100644 index 000000000..f5496484d --- /dev/null +++ b/tests/unittests/test_platform_libs.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +"""Vendor library loading prefers copies already mapped into the process.""" + +import io +import logging +import os + +import pytest +import torch + +from iris.host.platform import libs + + +def _fake_maps(monkeypatch, paths): + text = "".join(f"7f0000000000-7f0000001000 r-xp 00000000 08:01 1 {p}\n" for p in paths) + text += "7ffd00000000-7ffd00021000 rw-p 00000000 00:00 0 [stack]\n" + text += "7ffe00000000-7ffe00001000 rw-p 00000000 00:00 0\n" + monkeypatch.setattr(libs, "open", lambda path: io.StringIO(text), raising=False) + + +class TestMappedLibraryPaths: + def test_matches_versioned_names_and_dedupes_hardlinks(self, monkeypatch, tmp_path): + core = tmp_path / "core" / "libamdhip64.so.7" + core.parent.mkdir() + core.write_bytes(b"hip") + devel = tmp_path / "devel" / "libamdhip64.so" + devel.parent.mkdir() + os.link(core, devel) + system = tmp_path / "system" / "libamdhip64.so.7.17.0" + system.parent.mkdir() + system.write_bytes(b"other hip") + other = tmp_path / "core" / "libamdhip64_helper.so" + other.write_bytes(b"not hip") + + _fake_maps(monkeypatch, [core, core, devel, other, system]) + + assert libs._mapped_library_paths("libamdhip64") == [str(core), str(system)] + + def test_unreadable_maps(self, monkeypatch): + def fail(path): + raise OSError("no procfs") + + monkeypatch.setattr(libs, "open", fail, raising=False) + assert libs._mapped_library_paths("libamdhip64") == [] + + @pytest.mark.skipif(torch.version.hip is None, reason="needs a ROCm build of torch") + def test_finds_torch_hip_runtime(self): + assert libs._mapped_library_paths("libamdhip64") + + +class TestLoadVendorLibrary: + @pytest.fixture(autouse=True) + def _fresh_reports(self, monkeypatch): + monkeypatch.setattr(libs, "_reported", set()) + + def test_prefers_mapped_copy(self, monkeypatch): + opened = [] + monkeypatch.setattr(libs, "_mapped_library_paths", lambda stem: ["/sdk/libamdhip64.so.7"]) + monkeypatch.setattr(libs.ctypes, "CDLL", lambda name: opened.append(name) or name) + + assert libs.load_vendor_library("libamdhip64", "libamdhip64.so") == "/sdk/libamdhip64.so.7" + assert opened == ["/sdk/libamdhip64.so.7"] + + def test_falls_back_in_order(self, monkeypatch): + opened = [] + + def cdll(name): + opened.append(name) + if name != "/opt/rocm/lib/libamdhip64.so": + raise OSError(name) + return name + + monkeypatch.setattr(libs, "_mapped_library_paths", lambda stem: []) + monkeypatch.setattr(libs.ctypes, "CDLL", cdll) + + result = libs.load_vendor_library("libamdhip64", "libamdhip64.so", None, "/opt/rocm/lib/libamdhip64.so") + assert result == "/opt/rocm/lib/libamdhip64.so" + assert opened == ["libamdhip64.so", "/opt/rocm/lib/libamdhip64.so"] + + def test_returns_none_when_nothing_loads(self, monkeypatch): + def cdll(name): + raise OSError(name) + + monkeypatch.setattr(libs, "_mapped_library_paths", lambda stem: []) + monkeypatch.setattr(libs.ctypes, "CDLL", cdll) + + assert libs.load_vendor_library("libamdhip64", "libamdhip64.so") is None + + def test_warns_when_a_second_copy_gets_loaded(self, monkeypatch, caplog): + calls = iter([[], ["/sdk/libamdhip64.so.7", "/opt/rocm/lib/libamdhip64.so.7"]]) + monkeypatch.setattr(libs, "_mapped_library_paths", lambda stem: next(calls)) + monkeypatch.setattr(libs.ctypes, "CDLL", lambda name: name) + caplog.set_level(logging.WARNING, logger="iris.platform") + + libs.load_vendor_library("libamdhip64", "libamdhip64.so") + + assert any("Multiple copies of libamdhip64" in r.message for r in caplog.records) + + def test_info_once_when_torch_copy_overrides_ld_library_path(self, monkeypatch, caplog, tmp_path): + sdk, system = tmp_path / "sdk", tmp_path / "system" + sdk.mkdir() + system.mkdir() + (sdk / "libamdhip64.so.7").write_bytes(b"sdk hip") + (system / "libamdhip64.so").write_bytes(b"system hip") + monkeypatch.setenv("LD_LIBRARY_PATH", f"{tmp_path / 'missing'}:{system}") + monkeypatch.setattr(libs, "_mapped_library_paths", lambda stem: [str(sdk / "libamdhip64.so.7")]) + monkeypatch.setattr(libs.ctypes, "CDLL", lambda name: name) + caplog.set_level(logging.INFO, logger="iris.platform") + + libs.load_vendor_library("libamdhip64", "libamdhip64.so") + libs.load_vendor_library("libamdhip64", "libamdhip64.so") + + infos = [r.getMessage() for r in caplog.records if r.levelno == logging.INFO] + assert infos == [ + f"Using libamdhip64 from {sdk / 'libamdhip64.so.7'}, not {system / 'libamdhip64.so'} from LD_LIBRARY_PATH" + ] + + def test_debug_only_when_ld_library_path_has_the_same_copy(self, monkeypatch, caplog, tmp_path): + lib = tmp_path / "libamdhip64.so.7" + lib.write_bytes(b"hip") + os.symlink(lib, tmp_path / "libamdhip64.so") + monkeypatch.setenv("LD_LIBRARY_PATH", str(tmp_path)) + monkeypatch.setattr(libs, "_mapped_library_paths", lambda stem: [str(lib)]) + monkeypatch.setattr(libs.ctypes, "CDLL", lambda name: name) + caplog.set_level(logging.INFO, logger="iris.platform") + + libs.load_vendor_library("libamdhip64", "libamdhip64.so") + + assert not caplog.records + + def test_unmapped_library_comes_from_beside_directory(self, monkeypatch, tmp_path): + sdk = tmp_path / "sdk" + sdk.mkdir() + (sdk / "libamd_smi.so.27").write_bytes(b"smi") + mapped = {"libamdhip64": [str(sdk / "libamdhip64.so.7")], "libamd_smi": []} + opened = [] + monkeypatch.setattr(libs, "_mapped_library_paths", lambda stem: mapped[stem]) + monkeypatch.setattr(libs.ctypes, "CDLL", lambda name: opened.append(name) or name) + + libs.load_vendor_library("libamd_smi", "libamd_smi.so", beside="libamdhip64") + + assert opened == [str(sdk / "libamd_smi.so.27")] + + def test_beside_falls_back_when_directory_has_no_copy(self, monkeypatch, tmp_path): + mapped = {"libamdhip64": [str(tmp_path / "libamdhip64.so.7")], "libamd_smi": []} + opened = [] + monkeypatch.setattr(libs, "_mapped_library_paths", lambda stem: mapped[stem]) + monkeypatch.setattr(libs.ctypes, "CDLL", lambda name: opened.append(name) or name) + + libs.load_vendor_library("libamd_smi", "libamd_smi.so", beside="libamdhip64") + + assert opened == ["libamd_smi.so"] diff --git a/tests/unittests/test_topology.py b/tests/unittests/test_topology.py index cfd39f57d..901e92cc6 100644 --- a/tests/unittests/test_topology.py +++ b/tests/unittests/test_topology.py @@ -627,6 +627,165 @@ def test_discover_unknown_vendor_uses_noop_outer_lifecycle(self, monkeypatch): assert topo.nodes["test-host"].num_gpus == 1 +def _block_vendor_imports(monkeypatch): + real_import = builtins.__import__ + + def blocked_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in {"pynvml", "amdsmi"}: + raise ImportError(f"{name} unavailable") + return real_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", blocked_import) + + +def _fake_torch_pci(monkeypatch, bus_by_device): + def props(device_idx): + return types.SimpleNamespace(pci_domain_id=1, pci_bus_id=bus_by_device[device_idx], pci_device_id=0) + + monkeypatch.setattr(torch.cuda, "get_device_properties", props) + + +class TestAmdPciBusIdFallback: + def test_falls_back_to_torch_without_amdsmi_bindings(self, monkeypatch): + _block_vendor_imports(monkeypatch) + _fake_torch_pci(monkeypatch, {0: 0x01, 1: 0x21}) + + assert topology._get_pci_bus_id(0, "amd") == "0001:01:00.0" + assert topology._get_pci_bus_id(1, "amd") == "0001:21:00.0" + + def test_amdsmi_bindings_still_preferred(self, monkeypatch): + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising=False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising=False) + + monkeypatch.setitem(sys.modules, "amdsmi", _make_fake_amdsmi()) + _fake_torch_pci(monkeypatch, {1: 0x99}) + + assert topology._get_pci_bus_id(1, "amd") == "0000:42:00.0" + + def test_unknown_records_reason_without_warning(self, monkeypatch, caplog): + _block_vendor_imports(monkeypatch) + + def no_props(device_idx): + raise RuntimeError("no device") + + monkeypatch.setattr(torch.cuda, "get_device_properties", no_props) + monkeypatch.setattr(topology, "_fabric_failure", None) + caplog.set_level(logging.WARNING, logger="iris.topology") + + assert topology._get_pci_bus_id(0, "amd") == "unknown" + assert "no PCI bus ID for GPU 0" in topology.fabric_failure_reason() + assert not caplog.records + + +class TestAmdFabricFailureReason: + """Each way AMD fabric discovery comes back empty leaves a reason for the multi-host error.""" + + @staticmethod + def _fake_amdsmi(monkeypatch, init=0, handle=0, fabric=0, ppod=b"\xab" * 16, vpod=9): + class call: + def __init__(self, fn): + self.fn = fn + self.argtypes = None + self.restype = None + + def __call__(self, *args): + return self.fn(*args) + + def get_handle(bdf, processor_ptr): + processor_ptr._obj.value = 0xABCD + return handle + + def get_fabric(processor, info_ptr): + v1 = info_ptr._obj.fabric_info.fabric_version.v1 + for idx, value in enumerate(ppod): + v1.ppod_id[idx] = value + v1.vpod_id = vpod + return fabric + + lib = types.SimpleNamespace( + amdsmi_init=call(lambda flags: init), + amdsmi_get_processor_handle_from_bdf=call(get_handle), + amdsmi_get_gpu_fabric_info=call(get_fabric), + ) + monkeypatch.setattr(topology, "load_vendor_library", lambda *a, **k: lib) + monkeypatch.setattr(topology, "_fabric_failure", None) + + @pytest.mark.parametrize( + "kwargs, expected", + [ + ({"init": 8}, "amdsmi_init() failed with status 8"), + ({"handle": 2}, "amdsmi_get_processor_handle_from_bdf(0000:41:00.0) failed for GPU 0 with status 2"), + ({"fabric": 2}, "amdsmi_get_gpu_fabric_info failed for GPU 0 with status 2"), + ({"ppod": b"\x99" * 16}, "amdsmi reports no pod membership for GPU 0"), + ], + ) + def test_failure_reason(self, monkeypatch, kwargs, expected): + self._fake_amdsmi(monkeypatch, **kwargs) + + assert topology._amd_get_gpu_fabric_info(0, "0000:41:00.0") == FabricInfo() + assert expected in topology.fabric_failure_reason() + + def test_missing_library(self, monkeypatch): + monkeypatch.setattr(topology, "load_vendor_library", lambda *a, **k: None) + monkeypatch.setattr(topology, "_fabric_failure", None) + + assert topology._amd_get_gpu_fabric_info(0, "0000:41:00.0") == FabricInfo() + assert topology.fabric_failure_reason() == "libamd_smi.so could not be loaded" + + def test_success_leaves_no_reason(self, monkeypatch): + self._fake_amdsmi(monkeypatch) + + assert topology._amd_get_gpu_fabric_info(0, "0000:41:00.0").domain_key == f"{'ab' * 16}:9" + assert topology.fabric_failure_reason() is None + + def test_two_host_discovery_finds_fabric_domain_without_amdsmi_bindings(self, monkeypatch, caplog): + """Without the bindings, two hosts in one pod still land in the same fabric domain.""" + _block_vendor_imports(monkeypatch) + _fake_torch_pci(monkeypatch, {0: 0x01}) + monkeypatch.setattr(topology, "_detect_vendor", lambda: "amd") + monkeypatch.setattr(socket, "gethostname", lambda: "tray-0") + monkeypatch.setattr(torch.cuda, "get_device_name", lambda gpu_id: "Fake MI455X") + monkeypatch.setattr(topology, "_get_total_memory_mb", lambda gpu_id: 1024) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 1) + monkeypatch.setattr(topology, "_get_numa_node", lambda pci_bus_id: 0) + monkeypatch.setattr(topology, "_detect_infiniband", lambda: (False, [])) + monkeypatch.setattr(topology, "_detect_intra_node_topology", lambda pci_bus_ids, vendor: (None, None)) + monkeypatch.setattr(topology, "_outer_init_vendor_lib", lambda vendor: None) + monkeypatch.setattr(topology, "_outer_shutdown_vendor_lib", lambda vendor: None) + + def fabric(gpu_id, pci_bus_id=""): + if pci_bus_id in ("", "unknown"): + return FabricInfo() + return FabricInfo(cluster_uuid="3e568d58317a4bff84fcb661a751878f", clique_id=2) + + monkeypatch.setattr(topology, "_amd_get_gpu_fabric_info", fabric) + + def gather(payload, world_size): + data = json.loads(payload) + if "global_rank" not in data: + return [payload, payload] + peer = dict(data, global_rank=1, hostname="tray-1") + return [payload, json.dumps(peer)] + + monkeypatch.setattr(topology, "_all_gather_strings", gather) + caplog.set_level(logging.INFO, logger="iris.topology") + + discovery = TopologyDiscovery.__new__(TopologyDiscovery) + discovery.rank = 0 + discovery.world_size = 2 + discovery.gpu_id = 0 + discovery._topology = None + topo = discovery.discover() + + assert topo.gpu_info[0].pci_bus_id == "0001:01:00.0" + assert topo.gpu_info[0].fabric_info.domain_key == "3e568d58317a4bff84fcb661a751878f:2" + assert topo.get_interconnect_level(0, 1) == InterconnectLevel.INTRA_RACK_FABRIC + assert any( + "pci_bus_id=0001:01:00.0 fabric_domain=3e568d58317a4bff84fcb661a751878f:2" in r.getMessage() + for r in caplog.records + ) + + class TestNodeInfo: """Tests for NodeInfo safe accessors."""