From cf4fb50f34bea28c89223a94c57bbd23054af96f Mon Sep 17 00:00:00 2001
From: RunguoLi <li19107254665@gmail.com>
Date: Mon, 12 Oct 2026 15:04:33 -0500
Subject: [PATCH] [Fix][Relax][VM] Keep the call_py_func registry alive until
 process exit

The Python function registry used by `R.call_py_func` was a static
std::unordered_map. Its destructor runs during static destruction, after
the Python interpreter has been finalized, and releases the Python
callables it still holds. Any process that exits with a function still
registered therefore crashes:

  import tvm
  tvm.get_global_func("vm.builtin.register_py_func")("f", lambda x: x)
  # exits with SIGSEGV

Make the registry a heap-allocated singleton that is never destroyed, so
nothing is released after finalization, and add a test that a process
exiting with a function still registered returns 0.
---
 src/runtime/vm/builtin.cc                  | 26 +++++++++++++++-------
 tests/python/relax/test_relax_operators.py | 12 ++++++++++
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/src/runtime/vm/builtin.cc b/src/runtime/vm/builtin.cc
index 334e1d376081..3c2697d5e9b4 100644
--- a/src/runtime/vm/builtin.cc
+++ b/src/runtime/vm/builtin.cc
@@ -470,26 +470,35 @@ TVM_FFI_STATIC_INIT_BLOCK() {
 //  Python function call support
 //-------------------------------------
 
-// Global registry for Python functions
-static std::unordered_map<std::string, ffi::Function> py_func_registry;
+/*!
+ * \brief Global registry for Python functions.
+ *
+ * The registry is heap-allocated and intentionally never destroyed. Its entries hold Python
+ * callables, and releasing them from a static destructor after the Python interpreter has
+ * been finalized corrupts the heap at process exit.
+ */
+static std::unordered_map<std::string, ffi::Function>& PyFuncRegistry() {
+  static auto* registry = new std::unordered_map<std::string, ffi::Function>();
+  return *registry;
+}
 
 /*!
- * \brief Clear the Python function registry on shutdown
+ * \brief Clear the Python function registry
  */
-void ClearPyFuncRegistry() { py_func_registry.clear(); }
+void ClearPyFuncRegistry() { PyFuncRegistry().clear(); }
 
 /*!
  * \brief Register a Python function for call_py_func
  * \param name The function name
  * \param func The Python function wrapped as ffi::Function
  */
-void RegisterPyFunc(const std::string& name, ffi::Function func) { py_func_registry[name] = func; }
+void RegisterPyFunc(const std::string& name, ffi::Function func) { PyFuncRegistry()[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); }
+void UnregisterPyFunc(const std::string& name) { PyFuncRegistry().erase(name); }
 
 /*!
  * \brief Get a registered Python function
@@ -497,8 +506,9 @@ void UnregisterPyFunc(const std::string& name) { py_func_registry.erase(name); }
  * \return The Python function
  */
 ffi::Function GetPyFunc(const std::string& name) {
-  auto it = py_func_registry.find(name);
-  if (it == py_func_registry.end()) {
+  auto& registry = PyFuncRegistry();
+  auto it = registry.find(name);
+  if (it == registry.end()) {
     TVM_FFI_THROW(InternalError) << "Python function '" << name << "' not found in registry";
   }
   return it->second;
diff --git a/tests/python/relax/test_relax_operators.py b/tests/python/relax/test_relax_operators.py
index ccde7417343e..aa263f84a96e 100644
--- a/tests/python/relax/test_relax_operators.py
+++ b/tests/python/relax/test_relax_operators.py
@@ -731,6 +731,18 @@ def test_py_func_registry_is_scoped_to_its_module():
     tvm_ffi.get_global_func("vm.builtin.unregister_py_func")("registry_probe")
 
 
+def test_py_func_registry_exit_with_registered_function():
+    """Exiting with a Python function still registered must not crash the process."""
+    import subprocess
+
+    code = (
+        "import tvm\n"
+        "tvm.get_global_func('vm.builtin.register_py_func')('py_identity', lambda x: x)\n"
+    )
+    result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True)
+    assert result.returncode == 0, result.stderr
+
+
 def test_op_to_device():
     @tvm_tirx.script.ir_module
     class CallToDevice: