This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 60c250be20 [Fix][Relax] Scope BasePyModule's Python function registry
entries to the module (#20417)
60c250be20 is described below
commit 60c250be20c1109d1e54eb12a91b342c500ec0dc
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Thu Sep 24 11:38:16 2026 -0700
[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.
---
python/tvm/relax/base_py_module.py | 57 +++++++++++++++++++++++++-----
src/runtime/vm/builtin.cc | 7 ++++
tests/python/relax/test_relax_operators.py | 31 ++++++++++++++--
3 files changed, 85 insertions(+), 10 deletions(-)
diff --git a/python/tvm/relax/base_py_module.py
b/python/tvm/relax/base_py_module.py
index cba22b5dba..284201acb8 100644
--- a/python/tvm/relax/base_py_module.py
+++ b/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 @@ except ImportError:
_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 @@ class BasePyModule:
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."""
diff --git a/src/runtime/vm/builtin.cc b/src/runtime/vm/builtin.cc
index fa5f5ce7ce..9c2651e7c8 100644
--- a/src/runtime/vm/builtin.cc
+++ b/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);
}
diff --git a/tests/python/relax/test_relax_operators.py
b/tests/python/relax/test_relax_operators.py
index 37243843e3..639ccbe67a 100644
--- a/tests/python/relax/test_relax_operators.py
+++ b/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 test_op_call_py_func(exec_mode):
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):
