diff --git a/iris/drivers/__init__.py b/iris/drivers/__init__.py index a46501665..68f4dcec9 100644 --- a/iris/drivers/__init__.py +++ b/iris/drivers/__init__.py @@ -10,9 +10,9 @@ from dataclasses import dataclass from typing import Optional -from iris.drivers.base import BaseDriver +from iris.drivers.base import BaseDriver, CleanupTarget, ExportableMemory, MappingPlacement -__all__ = ["DriverStack"] +__all__ = ["DriverStack", "CleanupTarget", "ExportableMemory", "MappingPlacement"] @dataclass diff --git a/iris/drivers/base.py b/iris/drivers/base.py index 272fe0af4..f276c7e00 100644 --- a/iris/drivers/base.py +++ b/iris/drivers/base.py @@ -16,6 +16,9 @@ __all__ = [ "PeerMapping", "LocalAllocation", + "CleanupTarget", + "ExportableMemory", + "MappingPlacement", "BaseDriver", "DriverError", "DriverNotSupported", @@ -43,6 +46,40 @@ class LocalAllocation: _va_owned: bool = True +CleanupTarget = LocalAllocation | PeerMapping + + +@dataclass(frozen=True) +class ExportableMemory: + """A local memory range that can be exported to peers.""" + + va: int + size: int + allocation: Optional[LocalAllocation] = None + + +@dataclass(frozen=True) +class MappingPlacement: + """Caller-reserved virtual address placement for a VMM mapping. + + If arena_base is set, the mapping belongs to a larger reserved VA arena. + Drivers that need cumulative access programming can use that arena context + without exposing access-range details in the public API. + """ + + va: int + arena_base: Optional[int] = None + + def access_range(self, mapped_size: int, cumulative: bool = False) -> tuple[int, int]: + if not cumulative or self.arena_base is None: + return int(self.va), int(mapped_size) + + offset = int(self.va) - int(self.arena_base) + if offset < 0: + raise ValueError(f"placement va {self.va} is below arena base {self.arena_base}") + return int(self.arena_base), offset + int(mapped_size) + + class DriverError(RuntimeError): """Base exception for driver operations.""" @@ -62,16 +99,13 @@ def initialize(self, device_ordinal: int) -> None: def allocate_exportable( self, size: int, - va: Optional[int] = None, - *, - access_va: Optional[int] = None, - access_size: Optional[int] = None, + placement: Optional[MappingPlacement] = None, ) -> LocalAllocation: """Allocate exportable memory, optionally mapping it at a caller-reserved VA.""" @abstractmethod - def export_handle(self, allocation: LocalAllocation) -> bytes: - """Export a transport-specific handle for a local allocation.""" + def export_handle(self, memory: ExportableMemory) -> bytes: + """Export a transport-specific handle for a local memory range.""" @abstractmethod def import_and_map( @@ -79,20 +113,13 @@ def import_and_map( peer_rank: int, handle_bytes: bytes, size: int, - va: Optional[int] = None, - *, - access_va: Optional[int] = None, - access_size: Optional[int] = None, + placement: Optional[MappingPlacement] = None, ) -> PeerMapping: """Import a peer handle and map it into the local virtual address space.""" @abstractmethod - def cleanup_import(self, mapping: PeerMapping) -> None: - """Release a mapped peer allocation.""" - - @abstractmethod - def cleanup_local(self, allocation: LocalAllocation) -> None: - """Release a locally-exported allocation.""" + def cleanup(self, target: CleanupTarget) -> None: + """Release a local allocation or imported peer mapping.""" @abstractmethod def get_minimum_granularity(self) -> int: @@ -109,7 +136,3 @@ def free_va(self, va: int, size: int) -> None: def get_address_range(self, ptr: int) -> tuple[int, int]: """Return the base VA and size of the allocation containing ptr.""" raise DriverNotSupported(f"{type(self).__name__} does not support get_address_range") - - def export_pointer_handle(self, ptr: int, size: int) -> bytes: - """Export a peer handle for an arbitrary device pointer.""" - raise DriverNotSupported(f"{type(self).__name__} does not support export_pointer_handle") diff --git a/iris/drivers/fabric/amd.py b/iris/drivers/fabric/amd.py index a7af84497..9d1ed4c82 100644 --- a/iris/drivers/fabric/amd.py +++ b/iris/drivers/fabric/amd.py @@ -12,7 +12,9 @@ from iris.drivers.base import ( BaseDriver, DriverNotSupported, + ExportableMemory, LocalAllocation, + MappingPlacement, PeerMapping, ) @@ -30,14 +32,11 @@ def initialize(self, device_ordinal: int) -> None: def allocate_exportable( self, size: int, - va: Optional[int] = None, - *, - access_va: Optional[int] = None, - access_size: Optional[int] = None, + placement: Optional[MappingPlacement] = None, ) -> LocalAllocation: raise DriverNotSupported(_NOT_IMPLEMENTED_MESSAGE) - def export_handle(self, allocation: LocalAllocation) -> bytes: + def export_handle(self, memory: ExportableMemory) -> bytes: raise DriverNotSupported(_NOT_IMPLEMENTED_MESSAGE) def import_and_map( @@ -45,17 +44,11 @@ def import_and_map( peer_rank: int, handle_bytes: bytes, size: int, - va: Optional[int] = None, - *, - access_va: Optional[int] = None, - access_size: Optional[int] = None, + placement: Optional[MappingPlacement] = None, ) -> PeerMapping: raise DriverNotSupported(_NOT_IMPLEMENTED_MESSAGE) - def cleanup_import(self, mapping: PeerMapping) -> None: - raise DriverNotSupported(_NOT_IMPLEMENTED_MESSAGE) - - def cleanup_local(self, allocation: LocalAllocation) -> None: + def cleanup(self, target: LocalAllocation | PeerMapping) -> None: raise DriverNotSupported(_NOT_IMPLEMENTED_MESSAGE) def get_minimum_granularity(self) -> int: diff --git a/iris/drivers/fabric/nvidia.py b/iris/drivers/fabric/nvidia.py index fc948f17a..685a3a5c9 100644 --- a/iris/drivers/fabric/nvidia.py +++ b/iris/drivers/fabric/nvidia.py @@ -18,7 +18,9 @@ BaseDriver, DriverError, DriverNotSupported, + ExportableMemory, LocalAllocation, + MappingPlacement, PeerMapping, ) from iris.host.distributed.topology import InterconnectLevel @@ -339,20 +341,15 @@ def _check_initialized(self) -> None: def allocate_exportable( self, size: int, - va: Optional[int] = None, - *, - access_va: Optional[int] = None, - access_size: Optional[int] = None, + placement: Optional[MappingPlacement] = None, ) -> LocalAllocation: self._check_initialized() - if (access_va is None) != (access_size is None): - raise CudaFabricError("access_va and access_size must be provided together") props = self._make_alloc_props() granularity = self._get_granularity() alloc_size = _round_up(size, granularity) - reserved_va = va is None - mapped_va = int(va) if va is not None else 0 + reserved_va = placement is None + mapped_va = int(placement.va) if placement is not None else 0 handle = ctypes.c_uint64() mapped = False @@ -373,10 +370,10 @@ def allocate_exportable( "cuMemMap", ) mapped = True - self._mem_set_access( - int(access_va) if access_va is not None else mapped_va, - int(access_size) if access_size is not None else alloc_size, + access_base, access_bytes = ( + placement.access_range(alloc_size) if placement is not None else (mapped_va, alloc_size) ) + self._mem_set_access(access_base, access_bytes) return LocalAllocation( va=mapped_va, size=alloc_size, @@ -404,13 +401,16 @@ def allocate_exportable( pass raise - def export_handle(self, allocation: LocalAllocation) -> bytes: + def export_handle(self, memory: ExportableMemory) -> bytes: self._check_initialized() + if memory.allocation is None: + raise CudaFabricNotSupported("NVIDIA fabric driver can only export driver-created VMM allocations") + raw = (ctypes.c_ubyte * FABRIC_HANDLE_BYTES)() _cuda_try( _cuda_driver.cuMemExportToShareableHandle( ctypes.byref(raw), - int(allocation.handle), + int(memory.allocation.handle), _CU_MEM_HANDLE_TYPE_FABRIC, 0, ), @@ -437,19 +437,14 @@ def import_and_map( peer_rank: int, handle_bytes: bytes, size: int, - va: Optional[int] = None, - *, - access_va: Optional[int] = None, - access_size: Optional[int] = None, + placement: Optional[MappingPlacement] = None, ) -> PeerMapping: self._check_initialized() - if (access_va is None) != (access_size is None): - raise CudaFabricError("access_va and access_size must be provided together") imported_handle = self._import_handle(handle_bytes) granularity = self._get_granularity() - va_owned = va is None - mapped_va = int(va) if va is not None else 0 + va_owned = placement is None + mapped_va = int(placement.va) if placement is not None else 0 mapped = False try: if va_owned: @@ -464,10 +459,8 @@ def import_and_map( "cuMemMap", ) mapped = True - self._mem_set_access( - int(access_va) if access_va is not None else mapped_va, - int(access_size) if access_size is not None else size, - ) + access_base, access_bytes = placement.access_range(size) if placement is not None else (mapped_va, size) + self._mem_set_access(access_base, access_bytes) except Exception: if mapped: try: @@ -497,7 +490,16 @@ def import_and_map( _driver_handle=(tag, imported_handle), ) - def cleanup_import(self, mapping: PeerMapping) -> None: + def cleanup(self, target: LocalAllocation | PeerMapping) -> None: + if isinstance(target, LocalAllocation): + self._cleanup_local(target) + return + if isinstance(target, PeerMapping): + self._cleanup_import(target) + return + raise CudaFabricError(f"Unsupported cleanup target: {type(target).__name__}") + + def _cleanup_import(self, mapping: PeerMapping) -> None: self._check_initialized() if isinstance(mapping._driver_handle, tuple) and len(mapping._driver_handle) == 2: tag, imported_handle = mapping._driver_handle @@ -533,7 +535,7 @@ def cleanup_import(self, mapping: PeerMapping) -> None: ) _run_cleanup_steps(*steps) - def cleanup_local(self, allocation: LocalAllocation) -> None: + def _cleanup_local(self, allocation: LocalAllocation) -> None: self._check_initialized() steps = [ ( diff --git a/iris/drivers/local/amd.py b/iris/drivers/local/amd.py index 6ce10bb82..ff346f942 100644 --- a/iris/drivers/local/amd.py +++ b/iris/drivers/local/amd.py @@ -16,7 +16,9 @@ BaseDriver, DriverError, DriverNotSupported, + ExportableMemory, LocalAllocation, + MappingPlacement, PeerMapping, ) from iris.host.distributed.topology import InterconnectLevel @@ -368,26 +370,22 @@ def initialize(self, device_ordinal: int) -> None: def allocate_exportable( self, size: int, - va: Optional[int] = None, - *, - access_va: Optional[int] = None, - access_size: Optional[int] = None, + placement: Optional[MappingPlacement] = None, ) -> LocalAllocation: """ Allocate HIP VMem exportable as a DMA-BUF. - If va is supplied, the caller must already own a sufficiently large, - granularity-aligned VA range containing [va, va + size). + If placement is supplied, the caller must already own a sufficiently + large, granularity-aligned VA range containing [placement.va, + placement.va + size). """ self._check_initialized() - if (access_va is None) != (access_size is None): - raise LocalHipError("access_va and access_size must be provided together") props = self._make_alloc_props() granularity = self._get_granularity() alloc_size = _round_up(size, granularity) - reserved_va = va is None - mapped_va = int(va) if va is not None else 0 + reserved_va = placement is None + mapped_va = int(placement.va) if placement is not None else 0 handle = hipMemGenericAllocationHandle_t() mapped = False @@ -409,10 +407,12 @@ def allocate_exportable( "hipMemMap", ) mapped = True - self._mem_set_access( - int(access_va) if access_va is not None else mapped_va, - int(access_size) if access_size is not None else alloc_size, + access_base, access_bytes = ( + placement.access_range(alloc_size, cumulative=True) + if placement is not None + else (mapped_va, alloc_size) ) + self._mem_set_access(access_base, access_bytes) return LocalAllocation( va=mapped_va, size=alloc_size, @@ -489,25 +489,20 @@ def _export_range(self, va: int, size: int) -> bytes: pass raise - def export_handle(self, allocation: LocalAllocation) -> bytes: - """Export a 20-byte DMA-BUF descriptor for a local HIP allocation.""" + def export_handle(self, memory: ExportableMemory) -> bytes: + """Export a 20-byte DMA-BUF descriptor for a local HIP memory range.""" self._check_initialized() - return self._export_range(allocation.va, allocation.size) + return self._export_range(memory.va, memory.size) def import_and_map( self, peer_rank: int, handle_bytes: bytes, size: int, - va: Optional[int] = None, - *, - access_va: Optional[int] = None, - access_size: Optional[int] = None, + placement: Optional[MappingPlacement] = None, ) -> PeerMapping: """Import a DMA-BUF descriptor and map it into local GPU address space.""" self._check_initialized() - if (access_va is None) != (access_size is None): - raise LocalHipError("access_va and access_size must be provided together") if len(handle_bytes) != _AMD_HANDLE_BYTES: raise LocalHipError(f"AMD local handle must be {_AMD_HANDLE_BYTES} bytes, got {len(handle_bytes)}") @@ -515,8 +510,8 @@ def import_and_map( if size > base_size - offset: raise LocalHipError(f"Requested map size {size} exceeds imported base range {base_size} at offset {offset}") - if va is not None: - mapped_va = int(va) + if placement is not None: + mapped_va = int(placement.va) imported_handle = hipMemGenericAllocationHandle_t() mapped = False fd_open = True @@ -537,10 +532,8 @@ def import_and_map( "hipMemMap", ) mapped = True - self._mem_set_access( - int(access_va) if access_va is not None else mapped_va, - int(access_size) if access_size is not None else size, - ) + access_base, access_bytes = placement.access_range(size, cumulative=True) + self._mem_set_access(access_base, access_bytes) return PeerMapping( peer_rank=peer_rank, transport=InterconnectLevel.INTRA_NODE, @@ -628,7 +621,17 @@ def import_and_map( ) raise - def cleanup_import(self, mapping: PeerMapping) -> None: + def cleanup(self, target: LocalAllocation | PeerMapping) -> None: + """Release a local HIP allocation or imported HIP mapping.""" + if isinstance(target, LocalAllocation): + self._cleanup_local(target) + return + if isinstance(target, PeerMapping): + self._cleanup_import(target) + return + raise LocalHipError(f"Unsupported cleanup target: {type(target).__name__}") + + def _cleanup_import(self, mapping: PeerMapping) -> None: """Release an imported HIP external-memory mapping.""" self._check_initialized() if ( @@ -662,7 +665,7 @@ def cleanup_import(self, mapping: PeerMapping) -> None: logger.warning("hipDestroyExternalMemory failed during import cleanup", exc_info=True) raise - def cleanup_local(self, allocation: LocalAllocation) -> None: + def _cleanup_local(self, allocation: LocalAllocation) -> None: """Unmap, release, and free a local HIP VMem allocation.""" self._check_initialized() steps = [ @@ -734,8 +737,3 @@ def get_address_range(self, ptr: int) -> tuple[int, int]: if base_ptr.value is None: raise LocalHipError("hipMemGetAddressRange returned a null base pointer") return int(base_ptr.value), int(base_size.value) - - def export_pointer_handle(self, ptr: int, size: int) -> bytes: - """Export a DMA-BUF descriptor for an arbitrary HIP device pointer range.""" - self._check_initialized() - return self._export_range(int(ptr), int(size)) diff --git a/iris/drivers/local/nvidia.py b/iris/drivers/local/nvidia.py index 3f4530012..f6f49dde2 100644 --- a/iris/drivers/local/nvidia.py +++ b/iris/drivers/local/nvidia.py @@ -16,7 +16,9 @@ BaseDriver, DriverError, DriverNotSupported, + ExportableMemory, LocalAllocation, + MappingPlacement, PeerMapping, ) from iris.host.distributed.topology import InterconnectLevel @@ -375,26 +377,22 @@ def initialize(self, device_ordinal: int) -> None: def allocate_exportable( self, size: int, - va: Optional[int] = None, - *, - access_va: Optional[int] = None, - access_size: Optional[int] = None, + placement: Optional[MappingPlacement] = None, ) -> LocalAllocation: """ Allocate CUDA VMM memory exportable as a POSIX FD. - If va is supplied, the caller must already own a sufficiently large, - granularity-aligned VA range containing [va, va + size). + If placement is supplied, the caller must already own a sufficiently + large, granularity-aligned VA range containing [placement.va, + placement.va + size). """ self._check_initialized() - if (access_va is None) != (access_size is None): - raise LocalCudaError("access_va and access_size must be provided together") props = self._make_alloc_props() granularity = self._get_granularity() alloc_size = _round_up(size, granularity) - reserved_va = va is None - mapped_va = int(va) if va is not None else 0 + reserved_va = placement is None + mapped_va = int(placement.va) if placement is not None else 0 handle = ctypes.c_uint64() mapped = False @@ -415,10 +413,10 @@ def allocate_exportable( "cuMemMap", ) mapped = True - self._mem_set_access( - int(access_va) if access_va is not None else mapped_va, - int(access_size) if access_size is not None else alloc_size, + access_base, access_bytes = ( + placement.access_range(alloc_size) if placement is not None else (mapped_va, alloc_size) ) + self._mem_set_access(access_base, access_bytes) return LocalAllocation( va=mapped_va, size=alloc_size, @@ -454,14 +452,12 @@ def allocate_exportable( _cleanup_after_failure(*steps) raise - def export_handle(self, allocation: LocalAllocation) -> bytes: - """Export a 4-byte native-endian POSIX-FD descriptor for a local allocation.""" - self._check_initialized() + def _export_allocation_handle(self, handle: int) -> bytes: fd = ctypes.c_int(-1) _cuda_try( _cuda_driver.cuMemExportToShareableHandle( ctypes.byref(fd), - int(allocation.handle), + int(handle), _CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0, ), @@ -469,9 +465,11 @@ def export_handle(self, allocation: LocalAllocation) -> bytes: ) return struct.pack(_CUDA_HANDLE_FMT, int(fd.value)) - def export_pointer_handle(self, ptr: int, size: int) -> bytes: - """Export the VMM allocation containing ptr as a 4-byte native-endian POSIX FD.""" + def export_handle(self, memory: ExportableMemory) -> bytes: + """Export a 4-byte native-endian POSIX-FD descriptor for a local memory range.""" self._check_initialized() + if memory.allocation is not None: + return self._export_allocation_handle(int(memory.allocation.handle)) retain_handle = getattr(_cuda_driver, "cuMemRetainAllocationHandle", None) if retain_handle is None: @@ -480,7 +478,7 @@ def export_pointer_handle(self, ptr: int, size: int) -> bytes: handle = ctypes.c_uint64() try: _cuda_try( - retain_handle(ctypes.byref(handle), ctypes.c_void_p(ptr)), + retain_handle(ctypes.byref(handle), ctypes.c_void_p(memory.va)), "cuMemRetainAllocationHandle", ) except LocalCudaError as exc: @@ -489,20 +487,10 @@ def export_pointer_handle(self, ptr: int, size: int) -> bytes: ) from exc try: - fd = ctypes.c_int(-1) try: - _cuda_try( - _cuda_driver.cuMemExportToShareableHandle( - ctypes.byref(fd), - handle.value, - _CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, - 0, - ), - "cuMemExportToShareableHandle", - ) + return self._export_allocation_handle(handle.value) except LocalCudaError as exc: raise LocalCudaNotSupported("CUDA could not export the retained allocation handle") from exc - return struct.pack(_CUDA_HANDLE_FMT, int(fd.value)) finally: _cuda_try(_cuda_driver.cuMemRelease(handle.value), "cuMemRelease") @@ -529,20 +517,15 @@ def import_and_map( peer_rank: int, handle_bytes: bytes, size: int, - va: Optional[int] = None, - *, - access_va: Optional[int] = None, - access_size: Optional[int] = None, + placement: Optional[MappingPlacement] = None, ) -> PeerMapping: """Import a POSIX-FD handle and map it into local CUDA VMM VA space.""" self._check_initialized() - if (access_va is None) != (access_size is None): - raise LocalCudaError("access_va and access_size must be provided together") imported_handle = self._import_handle(handle_bytes) granularity = self._get_granularity() - va_owned = va is None - mapped_va = int(va) if va is not None else 0 + va_owned = placement is None + mapped_va = int(placement.va) if placement is not None else 0 mapped = False try: if va_owned: @@ -557,10 +540,8 @@ def import_and_map( "cuMemMap", ) mapped = True - self._mem_set_access( - int(access_va) if access_va is not None else mapped_va, - int(access_size) if access_size is not None else size, - ) + access_base, access_bytes = placement.access_range(size) if placement is not None else (mapped_va, size) + self._mem_set_access(access_base, access_bytes) except Exception: steps: list[tuple[str, Callable[[], None]]] = [] if mapped: @@ -598,7 +579,17 @@ def import_and_map( _driver_handle=(tag, imported_handle), ) - def cleanup_import(self, mapping: PeerMapping) -> None: + def cleanup(self, target: LocalAllocation | PeerMapping) -> None: + """Release a local CUDA allocation or imported CUDA mapping.""" + if isinstance(target, LocalAllocation): + self._cleanup_local(target) + return + if isinstance(target, PeerMapping): + self._cleanup_import(target) + return + raise LocalCudaError(f"Unsupported cleanup target: {type(target).__name__}") + + def _cleanup_import(self, mapping: PeerMapping) -> None: """Unmap, release, and free an imported CUDA VMM mapping.""" self._check_initialized() if isinstance(mapping._driver_handle, tuple) and len(mapping._driver_handle) == 2: @@ -632,7 +623,7 @@ def cleanup_import(self, mapping: PeerMapping) -> None: ) _run_cleanup_steps(*steps) - def cleanup_local(self, allocation: LocalAllocation) -> None: + def _cleanup_local(self, allocation: LocalAllocation) -> None: """Unmap, release, and conditionally free a local CUDA VMM allocation.""" self._check_initialized() steps = [ diff --git a/iris/host/distributed/topology.py b/iris/host/distributed/topology.py index 20da03e94..b0e7256c9 100644 --- a/iris/host/distributed/topology.py +++ b/iris/host/distributed/topology.py @@ -8,7 +8,8 @@ import os import re import socket -import hashlib +import ctypes +import ctypes.util from dataclasses import dataclass, field from enum import IntEnum from typing import Any, Dict, List, Optional, Set, Tuple @@ -242,6 +243,72 @@ def _amd_get_gpu_fabric_info(gpu_id: int, pci_bus_id: str = "") -> FabricInfo: return FabricInfo() +class _NvmlGpuFabricInfoV2(ctypes.Structure): + _fields_ = [ + ("version", ctypes.c_uint), + ("clusterUuid", ctypes.c_ubyte * 16), + ("status", ctypes.c_int), + ("cliqueId", ctypes.c_uint), + ("state", ctypes.c_ubyte), + ("healthMask", ctypes.c_uint), + ] + + def __init__(self) -> None: + super().__init__() + self.version = (2 << 24) | ctypes.sizeof(type(self)) + + +class _NvmlGpuFabricInfoV3(ctypes.Structure): + _fields_ = [ + ("version", ctypes.c_uint), + ("clusterUuid", ctypes.c_ubyte * 16), + ("status", ctypes.c_int), + ("cliqueId", ctypes.c_uint), + ("state", ctypes.c_ubyte), + ("healthMask", ctypes.c_uint), + ("healthSummary", ctypes.c_ubyte), + ] + + def __init__(self) -> None: + super().__init__() + self.version = (3 << 24) | ctypes.sizeof(type(self)) + + +def _fabric_info_from_nvml_struct(fabric_info) -> FabricInfo: + # Check registration state — must be COMPLETED (value 3) + state = getattr(fabric_info, "state", None) + if state is not None and int(state) != 3: + return FabricInfo() + + # Check status — must be SUCCESS (value 0) + status = getattr(fabric_info, "status", None) + if status is not None and int(status) != 0: + return FabricInfo() + + cluster_uuid_raw = getattr(fabric_info, "clusterUuid", None) + if cluster_uuid_raw is None: + return FabricInfo() + + if isinstance(cluster_uuid_raw, bytes): + cluster_uuid_hex = cluster_uuid_raw.hex() + elif isinstance(cluster_uuid_raw, (list, tuple)): + cluster_uuid_hex = bytes(cluster_uuid_raw).hex() + else: + try: + cluster_uuid_hex = bytes(cluster_uuid_raw).hex() + except TypeError: + cluster_uuid_hex = str(cluster_uuid_raw) + + if all(c == "0" for c in cluster_uuid_hex): + return FabricInfo() + + clique_id = getattr(fabric_info, "cliqueId", 0) + return FabricInfo( + cluster_uuid=cluster_uuid_hex, + clique_id=int(clique_id), + ) + + def _nvidia_get_gpu_fabric_info(gpu_id: int, pci_bus_id: str = "") -> FabricInfo: """ Get GPU fabric info from NVIDIA's NVML library. @@ -266,49 +333,20 @@ def _nvidia_get_gpu_fabric_info(gpu_id: int, pci_bus_id: str = "") -> FabricInfo physical_idx = _logical_to_physical_gpu_index(gpu_id, "nvidia") handle = pynvml.nvmlDeviceGetHandleByIndex(physical_idx) - fabric_info = None - try: - info_struct = pynvml.c_nvmlGpuFabricInfo_v2_t() - pynvml.nvmlDeviceGetGpuFabricInfoV(handle, info_struct) - fabric_info = info_struct - except (AttributeError, TypeError, pynvml.NVMLError): - # GPU doesn't support fabric - return FabricInfo() - - if fabric_info is None: - return FabricInfo() - - # Check registration state — must be COMPLETED (value 3) - state = getattr(fabric_info, "state", None) - if state is not None and state != 3: - return FabricInfo() - - # Check status — must be SUCCESS (value 0) - status = getattr(fabric_info, "status", None) - if status is not None and status != 0: - return FabricInfo() - - # Extract clusterUuid - cluster_uuid_raw = getattr(fabric_info, "clusterUuid", None) - if cluster_uuid_raw is None: - return FabricInfo() - - if isinstance(cluster_uuid_raw, bytes): - cluster_uuid_hex = cluster_uuid_raw.hex() - elif isinstance(cluster_uuid_raw, (list, tuple)): - cluster_uuid_hex = bytes(cluster_uuid_raw).hex() - else: - cluster_uuid_hex = str(cluster_uuid_raw) - - if all(c == "0" for c in cluster_uuid_hex): - return FabricInfo() - - clique_id = getattr(fabric_info, "cliqueId", 0) + nvml_path = ctypes.util.find_library("nvidia-ml") or "libnvidia-ml.so.1" + nvml = ctypes.CDLL(nvml_path) + get_fabric_info = nvml.nvmlDeviceGetGpuFabricInfoV + get_fabric_info.restype = ctypes.c_int - return FabricInfo( - cluster_uuid=cluster_uuid_hex, - clique_id=int(clique_id), - ) + for info_type in (_NvmlGpuFabricInfoV3, _NvmlGpuFabricInfoV2): + info = info_type() + get_fabric_info.argtypes = [ctypes.c_void_p, ctypes.POINTER(info_type)] + ret = get_fabric_info(handle, ctypes.byref(info)) + if ret != 0: + continue + fabric_info = _fabric_info_from_nvml_struct(info) + if fabric_info.is_valid: + return fabric_info except ImportError: logger.debug("pynvml not available, skipping NVML fabric info") except Exception as e: @@ -344,111 +382,6 @@ def _get_gpu_fabric_info(gpu_id: int, vendor: str, pci_bus_id: str = "") -> Fabr return _nvidia_get_gpu_fabric_info(gpu_id, pci_bus_id=pci_bus_id) -def _probe_nvidia_fabric_connectivity(gpu_id: int, rank: int, world_size: int) -> Optional[List[List[bool]]]: - """ - Probe NVIDIA fabric reachability with CUDA fabric memory handles. - - Some GB200/MNNVL environments expose working fabric handles while NVML - reports an empty GPU fabric UUID. This collective probe uses the same - driver interface as Iris memory sharing: each rank exports a tiny fabric - allocation and every other rank tries to import it. A successful symmetric - import means the ranks share an NVIDIA fabric memory domain. - """ - if not dist.is_initialized() or world_size <= 1: - return None - - local_record: dict[str, Any] = { - "rank": rank, - "ok": False, - "handle": b"", - "size": 0, - } - driver = None - allocation = None - imported_mappings = [] - - try: - from iris.drivers.fabric.nvidia import NvidiaFabricDriver - - driver = NvidiaFabricDriver() - driver.initialize(gpu_id) - size = driver.get_minimum_granularity() - allocation = driver.allocate_exportable(size) - local_record = { - "rank": rank, - "ok": True, - "handle": driver.export_handle(allocation), - "size": allocation.size, - } - except Exception as exc: - logger.debug("[Rank %d] CUDA fabric handle export probe failed: %s", rank, exc) - - records: List[Optional[dict[str, Any]]] = [None] * world_size - dist.all_gather_object(records, local_record) - - local_row = [False] * world_size - for record in records: - if not record or not record.get("ok"): - continue - peer_rank = int(record["rank"]) - if peer_rank == rank: - local_row[peer_rank] = True - continue - if driver is None: - continue - try: - mapping = driver.import_and_map(peer_rank, record["handle"], int(record["size"])) - imported_mappings.append(mapping) - local_row[peer_rank] = True - except Exception as exc: - logger.debug("[Rank %d] CUDA fabric handle import from rank %d failed: %s", rank, peer_rank, exc) - - rows: List[Optional[List[bool]]] = [None] * world_size - dist.all_gather_object(rows, local_row) - - if driver is not None: - for mapping in imported_mappings: - try: - driver.cleanup_import(mapping) - except Exception as exc: - logger.debug("[Rank %d] CUDA fabric probe import cleanup failed: %s", rank, exc) - if allocation is not None: - try: - driver.cleanup_local(allocation) - except Exception as exc: - logger.debug("[Rank %d] CUDA fabric probe local cleanup failed: %s", rank, exc) - - if any(row is None for row in rows): - return None - return [list(row) for row in rows if row is not None] - - -def _fabric_components_from_connectivity(connectivity: List[List[bool]]) -> List[Set[int]]: - """Return bidirectionally reachable components from a fabric probe matrix.""" - world_size = len(connectivity) - visited: Set[int] = set() - components: List[Set[int]] = [] - - for start in range(world_size): - if start in visited: - continue - stack = [start] - component: Set[int] = set() - visited.add(start) - while stack: - rank = stack.pop() - component.add(rank) - for peer in range(world_size): - if peer in visited: - continue - if connectivity[rank][peer] and connectivity[peer][rank]: - visited.add(peer) - stack.append(peer) - components.append(component) - - return components - - def _normalize_pci_bus_id(bus_id: str) -> str: """ Normalize a PCI bus ID to a canonical lowercase form for comparison. @@ -1331,30 +1264,6 @@ def discover(self) -> TopologyMap: info = GPUInfo.from_dict(json.loads(gpu_json)) gpu_info_map[info.global_rank] = info - if ( - vendor == "nvidia" - and self.world_size > 1 - and any(not info.fabric_info.domain_key for info in gpu_info_map.values()) - ): - connectivity = _probe_nvidia_fabric_connectivity(self.gpu_id, self.rank, self.world_size) - if connectivity is not None: - for component in _fabric_components_from_connectivity(connectivity): - if len(component) <= 1: - continue - component_uuids = sorted(gpu_info_map[r].uuid for r in component) - cluster_uuid = "cuda-probe-" + hashlib.sha1(",".join(component_uuids).encode()).hexdigest()[:16] - for component_rank in component: - if not gpu_info_map[component_rank].fabric_info.domain_key: - gpu_info_map[component_rank].fabric_info = FabricInfo( - cluster_uuid=cluster_uuid, - clique_id=0, - ) - logger.info( - "Detected NVIDIA fabric domain via CUDA fabric handle probe: ranks=%s domain=%s", - sorted(component), - cluster_uuid, - ) - all_node_infos = [json.loads(s) for s in all_node_jsons] # Group ranks by hostname diff --git a/iris/host/memory/allocators/vmem_chunked_allocator.py b/iris/host/memory/allocators/vmem_chunked_allocator.py index 62e922902..e66db02df 100644 --- a/iris/host/memory/allocators/vmem_chunked_allocator.py +++ b/iris/host/memory/allocators/vmem_chunked_allocator.py @@ -24,14 +24,13 @@ import logging import weakref from collections import defaultdict, deque -from dataclasses import dataclass from threading import Lock from typing import List, Optional import torch from .base import BaseAllocator -from iris.drivers.base import DriverNotSupported, LocalAllocation, PeerMapping +from iris.drivers.base import DriverNotSupported, ExportableMemory, LocalAllocation, MappingPlacement, PeerMapping from iris.drivers.factory import DriverFactory from iris.host.distributed.topology import ( InterconnectLevel, @@ -81,15 +80,6 @@ def _is_power_of_two(n: int) -> bool: return n > 0 and (n & (n - 1)) == 0 -@dataclass -class _SharedRegion: - """Exported heap region tracked for peer refresh.""" - - va: int - size: int - allocation: Optional[LocalAllocation] = None - - class VMemChunkedAllocator(BaseAllocator): """ Chunked VMem allocator with power-of-two free lists. @@ -160,7 +150,7 @@ def __init__( self.alloc_sizes = {} self._pending_free = deque() self.chunks: List[LocalAllocation] = [] - self._shared_regions: List[_SharedRegion] = [] + self._shared_regions: List[ExportableMemory] = [] self._peer_mappings: List[PeerMapping] = [] self._imported_heap_mappings: List[PeerMapping] = [] self.mapped_extent = 0 @@ -259,15 +249,10 @@ def _grow_chunk(self): ) target_va = self.base_va + self.mapped_extent - alloc_kwargs = {} - if self.driver.__class__.__name__ == "LocalHipDriver": - alloc_kwargs = { - "access_va": self.base_va, - "access_size": self.mapped_extent + self.chunk_size, - } - allocation = self.driver.allocate_exportable(self.chunk_size, va=target_va, **alloc_kwargs) + placement = MappingPlacement(target_va, arena_base=self.base_va) + allocation = self.driver.allocate_exportable(self.chunk_size, placement) self.chunks.append(allocation) - self._shared_regions.append(_SharedRegion(va=allocation.va, size=allocation.size, allocation=allocation)) + self._shared_regions.append(ExportableMemory(va=allocation.va, size=allocation.size, allocation=allocation)) self.mapped_extent += self.chunk_size def _process_pending_frees(self): @@ -437,10 +422,7 @@ def get_allocation_chunks_since(self, start_index: int): for i in range(start_index, len(self._shared_regions)): region = self._shared_regions[i] offset = region.va - self.base_va - if region.allocation is not None: - handle_bytes = self.driver.export_handle(region.allocation) - else: - handle_bytes = self.driver.export_pointer_handle(region.va, region.size) + handle_bytes = self.driver.export_handle(region) result.append((i, offset, region.size, handle_bytes)) return result @@ -503,27 +485,21 @@ def import_external_tensor(self, external_tensor: torch.Tensor) -> torch.Tensor: target_base_va = self.base_va + target_offset try: - handle_bytes = self.driver.export_pointer_handle(alloc_base, alloc_size) + handle_bytes = self.driver.export_handle(ExportableMemory(alloc_base, alloc_size)) except DriverNotSupported: if self._can_return_external_tensor_alias(): return self._external_tensor_alias(external_tensor) raise - import_kwargs = {} - if self.driver.__class__.__name__ == "LocalHipDriver": - import_kwargs = { - "access_va": self.base_va, - "access_size": target_offset + aligned_alloc_size, - } + placement = MappingPlacement(target_base_va, arena_base=self.base_va) mapping = self.driver.import_and_map( self.cur_rank, handle_bytes, aligned_alloc_size, - va=target_base_va, - **import_kwargs, + placement, ) self._imported_heap_mappings.append(mapping) - self._shared_regions.append(_SharedRegion(va=target_base_va, size=aligned_alloc_size, allocation=None)) + self._shared_regions.append(ExportableMemory(va=target_base_va, size=aligned_alloc_size)) self.mapped_extent = target_offset + aligned_alloc_size self.bump = max(self.bump, self.mapped_extent) @@ -539,7 +515,7 @@ def _import_release_callback(self, mapping: PeerMapping) -> None: combined with using self._peer_mappings.remove(mapping) as the gate, is what makes cleanup race-free against close() and release_peer_chunk: only one code path can successfully remove a given mapping from the - list, and that code path owns the cleanup_import call. + list, and that code path owns the cleanup call. The self._closed check before the lock is a fast-path optimization only -- it is NOT a correctness gate. On weakly-ordered architectures the @@ -563,9 +539,9 @@ def _import_release_callback(self, mapping: PeerMapping) -> None: except ValueError: return try: - self.driver.cleanup_import(mapping) + self.driver.cleanup(mapping) except Exception as exc: - logger.warning("cleanup_import failed in finalizer: %s", exc) + logger.warning("cleanup failed in finalizer: %s", exc) def import_peer_chunk(self, peer_rank: int, handle_bytes: bytes, size: int) -> int: """ @@ -575,7 +551,7 @@ def import_peer_chunk(self, peer_rank: int, handle_bytes: bytes, size: int) -> i The caller is responsible for calling release_peer_chunk when done. """ with self.lock: - mapping = self.driver.import_and_map(peer_rank, handle_bytes, size, va=None) + mapping = self.driver.import_and_map(peer_rank, handle_bytes, size) self._peer_mappings.append(mapping) return mapping.remote_va @@ -592,7 +568,7 @@ def release_peer_chunk(self, remote_va: int) -> None: for i, mapping in enumerate(self._peer_mappings): if mapping.remote_va == remote_va: try: - self.driver.cleanup_import(mapping) + self.driver.cleanup(mapping) finally: self._peer_mappings.pop(i) return @@ -640,7 +616,7 @@ def close(self): # Release imported peer mappings for mapping in self._peer_mappings: try: - self.driver.cleanup_import(mapping) + self.driver.cleanup(mapping) except Exception: pass self._peer_mappings.clear() @@ -649,18 +625,18 @@ def close(self): # heap VA layout via import_external_tensor. for mapping in self._imported_heap_mappings: try: - self.driver.cleanup_import(mapping) + self.driver.cleanup(mapping) except Exception: pass self._imported_heap_mappings.clear() self._shared_regions.clear() # Release locally-mapped chunks. Each chunk has _va_owned=False, - # so cleanup_local will unmap and release the physical handle + # so cleanup will unmap and release the physical handle # but will NOT free VA for alloc in self.chunks: try: - self.driver.cleanup_local(alloc) + self.driver.cleanup(alloc) except Exception: pass self.chunks.clear() diff --git a/iris/host/memory/symmetric_heap.py b/iris/host/memory/symmetric_heap.py index 28165288b..2620c18e7 100644 --- a/iris/host/memory/symmetric_heap.py +++ b/iris/host/memory/symmetric_heap.py @@ -17,7 +17,7 @@ from iris.host.logging.logging import _log_rank from iris.host.memory.allocators import TorchAllocator, VMemAllocator, VMemChunkedAllocator -from iris.drivers.base import PeerMapping +from iris.drivers.base import MappingPlacement, PeerMapping from iris.host.distributed.fd_passing import setup_fd_infrastructure from iris.host.distributed.helpers import distributed_allgather from iris.host.platform.utils import is_simulation_env @@ -522,18 +522,12 @@ def _refresh_peer_access_chunked(self, dist): # cleanup path below only closes never-consumed FDs. pending_cloned_fds.pop(0) reconstructed_handle = _replace_fd_in_local_handle(peer_handle_bytes, cloned_fd) - import_kwargs = {} - if len(peer_handle_bytes) == _LOCAL_HIP_HANDLE_BYTES: - import_kwargs = { - "access_va": peer_va_base, - "access_size": peer_offset + peer_size, - } + placement = MappingPlacement(peer_va_base + peer_offset, arena_base=peer_va_base) mapping = self.allocator.driver.import_and_map( peer, reconstructed_handle, peer_size, - va=peer_va_base + peer_offset, - **import_kwargs, + placement, ) self._peer_imported_mappings[peer].append(mapping) @@ -638,7 +632,7 @@ def _refresh_peer_access_fabric(self, dist): peer, handle_bytes, chunk_size, - va=peer_va_base + chunk_offset, + MappingPlacement(peer_va_base + chunk_offset, arena_base=peer_va_base), ) self._peer_imported_mappings[peer].append(mapping) @@ -796,7 +790,7 @@ def close(self): for entry in mappings: try: if has_driver and isinstance(entry, PeerMapping): - self.allocator.driver.cleanup_import(entry) + self.allocator.driver.cleanup(entry) else: from iris.host.platform.hip import mem_release, mem_unmap diff --git a/tests/unittests/test_drivers.py b/tests/unittests/test_drivers.py index 428063e63..4710286ea 100644 --- a/tests/unittests/test_drivers.py +++ b/tests/unittests/test_drivers.py @@ -15,6 +15,7 @@ from iris.drivers.base import ( DriverError, DriverNotSupported, + ExportableMemory, LocalAllocation, PeerMapping, ) @@ -94,10 +95,10 @@ class TestNvidiaFabricDriver: ("method_name", "args"), [ ("allocate_exportable", (4096,)), - ("export_handle", (LocalAllocation(va=0, size=0, handle=0),)), + ("export_handle", (ExportableMemory(va=0, size=0, allocation=LocalAllocation(va=0, size=0, handle=0)),)), ("import_and_map", (0, b"\x00" * FABRIC_HANDLE_BYTES, 4096)), ( - "cleanup_import", + "cleanup", ( PeerMapping( peer_rank=0, @@ -107,7 +108,7 @@ class TestNvidiaFabricDriver: ), ), ), - ("cleanup_local", (LocalAllocation(va=0, size=0, handle=0),)), + ("cleanup", (LocalAllocation(va=0, size=0, handle=0),)), ], ) def test_public_methods_require_initialize(self, method_name, args): @@ -131,7 +132,7 @@ class IncompleteCudaDriver: with pytest.raises(CudaFabricNotSupported, match="missing required VMM symbol: cuInit"): driver.initialize(0) - def test_cleanup_import_attempts_all_cleanup_steps(self, monkeypatch): + def test_cleanup_import_target_attempts_all_cleanup_steps(self, monkeypatch): calls = [] class FakeCudaDriver: @@ -159,7 +160,7 @@ def cuMemAddressFree(self, remote_va, size): ) with pytest.raises(CudaFabricError, match="cuMemUnmap"): - driver.cleanup_import(mapping) + driver.cleanup(mapping) assert calls == [ ("unmap", 0x2000, 4096), @@ -167,7 +168,7 @@ def cuMemAddressFree(self, remote_va, size): ("free", 0x2000, 4096), ] - def test_cleanup_local_attempts_all_cleanup_steps(self, monkeypatch): + def test_cleanup_local_target_attempts_all_cleanup_steps(self, monkeypatch): calls = [] class FakeCudaDriver: @@ -189,7 +190,7 @@ def cuMemAddressFree(self, va, size): allocation = LocalAllocation(va=0x1000, size=8192, handle=77) with pytest.raises(CudaFabricError, match="cuMemUnmap"): - driver.cleanup_local(allocation) + driver.cleanup(allocation) assert calls == [ ("unmap", 0x1000, 8192), diff --git a/tests/unittests/test_topology.py b/tests/unittests/test_topology.py index 5ec3062f9..b12b35d9b 100644 --- a/tests/unittests/test_topology.py +++ b/tests/unittests/test_topology.py @@ -11,6 +11,7 @@ import json import builtins +import ctypes import logging import socket import sys @@ -365,12 +366,65 @@ def test_get_gpu_fabric_info_uses_existing_vendor_lifecycle(self, monkeypatch): fake_pynvml = _make_fake_pynvml() monkeypatch.setitem(sys.modules, "pynvml", fake_pynvml) + class FakeNvmlFunction: + def __init__(self): + self.restype = None + self.argtypes = None + + def __call__(self, handle, info_ptr): + assert handle == "gpu0" + info = info_ptr._obj + for idx, value in enumerate(bytes.fromhex("01" * 16)): + info.clusterUuid[idx] = value + info.status = 0 + info.cliqueId = 7 + info.state = 3 + return 0 + + class FakeNvmlLib: + def __init__(self): + self.nvmlDeviceGetGpuFabricInfoV = FakeNvmlFunction() + + monkeypatch.setattr(topology.ctypes.util, "find_library", lambda name: "fake-nvml") + monkeypatch.setattr(topology.ctypes, "CDLL", lambda path: FakeNvmlLib()) + fabric = topology._get_gpu_fabric_info(0, "nvidia", pci_bus_id="0000:41:00.0") assert fabric.domain_key == f"{'01' * 16}:7" assert fake_pynvml.init_calls == 0 assert fake_pynvml.shutdown_calls == 0 + def test_nvidia_fabric_info_uses_direct_nvml_byref(self, monkeypatch): + fake_pynvml = _make_fake_pynvml() + monkeypatch.setitem(sys.modules, "pynvml", fake_pynvml) + + class FakeNvmlFunction: + def __init__(self): + self.restype = None + self.argtypes = None + + def __call__(self, handle, info_ptr): + assert handle == "gpu0" + assert self.argtypes[0] is ctypes.c_void_p + info = info_ptr._obj + for idx, value in enumerate(bytes.fromhex("ab" * 16)): + info.clusterUuid[idx] = value + info.status = 0 + info.cliqueId = 32766 + info.state = 3 + return 0 + + class FakeNvmlLib: + def __init__(self): + self.nvmlDeviceGetGpuFabricInfoV = FakeNvmlFunction() + + monkeypatch.setattr(topology.ctypes.util, "find_library", lambda name: "fake-nvml") + monkeypatch.setattr(topology.ctypes, "CDLL", lambda path: FakeNvmlLib()) + + fabric = topology._get_gpu_fabric_info(0, "nvidia", pci_bus_id="0000:41:00.0") + + assert fabric.domain_key == f"{'ab' * 16}:32766" + def test_missing_library_fallbacks_preserve_logs(self, monkeypatch, caplog): caplog.set_level(logging.DEBUG, logger="iris.topology") real_import = builtins.__import__