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 <thiago.crepaldi@microsoft.com>

Co-authored-by: Thiago Crepaldi <thiago.crepaldi@microsoft.com>

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")

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);

orttraining/orttraining/eager/ort_backends.h

@@ -43,8 +43,6 @@ private:
   std::unordered_map<at::DeviceIndex, ProviderInfoMap> device_ep_info_;
 };
 
-ORTBackendsManager& GetORTBackendsManager();
-
 onnxruntime::ORTInvoker& GetORTInvoker(const at::Device device);
 
 } // namespace eager

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<Tensor>();
-  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<std::string, std::string>& 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");

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<ORTBackendsManager> 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<ORTBackendsManager>(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<std::string>& 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
   }));
 }
 

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();
+
+}
+}