LiRunGuo opened a new pull request, #20413:
URL: https://github.com/apache/tvm/pull/20413
`tests/python/relax/test_relax_operators.py::test_op_call_py_func` fails
intermittently in CI with
```
tvm.error.InternalError: Python function 'torch_relu' not found in registry
```
(for example in the `cpu/pr-head` run of #20406, where it was the only
failure out of 11,464 tests). The cause is `BasePyModule.__del__`, which clears
the entire global `call_py_func` registry. `BasePyModule` instances sit in
reference cycles, because the wrappers they attach capture `self`, so they are
freed by the cyclic garbage collector at an arbitrary later point. When a
module left over from an earlier test in the same pytest-xdist worker is
collected between `register_py_func("torch_relu", ...)` and the VM call, the
function the test just registered disappears. The same thing can happen to
users: dropping any `BasePyModule` silently removes Python functions registered
by unrelated code.
This reproduces deterministically:
```python
import gc, numpy as np, tvm
from tvm.relax.base_py_module import BasePyModule
from tvm.script import relax as R
@tvm.script.ir_module
class Mod:
@R.function
def main(x: R.Tensor((3,), "float32")):
return R.call_py_func(R.str("f"), (x,), out_ty=R.Tensor((3,),
"float32"))
vm = tvm.relax.VirtualMachine(tvm.compile(Mod, target="llvm"), tvm.cpu())
gc.disable()
BasePyModule(tvm.IRModule({}), device=tvm.cpu(0), target="llvm")
tvm.get_global_func("vm.builtin.register_py_func")("f", lambda x: x[0])
gc.collect()
vm["main"](tvm.runtime.tensor(np.zeros(3, "float32"))) # "Python function
'f' not found in registry"
```
The global clear cannot remove the module's own entries in any case, because
the registered wrappers keep the module alive until they are removed. It
appears to exist to avoid a crash at exit. The registry was a static
`std::unordered_map`, and destroying it after the interpreter has been
finalized releases the Python callables it holds. Any process that exits with a
function still registered aborts with `double free or corruption` (exit code
134); this is reproducible with just `import tvm` and one `register_py_func`
call.
This PR fixes both problems at the source. The registry becomes a
heap-allocated singleton that is never destroyed, which is the usual way to
avoid static destruction order issues, so nothing is released after
finalization. `BasePyModule.__del__`, whose only job was the global clear, is
removed.
The new tests check that collecting an unrelated `BasePyModule` keeps a
registered function callable from the VM in both bytecode and compiled mode,
and that a subprocess exiting with a function still registered returns 0. All
three fail without this change. The existing `call_py_func`, `BasePyModule`,
PyTorch integration and DLPack tests pass.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
