From 4aa7cee0d8baf7485932b64370b4fc7dc1611874 Mon Sep 17 00:00:00 2001 From: Abhishek Jindal Date: Wed, 19 Jan 2022 14:20:09 -0800 Subject: [PATCH] Abjindal/clean eager backend (#10055) * clearing map for eager mode backends * clearing map for eager mode backends manager * making OrtBackendsManager an extern variable and trying to delete it * cleaning backends manager when the python interpret exits * adding ifdef for eager mode code * disabling warning for pybind state file * disabling warning for python module file * running clang auto format and reducing redundancy * remove new line * moving declaration to a new header file * adding the header file for eager mode for python module * removing source files for eager mode * add source file for python module in eager mode * Update orttraining/orttraining/python/orttraining_python_module_eager.h Co-authored-by: Thiago Crepaldi Co-authored-by: Thiago Crepaldi --- cmake/onnxruntime_python.cmake | 1 + .../orttraining/eager/ort_backends.cpp | 8 ++--- orttraining/orttraining/eager/ort_backends.h | 2 -- orttraining/orttraining/eager/ort_eager.cpp | 9 ++--- .../python/orttraining_python_module.cc | 33 +++++++++++++++++++ .../python/orttraining_python_module_eager.h | 9 +++++ 6 files changed, 50 insertions(+), 12 deletions(-) create mode 100644 orttraining/orttraining/python/orttraining_python_module_eager.h diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index 4a1a6a360b..824d14e302 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -126,6 +126,7 @@ if (onnxruntime_ENABLE_EAGER_MODE) set_source_files_properties("${ORTTRAINING_ROOT}/orttraining/eager/ort_hooks.cpp" PROPERTIES COMPILE_FLAGS -Wno-unused-parameter) set_source_files_properties("${ORTTRAINING_ROOT}/orttraining/eager/ort_ops.cpp" PROPERTIES COMPILE_FLAGS -Wno-unused-parameter) set_source_files_properties("${ORTTRAINING_ROOT}/orttraining/eager/ort_util.cpp" PROPERTIES COMPILE_FLAGS -Wno-unused-parameter) + set_source_files_properties("${ORTTRAINING_ROOT}/orttraining/python/orttraining_python_module.cc" PROPERTIES COMPILE_FLAGS -Wno-unused-parameter) endif() if (MSVC) target_compile_options(onnxruntime_pybind11_state PRIVATE "/wd4100" "/wd4324" "/wd4458" "/wd4127" "/wd4193" "/wd4624" "/wd4702") diff --git a/orttraining/orttraining/eager/ort_backends.cpp b/orttraining/orttraining/eager/ort_backends.cpp index e14ba9d2a1..6e3c2080a6 100644 --- a/orttraining/orttraining/eager/ort_backends.cpp +++ b/orttraining/orttraining/eager/ort_backends.cpp @@ -8,6 +8,7 @@ #include "ort_backends.h" #include "ort_log.h" #include "core/platform/env.h" +#include "orttraining/python/orttraining_python_module_eager.h" //use the environment from python module @@ -24,12 +25,7 @@ namespace torch_ort { namespace eager { using namespace onnxruntime; - -ORTBackendsManager& GetORTBackendsManager() { - auto& env = onnxruntime::python::GetTrainingORTEnv(); - static ORTBackendsManager instance {env.GetLoggingManager()->DefaultLogger()}; - return instance; -} +using namespace onnxruntime::python; onnxruntime::ORTInvoker& GetORTInvoker(const at::Device device) { return GetORTBackendsManager().GetInvoker(device); diff --git a/orttraining/orttraining/eager/ort_backends.h b/orttraining/orttraining/eager/ort_backends.h index 4f3d885a5e..484ab1b192 100644 --- a/orttraining/orttraining/eager/ort_backends.h +++ b/orttraining/orttraining/eager/ort_backends.h @@ -43,8 +43,6 @@ private: std::unordered_map device_ep_info_; }; -ORTBackendsManager& GetORTBackendsManager(); - onnxruntime::ORTInvoker& GetORTInvoker(const at::Device device); } // namespace eager diff --git a/orttraining/orttraining/eager/ort_eager.cpp b/orttraining/orttraining/eager/ort_eager.cpp index f679335258..f7d3d730c3 100644 --- a/orttraining/orttraining/eager/ort_eager.cpp +++ b/orttraining/orttraining/eager/ort_eager.cpp @@ -12,6 +12,7 @@ #include "ort_customops.h" #include "torch/csrc/autograd/python_variable.h" #include "core/framework/tensor.h" +#include "orttraining/python/orttraining_python_module_eager.h" namespace onnxruntime{ namespace python{ @@ -49,7 +50,7 @@ OrtValue ORTTensor_toORTValue(const at::Tensor& data) at::Tensor OrtValue_To_ATen_Tensor(OrtValue& ortvalue) { auto& ort_tensor = ortvalue.Get(); - size_t ort_device_idx = torch_ort::eager::GetORTBackendsManager().GetOrtDeviceIndex(ort_tensor.Location()); + size_t ort_device_idx = GetORTBackendsManager().GetOrtDeviceIndex(ort_tensor.Location()); return torch_ort::eager::aten_tensor_from_ort( std::move(ortvalue), at::TensorOptions() @@ -78,15 +79,15 @@ void addObjectMethodsForEager(py::module& m){ m.def("set_device", [](size_t device_index, const std::string& provider_type, const std::unordered_map& arguments){ - auto status = torch_ort::eager::GetORTBackendsManager().set_device(device_index, provider_type, arguments); + auto status = GetORTBackendsManager().set_device(device_index, provider_type, arguments); if (!status.IsOK()) throw std::runtime_error(status.ErrorMessage()); }); m.def("get_ort_device", [](size_t torch_device_index){ - return torch_ort::eager::GetORTBackendsManager().GetOrtDeviceInfo(torch_device_index); + return GetORTBackendsManager().GetOrtDeviceInfo(torch_device_index); }); m.def("get_ort_device_provider_info", [](size_t torch_device_index){ - return torch_ort::eager::GetORTBackendsManager().GetOrtDeviceProviderInfo(torch_device_index); + return GetORTBackendsManager().GetOrtDeviceProviderInfo(torch_device_index); }); auto customop_module = m.def_submodule("custom_ops"); diff --git a/orttraining/orttraining/python/orttraining_python_module.cc b/orttraining/orttraining/python/orttraining_python_module.cc index ec998bb5d6..3096d1c2cd 100644 --- a/orttraining/orttraining/python/orttraining_python_module.cc +++ b/orttraining/orttraining/python/orttraining_python_module.cc @@ -8,6 +8,10 @@ #include "core/providers/get_execution_providers.h" #include "core/session/provider_bridge_ort.h" +#ifdef ENABLE_EAGER_MODE +#include "orttraining/python/orttraining_python_module_eager.h" +#endif + namespace onnxruntime { namespace python { namespace py = pybind11; @@ -208,6 +212,32 @@ Environment& GetTrainingORTEnv() { return ort_training_env->GetORTEnv(); } +#ifdef ENABLE_EAGER_MODE +using namespace torch_ort::eager; +static std::unique_ptr ort_backends_manager_instance; + +void InitializeBackendsManager() { + auto initialize = [&]() { + static bool initialized = false; + if (initialized) { + return; + } + // Initialization of the module + auto& env = onnxruntime::python::GetTrainingORTEnv(); + ort_backends_manager_instance = std::make_unique(env.GetLoggingManager()->DefaultLogger()); + initialized = true; + }; + initialize(); +} + +ORTBackendsManager& GetORTBackendsManager() { + if (!ort_backends_manager_instance) { + InitializeBackendsManager(); + } + return *ort_backends_manager_instance; +} +#endif + void ResolveExtraProviderOptions(const std::vector& provider_types, const ProviderOptionsMap& original_provider_options_map, ProviderOptionsMap& merged_options){ @@ -325,6 +355,9 @@ PYBIND11_MODULE(onnxruntime_pybind11_state, m) { auto atexit = py::module_::import("atexit"); atexit.attr("register")(py::cpp_function([]() { ort_training_env = nullptr; +#ifdef ENABLE_EAGER_MODE + ort_backends_manager_instance = nullptr; +#endif })); } diff --git a/orttraining/orttraining/python/orttraining_python_module_eager.h b/orttraining/orttraining/python/orttraining_python_module_eager.h new file mode 100644 index 0000000000..ef8d511d84 --- /dev/null +++ b/orttraining/orttraining/python/orttraining_python_module_eager.h @@ -0,0 +1,9 @@ +#include "orttraining/eager/ort_backends.h" + +namespace onnxruntime { +namespace python { + +torch_ort::eager::ORTBackendsManager& GetORTBackendsManager(); + +} +}