fix: remove unused unload API from Module and decouple from Context to fix weakref lifetime issue

Author: cpcloud

numba_cuda/numba/cuda/cudadrv/driver.py

@@ -1321,7 +1321,7 @@ def create_module_image(
         )
         key = module.handle
         self.modules[key] = module
-        return weakref.proxy(module)
+        return module
 
     def unload_module(self, module):
         key = module.handle
@@ -1432,7 +1432,6 @@ def load_module_image_ctypes(
     info_log = jitinfo.value
 
     return CtypesModule(
-        weakref.proxy(context),
         handle,
         info_log,
         _module_finalizer(context, handle),
@@ -1476,7 +1475,6 @@ def load_module_image_cuda_python(
     info_log = jitinfo.decode("utf-8")
 
     return CudaPythonModule(
-        weakref.proxy(context),
         handle,
         info_log,
         _module_finalizer(context, handle),
@@ -2251,14 +2249,12 @@ class Module(metaclass=ABCMeta):
 
     def __init__(
         self,
-        context,
         handle,
         info_log,
         finalizer=None,
         setup_callbacks=None,
         teardown_callbacks=None,
     ):
-        self.context = context
         self.handle = handle
         self.info_log = info_log
         if finalizer is not None:
@@ -2270,10 +2266,6 @@ def __init__(
 
         self._set_finalizers()
 
-    def unload(self):
-        """Unload this module from the context"""
-        self.context.unload_module(self)
-
     @abstractmethod
     def get_function(self, name):
         """Returns a Function object encapsulating the named function"""

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

@@ -119,8 +119,6 @@ def test_cuda_driver_basic(self):
         for i, v in enumerate(array):
             self.assertEqual(i, v)
 
-        module.unload()
-
     def test_cuda_driver_stream_operations(self):
         module = self.context.create_module_ptx(self.ptx)
         function = module.get_function("_Z10helloworldPi")

numba_cuda/numba/cuda/tests/nrt/test_module.py

@@ -0,0 +1,42 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import pytest
+
+from numba.cuda.cudadrv import devices
+from numba.cuda.memory_management.nrt import rtsys
+
+
+@pytest.fixture
+def alloc_init():
+    rtsys.ensure_allocated()
+    rtsys.ensure_initialized()
+
+
+@pytest.fixture
+def ctx(alloc_init):
+    ctx = devices.get_context()
+    yield ctx
+    ctx.reset()
+
+    # this fails with a ReferenceError
+    str(rtsys._memsys_module)
+
+
+def test_nothing(ctx):
+    """
+    This test _must_ remain in a separate module, otherwise it is not reproducible.
+
+    The sequence that previously caused the failure is:
+
+    1. the alloc_init fixture creates a weakref proxy to a CudaPythonModule
+       (rtsys._memsys_module) from devices.get_context()
+    2. ctx.reset() calls dict.clear() on the context's module dict
+    3. since there's only a single strong ref to `rtsys._memsys_module` it gets collected
+       during the clear
+    4. All ops on the module (a weakref proxy) fail, because it is now
+       referencing a dead object
+
+    The solution was to decouple `Context` from `Module`, allowing a strong
+    reference to `Module` to exist, cleaned up with the usual mechanisms.
+    """