Simplify driver APIs and improve nvl72 fabric network detection (#540)

Co-authored-by: ayildiri <ayildiri@amd.com>

Author: artulab

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

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,37 +99,27 @@ 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(
         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 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")

iris/drivers/fabric/amd.py

@@ -12,7 +12,9 @@
 from iris.drivers.base import (
     BaseDriver,
     DriverNotSupported,
+    ExportableMemory,
     LocalAllocation,
+    MappingPlacement,
     PeerMapping,
 )
 
@@ -30,32 +32,23 @@ 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(
         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:
         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:

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 = [
             (