fix: load the CUDA runtime on Windows, and never let the P2P probe be fatal

p2p_registry.py

@@ -5,19 +5,49 @@
 """
 
 import ctypes
+import glob
 import logging
+import os
 import torch
 
 logger = logging.getLogger("MultiGPU")
 
 _libcudart = None
 
 
+def _cudart_candidates():
+    """CUDA runtime library names to try, most-specific first.
+
+    On Windows the runtime is cudart64_<major>.dll, not libcudart.so. Prefer the
+    copy PyTorch ships in torch/lib so the version always matches the build in
+    use; then PATH-resolved names; then the POSIX sonames.
+    """
+    names = []
+    if os.name == "nt":
+        try:
+            libdir = os.path.join(os.path.dirname(torch.__file__), "lib")
+            names.extend(sorted(glob.glob(os.path.join(libdir, "cudart64_*.dll")), reverse=True))
+        except Exception:
+            pass
+        names.extend(["cudart64_13.dll", "cudart64_12.dll", "cudart64_110.dll", "cudart64_101.dll"])
+    names.extend(["libcudart.so", "libcudart.dylib"])
+    return names
+
+
 def _get_libcudart():
-    """Load libcudart.so once and cache the handle."""
+    """Load the CUDA runtime once and cache the handle."""
     global _libcudart
     if _libcudart is None:
-        _libcudart = ctypes.CDLL("libcudart.so")
+        last_err = None
+        for name in _cudart_candidates():
+            try:
+                _libcudart = ctypes.CDLL(name)
+                logger.debug(f"[MultiGPU P2P] loaded CUDA runtime: {name}")
+                break
+            except OSError as e:
+                last_err = e
+        if _libcudart is None:
+            raise OSError(f"could not load CUDA runtime (tried {_cudart_candidates()})") from last_err
     return _libcudart
 
 
@@ -34,17 +64,44 @@ def __init__(self):
 
     @staticmethod
     def _raw_can_access_peer(device_a: int, device_b: int) -> bool:
-        """Call cudaDeviceCanAccessPeer via ctypes. Returns True if P2P is available."""
-        lib = _get_libcudart()
-        can_access = ctypes.c_int(0)
-        result = lib.cudaDeviceCanAccessPeer(ctypes.byref(can_access), device_a, device_b)
-        if result != 0:
+        """Return True if P2P is available between the two devices.
+
+        Prefers torch.cuda.can_device_access_peer, which does exist in torch 2.x
+        (the previous comment referred to "can_access_peer", which is not the
+        API name -- so this always fell through to ctypes).
+
+        This is an OPTIMIZATION probe: a False answer only means transfers get
+        staged through host memory. Every failure mode therefore degrades to
+        False rather than propagating -- an unloadable CUDA runtime must never
+        abort a render.
+        """
+        fn = getattr(torch.cuda, "can_device_access_peer", None)
+        if fn is not None:
+            try:
+                return bool(fn(device_a, device_b))
+            except Exception as e:
+                logger.warning(
+                    f"[MultiGPU P2P] torch.cuda.can_device_access_peer({device_a}, {device_b}) "
+                    f"failed ({e}); falling back to the CUDA runtime"
+                )
+
+        try:
+            lib = _get_libcudart()
+            can_access = ctypes.c_int(0)
+            result = lib.cudaDeviceCanAccessPeer(ctypes.byref(can_access), device_a, device_b)
+            if result != 0:
+                logger.warning(
+                    f"[MultiGPU P2P] cudaDeviceCanAccessPeer({device_a}, {device_b}) "
+                    f"returned error code {result}, assuming no P2P"
+                )
+                return False
+            return bool(can_access.value)
+        except Exception as e:
             logger.warning(
-                f"[MultiGPU P2P] cudaDeviceCanAccessPeer({device_a}, {device_b}) "
-                f"returned error code {result}, assuming no P2P"
+                f"[MultiGPU P2P] could not probe P2P for ({device_a}, {device_b}): {e}. "
+                f"Assuming no P2P; transfers will be staged through host memory."
             )
             return False
-        return bool(can_access.value)
 
     def can_access_peer(self, src_device: int, dst_device: int) -> bool:
         """Check if src_device can access dst_device memory via P2P.