Restore v0.3.x MemoryResource behaviors + MR minor follow-ups (NVIDIA#1388)

* add back missing MR attributes to restore/mimick v0.3.x behavior

* update tests to avoid deprecation warnings

* consolidate common mempool attributes + protect multi-threading warnings

* fix linter

Author: leofang

cuda_core/cuda/core/experimental/_memory/_buffer.pyx

@@ -16,7 +16,6 @@ from cuda.core.experimental._memory cimport _ipc
 from cuda.core.experimental._stream cimport Stream_accept, Stream
 from cuda.core.experimental._utils.cuda_utils cimport HANDLE_RETURN
 
-import abc
 import sys
 from typing import TypeVar, Union
 
@@ -492,7 +491,6 @@ cdef class MemoryResource:
     resource's respective property.)
     """
 
-    @abc.abstractmethod
     def allocate(self, size_t size, stream: Stream | GraphBuilder | None = None) -> Buffer:
         """Allocate a buffer of the requested size.
 
@@ -511,9 +509,8 @@ cdef class MemoryResource:
             The allocated buffer object, which can be used for device or host operations
             depending on the resource's properties.
         """
-        ...
+        raise TypeError("MemoryResource.allocate must be implemented by subclasses.")
 
-    @abc.abstractmethod
     def deallocate(self, ptr: DevicePointerT, size_t size, stream: Stream | GraphBuilder | None = None):
         """Deallocate a buffer previously allocated by this resource.
 
@@ -528,4 +525,19 @@ cdef class MemoryResource:
             If None, it is up to each memory resource implementation to decide
             and document the behavior.
         """
-        ...
+        raise TypeError("MemoryResource.deallocate must be implemented by subclasses.")
+
+    @property
+    def is_device_accessible(self) -> bool:
+        """Whether buffers allocated by this resource are device-accessible."""
+        raise TypeError("MemoryResource.is_device_accessible must be implemented by subclasses.")
+
+    @property
+    def is_host_accessible(self) -> bool:
+        """Whether buffers allocated by this resource are host-accessible."""
+        raise TypeError("MemoryResource.is_host_accessible must be implemented by subclasses.")
+
+    @property
+    def device_id(self) -> int:
+        """Device ID associated with this memory resource, or -1 if not applicable."""
+        raise TypeError("MemoryResource.device_id must be implemented by subclasses.")

cuda_core/cuda/core/experimental/_memory/_device_memory_resource.pyx

@@ -15,7 +15,7 @@ from cuda.core.experimental._utils.cuda_utils cimport (
 
 from dataclasses import dataclass
 import multiprocessing
-from typing import Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING
 import platform  # no-cython-lint
 import uuid
 
@@ -228,27 +228,6 @@ cdef class DeviceMemoryResource(_MemPool):
         """Return False. This memory resource does not provide host-accessible buffers."""
         return False
 
-    @property
-    def is_ipc_enabled(self) -> bool:
-        """Whether this memory resource has IPC enabled."""
-        return self._ipc_data is not None
-
-    @property
-    def is_mapped(self) -> bool:
-        """
-        Whether this is a mapping of an IPC-enabled memory resource from
-        another process.  If True, allocation is not permitted.
-        """
-        return self._ipc_data is not None and self._ipc_data._is_mapped
-
-    @property
-    def uuid(self) -> Optional[uuid.UUID]:
-        """
-        A universally unique identifier for this memory resource. Meaningful
-        only for IPC-enabled memory resources.
-        """
-        return getattr(self._ipc_data, 'uuid', None)
-
 
 # Note: this is referenced in instructions to debug nvbug 5698116.
 cpdef DMR_mempool_get_access(DeviceMemoryResource dmr, int device_id):

cuda_core/cuda/core/experimental/_memory/_managed_memory_resource.pyx

@@ -104,16 +104,3 @@ cdef class ManagedMemoryResource(_MemPool):
     def is_host_accessible(self) -> bool:
         """Return True. This memory resource provides host-accessible buffers."""
         return True
-
-    @property
-    def is_ipc_enabled(self) -> bool:
-        """Whether this memory resource has IPC enabled."""
-        return self._ipc_data is not None
-
-    @property
-    def is_mapped(self) -> bool:
-        """
-        Whether this is a mapping of an IPC-enabled memory resource from
-        another process.  If True, allocation is not permitted.
-        """
-        return self._ipc_data is not None and self._ipc_data._is_mapped

cuda_core/cuda/core/experimental/_memory/_memory_pool.pyx

@@ -278,6 +278,27 @@ cdef class _MemPool(MemoryResource):
 
             self._peer_accessible_by = tuple(target_ids)
 
+    @property
+    def is_ipc_enabled(self) -> bool:
+        """Whether this memory resource has IPC enabled."""
+        return self._ipc_data is not None
+
+    @property
+    def is_mapped(self) -> bool:
+        """
+        Whether this is a mapping of an IPC-enabled memory resource from
+        another process.  If True, allocation is not permitted.
+        """
+        return self._ipc_data is not None and self._ipc_data._is_mapped
+
+    @property
+    def uuid(self) -> Optional[uuid.UUID]:
+        """
+        A universally unique identifier for this memory resource. Meaningful
+        only for IPC-enabled memory resources.
+        """
+        return getattr(self._ipc_data, 'uuid', None)
+
 
 # _MemPool Implementation
 # -----------------------
@@ -406,7 +427,7 @@ cdef inline Buffer _MP_allocate(_MemPool self, size_t size, Stream stream):
 
 cdef inline void _MP_deallocate(
     _MemPool self, uintptr_t ptr, size_t size, Stream stream
-) noexcept:
+) noexcept nogil:
     cdef cydriver.CUstream s = stream._handle
     cdef cydriver.CUdeviceptr devptr = <cydriver.CUdeviceptr>ptr
     cdef cydriver.CUresult r

cuda_core/cuda/core/experimental/_memory/_pinned_memory_resource.pyx

@@ -14,10 +14,11 @@ from cuda.core.experimental._utils.cuda_utils cimport (
 )
 
 from dataclasses import dataclass
-from typing import Optional
 import multiprocessing
+import os
 import platform  # no-cython-lint
 import subprocess
+import threading
 import uuid
 import warnings
 
@@ -26,6 +27,7 @@ from cuda.core.experimental._utils.cuda_utils import check_multiprocessing_start
 
 # Cache to ensure NUMA warning is only raised once per process
 cdef bint _numa_warning_shown = False
+cdef object _lock = threading.Lock()
 
 
 def _check_numa_nodes():
@@ -34,49 +36,53 @@ def _check_numa_nodes():
     if _numa_warning_shown:
         return
 
-    if platform.system() != "Linux":
-        return
+    with _lock:
+        if _numa_warning_shown:
+            return
+
+        if platform.system() != "Linux":
+            _numa_warning_shown = True
+            return
 
-    numa_count = None
-
-    # Try /sys filesystem first (most reliable and doesn't require external tools)
-    try:
-        import os
-        node_path = "/sys/devices/system/node"
-        if os.path.exists(node_path):
-            # Count directories named "node[0-9]+"
-            nodes = [d for d in os.listdir(node_path) if d.startswith("node") and d[4:].isdigit()]
-            numa_count = len(nodes)
-    except (OSError, PermissionError):
-        pass
-
-    # Fallback to lscpu if /sys check didn't work
-    if numa_count is None:
+        numa_count = None
+
+        # Try /sys filesystem first (most reliable and doesn't require external tools)
         try:
-            result = subprocess.run(
-                ["lscpu"],
-                capture_output=True,
-                text=True,
-                timeout=1
-            )
-            for line in result.stdout.splitlines():
-                if line.startswith("NUMA node(s):"):
-                    numa_count = int(line.split(":")[1].strip())
-                    break
-        except (subprocess.SubprocessError, ValueError, FileNotFoundError):
+            node_path = "/sys/devices/system/node"
+            if os.path.exists(node_path):
+                # Count directories named "node[0-9]+"
+                nodes = [d for d in os.listdir(node_path) if d.startswith("node") and d[4:].isdigit()]
+                numa_count = len(nodes)
+        except (OSError, PermissionError):
             pass
 
-    # Warn if multiple NUMA nodes detected
-    if numa_count is not None and numa_count > 1:
-        warnings.warn(
-            f"System has {numa_count} NUMA nodes. IPC-enabled pinned memory "
-            f"uses location ID 0, which may not work correctly with multiple "
-            f"NUMA nodes.",
-            UserWarning,
-            stacklevel=3
-        )
+        # Fallback to lscpu if /sys check didn't work
+        if numa_count is None:
+            try:
+                result = subprocess.run(
+                    ["lscpu"],
+                    capture_output=True,
+                    text=True,
+                    timeout=1
+                )
+                for line in result.stdout.splitlines():
+                    if line.startswith("NUMA node(s):"):
+                        numa_count = int(line.split(":")[1].strip())
+                        break
+            except (subprocess.SubprocessError, ValueError, FileNotFoundError):
+                pass
+
+        # Warn if multiple NUMA nodes detected
+        if numa_count is not None and numa_count > 1:
+            warnings.warn(
+                f"System has {numa_count} NUMA nodes. IPC-enabled pinned memory "
+                f"uses location ID 0, which may not work correctly with multiple "
+                f"NUMA nodes.",
+                UserWarning,
+                stacklevel=3
+            )
 
-    _numa_warning_shown = True
+        _numa_warning_shown = True
 
 
 __all__ = ['PinnedMemoryResource', 'PinnedMemoryResourceOptions']
@@ -243,27 +249,6 @@ cdef class PinnedMemoryResource(_MemPool):
         """Return True. This memory resource provides host-accessible buffers."""
         return True
 
-    @property
-    def is_ipc_enabled(self) -> bool:
-        """Whether this memory resource has IPC enabled."""
-        return self._ipc_data is not None
-
-    @property
-    def is_mapped(self) -> bool:
-        """
-        Whether this is a mapping of an IPC-enabled memory resource from
-        another process.  If True, allocation is not permitted.
-        """
-        return self._ipc_data is not None and self._ipc_data._is_mapped
-
-    @property
-    def uuid(self) -> Optional[uuid.UUID]:
-        """
-        A universally unique identifier for this memory resource. Meaningful
-        only for IPC-enabled memory resources.
-        """
-        return getattr(self._ipc_data, 'uuid', None)
-
 
 def _deep_reduce_pinned_memory_resource(mr):
     check_multiprocessing_start_method()

cuda_core/tests/test_memory.py

@@ -172,6 +172,8 @@ def test_buffer_initialization():
     buffer_initialization(DummyHostMemoryResource())
     buffer_initialization(DummyUnifiedMemoryResource(device))
     buffer_initialization(DummyPinnedMemoryResource(device))
+    with pytest.raises(TypeError):
+        buffer_initialization(MemoryResource())
 
 
 def buffer_copy_to(dummy_mr: MemoryResource, device: Device, check=False):
@@ -365,7 +367,7 @@ def test_buffer_external_host():
 
 @pytest.mark.parametrize("change_device", [True, False])
 def test_buffer_external_device(change_device):
-    n = ccx_system.num_devices
+    n = ccx_system.get_num_devices()
     if n < 1:
         pytest.skip("No devices found")
     dev_id = n - 1
@@ -389,7 +391,7 @@ def test_buffer_external_device(change_device):
 
 @pytest.mark.parametrize("change_device", [True, False])
 def test_buffer_external_pinned_alloc(change_device):
-    n = ccx_system.num_devices
+    n = ccx_system.get_num_devices()
     if n < 1:
         pytest.skip("No devices found")
     dev_id = n - 1
@@ -414,7 +416,7 @@ def test_buffer_external_pinned_alloc(change_device):
 
 @pytest.mark.parametrize("change_device", [True, False])
 def test_buffer_external_pinned_registered(change_device):
-    n = ccx_system.num_devices
+    n = ccx_system.get_num_devices()
     if n < 1:
         pytest.skip("No devices found")
     dev_id = n - 1
@@ -447,7 +449,7 @@ def test_buffer_external_pinned_registered(change_device):
 
 @pytest.mark.parametrize("change_device", [True, False])
 def test_buffer_external_managed(change_device):
-    n = ccx_system.num_devices
+    n = ccx_system.get_num_devices()
     if n < 1:
         pytest.skip("No devices found")
     dev_id = n - 1