diff --git a/pymtml.py b/pymtml.py index aa7277c..75c6bc3 100644 --- a/pymtml.py +++ b/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 diff --git a/setup.py b/setup.py index f6ba8b7..56c301b 100644 --- a/setup.py +++ b/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', diff --git a/test_mtlink_visibility.py b/test_mtlink_visibility.py new file mode 100644 index 0000000..098ffe2 --- /dev/null +++ b/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