[Fix][Relax] Scope BasePyModule's Python function registry entries to the module (#20417)

`BasePyModule.__del__` calls `vm.builtin.clear_py_func_registry`, which
empties the process-wide registry rather than removing the module's own
entries. Collecting any module therefore drops every registration in the
process, and a later `R.call_py_func` fails with `Python function
'<name>' not found in registry`.

Whether it happens depends on when the collector runs, so it shows up as
an unrelated test failing. The module built in
`tests/python/relax/test_pytorch_integration.py` registers nothing
itself, but its finalizer wipes the names registered by
`tests/python/relax/test_relax_operators.py`:

```python
tvm.get_global_func("vm.builtin.register_py_func")("victim", tvm.runtime.convert(lambda x: x))
bystander = BasePyModule.__new__(BasePyModule)   # registers nothing
del bystander
gc.collect()
tvm.get_global_func("vm.builtin.get_py_func")("victim")
# InternalError: Python function 'victim' not found in registry
```

This adds `vm.builtin.unregister_py_func` and records the names each
module registers, so the finalizer removes only its own entries.
Registering a name again transfers ownership to the newer module, which
keeps a stale module from dropping a live one's function.

The per-module cleanup was also unreachable before this change. The
wrapper stored in the registry captured `self`, so the registry held a
strong reference to every module that registered a function and none of
them were ever collected — the finalizer only ever ran for modules that
had registered nothing, which are exactly the ones that then cleared
everyone else's entries. The wrapper now holds a weak reference to its
module, and reports a clear error if it is called after the module is
gone:

```
RuntimeError: Python function 'f' belongs to a BasePyModule that has been destroyed;
keep the module alive while calling into it.
```

`vm.builtin.clear_py_func_registry` is left in place; it simply has no
callers now.

Author: akaashrp

python/tvm/relax/base_py_module.py

@@ -19,6 +19,8 @@
 
 import inspect
 import os
+import sys
+import weakref
 from typing import Any, Optional, Union
 
 import numpy as np
@@ -43,6 +45,11 @@
     _FASTER_DLPACK_EXTENSION = None
 
 
+# Maps a registered Python function name to the id of the BasePyModule that registered it last.
+# A module only unregisters names it still owns, so re-registering a name transfers ownership.
+_PY_FUNC_OWNERS: dict[str, int] = {}
+
+
 class BasePyModule:
     """Base class that allows Python functions in IRModule with DLPack conversion.
 
@@ -56,12 +63,31 @@ class BasePyModule:
     subclass with ``R.py_module`` adds this executable runtime interface.
     """
 
+    # Names this instance registered with the VM's Python function registry. The registry is
+    # global, so a module must only remove entries it still owns.
+    _registered_py_funcs: tuple[str, ...] = ()
+
     def __del__(self):
-        """Clean up registered Python functions on module destruction."""
+        """Unregister the Python functions this module still owns."""
+        registered = getattr(self, "_registered_py_funcs", None)
+        if not registered:
+            return
         try:
