Remove call to set_current in Device reconstruction. Add device set-up to tests.

Author: Andy-Jost

cuda_core/cuda/core/experimental/_device.py

@@ -2,7 +2,6 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-import multiprocessing
 import threading
 from typing import Optional, Union
 
@@ -1161,6 +1160,9 @@ def __int__(self):
     def __repr__(self):
         return f"<Device {self._id} ({self.name})>"
 
+    def __reduce__(self):
+        return Device, (self.device_id,)
+
     def set_current(self, ctx: Context = None) -> Union[Context, None]:
         """Set device to be used for GPU executions.
 
@@ -1335,17 +1337,3 @@ def create_graph_builder(self) -> GraphBuilder:
         """
         self._check_context_initialized()
         return GraphBuilder._init(stream=self.create_stream(), is_stream_owner=True)
-
-
-def _reduce_device(device):
-    return _reconstruct_device, (device.device_id,)
-
-
-def _reconstruct_device(device_id):
-    device = Device(device_id)
-    if not device._has_inited:
-        device.set_current()
-    return device
-
-
-multiprocessing.reduction.register(Device, _reduce_device)

cuda_core/cuda/core/experimental/_memory.pyx

@@ -923,7 +923,7 @@ class DeviceMemoryResource(MemoryResource):
 
 
 def _deep_reduce_device_memory_resource(mr):
-    from ._device import Device
+    from . import Device
     device = Device(mr.device_id)
     alloc_handle = mr.get_allocation_handle()
     return mr.from_allocation_handle, (device, alloc_handle)

cuda_core/tests/memory_ipc/test_errors.py

@@ -43,6 +43,7 @@ def test_main(self, ipc_device, ipc_memory_resource):
     def child_main(self, pipe, device, mr):
         """Child process that pushes IPC errors to a shared pipe for testing."""
         self.device = device
+        self.device.set_current()
         self.mr = mr
         try:
             self.CHILD_ACTION(pipe[0])
@@ -129,13 +130,13 @@ def PARENT_ACTION(self, queue):
         options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True)
         mr2 = DeviceMemoryResource(self.device, options=options)
         self.buffer = mr2.allocate(NBYTES)
-        buffer_s = pickle.dumps(self.buffer)
+        buffer_s = pickle.dumps(self.buffer)  # noqa: S301
         queue.put(buffer_s)  # Note: mr2 not sent
 
     def CHILD_ACTION(self, queue):
         Device().set_current()
         buffer_s = queue.get(timeout=CHILD_TIMEOUT_SEC)
-        pickle.loads(buffer_s)
+        pickle.loads(buffer_s)  # noqa: S301
 
     def ASSERT(self, exc_type, exc_msg):
         assert exc_type is RuntimeError

cuda_core/tests/memory_ipc/test_leaks.py

@@ -33,7 +33,7 @@ def test_alloc_handle(ipc_memory_resource):
         [mr.get_allocation_handle() for _ in range(10)]
 
 
-def exec_with_object(obj, number=1):
+def exec_success(obj, number=1):
     """Succesfully run a child process."""
     for _ in range(number):
         process = mp.Process(target=child_main, args=(obj,))
@@ -92,7 +92,7 @@ def __reduce__(self):
     ],
     ids=["alloc_handle", "mr", "buffer", "buffer_desc"],
 )
-@pytest.mark.parametrize("launcher", [exec_with_object, exec_launch_failure, exec_reduce_failure])
+@pytest.mark.parametrize("launcher", [exec_success, exec_launch_failure, exec_reduce_failure])
 def test_pass_object(ipc_memory_resource, launcher, getobject):
     """Check for fd leaks when an object is sent as a subprocess argument."""
     mr = ipc_memory_resource

cuda_core/tests/memory_ipc/test_memory_ipc.py

@@ -3,7 +3,7 @@
 
 import multiprocessing as mp
 
-from cuda.core.experimental import Buffer, Device, DeviceMemoryResource
+from cuda.core.experimental import Buffer, DeviceMemoryResource
 from utility import IPCBufferTestHelper
 
 CHILD_TIMEOUT_SEC = 20
