Repository: numba/numba

Issue:
RAPIDS RMM + Numba Integration: `context.deallocations` is `NoneType`
We're getting the following at the exit of a process with Numba 0.45:
```
Traceback (most recent call last):
  File "/conda/envs/gdf/lib/python3.6/site-packages/numba/utils.py", line 745, in _exitfunc
    f()
  File "/conda/envs/gdf/lib/python3.6/site-packages/numba/utils.py", line 669, in __call__
    return info.func(*info.args, **(info.kwargs or {}))
  File "/conda/envs/gdf/lib/python3.6/site-packages/numba/cuda/cudadrv/driver.py", line 1037, in core
    dealloc.add_item(module_unload, handle)
AttributeError: 'NoneType' object has no attribute 'add_item'
```

This didn't happen with previous versions of Numba where I'm wondering if something changed which broke the backdoor integration with RMM and Numba.

I apologize I couldn't give a clearer title, so please feel free to change it if appropriate.
RAPIDS RMM + Numba Integration: `context.deallocations` is `NoneType`
We're getting the following at the exit of a process with Numba 0.45:
```
Traceback (most recent call last):
  File "/conda/envs/gdf/lib/python3.6/site-packages/numba/utils.py", line 745, in _exitfunc
    f()
  File "/conda/envs/gdf/lib/python3.6/site-packages/numba/utils.py", line 669, in __call__
    return info.func(*info.args, **(info.kwargs or {}))
  File "/conda/envs/gdf/lib/python3.6/site-packages/numba/cuda/cudadrv/driver.py", line 1037, in core
    dealloc.add_item(module_unload, handle)
AttributeError: 'NoneType' object has no attribute 'add_item'
```

This didn't happen with previous versions of Numba where I'm wondering if something changed which broke the backdoor integration with RMM and Numba.

I apologize I couldn't give a clearer title, so please feel free to change it if appropriate.


Discussion:
This is due to CPython's attempt to free cyclic objects by replacing them with `None` in unpredictable manner during interpreter shutdown.  I usually resolve such issues by adding a `if something is not None` for all offending lines reported, or trying to detect if the CPython is shutting down and do things differently.

I can push a patch for you to test.
@sklam Are you suggesting a patch for Numba or a patch for RMM? Either way I'm more than happy to test.
Patch in numba, but it will need testing against RMM.  This sort of error is very random and depends on the order in which modules are imported.
This is due to CPython's attempt to free cyclic objects by replacing them with `None` in unpredictable manner during interpreter shutdown.  I usually resolve such issues by adding a `if something is not None` for all offending lines reported, or trying to detect if the CPython is shutting down and do things differently.

I can push a patch for you to test.
@sklam Are you suggesting a patch for Numba or a patch for RMM? Either way I'm more than happy to test.
Patch in numba, but it will need testing against RMM.  This sort of error is very random and depends on the order in which modules are imported.

Resolving patch:
diff --git a/numba/cuda/cudadrv/devices.py b/numba/cuda/cudadrv/devices.py
--- a/numba/cuda/cudadrv/devices.py
+++ b/numba/cuda/cudadrv/devices.py
@@ -160,6 +160,8 @@ def _get_or_create_context_uncached(self, devnum):
                         msg = ('Numba cannot operate on non-primary'
                                ' CUDA context {:x}')
                         raise RuntimeError(msg.format(ac.context_handle.value))
+                    # Ensure the context is ready
+                    ctx.prepare_for_use()
                 return ctx
 
     def _activate_context_for(self, devnum):
diff --git a/numba/cuda/cudadrv/driver.py b/numba/cuda/cudadrv/driver.py
--- a/numba/cuda/cudadrv/driver.py
+++ b/numba/cuda/cudadrv/driver.py
@@ -708,14 +708,20 @@ def get_max_potential_block_size(self, func, b2d_func, memsize, blocksizelimit,
                                                              memsize, blocksizelimit, flags)
         return (gridsize.value, blocksize.value)
 
+    def prepare_for_use(self):
+        """Initialize the context for use.
+        It's safe to be called multiple times.
+        """
+        # setup *deallocations* as the context becomes active for the first time
+        if self.deallocations is None:
+            self.deallocations = _PendingDeallocs(self.get_memory_info().total)
+
     def push(self):
         """
         Pushes this context on the current CPU Thread.
         """
         driver.cuCtxPushCurrent(self.handle)
-        # setup *deallocations* as the context becomes active for the first time
-        if self.deallocations is None:
-            self.deallocations = _PendingDeallocs(self.get_memory_info().total)
+        self.prepare_for_use()
 
     def pop(self):
         """


Test patch:
diff --git a/numba/cuda/tests/cudadrv/test_context_stack.py b/numba/cuda/tests/cudadrv/test_context_stack.py
--- a/numba/cuda/tests/cudadrv/test_context_stack.py
+++ b/numba/cuda/tests/cudadrv/test_context_stack.py
@@ -74,7 +74,7 @@ class Test3rdPartyContext(SerialMixin, unittest.TestCase):
     def tearDown(self):
         cuda.close()
 
-    def test_attached_primary(self):
+    def test_attached_primary(self, extra_work=lambda: None):
         # Emulate primary context creation by 3rd party
         the_driver = driver.driver
         hctx = driver.drvapi.cu_context()
@@ -86,6 +86,8 @@ def test_attached_primary(self):
             # context.
             my_ctx = cuda.current_context()
             self.assertEqual(my_ctx.handle.value, ctx.handle.value)
+
+            extra_work()
         finally:
             ctx.pop()
             the_driver.cuDevicePrimaryCtxRelease(0)
@@ -106,6 +108,21 @@ def test_attached_non_primary(self):
         finally:
             the_driver.cuCtxDestroy(hctx)
 
+    def test_cudajit_in_attached_primary_context(self):
+        def do():
+            from numba import cuda
+
+            @cuda.jit
+            def foo(a):
+                for i in range(a.size):
+                    a[i] = i
+
+            a = cuda.device_array(10)
+            foo(a)
+            self.assertEqual(list(a.copy_to_host()), list(range(10)))
+
+        self.test_attached_primary(do)
+
 
 if __name__ == '__main__':
     unittest.main()