Fix MtLink P2P detection with partially visible GPUs

pymtml.py

@@ -2492,6 +2492,21 @@ def nvmlDeviceGetP2PStatus(device1, device2, p2pIndex):
             except MTMLError:
                 pass
 
+            # Prefer the pairwise layout API.  A process or container may expose
+            # only a subset of the GPUs on a host while MtLinkSpec still reports
+            # every physical link.  Walking those links asks MTML to materialize
+            # handles for hidden peers and legitimately returns NOT_FOUND.  The
+            # layout query operates only on the two visible handles supplied by
+            # the caller and avoids that ambiguity.
+            try:
+                link_count = mtmlDeviceCountMtLinkLayouts(device1, device2)
+                if link_count > 0:
+                    return NVML_P2P_STATUS_OK
+            except MTMLError:
+                # Older runtimes may not implement the layout API.  Retain the
+                # per-link scan below as a compatibility fallback.
+                pass
+
             # Check MtLink connectivity by iterating through links
             # This is the detailed check: for each link on device1, check if
             # the remote device matches device2's UUID
@@ -2517,14 +2532,6 @@ def nvmlDeviceGetP2PStatus(device1, device2, p2pIndex):
             except MTMLError:
                 pass
 
-            # Fallback: try mtmlDeviceCountMtLinkLayouts as a secondary check
-            try:
-                link_count = mtmlDeviceCountMtLinkLayouts(device1, device2)
-                if link_count > 0:
-                    return NVML_P2P_STATUS_OK
-            except MTMLError:
-                pass
-
             return NVML_P2P_STATUS_NOT_SUPPORTED
         else:
             # For other P2P caps, use MTML P2P status

setup.py

@@ -47,7 +47,7 @@
 
 
 setup(name=_package_name,
-      version='2.4.1',
+      version='2.4.2',
       description='Python Bindings for the Moore Threads GPU Management Library',
       long_description=long_description,
       long_description_content_type='text/markdown',

test_mtlink_visibility.py

@@ -0,0 +1,73 @@
+from types import SimpleNamespace
+
+import pymtml
+
+
+def test_nvlink_p2p_prefers_pairwise_layout_query(monkeypatch):
+    local = object()
+    peer = object()
+
+    monkeypatch.setattr(
+        pymtml,
+        "mtmlDeviceGetTopologyLevel",
+        lambda device1, device2: pymtml.MTML_TOPOLOGY_SINGLE,
+    )
+    monkeypatch.setattr(
+        pymtml, "mtmlDeviceCountMtLinkLayouts", lambda device1, device2: 2
+    )
+
+    def unexpected_link_scan(device):
+        raise AssertionError("pairwise layout detection should avoid link scanning")
+
+    monkeypatch.setattr(pymtml, "mtmlDeviceGetMtLinkSpec", unexpected_link_scan)
+
+    status = pymtml.nvmlDeviceGetP2PStatus(
+        local, peer, pymtml.NVML_P2P_CAPS_INDEX_NVLINK
+    )
+
+    assert status == pymtml.NVML_P2P_STATUS_OK
+
+
+def test_nvlink_p2p_skips_hidden_peers_when_layout_query_is_unavailable(monkeypatch):
+    local = object()
+    peer = object()
+    visible_peer = object()
+
+    monkeypatch.setattr(
+        pymtml,
+        "mtmlDeviceGetTopologyLevel",
+        lambda device1, device2: pymtml.MTML_TOPOLOGY_SINGLE,
+    )
+
+    def unsupported_layout_query(device1, device2):
+        raise pymtml.MTMLError_NotSupported(pymtml.MTML_ERROR_NOT_SUPPORTED)
+
+    monkeypatch.setattr(
+        pymtml, "mtmlDeviceCountMtLinkLayouts", unsupported_layout_query
+    )
+    monkeypatch.setattr(
+        pymtml, "mtmlDeviceGetMtLinkSpec", lambda device: SimpleNamespace(linkNum=2)
+    )
+    monkeypatch.setattr(
+        pymtml,
+        "mtmlDeviceGetMtLinkState",
+        lambda device, link: pymtml.MTML_MTLINK_STATE_UP,
+    )
+
+    def get_remote(device, link):
+        if link == 0:
+            raise pymtml.MTMLError_NotFound(pymtml.MTML_ERROR_NOT_FOUND)
+        return visible_peer
+
+    monkeypatch.setattr(pymtml, "mtmlDeviceGetMtLinkRemoteDevice", get_remote)
+    monkeypatch.setattr(
+        pymtml,
+        "mtmlDeviceGetUUID",
+        lambda device: b"visible-peer" if device in (peer, visible_peer) else b"hidden",
+    )
+
+    status = pymtml.nvmlDeviceGetP2PStatus(
+        local, peer, pymtml.NVML_P2P_CAPS_INDEX_NVLINK
+    )
+
+    assert status == pymtml.NVML_P2P_STATUS_OK