@@ -21,7 +21,7 @@ def test_main(self, ipc_device, ipc_memory_resource):
 
         # Start the child process.
         queue = mp.Queue()
-        process = mp.Process(target=self.child_main, args=(mr, queue))
+        process = mp.Process(target=self.child_main, args=(device, mr, queue))
         process.start()
 
         # Allocate and fill memory.
@@ -39,8 +39,8 @@ def test_main(self, ipc_device, ipc_memory_resource):
         # Verify that the buffer was modified.
         helper.verify_buffer(flipped=True)
 
-    def child_main(self, mr, queue):
-        device = Device()
+    def child_main(self, device, mr, queue):
+        device.set_current()
         buffer = queue.get(timeout=CHILD_TIMEOUT_SEC)
         helper = IPCBufferTestHelper(device, buffer)
         helper.verify_buffer(flipped=False)
@@ -64,8 +64,8 @@ def test_main(self, ipc_device, ipc_memory_resource):
         q2.put(buffer2)
 
         # Start the child processes.
-        p1 = mp.Process(target=self.child_main, args=(mr, 1, q1))
-        p2 = mp.Process(target=self.child_main, args=(mr, 2, q2))
+        p1 = mp.Process(target=self.child_main, args=(device, mr, 1, q1))
+        p2 = mp.Process(target=self.child_main, args=(device, mr, 2, q2))
         p1.start()
         p2.start()
 
@@ -79,10 +79,10 @@ def test_main(self, ipc_device, ipc_memory_resource):
         IPCBufferTestHelper(device, buffer1).verify_buffer(flipped=False)
         IPCBufferTestHelper(device, buffer2).verify_buffer(flipped=True)
 
-    def child_main(self, mr, idx, queue):
+    def child_main(self, device, mr, idx, queue):
         # Note: passing the mr registers it so that buffers can be passed
         # directly.
-        device = Device()
+        device.set_current()
         buffer1 = queue.get(timeout=CHILD_TIMEOUT_SEC)
         buffer2 = queue.get(timeout=CHILD_TIMEOUT_SEC)
         if idx == 1:
@@ -104,8 +104,8 @@ def test_main(self, ipc_device, ipc_memory_resource):
 
         # Start children.
         q1, q2 = (mp.Queue() for _ in range(2))
-        p1 = mp.Process(target=self.child_main, args=(alloc_handle, 1, q1))
-        p2 = mp.Process(target=self.child_main, args=(alloc_handle, 2, q2))
+        p1 = mp.Process(target=self.child_main, args=(device, alloc_handle, 1, q1))
+        p2 = mp.Process(target=self.child_main, args=(device, alloc_handle, 2, q2))
         p1.start()
         p2.start()
 
@@ -125,11 +125,10 @@ def test_main(self, ipc_device, ipc_memory_resource):
         IPCBufferTestHelper(device, buf1).verify_buffer(starting_from=1)
         IPCBufferTestHelper(device, buf2).verify_buffer(starting_from=2)
 
-    def child_main(self, alloc_handle, idx, queue):
+    def child_main(self, device, alloc_handle, idx, queue):
         """Fills a shared memory buffer."""
         # In this case, the device needs to be set up (passing the mr does it
         # implicitly in other tests).
-        device = Device()
         device.set_current()
         mr = DeviceMemoryResource.from_allocation_handle(device, alloc_handle)
         buffer_descriptor = queue.get(timeout=CHILD_TIMEOUT_SEC)
@@ -149,8 +148,8 @@ def test_main(self, ipc_device, ipc_memory_resource):
 
         # Start children.
         q1, q2 = (mp.Queue() for _ in range(2))
-        p1 = mp.Process(target=self.child_main, args=(alloc_handle, 1, q1))
-        p2 = mp.Process(target=self.child_main, args=(alloc_handle, 2, q2))
+        p1 = mp.Process(target=self.child_main, args=(device, alloc_handle, 1, q1))
+        p2 = mp.Process(target=self.child_main, args=(device, alloc_handle, 2, q2))
         p1.start()
         p2.start()
 
