Try addressing issues that may be causing CI failures

Author: shwina

cuda_core/examples/memory_ops.py

@@ -2,6 +2,18 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+# ################################################################################
+#
+# This demo illustrates:
+#
+#   1. How to use different memory resources to allocate and manage memory
+#   2. How to copy data between different memory types
+#   3. How to use DLPack to interoperate with other libraries
+#
+# ################################################################################
+
+import sys
+
 import cupy as cp
 import numpy as np
 
@@ -15,6 +27,10 @@
     launch,
 )
 
+if np.__version__ < "2.1.0":
+    print("This example requires NumPy 2.1.0 or later", file=sys.stderr)
+    sys.exit(0)
+
 # Kernel for memory operations
 code = """
 extern "C"

cuda_core/tests/test_launcher.py

@@ -5,6 +5,7 @@
 import os
 import pathlib
 
+import cupy as cp
 import numpy as np
 import pytest
 from conftest import skipif_need_cuda_headers
@@ -219,6 +220,8 @@ def test_launch_with_buffers_allocated_by_memory_resource(init_cuda, memory_reso
     dev = Device()
     dev.set_current()
     stream = dev.create_stream()
+    # tell CuPy to use our stream as the current stream:
+    cp.cuda.ExternalStream(int(stream.handle)).use()
 
     # Kernel that operates on memory
     code = """
@@ -258,18 +261,13 @@ def test_launch_with_buffers_allocated_by_memory_resource(init_cuda, memory_reso
         # For pinned memory, use numpy
         array = np.from_dlpack(buffer).view(dtype=dtype)
     else:
-        # For device memory, use cupy
-        import cupy as cp
-
         array = cp.from_dlpack(buffer).view(dtype=dtype)
 
     # Initialize data with random values
     if mr.is_host_accessible:
         rng = np.random.default_rng()
         array[:] = rng.random(size, dtype=dtype)
     else:
-        import cupy as cp
-
         rng = cp.random.default_rng()
         array[:] = rng.random(size, dtype=dtype)
 
@@ -288,16 +286,13 @@ def test_launch_with_buffers_allocated_by_memory_resource(init_cuda, memory_reso
     stream.sync()
 
     # Verify kernel operations
-    if mr.is_host_accessible:
-        assert np.allclose(array, original * 3.0), f"{memory_resource_class.__name__} operation failed"
-    else:
-        import cupy as cp
-
-        assert cp.allclose(array, original * 3.0), f"{memory_resource_class.__name__} operation failed"
+    assert cp.allclose(array, original * 3.0), f"{memory_resource_class.__name__} operation failed"
 
     # Clean up
     buffer.close(stream)
     stream.close()
 
+    cp.cuda.Stream.null.use()  # reset CuPy's current stream to the null stream
+
     # Verify buffer is properly closed
     assert buffer.handle == 0, f"{memory_resource_class.__name__} buffer should be closed"