Fix segment fault for custom function (#8331)

* unregister registered python functions upon normal interpreter termination
* atexit.register(unregister_python_functions) should be called by __init__.py
* minor fix

onnxruntime/core/language_interop_ops/torch/custom_function_register.cc

@@ -109,7 +109,7 @@ static void RegisterEntry(
   // Own the Python object.
   Py_INCREF(obj);
   PythonObjectPtr ptr(obj, PythonObjectDeleter);
-  
+
   // If an obj has been registered, this old ownership is automatically released
   // after this move-assignment. Then, the "storage" owns the new object.
   storage = std::move(ptr);
@@ -184,6 +184,15 @@ PyObject* OrtTorchFunctionPool::GetContext(int64_t context_index) {
   ORT_ENFORCE(iter != func_context_pool_.end(), "No context registered for ", context_index);
   return iter->second.get();
 }
+
+void OrtTorchFunctionPool::UnRegisterFunctions() {
+  forward_runner_.reset();
+  backward_runner_.reset();
+  forward_core_pool_.clear();
+  backward_core_pool_.clear();
+  func_context_pool_.clear();
+}
+
 }  // namespace torch
 }  // namespace language_interop_ops
 }  // namespace onnxruntime

onnxruntime/core/language_interop_ops/torch/custom_function_register.h

@@ -23,12 +23,12 @@ class OrtTorchFunctionPool final {
   // AutogradFunction includes ForwardCore and BackwardCore.
   // ForwardCore is the apply() function pointer.
   // BackwardCore is the backward() function pointer.
-  // RegisterTorchAutogradFunction owns the input "obj" and will release its ownership only in its destructor. 
+  // RegisterTorchAutogradFunction owns the input "obj" and will release its ownership only in its destructor.
   void RegisterTorchAutogradFunction(const std::string& key, PyObject* obj);
   // Return a borrowed reference to the stored Python function. Thus,
   //  1. The returned value doesn't own its Python function.
   //  2. Caller of GetForwardCore should not decrease the reference count of the returned object.
-  PyObject* GetForwardCore(const std::string& key);   // The "key" is the "name" attribute in PythonOp.
+  PyObject* GetForwardCore(const std::string& key);  // The "key" is the "name" attribute in PythonOp.
   // Return a borrowed reference to the stored Python function. Thus,
   //  1. The returned value doesn't own its Python function.
   //  2. Caller of GetBackwardCore should not decrease the reference count of the returned object.
@@ -59,8 +59,21 @@ class OrtTorchFunctionPool final {
   // is responsible for executing autograd.Function.apply.
   PyObject* GetBackwardRunner();
 
+  // The reason we provide this unregister api is:
+  //   A static OrtTorchFunctionPool instance will be destructed after
+  //   Python modules/functions are released. Once we own func pointers
+  //   by increasing ref count for the functions, we need decrease the
+  //   ref count in ~OrtTorchFunctionPool, but at that time some properties
+  //   of the python function object, for example co_consts
+  //   (tuple type, https://github.com/python/cpython/blob/3.7/Objects/funcobject.c#L38)
+  //   already released, there will be a segment fault.
+
+  //   Calling this function upon normal interpreter termination helps release the
+  //   registered functions earlier, avoiding segment fault.
+  void UnRegisterFunctions();
+
  private:
-  OrtTorchFunctionPool() {};
+  OrtTorchFunctionPool(){};
   ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(OrtTorchFunctionPool);
 
   PythonObjectPtr forward_runner_;

onnxruntime/python/onnxruntime_pybind_state.cc

@@ -849,6 +849,13 @@ void addGlobalMethods(py::module& m, Environment& env) {
 #else
         ORT_UNUSED_PARAMETER(key);
         ORT_UNUSED_PARAMETER(obj);
+#endif
+  });
+  m.def("unregister_python_functions", []() -> void {
+#ifdef ENABLE_TRAINING_TORCH_INTEROP
+    // Release all custom python functions registered.
+    auto& pool = onnxruntime::language_interop_ops::torch::OrtTorchFunctionPool::GetInstance();
+    pool.UnRegisterFunctions();
 #endif
   });
 #endif

orttraining/orttraining/core/agent/training_agent.h

@@ -26,16 +26,14 @@ class TrainingAgent {
   // For ORTModule.forward()
   common::Status RunForward(const std::vector<OrtValue>& feeds, std::vector<OrtValue>& fetches,
                             PartialGraphExecutionState& state) ORT_MUST_USE_RESULT;
-  ;
+
   // For ORTModule.backward()
   common::Status RunBackward(const std::vector<OrtValue>& feeds, std::vector<OrtValue>& fetches,
                              PartialGraphExecutionState& state) ORT_MUST_USE_RESULT;
-  ;
 
   common::Status RunCore(const std::vector<OrtValue>& feeds, std::vector<OrtValue>& fetches,
                          PartialGraphExecutionState& state, FeedsFetchesManager& feeds_fetches_manager)
       ORT_MUST_USE_RESULT;
-  ;
 
   void CreateAndInitializeFeedsFetchesManager(const SessionState& session_state,
                                               const std::vector<std::string>& feed_names,

orttraining/orttraining/python/training/ortmodule/_custom_autograd_function.py

@@ -5,12 +5,16 @@
 
 # Initialize static objects needed to run custom autograd.Function's.
 def enable_custom_autograd_support():
-    from onnxruntime.capi._pybind_state import register_forward_runner, register_backward_runner
+    from onnxruntime.capi._pybind_state import register_forward_runner, register_backward_runner, unregister_python_functions
     from torch.onnx import register_custom_op_symbolic
     from ._custom_autograd_function_exporter import _export
     from ._custom_autograd_function_runner import call_python_forward_function, call_python_backward_function
+    import atexit
 
     register_forward_runner(call_python_forward_function)
     register_backward_runner(call_python_backward_function)
 
+    # Unregister all python functions automatically upon normal interpreter termination.
+    atexit.register(unregister_python_functions)
+
     register_custom_op_symbolic('::prim_PythonOp', _export, 1)

orttraining/orttraining/test/python/_test_helpers.py

@@ -188,7 +188,7 @@ def assert_values_are_close(input, other, rtol=1e-05, atol=1e-06):
 
 def enable_custom_autograd_function(module):
     for mode in [True, False]:
-        module._execution_manager(mode)._enable_custom_autograd_function = True
+        module._torch_module._execution_manager(mode)._enable_custom_autograd_function = True
 
 def run_with_pytorch_on_device(device, model, input_list, label_input, is_eval_mode=False):
     model.to(device)