@@ -170,9 +169,8 @@ def test_main(self, ipc_device, ipc_memory_resource):
         IPCBufferTestHelper(device, buf1).verify_buffer(starting_from=1)
         IPCBufferTestHelper(device, buf2).verify_buffer(starting_from=2)
 
-    def child_main(self, alloc_handle, idx, queue):
+    def child_main(self, device, alloc_handle, idx, queue):
         """Fills a shared memory buffer."""
-        device = Device()
         device.set_current()
 
         # Register the memory resource.

cuda_core/tests/memory_ipc/test_send_buffers.py

@@ -1,10 +1,11 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-import multiprocessing
+import multiprocessing as mp
 from itertools import cycle
 
-from cuda.core.experimental import Device, DeviceMemoryResource, DeviceMemoryResourceOptions
+import pytest
+from cuda.core.experimental import DeviceMemoryResource, DeviceMemoryResourceOptions
 from utility import IPCBufferTestHelper
 
 CHILD_TIMEOUT_SEC = 20
@@ -14,37 +15,13 @@
 POOL_SIZE = 2097152
 
 
-def test_ipc_send_buffers(ipc_device, ipc_memory_resource):
-    """Test passing buffers directly to a child separately from a memory resource."""
-    device = ipc_device
-    mr = ipc_memory_resource
-
-    # Allocate and fill memory.
-    buffers = [mr.allocate(NBYTES) for _ in range(NTASKS)]
-    for buffer in buffers:
-        helper = IPCBufferTestHelper(device, buffer)
-        helper.fill_buffer(flipped=False)
-
-    # Start the child process. Send the buffer directly.
-    process = multiprocessing.Process(target=child_main, args=(buffers,))
-    process.start()
-
-    # Wait for the child process.
-    process.join(timeout=CHILD_TIMEOUT_SEC)
-    assert process.exitcode == 0
-
-    # Verify that the buffers were modified.
-    for buffer in buffers:
-        helper = IPCBufferTestHelper(device, buffer)
-        helper.verify_buffer(flipped=True)
-
-
-def test_ipc_send_buffers_multi(ipc_device, ipc_memory_resource):
+@pytest.mark.parametrize("nmrs", (1, NMRS))
+def test_ipc_send_buffers(ipc_device, nmrs):
     """Test passing buffers sourced from multiple memory resources."""
     # Set up several IPC-enabled memory pools.
     device = ipc_device
     options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True)
-    mrs = [ipc_memory_resource] + [DeviceMemoryResource(device, options=options) for _ in range(NMRS - 1)]
+    mrs = [DeviceMemoryResource(device, options=options) for _ in range(NMRS)]
 
     # Allocate and fill memory.
     buffers = [mr.allocate(NBYTES) for mr, _ in zip(cycle(mrs), range(NTASKS))]
@@ -53,7 +30,13 @@ def test_ipc_send_buffers_multi(ipc_device, ipc_memory_resource):
         helper.fill_buffer(flipped=False)
 
     # Start the child process.
-    process = multiprocessing.Process(target=child_main, args=(buffers,))
+    process = mp.Process(
+        target=child_main,
+        args=(
+            device,
+            buffers,
+        ),
+    )
     process.start()
 
     # Wait for the child process.
@@ -66,8 +49,8 @@ def test_ipc_send_buffers_multi(ipc_device, ipc_memory_resource):
         helper.verify_buffer(flipped=True)
 
 
-def child_main(buffers):
-    device = Device()
+def child_main(device, buffers):
+    device.set_current()
     for buffer in buffers:
         helper = IPCBufferTestHelper(device, buffer)
         helper.verify_buffer(flipped=False)

cuda_core/tests/memory_ipc/test_serialize.py

@@ -33,7 +33,6 @@ def test_main(self, ipc_device, ipc_memory_resource):
         # Send a memory resource by allocation handle.
         alloc_handle = mr.get_allocation_handle()
         mp.reduction.send_handle(parent_conn, alloc_handle.handle, process.pid)
-        parent_conn.send(mr.uuid)
 
         # Send a buffer.
         buffer1 = mr.allocate(NBYTES)
