Wrap IPC memory resource cleanup in try/finally

Ensure MR.close() runs even when tests fail (timeout, assertion,
queue error) so pools don't leak on the unhappy path either.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: cpcloud

cuda_core/tests/memory_ipc/test_errors.py

@@ -28,24 +28,25 @@ def test_main(self, ipc_device, ipc_memory_resource):
         self.mr = ipc_memory_resource
         self._extra_mrs = []
 
-        # Start a child process to generate error info.
-        pipe = [multiprocessing.Queue() for _ in range(2)]
-        process = multiprocessing.Process(target=self.child_main, args=(pipe, self.device, self.mr))
-        process.start()
-
-        # Interact.
-        self.PARENT_ACTION(pipe[0])
-
-        # Check the error.
-        exc_type, exc_msg = pipe[1].get(timeout=CHILD_TIMEOUT_SEC)
-        self.ASSERT(exc_type, exc_msg)
-
-        # Wait for the child process.
-        process.join(timeout=CHILD_TIMEOUT_SEC)
-        assert process.exitcode == 0
-
-        for mr in self._extra_mrs:
-            mr.close()
+        try:
+            # Start a child process to generate error info.
+            pipe = [multiprocessing.Queue() for _ in range(2)]
+            process = multiprocessing.Process(target=self.child_main, args=(pipe, self.device, self.mr))
+            process.start()
+
+            # Interact.
+            self.PARENT_ACTION(pipe[0])
+
+            # Check the error.
+            exc_type, exc_msg = pipe[1].get(timeout=CHILD_TIMEOUT_SEC)
+            self.ASSERT(exc_type, exc_msg)
+
+            # Wait for the child process.
+            process.join(timeout=CHILD_TIMEOUT_SEC)
+            assert process.exitcode == 0
+        finally:
+            for mr in self._extra_mrs:
+                mr.close()
 
     def child_main(self, pipe, device, mr):
         """Child process that pushes IPC errors to a shared pipe for testing."""

cuda_core/tests/memory_ipc/test_send_buffers.py

@@ -26,27 +26,29 @@ def test_main(self, ipc_device, nmrs):
         options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True)
         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))]
-        pgen = PatternGen(device, NBYTES)
-        for buffer in buffers:
-            pgen.fill_buffer(buffer, seed=False)
-
-        # Start the child process.
-        process = mp.Process(target=self.child_main, args=(device, buffers))
-        process.start()
-
-        # Wait for the child process.
-        process.join(timeout=CHILD_TIMEOUT_SEC)
-        assert process.exitcode == 0
-
-        # Verify that the buffers were modified.
-        pgen = PatternGen(device, NBYTES)
-        for buffer in buffers:
-            pgen.verify_buffer(buffer, seed=True)
-            buffer.close()
-        for mr in mrs:
-            mr.close()
+        try:
+            # Allocate and fill memory.
+            buffers = [mr.allocate(NBYTES) for mr, _ in zip(cycle(mrs), range(NTASKS))]
+            pgen = PatternGen(device, NBYTES)
+            for buffer in buffers:
+                pgen.fill_buffer(buffer, seed=False)
+
+            # Start the child process.
+            process = mp.Process(target=self.child_main, args=(device, buffers))
+            process.start()
+
+            # Wait for the child process.
+            process.join(timeout=CHILD_TIMEOUT_SEC)
+            assert process.exitcode == 0
+
+            # Verify that the buffers were modified.
+            pgen = PatternGen(device, NBYTES)
+            for buffer in buffers:
+                pgen.verify_buffer(buffer, seed=True)
+                buffer.close()
+        finally:
+            for mr in mrs:
+                mr.close()
 
     def child_main(self, device, buffers):
         device.set_current()

cuda_core/tests/memory_ipc/test_workerpool.py

@@ -33,17 +33,20 @@ def test_main(self, ipc_device, nmrs):
         device = ipc_device
         options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True)
         mrs = [DeviceMemoryResource(device, options=options) for _ in range(nmrs)]
-        buffers = [mr.allocate(NBYTES) for mr, _ in zip(cycle(mrs), range(NTASKS))]
 
-        with mp.Pool(NWORKERS) as pool:
-            pool.map(self.process_buffer, buffers)
+        try:
+            buffers = [mr.allocate(NBYTES) for mr, _ in zip(cycle(mrs), range(NTASKS))]
 
-        pgen = PatternGen(device, NBYTES)
-        for buffer in buffers:
-            pgen.verify_buffer(buffer, seed=True)
-            buffer.close()
-        for mr in mrs:
-            mr.close()
+            with mp.Pool(NWORKERS) as pool:
+                pool.map(self.process_buffer, buffers)
+
+            pgen = PatternGen(device, NBYTES)
+            for buffer in buffers:
+                pgen.verify_buffer(buffer, seed=True)
+                buffer.close()
+        finally:
+            for mr in mrs:
+                mr.close()
 
     def process_buffer(self, buffer):
         device = Device(buffer.memory_resource.device_id)
@@ -72,20 +75,23 @@ def test_main(self, ipc_device, nmrs):
         device = ipc_device
         options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True)
         mrs = [DeviceMemoryResource(device, options=options) for _ in range(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.starmap(
-                self.process_buffer,
-                [(mrs.index(buffer.memory_resource), buffer.get_ipc_descriptor()) for buffer in buffers],
-            )
+        try:
+            buffers = [mr.allocate(NBYTES) for mr, _ in zip(cycle(mrs), range(NTASKS))]
 
-        pgen = PatternGen(device, NBYTES)
-        for buffer in buffers:
-            pgen.verify_buffer(buffer, seed=True)
-            buffer.close()
-        for mr in mrs:
-            mr.close()
+            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],
+                )
+
+            pgen = PatternGen(device, NBYTES)
+            for buffer in buffers:
+                pgen.verify_buffer(buffer, seed=True)
+                buffer.close()
+        finally:
+            for mr in mrs:
+                mr.close()
 
     def process_buffer(self, mr_idx, buffer_desc):
         mr = self.mrs[mr_idx]
@@ -119,17 +125,20 @@ def test_main(self, ipc_device, nmrs):
         device = ipc_device
         options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True)
         mrs = [DeviceMemoryResource(device, options=options) for _ in range(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.starmap(self.process_buffer, [(device, pickle.dumps(buffer)) for buffer in buffers])
+        try:
+            buffers = [mr.allocate(NBYTES) for mr, _ in zip(cycle(mrs), range(NTASKS))]
 
-        pgen = PatternGen(device, NBYTES)
-        for buffer in buffers:
-            pgen.verify_buffer(buffer, seed=True)
-            buffer.close()
-        for mr in mrs:
-            mr.close()
+            with mp.Pool(NWORKERS, initializer=self.init_worker, initargs=(mrs,)) as pool:
+                pool.starmap(self.process_buffer, [(device, pickle.dumps(buffer)) for buffer in buffers])
+
+            pgen = PatternGen(device, NBYTES)
+            for buffer in buffers:
+                pgen.verify_buffer(buffer, seed=True)
+                buffer.close()
+        finally:
+            for mr in mrs:
+                mr.close()
 
     def process_buffer(self, device, buffer_s):
         device.set_current()