add proper cleanup to test files

Author: keenan-simpson

cuda_core/tests/test_memory.py

@@ -301,31 +301,41 @@ def test_mempool_properties(property_name, expected_type):
     pool_size = 2097152  # 2MB size
     mr = Mempool.create(device.device_id, pool_size, enable_ipc=False)
 
-    # Get the property value
-    value = getattr(mr, property_name)
-
-    # Test type
-    assert isinstance(value, expected_type), f"{property_name} should return {expected_type}, got {type(value)}"
-
-    # Test value constraints
-    if expected_type is int:
-        assert value >= 0, f"{property_name} should be non-negative"
-
-    # Test memory usage properties with actual allocations
-    if property_name in ["reserved_mem_current", "used_mem_current"]:
-        # Allocate some memory and check if values increase
-        initial_value = value
-        buffer = mr.allocate(1024)
-        new_value = getattr(mr, property_name)
-        assert new_value >= initial_value, f"{property_name} should increase or stay same after allocation"
-        buffer.close()
-
-    # Test high watermark properties
-    if property_name in ["reserved_mem_high", "used_mem_high"]:
-        # High watermark should never be less than current
-        current_prop = property_name.replace("_high", "_current")
-        current_value = getattr(mr, current_prop)
-        assert value >= current_value, f"{property_name} should be >= {current_prop}"
+    try:
+        # Get the property value
+        value = getattr(mr, property_name)
+
+        # Test type
+        assert isinstance(value, expected_type), f"{property_name} should return {expected_type}, got {type(value)}"
+
+        # Test value constraints
+        if expected_type is int:
+            assert value >= 0, f"{property_name} should be non-negative"
+
+        # Test memory usage properties with actual allocations
+        if property_name in ["reserved_mem_current", "used_mem_current"]:
+            # Allocate some memory and check if values increase
+            initial_value = value
+            buffer = None
+            try:
+                buffer = mr.allocate(1024)
+                new_value = getattr(mr, property_name)
+                assert new_value >= initial_value, f"{property_name} should increase or stay same after allocation"
+            finally:
+                if buffer is not None:
+                    buffer.close()
+
+        # Test high watermark properties
+        if property_name in ["reserved_mem_high", "used_mem_high"]:
+            # High watermark should never be less than current
+            current_prop = property_name.replace("_high", "_current")
+            current_value = getattr(mr, current_prop)
+            assert value >= current_value, f"{property_name} should be >= {current_prop}"
+
+    finally:
+        # Ensure we allocate and deallocate a small buffer to flush any pending operations
+        flush_buffer = mr.allocate(64)
+        flush_buffer.close()
 
 
 def mempool_child_process(importer, queue):
@@ -387,62 +397,84 @@ def test_ipc_mempool():
     stream = device.create_stream()
     pool_size = 2097152  # 2MB size
     mr = Mempool.create(device.device_id, pool_size, enable_ipc=True)
-    shareable_handle = mr.get_shareable_handle()
-
-    # Allocate and export memory
-    buffer = mr.allocate(64)
-
-    # Fill buffer with test pattern using unified memory
-    unified_mr = DummyUnifiedMemoryResource(device)
-    src_buffer = unified_mr.allocate(64)
-    src_ptr = ctypes.cast(int(src_buffer.handle), ctypes.POINTER(ctypes.c_byte))
-    for i in range(64):
-        src_ptr[i] = ctypes.c_byte(i)
-
-    buffer.copy_from(src_buffer, stream=stream)
-    device.sync()
-    src_buffer.close()
-
-    # Export buffer for IPC
-    ipc_buffer = mr.export_buffer(buffer)
 
     # Create socket pair for handle transfer
     exporter, importer = socketpair(AF_UNIX, SOCK_DGRAM)
-
-    # Start child process
-    multiprocessing.set_start_method("spawn", force=True)
     queue = multiprocessing.Queue()
-    process = multiprocessing.Process(target=mempool_child_process, args=(importer, queue))
-    process.start()
-
-    # Send handles to child process
-    exporter.sendmsg([], [(SOL_SOCKET, SCM_RIGHTS, array.array("i", [shareable_handle]))])
-    queue.put(ipc_buffer)
-
-    # Wait for child process
-    process.join(timeout=10)
-    assert process.exitcode == 0
-
-    # Check for exceptions
-    if not queue.empty():
-        result = queue.get()
-        if isinstance(result, tuple):
-            exception, traceback_str = result
-            print("\nException in child process:")
-            print(traceback_str)
-            raise exception
-        assert result is True
-
-    # Verify child process wrote the inverted pattern using unified memory
-    verify_buffer = unified_mr.allocate(64)
-    verify_buffer.copy_from(buffer, stream=stream)
-    device.sync()
+    process = None
 
-    verify_ptr = ctypes.cast(int(verify_buffer.handle), ctypes.POINTER(ctypes.c_byte))
-    for i in range(64):
-        assert (
-            ctypes.c_byte(verify_ptr[i]).value == ctypes.c_byte(255 - i).value
-        ), f"Child process data not reflected in parent at index {i}"
-
-    verify_buffer.close()
-    buffer.close()
+    try:
+        shareable_handle = mr.get_shareable_handle()
+
+        # Allocate and export memory
+        buffer = mr.allocate(64)
+
+        try:
+            # Fill buffer with test pattern using unified memory
+            unified_mr = DummyUnifiedMemoryResource(device)
+            src_buffer = unified_mr.allocate(64)
+            try:
+                src_ptr = ctypes.cast(int(src_buffer.handle), ctypes.POINTER(ctypes.c_byte))
+                for i in range(64):
+                    src_ptr[i] = ctypes.c_byte(i)
+
+                buffer.copy_from(src_buffer, stream=stream)
+                device.sync()
+            finally:
+                src_buffer.close()
+
+            # Export buffer for IPC
+            ipc_buffer = mr.export_buffer(buffer)
+
+            # Start child process
+            multiprocessing.set_start_method("spawn", force=True)
+            process = multiprocessing.Process(target=mempool_child_process, args=(importer, queue))
+            process.start()
+
+            # Send handles to child process
+            exporter.sendmsg([], [(SOL_SOCKET, SCM_RIGHTS, array.array("i", [shareable_handle]))])
+            queue.put(ipc_buffer)
+
+            # Wait for child process
+            process.join(timeout=10)
+            assert process.exitcode == 0
+
+            # Check for exceptions
+            if not queue.empty():
+                result = queue.get()
+                if isinstance(result, tuple):
+                    exception, traceback_str = result
+                    print("\nException in child process:")
+                    print(traceback_str)
+                    raise exception
+                assert result is True
+
+            # Verify child process wrote the inverted pattern using unified memory
+            verify_buffer = unified_mr.allocate(64)
+            try:
+                verify_buffer.copy_from(buffer, stream=stream)
+                device.sync()
+
+                verify_ptr = ctypes.cast(int(verify_buffer.handle), ctypes.POINTER(ctypes.c_byte))
+                for i in range(64):
+                    assert (
+                        ctypes.c_byte(verify_ptr[i]).value == ctypes.c_byte(255 - i).value
+                    ), f"Child process data not reflected in parent at index {i}"
+            finally:
+                verify_buffer.close()
+
+        finally:
+            buffer.close()
+
+    finally:
+        # Clean up all resources
+        if process is not None and process.is_alive():
+            process.terminate()
+            process.join(timeout=1)
+        queue.close()
+        queue.join_thread()  # Ensure the queue's background thread is cleaned up
+        exporter.close()
+        importer.close()
+        # Flush any pending operations
+        flush_buffer = mr.allocate(64)
+        flush_buffer.close()