@@ -57,9 +56,7 @@ def child_main(self, conn):
 
         # Receive the memory resource.
         handle = mp.reduction.recv_handle(conn)
-        uuid = conn.recv()
         mr = DeviceMemoryResource.from_allocation_handle(device, handle)
-        mr.register(uuid)
         os.close(handle)
 
         # Receive the buffers.
@@ -135,15 +132,17 @@ def test_object_passing(ipc_device, ipc_memory_resource):
     helper.fill_buffer(flipped=False)
 
     # Start the child process.
-    process = mp.Process(target=child_main, args=(device, alloc_handle, mr, buffer_desc, buffer))
+    process = mp.Process(target=child_main, args=(alloc_handle, mr, buffer_desc, buffer))
     process.start()
     process.join(timeout=CHILD_TIMEOUT_SEC)
     assert process.exitcode == 0
 
     helper.verify_buffer(flipped=True)
 
 
-def child_main(device, alloc_handle, mr1, buffer_desc, buffer1):
+def child_main(alloc_handle, mr1, buffer_desc, buffer1):
+    device = Device()
+    device.set_current()
     mr2 = DeviceMemoryResource.from_allocation_handle(device, alloc_handle)
 
     # OK to build the buffer from either mr and the descriptor.

cuda_core/tests/memory_ipc/test_workerpool.py

@@ -1,11 +1,11 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from itertools import cycle
 import multiprocessing as mp
 import pickle
-import pytest
+from itertools import cycle
 
+import pytest
 from cuda.core.experimental import Buffer, Device, DeviceMemoryResource, DeviceMemoryResourceOptions
 from utility import IPCBufferTestHelper
 
@@ -40,7 +40,8 @@ def test_main(self, ipc_device, nmrs):
             IPCBufferTestHelper(device, buffer).verify_buffer(flipped=True)
 
     def process_buffer(self, buffer):
-        device = Device()
+        device = Device(buffer.memory_resource.device_id)
+        device.set_current()
         IPCBufferTestHelper(device, buffer).fill_buffer(flipped=True)
 
 
@@ -67,15 +68,17 @@ def test_main(self, ipc_device, nmrs):
         with mp.Pool(NWORKERS, initializer=self.init_worker, initargs=(mrs,)) as pool:
             pool.starmap(
                 self.process_buffer,
-                [(mrs.index(buffer.memory_resource), buffer.get_ipc_descriptor()) for buffer in buffers]
+                [(mrs.index(buffer.memory_resource), buffer.get_ipc_descriptor()) for buffer in buffers],
             )
 
         for buffer in buffers:
             IPCBufferTestHelper(device, buffer).verify_buffer(flipped=True)
 
     def process_buffer(self, mr_idx, buffer_desc):
-        device = Device()
-        buffer = Buffer.from_ipc_descriptor(self.mrs[mr_idx], buffer_desc)
+        mr = self.mrs[mr_idx]
+        device = Device(mr.device_id)
+        device.set_current()
+        buffer = Buffer.from_ipc_descriptor(mr, buffer_desc)
         IPCBufferTestHelper(device, buffer).fill_buffer(flipped=True)
 
 
@@ -103,14 +106,12 @@ def test_main(self, ipc_device, nmrs):
         buffers = [mr.allocate(NBYTES) for mr, _ in zip(cycle(mrs), range(NTASKS))]
 
         with mp.Pool(NWORKERS, initializer=self.init_worker, initargs=(mrs,)) as pool:
-            pool.map(self.process_buffer, [pickle.dumps(buffer) for buffer in buffers]
-            )
+            pool.starmap(self.process_buffer, [(device, pickle.dumps(buffer)) for buffer in buffers])
 
         for buffer in buffers:
             IPCBufferTestHelper(device, buffer).verify_buffer(flipped=True)
 
-    def process_buffer(self, buffer_s):
-        device = Device()
-        buffer = pickle.loads(buffer_s)
+    def process_buffer(self, device, buffer_s):
+        device.set_current()
+        buffer = pickle.loads(buffer_s)  # noqa: S301
         IPCBufferTestHelper(device, buffer).fill_buffer(flipped=True)
-