akaashrp opened a new pull request, #20417:
URL: https://github.com/apache/tvm/pull/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.
Tests: the `call_py_func` test unregisters the names it registered instead
of clearing the registry, and three tests cover an unrelated module's
finalizer, an entry being released together with its owning module, and two
modules registering the same name. `tests/python/relax` shows the same results
as main on this machine (112 pre-existing failures either way, unrelated to
Relax's Python functions).
--
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]
