diff --git a/onnxruntime/core/language_interop_ops/torch/custom_function_register.cc b/onnxruntime/core/language_interop_ops/torch/custom_function_register.cc index 2a3a8a7852..72af6a2356 100644 --- a/onnxruntime/core/language_interop_ops/torch/custom_function_register.cc +++ b/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 diff --git a/onnxruntime/core/language_interop_ops/torch/custom_function_register.h b/onnxruntime/core/language_interop_ops/torch/custom_function_register.h index 33e6e70be4..107b312bc9 100644 --- a/onnxruntime/core/language_interop_ops/torch/custom_function_register.h +++ b/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_; diff --git a/onnxruntime/python/onnxruntime_pybind_state.cc b/onnxruntime/python/onnxruntime_pybind_state.cc index 5070a00f85..d53fa87544 100644 --- a/onnxruntime/python/onnxruntime_pybind_state.cc +++ b/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 diff --git a/orttraining/orttraining/core/agent/training_agent.h b/orttraining/orttraining/core/agent/training_agent.h index 92361dccf5..63f4794a97 100644 --- a/orttraining/orttraining/core/agent/training_agent.h +++ b/orttraining/orttraining/core/agent/training_agent.h @@ -26,16 +26,14 @@ class TrainingAgent { // For ORTModule.forward() common::Status RunForward(const std::vector& feeds, std::vector& fetches, PartialGraphExecutionState& state) ORT_MUST_USE_RESULT; - ; + // For ORTModule.backward() common::Status RunBackward(const std::vector& feeds, std::vector& fetches, PartialGraphExecutionState& state) ORT_MUST_USE_RESULT; - ; common::Status RunCore(const std::vector& feeds, std::vector& fetches, PartialGraphExecutionState& state, FeedsFetchesManager& feeds_fetches_manager) ORT_MUST_USE_RESULT; - ; void CreateAndInitializeFeedsFetchesManager(const SessionState& session_state, const std::vector& feed_names, diff --git a/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function.py b/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function.py index 656f09db32..e6828982be 100644 --- a/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function.py +++ b/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) diff --git a/orttraining/orttraining/test/python/_test_helpers.py b/orttraining/orttraining/test/python/_test_helpers.py index c1d6192c95..98bfff95ae 100644 --- a/orttraining/orttraining/test/python/_test_helpers.py +++ b/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)