Simplify driver APIs and improve nvl72 fabric network detection#540
Merged
Conversation
github-actions
Bot
added
in-progress
Contributor
There was a problem hiding this comment.
Pull request overview
Note
Copilot couldn't run its full agentic review because no GitHub Actions runner was available. Make sure your repository has a runner available to run Copilot's review, or add a copilot-setup-steps.yml file specifying one with the runs-on attribute. See the docs for more details.
This PR simplifies the driver API for export/import/cleanup by introducing structured placement and export descriptors, and updates NVIDIA fabric discovery to query NVML directly via ctypes (including v3/v2 struct support) for improved NVL72 detection.
Changes:
- Replaced per-driver
va/access_va/access_sizeparameters withMappingPlacement, and unifiedcleanup_import/cleanup_localintocleanup(target). - Introduced
ExportableMemoryto standardize exporting either driver-created allocations or arbitrary local memory ranges (where supported). - Updated NVIDIA NVML fabric detection to use direct
ctypescalls and added unit tests for the new NVML path.
Reviewed changes
Copilot reviewed 11 out of 11 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unittests/test_topology.py | Adds tests to validate the direct-NVML (ctypes) fabric info path and byref usage. |
| tests/unittests/test_drivers.py | Updates tests to match the unified cleanup() API and ExportableMemory export API. |
| iris/host/memory/symmetric_heap.py | Switches peer mapping calls to MappingPlacement and unified cleanup(). |
| iris/host/memory/allocators/vmem_chunked_allocator.py | Refactors shared-region tracking to ExportableMemory, uses MappingPlacement, and updates cleanup/export calls. |
| iris/host/distributed/topology.py | Implements direct NVML fabric info query (v3/v2) and removes CUDA fabric-handle probing fallback. |
| iris/drivers/local/nvidia.py | Updates export/import/cleanup APIs to use ExportableMemory, MappingPlacement, and unified cleanup(). |
| iris/drivers/local/amd.py | Updates export/import/cleanup APIs similarly; removes export_pointer_handle. |
| iris/drivers/fabric/nvidia.py | Updates APIs to use ExportableMemory, MappingPlacement, and unified cleanup(). |
| iris/drivers/fabric/amd.py | Updates abstract API surface to match new base driver interfaces. |
| iris/drivers/base.py | Introduces ExportableMemory, MappingPlacement, and unified cleanup() in the base interface. |
| iris/drivers/init.py | Re-exports new driver API types. |
Comment on lines
+277
to
+309
| 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), | ||
| ) |
Comment on lines
+246
to
+254
| 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), | ||
| ] |
Comment on lines
+261
to
+270
| 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), | ||
| ] |
Comment on lines
+336
to
+339
| 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 |
Comment on lines
1267
to
1269
| all_node_infos = [json.loads(s) for s in all_node_jsons] | ||
|
|
||
| # Group ranks by hostname |
Comment on lines
+341
to
+349
| 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 |
mawad-amd
approved these changes
Jun 25, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Technical Details
Test Plan
Tested on NVL72 and AMD clusters with fabric and local network topologies.
Test Result
PASS
Submission Checklist