Remove some unnecessary uses of ContextResettingTestCase (#507)

Some cases that I think look safe to remove; a few more complicated ones
still remain.

Testing on CI for now - seems to pass locally for me.

---------

Co-authored-by: brandon-b-miller <[email protected]>

Author: gmarkall

numba_cuda/numba/cuda/testing.py

@@ -32,10 +32,7 @@
 @pytest.mark.usefixtures("initialize_from_pytest_config")
 class CUDATestCase(TestCase):
     """
-    For tests that use a CUDA device. Test methods in a CUDATestCase must not
-    be run out of module order, because the ContextResettingTestCase may reset
-    the context and destroy resources used by a normal CUDATestCase if any of
-    its tests are run between tests from a CUDATestCase.
+    For tests that use a CUDA device.
 
     Methods assertFileCheckAsm and assertFileCheckLLVM will inspect a
     CUDADispatcher and assert that the compilation artifacts match the
@@ -187,21 +184,6 @@ def assertFileCheckMatches(
             )
 
 
-class ContextResettingTestCase(CUDATestCase):
-    """
-    For tests where the context needs to be reset after each test. Typically
-    these inspect or modify parts of the context that would usually be expected
-    to be internal implementation details (such as the state of allocations and
-    deallocations, etc.).
-    """
-
-    def tearDown(self):
-        super().tearDown()
-        from numba.cuda.cudadrv.devices import reset
-
-        reset()
-
-
 def skip_on_cudasim(reason):
     """Skip this test if running on the CUDA simulator"""
     return unittest.skipIf(config.ENABLE_CUDASIM, reason)

numba_cuda/numba/cuda/tests/cudadrv/test_context_stack.py

@@ -14,12 +14,7 @@ class TestContextStack(CUDATestCase):
     def setUp(self):
         super().setUp()
         # Reset before testing
-        cuda.close()
-
-    def test_gpus_current(self):
-        self.assertIs(cuda.gpus.current, None)
-        with cuda.gpus[0]:
-            self.assertEqual(int(cuda.gpus.current.id), 0)
+        cuda.current_context().reset()
 
     def test_gpus_len(self):
         self.assertGreater(len(cuda.gpus), 0)
@@ -45,7 +40,7 @@ def test_gpus_cudevice_indexing(self):
 class TestContextAPI(CUDATestCase):
     def tearDown(self):
         super().tearDown()
-        cuda.close()
+        cuda.current_context().reset()
 
     def test_context_memory(self):
         try:
@@ -91,7 +86,7 @@ def switch_gpu():
 class Test3rdPartyContext(CUDATestCase):
     def tearDown(self):
         super().tearDown()
-        cuda.close()
+        cuda.current_context().reset()
 
     def test_attached_primary(self, extra_work=lambda: None):
         # Emulate primary context creation by 3rd party

numba_cuda/numba/cuda/tests/cudadrv/test_cuda_memory.py

@@ -6,17 +6,18 @@
 import numpy as np
 
 from numba.cuda.cudadrv import driver, drvapi, devices
-from numba.cuda.testing import unittest, ContextResettingTestCase
+from numba.cuda.testing import unittest, CUDATestCase
 from numba.cuda.testing import skip_on_cudasim
 
 
 @skip_on_cudasim("CUDA Memory API unsupported in the simulator")
-class TestCudaMemory(ContextResettingTestCase):
+class TestCudaMemory(CUDATestCase):
     def setUp(self):
         super().setUp()
         self.context = devices.get_context()
 
     def tearDown(self):
+        self.context.reset()
         del self.context
         super(TestCudaMemory, self).tearDown()
 
@@ -107,7 +108,7 @@ def dtor():
         self.assertEqual(dtor_invoked[0], 2)
 
 
-class TestCudaMemoryFunctions(ContextResettingTestCase):
+class TestCudaMemoryFunctions(CUDATestCase):
     def setUp(self):
         super().setUp()
         self.context = devices.get_context()
@@ -153,7 +154,7 @@ def test_d2d(self):
 
 
 @skip_on_cudasim("CUDA Memory API unsupported in the simulator")
-class TestMVExtent(ContextResettingTestCase):
+class TestMVExtent(CUDATestCase):
     def test_c_contiguous_array(self):
         ary = np.arange(100)
         arysz = ary.dtype.itemsize * ary.size

numba_cuda/numba/cuda/tests/cudadrv/test_emm_plugins.py

@@ -112,14 +112,16 @@ class TestDeviceOnlyEMMPlugin(CUDATestCase):
     def setUp(self):
         super().setUp()
         # Always start afresh with a new context and memory manager
-        cuda.close()
-        cuda.set_memory_manager(DeviceOnlyEMMPlugin)
+        ctx = cuda.current_context()
+        ctx.reset()
+        self._initial_memory_manager = ctx.memory_manager
+        ctx.memory_manager = DeviceOnlyEMMPlugin(context=ctx)
 
     def tearDown(self):
         super().tearDown()
-        # Unset the memory manager for subsequent tests
-        cuda.close()
-        cuda.cudadrv.driver._memory_manager = None
+        ctx = cuda.current_context()
+        ctx.reset()
+        ctx.memory_manager = self._initial_memory_manager
 
     def test_memalloc(self):
         mgr = cuda.current_context().memory_manager
@@ -129,6 +131,7 @@ def test_memalloc(self):
         arr_1 = np.arange(10)
         d_arr_1 = cuda.device_array_like(arr_1)
         self.assertTrue(mgr.memalloc_called)
+
         self.assertEqual(mgr.count, 1)
         self.assertEqual(mgr.allocations[1], arr_1.nbytes)
 

numba_cuda/numba/cuda/tests/cudadrv/test_host_alloc.py

@@ -4,10 +4,13 @@
 import numpy as np
 from numba.cuda.cudadrv import driver
 from numba import cuda
-from numba.cuda.testing import unittest, ContextResettingTestCase
+from numba.cuda.testing import unittest, CUDATestCase
 
 
-class TestHostAlloc(ContextResettingTestCase):
+class TestHostAlloc(CUDATestCase):
+    def tearDown(self):
+        cuda.current_context().reset()
+
     def test_host_alloc_driver(self):
         n = 32
         mem = cuda.current_context().memhostalloc(n, mapped=True)

numba_cuda/numba/cuda/tests/cudadrv/test_inline_ptx.py

@@ -4,12 +4,12 @@
 from llvmlite import ir
 
 from numba.cuda.cudadrv import nvvm
-from numba.cuda.testing import unittest, ContextResettingTestCase
+from numba.cuda.testing import unittest, CUDATestCase
 from numba.cuda.testing import skip_on_cudasim
 
 
 @skip_on_cudasim("Inline PTX cannot be used in the simulator")
-class TestCudaInlineAsm(ContextResettingTestCase):
+class TestCudaInlineAsm(CUDATestCase):
     def test_inline_rsqrt(self):
         mod = ir.Module(__name__)
         mod.triple = "nvptx64-nvidia-cuda"

numba_cuda/numba/cuda/tests/cudadrv/test_managed_alloc.py

@@ -5,15 +5,19 @@
 from ctypes import byref, c_size_t
 from numba.cuda.cudadrv.driver import device_memset, driver, USE_NV_BINDING
 from numba import cuda
-from numba.cuda.testing import unittest, ContextResettingTestCase
+from numba.cuda.testing import unittest, CUDATestCase
 from numba.cuda.testing import skip_on_cudasim, skip_on_arm
 from numba.cuda.tests.support import linux_only
 
 
 @skip_on_cudasim("CUDA Driver API unsupported in the simulator")
 @linux_only
 @skip_on_arm("Managed Alloc support is experimental/untested on ARM")
-class TestManagedAlloc(ContextResettingTestCase):
+class TestManagedAlloc(CUDATestCase):
+    def tearDown(self):
+        super().tearDown()
+        cuda.current_context().reset()
+
     def get_total_gpu_memory(self):
         # We use a driver function to directly get the total GPU memory because
         # an EMM plugin may report something different (or not implement

numba_cuda/numba/cuda/tests/cudadrv/test_module_callbacks.py

@@ -11,7 +11,6 @@
 from numba.cuda.cudadrv.linkable_code import CUSource
 from numba.cuda.testing import (
     CUDATestCase,
-    ContextResettingTestCase,
     skip_on_cudasim,
 )
 
@@ -42,7 +41,7 @@ def get_hashable_handle_value(handle):
 
 
 @skip_on_cudasim("Module loading not implemented in the simulator")
-class TestModuleCallbacksBasic(ContextResettingTestCase):
+class TestModuleCallbacksBasic(CUDATestCase):
     def test_basic(self):
         counter = 0
 

numba_cuda/numba/cuda/tests/cudadrv/test_pinned.py

@@ -5,10 +5,10 @@
 import platform
 
 from numba import cuda
-from numba.cuda.testing import unittest, ContextResettingTestCase
+from numba.cuda.testing import unittest, CUDATestCase
 
 
-class TestPinned(ContextResettingTestCase):
+class TestPinned(CUDATestCase):
     def _run_copies(self, A):
         A0 = np.copy(A)
 

numba_cuda/numba/cuda/tests/cudadrv/test_profiler.py

@@ -2,13 +2,13 @@
 # SPDX-License-Identifier: BSD-2-Clause
 
 import unittest
-from numba.cuda.testing import ContextResettingTestCase
+from numba.cuda.testing import CUDATestCase
 from numba import cuda
 from numba.cuda.testing import skip_on_cudasim
 
 
 @skip_on_cudasim("CUDA Profiler unsupported in the simulator")
-class TestProfiler(ContextResettingTestCase):
+class TestProfiler(CUDATestCase):
     def test_profiling(self):
         with cuda.profiling():
             a = cuda.device_array(10)