-            clear_func = tvm.get_global_func("vm.builtin.clear_py_func_registry")
-            clear_func()
-        except (ValueError, AttributeError):
+            # Release ownership first: if the call below fails, the map must not keep pointing at
+            # a destroyed module, whose id() a later module could reuse.
+            owned = [name for name in registered if _PY_FUNC_OWNERS.get(name) == id(self)]
+            for func_name in owned:
+                del _PY_FUNC_OWNERS[func_name]
+            # Once finalization starts the registry goes away with the process, and the module
+            # globals this needs may already be cleared.
+            if not owned or sys.is_finalizing():
+                return
+            unregister_py_func = tvm.get_global_func("vm.builtin.unregister_py_func")
+            for func_name in owned:
+                unregister_py_func(func_name)
+        except Exception:  # pylint: disable=broad-except
+            # A finalizer must not raise: either the interpreter is shutting down or the runtime
+            # is already unusable, and in both cases the registry no longer matters.
             pass
 
     def __init__(
@@ -203,21 +229,36 @@ def _register_python_functions(self):
         for func_name, py_func in self.ir_mod.__pyfuncs__.items():
 
             def create_py_func_wrapper(name, original_func):
+                # The registry owns the wrapper, so capture the module weakly. A strong capture
+                # would make every registering module immortal, keeping its entries registered
+                # for the lifetime of the process.
+                module_ref = weakref.ref(self)
+
                 def wrapper(*args, **kwargs):
-                    converted_args = [self._convert_tvm_to_pytorch(arg) for arg in args]
+                    module = module_ref()
+                    if module is None:
+                        raise RuntimeError(
+                            f"Python function '{name}' belongs to a BasePyModule that has been "
+                            "destroyed; keep the module alive while calling into it."
+                        )
+
+                    converted_args = [module._convert_tvm_to_pytorch(arg) for arg in args]
                     converted_kwargs = {
-                        k: self._convert_tvm_to_pytorch(v) for k, v in kwargs.items()
+                        k: module._convert_tvm_to_pytorch(v) for k, v in kwargs.items()
                     }
 
-                    result = original_func(self, *converted_args, **converted_kwargs)
+                    result = original_func(module, *converted_args, **converted_kwargs)
 
-                    return self._convert_pytorch_to_tvm(result)
+                    return module._convert_pytorch_to_tvm(result)
 
                 wrapper.__name__ = name
                 return wrapper
 
             wrapped_func = create_py_func_wrapper(func_name, py_func)
             register_py_func(func_name, wrapped_func)
+            _PY_FUNC_OWNERS[func_name] = id(self)
+            if func_name not in self._registered_py_funcs:
+                self._registered_py_funcs = (*self._registered_py_funcs, func_name)
 
     def call_tir(self, tir_func, args, out_ty):
         """Call a TIR function with PyTorch tensors."""

src/runtime/vm/builtin.cc

@@ -472,6 +472,12 @@ void ClearPyFuncRegistry() { py_func_registry.clear(); }
  */
 void RegisterPyFunc(const std::string& name, ffi::Function func) { py_func_registry[name] = func; }
 
+/*!
+ * \brief Unregister a Python function registered with RegisterPyFunc
+ * \param name The function name
+ */
+void UnregisterPyFunc(const std::string& name) { py_func_registry.erase(name); }
+
 /*!
  * \brief Get a registered Python function
  * \param name The function name
@@ -522,6 +528,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
   refl::GlobalDef()
       .def_packed("vm.builtin.call_py_func", CallPyFunc)
       .def("vm.builtin.register_py_func", RegisterPyFunc)
+      .def("vm.builtin.unregister_py_func", UnregisterPyFunc)
       .def("vm.builtin.get_py_func", GetPyFunc)
       .def("vm.builtin.clear_py_func_registry", ClearPyFuncRegistry);
 }

tests/python/relax/test_relax_operators.py

@@ -16,8 +16,11 @@
 # under the License.
 # ruff: noqa: E501, F841
 
+import gc
 import sys
 import tempfile
+import weakref
+from types import SimpleNamespace
 
 import numpy as np
 import pytest
@@ -478,8 +481,32 @@ def multiple_calls(x: R.Tensor((2,), "float32")):
     expected2 = 1.0 / (1.0 + np.exp(-np.maximum(y_data, 0.0)))
     assert (result2.numpy() == expected2).all()
 
-    clear_func = tvm.get_global_func("vm.builtin.clear_py_func_registry")
-    clear_func()
+    unregister_func = tvm.get_global_func("vm.builtin.unregister_py_func")
+    unregister_func("torch_relu")
+    unregister_func("torch_sigmoid")
+
+
+def test_py_func_registry_is_scoped_to_its_module():
+    """A module's finalizer must drop its own registrations and nothing else."""
+    from tvm.relax.base_py_module import BasePyModule
+
+    get_func = tvm.get_global_func("vm.builtin.get_py_func")
+    tvm.get_global_func("vm.builtin.register_py_func")("registry_probe", lambda x: x)
+
+    # __new__ skips __init__'s JIT compilation; only the registration matters here.
+    module = BasePyModule.__new__(BasePyModule)
+    module.ir_mod = SimpleNamespace(pyfuncs={"registry_owned": lambda self, x: x})
+    module._register_python_functions()
+    module_ref = weakref.ref(module)
+
+    del module
+    gc.collect()
+
+    assert module_ref() is None, "the registry must not keep the module alive"
+    assert get_func("registry_probe") is not None, "another owner's function was dropped"
+    with pytest.raises(tvm.error.InternalError, match="not found in registry"):
+        get_func("registry_owned")
+    tvm.get_global_func("vm.builtin.unregister_py_func")("registry_probe")
 
 
 def test_op_to_device(exec_mode):