Revert "Call pluggable EP's shutdown function in Environment::~Environment() (#11120)" (#11393)

This reverts commit 4983d6e5d6. We can't destroy OrtEnv through python's atexit function, because at that time there might be many other ORT python objects alive.

include/onnxruntime/core/session/environment.h

@@ -80,7 +80,6 @@ class Environment {
   Status UnregisterAllocator(const OrtMemoryInfo& mem_info);
 
   Environment() = default;
-  ~Environment();
 
  private:
   ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(Environment);

onnxruntime/core/session/environment.cc

@@ -12,7 +12,6 @@
 #include "onnx/defs/operator_sets_ml.h"
 #include "core/graph/contrib_ops/ms_opset.h"
 #include "core/graph/contrib_ops/onnx_deprecated_opset.h"
-#include "core/framework/provider_shutdown.h"
 #if defined(ENABLE_TRAINING) || defined(ENABLE_TRAINING_OPS)
 #include "onnx/defs/operator_sets_training.h"
 #endif
@@ -50,13 +49,6 @@ using namespace ONNX_NAMESPACE;
 
 std::once_flag schemaRegistrationOnceFlag;
 
-Environment::~Environment() {
-// We don't support any shared providers in the minimal build yet
-#if !defined(ORT_MINIMAL_BUILD)
-  UnloadSharedProviders();
-#endif
-}
-
 Status Environment::Create(std::unique_ptr<logging::LoggingManager> logging_manager,
                            std::unique_ptr<Environment>& environment,
                            const OrtThreadingOptions* tp_options,

onnxruntime/core/session/ort_env.cc

@@ -35,6 +35,13 @@ OrtEnv::OrtEnv(std::unique_ptr<onnxruntime::Environment> value1)
     : value_(std::move(value1)) {
 }
 
+OrtEnv::~OrtEnv() {
+// We don't support any shared providers in the minimal build yet
+#if !defined(ORT_MINIMAL_BUILD)
+  UnloadSharedProviders();
+#endif
+}
+
 OrtEnv* OrtEnv::GetInstance(const OrtEnv::LoggingManagerConstructionInfo& lm_info,
                             onnxruntime::common::Status& status,
                             const OrtThreadingOptions* tp_options) {

onnxruntime/core/session/ort_env.h

@@ -74,6 +74,7 @@ struct OrtEnv {
    */
   onnxruntime::common::Status UnregisterAllocator(const OrtMemoryInfo& mem_info);
   OrtEnv(std::unique_ptr<onnxruntime::Environment> value);
+  ~OrtEnv();
 
  private:
   static std::unique_ptr<OrtEnv> p_instance_;

onnxruntime/python/onnxruntime_pybind_state.cc

@@ -1563,9 +1563,6 @@ including arg name, arg type (contains both type and shape).)pbdoc")
       .export_values();
 }
 
-// static variable used to create inference session and training session.
-static std::unique_ptr<Environment> session_env;
-
 void CreateInferencePybindStateModule(py::module& m) {
   m.doc() = "pybind11 stateful interface to ONNX runtime";
   RegisterExceptions(m);
@@ -1596,10 +1593,6 @@ void CreateInferencePybindStateModule(py::module& m) {
   addOpSchemaSubmodule(m);
   addOpKernelSubmodule(m);
 #endif
-  auto atexit = py::module_::import("atexit");
-  atexit.attr("register")(py::cpp_function([]() {
-    session_env = nullptr;
-  }));
 }
 
 void InitArray() {
@@ -1609,6 +1602,9 @@ void InitArray() {
   })();
 }
 
+// static variable used to create inference session and training session.
+static std::unique_ptr<Environment> session_env;
+
 void InitializeEnv() {
   auto initialize = [&]() {
     // Initialization of the module

onnxruntime/test/framework/inference_session_test.cc

@@ -576,9 +576,7 @@ TEST(InferenceSessionTests, ModelMetadata) {
   }
 }
 #endif
-
-// TODO: move it to a new executable file
-TEST(InferenceSessionTests, DISABLED_CheckRunLogger) {
+TEST(InferenceSessionTests, CheckRunLogger) {
   SessionOptions so;
 
   so.session_logid = "CheckRunLogger";
@@ -776,8 +774,7 @@ TEST(InferenceSessionTests, PreAllocateOutputVector) {
   RunModel(session_object, run_options, is_preallocate_output_vec);
 }
 
-//TODO: move it to a new executable file
-TEST(InferenceSessionTests, DISABLED_ConfigureVerbosityLevel) {
+TEST(InferenceSessionTests, ConfigureVerbosityLevel) {
   SessionOptions so;
 
   so.session_logid = "ConfigureVerbosityLevel";
@@ -2367,9 +2364,8 @@ class InferenceSessionTestGlobalThreadPools : public InferenceSessionWrapper {
   }
 };
 
-// TODO: move it
 // Test 1: env created WITHOUT global tp / use per session tp (default case): in this case per session tps should be in use
-TEST(InferenceSessionTests, DISABLED_CheckIfPerSessionThreadPoolsAreBeingUsed) {
+TEST(InferenceSessionTests, CheckIfPerSessionThreadPoolsAreBeingUsed) {
   SessionOptions so;
   so.use_per_session_threads = true;
 
@@ -2408,7 +2404,7 @@ TEST(InferenceSessionTests, DISABLED_CheckIfPerSessionThreadPoolsAreBeingUsed) {
 }
 
 // Test 2: env created with global tp / DONT use per session tp: in this case global tps should be in use
-TEST(InferenceSessionTests, DISABLED_CheckIfGlobalThreadPoolsAreBeingUsed) {
+TEST(InferenceSessionTests, CheckIfGlobalThreadPoolsAreBeingUsed) {
   SessionOptions so;
   so.use_per_session_threads = false;
 
@@ -2446,7 +2442,7 @@ TEST(InferenceSessionTests, DISABLED_CheckIfGlobalThreadPoolsAreBeingUsed) {
 }
 
 // Test 3: env created with global tp / use per session tp: in this case per session tps should be in use
-TEST(InferenceSessionTests, DISABLED_CheckIfPerSessionThreadPoolsAreBeingUsed2) {
+TEST(InferenceSessionTests, CheckIfPerSessionThreadPoolsAreBeingUsed2) {
   SessionOptions so;
   so.use_per_session_threads = true;
 
@@ -2493,7 +2489,7 @@ TEST(InferenceSessionTests, DISABLED_CheckIfPerSessionThreadPoolsAreBeingUsed2)
 }
 
 // Test 4: env created WITHOUT global tp / DONT use per session tp --> this should throw an exception
-TEST(InferenceSessionTests, DISABLED_InvalidSessionEnvCombination) {
+TEST(InferenceSessionTests, InvalidSessionEnvCombination) {
   SessionOptions so;
   so.use_per_session_threads = false;
 
@@ -2530,7 +2526,7 @@ class InferenceSessionTestSharingAllocator : public InferenceSessionWrapper {
 };
 
 // Ensure sessions use the same allocator. It uses ORT created allocator.
-TEST(InferenceSessionTests, DISABLED_AllocatorSharing_EnsureSessionsUseSameOrtCreatedAllocator) {
+TEST(InferenceSessionTests, AllocatorSharing_EnsureSessionsUseSameOrtCreatedAllocator) {
   auto logging_manager = std::make_unique<logging::LoggingManager>(
       std::unique_ptr<ISink>(new CLogSink()), logging::Severity::kVERBOSE, false,
       LoggingManager::InstanceType::Temporal);
@@ -2575,7 +2571,7 @@ TEST(InferenceSessionTests, DISABLED_AllocatorSharing_EnsureSessionsUseSameOrtCr
 }
 
 // Ensure sessions don't use the same allocator. It uses ORT created allocator.
-TEST(InferenceSessionTests, DISABLED_AllocatorSharing_EnsureSessionsDontUseSameOrtCreatedAllocator) {
+TEST(InferenceSessionTests, AllocatorSharing_EnsureSessionsDontUseSameOrtCreatedAllocator) {
   auto logging_manager = std::make_unique<logging::LoggingManager>(
       std::unique_ptr<ISink>(new CLogSink()), logging::Severity::kVERBOSE, false,
       LoggingManager::InstanceType::Temporal);
@@ -2627,7 +2623,7 @@ class InferenceSessionTestSharingInitializer : public InferenceSessionWrapper {
   }
 };
 
-TEST(InferenceSessionTests, DISABLED_InitializerSharing_EnsureSessionsUseUserAddedInitializer) {
+TEST(InferenceSessionTests, InitializerSharing_EnsureSessionsUseUserAddedInitializer) {
   auto logging_manager = std::make_unique<logging::LoggingManager>(
       std::unique_ptr<ISink>(new CLogSink()), logging::Severity::kVERBOSE, false,
       LoggingManager::InstanceType::Temporal);
@@ -2759,7 +2755,7 @@ void VerifyThreadPoolWithDenormalAsZero(onnxruntime::concurrency::ThreadPool* tp
 }
 
 // test global thread pool with setting denormal as zero
-TEST(InferenceSessionTests, DISABLED_GlobalThreadPoolWithDenormalAsZero) {
+TEST(InferenceSessionTests, GlobalThreadPoolWithDenormalAsZero) {
   // test if denormal-as-zero mode is supported
   if (!SetDenormalAsZero(false)) {
     return;
@@ -2807,7 +2803,7 @@ TEST(InferenceSessionTests, DISABLED_GlobalThreadPoolWithDenormalAsZero) {
 }
 
 // test inter thread pool with setting denormal as zero
-TEST(InferenceSessionTests, DISABLED_InterThreadPoolWithDenormalAsZero) {
+TEST(InferenceSessionTests, InterThreadPoolWithDenormalAsZero) {
   // test if denormal-as-zero mode is supported
   if (!SetDenormalAsZero(false)) {
     return;

onnxruntime/test/providers/provider_test_utils.cc

@@ -1134,7 +1134,6 @@ void OpTester::Run(
               provider_type == onnxruntime::kDnnlExecutionProvider)
             continue;
           auto reg = execution_provider->GetKernelRegistry();
-          ASSERT_NE(nullptr, reg);
           if (!KernelRegistry::HasImplementationOf(*reg, node, execution_provider->Type())) {
             valid = false;
             for (auto& custom_session_registry : custom_session_registries_) {

orttraining/orttraining/test/graph/gradient_graph_builder_test.cc

@@ -6,7 +6,6 @@
 #include "orttraining/core/optimizer/gist_encode_decode.h"
 #include "test/providers/provider_test_utils.h"
 #include "test/framework/test_utils.h"
-#include "test/test_environment.h"
 #include "test/util/include/default_providers.h"
 #include "core/common/path_utils.h"
 #include "core/providers/cpu/cpu_execution_provider.h"
@@ -38,10 +37,11 @@ static Status BuildBackPropGraph(
     const PathString& forward_model_file,
     const TrainingSession::TrainingConfiguration& config,
     PathString& backward_model_file) {
-  const Environment& env = GetEnvironment();
+  std::unique_ptr<Environment> env;
+  ORT_RETURN_IF_ERROR(Environment::Create(nullptr, env));
 
   SessionOptions so{};
-  TrainingSession training_session{so, env};
+  TrainingSession training_session{so, *env};
 
   std::cout << "Loading source model file = " << ToUTF8String(forward_model_file) << "\n";
 
@@ -63,9 +63,10 @@ static Status BuildBackPropGraph(
  */
 static std::unique_ptr<TrainingSession> RunTrainingSessionWithChecks(
     const SessionOptions& so, const PathString& backprop_model_file) {
-  const Environment& env = GetEnvironment();
+  std::unique_ptr<Environment> env;
+  ORT_THROW_IF_ERROR(Environment::Create(nullptr, env));
 
-  std::unique_ptr<TrainingSession> training_session = std::make_unique<TrainingSession>(so, env);
+  std::unique_ptr<TrainingSession> training_session = std::make_unique<TrainingSession>(so, *env);
 
   ORT_THROW_IF_ERROR(training_session->Load(backprop_model_file));
 
@@ -223,10 +224,11 @@ TEST(GradientGraphBuilderTest, TrainingSession_WithGist) {
 
   // Create backward graph with gist transformations
   const PathString& forward_model_file = ORIGINAL_MODEL_PATH;
-  const Environment& env = GetEnvironment();
+  std::unique_ptr<Environment> env;
+  ORT_THROW_IF_ERROR(Environment::Create(nullptr, env));
 
   SessionOptions so{};
-  TrainingSession training_session{so, env};
+  TrainingSession training_session{so, *env};
 
   std::cout << "Loading source model file = " << ToUTF8String(forward_model_file) << "\n";
 
@@ -385,9 +387,10 @@ TEST(GradientGraphBuilderTest, TrainingSession_WithProfiler) {
 static void RunBertTrainingWithChecks(
     const SessionOptions& so,
     const PathString& backprop_model_file) {
-  const Environment& env = GetEnvironment();
+  std::unique_ptr<Environment> env;
+  ASSERT_STATUS_OK(Environment::Create(nullptr, env));
 
-  std::unique_ptr<TrainingSession> training_session = std::make_unique<TrainingSession>(so, env);
+  std::unique_ptr<TrainingSession> training_session = std::make_unique<TrainingSession>(so, *env);
 
   ASSERT_STATUS_OK(training_session->Load(backprop_model_file));
 
@@ -1719,7 +1722,8 @@ TEST(GradientGraphBuilderTest, TrainingSession_WithPipeline) {
   splitter.Split(backprop_model_file, sub_model_files, cuts);
 
   // create training sessions
-  const Environment& env = GetEnvironment();
+  std::unique_ptr<Environment> env;
+  ASSERT_STATUS_OK(Environment::Create(nullptr, env));
 
   struct SubSession {
     std::unique_ptr<TrainingSession> sess;
@@ -1737,7 +1741,7 @@ TEST(GradientGraphBuilderTest, TrainingSession_WithPipeline) {
     sub_sess.run_options.run_tag = sub_sess.so.session_logid;
     sub_sess.run_options.training_mode = true;
 
-    sub_sess.sess = std::make_unique<TrainingSession>(sub_sess.so, env);
+    sub_sess.sess = std::make_unique<TrainingSession>(sub_sess.so, *env);
     ASSERT_STATUS_OK(sub_sess.sess->Load(sub_model_files[sub_id]));
     ASSERT_STATUS_OK(sub_sess.sess->Initialize());
   }

orttraining/orttraining/test/model/training_runner_test.cc

@@ -3,7 +3,7 @@
 
 // ReduceAllL2 op only have the CUDA implementation
 #ifdef USE_CUDA
-#include "test/test_environment.h"
+
 #include "orttraining/models/runner/training_runner.h"
 
 #include "gtest/gtest.h"
@@ -40,9 +40,10 @@ TEST(TrainingRunnerTest, Basic) {
   params.fetch_names = {"predictions"};
   params.loss_func_info = LossFunctionInfo(OpDef("MeanSquaredError"), "loss", {"predictions", "labels"});
 
-  const Environment& env = onnxruntime::test::GetEnvironment();
+  std::unique_ptr<Environment> env;
+  ASSERT_TRUE(Environment::Create(nullptr, env).IsOK());
 
-  TrainingRunner runner{params, env};
+  TrainingRunner runner{params, *env};
 
   ASSERT_TRUE(runner.GetSession().RegisterExecutionProvider(onnxruntime::test::DefaultCudaExecutionProvider()).IsOK());
 
@@ -92,8 +93,9 @@ TEST(TrainingRunnerTest, DISABLED_PipelineRun) {
   params.pipeline_partition_cut_list.emplace_back(cut0);
   params.pipeline_partition_cut_list.emplace_back(cut1);
 
-  const Environment& env = onnxruntime::test::GetEnvironment();
-  TrainingRunner runner{params, env};
+  std::unique_ptr<Environment> env;
+  ASSERT_TRUE(Environment::Create(nullptr, env).IsOK());
+  TrainingRunner runner{params, *env};
 
   auto status = runner.Initialize();
   ASSERT_TRUE(status.IsOK()) << status.ErrorMessage();

orttraining/orttraining/test/session/training_session_test_utils.cc

@@ -162,9 +162,10 @@ TrainingSession::TrainingConfiguration MakeBasicTrainingConfig() {
 std::unique_ptr<TrainingSession> BuildAndRunTrainingSessionWithChecks(
     const SessionOptions& so, const PathString& forward_model_file,
     const TrainingSession::TrainingConfiguration& config) {
-  const Environment& env = onnxruntime::test::GetEnvironment();
+  std::unique_ptr<Environment> env;
+  ORT_THROW_IF_ERROR(Environment::Create(nullptr, env));
 
-  std::unique_ptr<TrainingSession> training_session = std::make_unique<TrainingSession>(so, env);
+  std::unique_ptr<TrainingSession> training_session = std::make_unique<TrainingSession>(so, *env);
 
   std::cout << "Loading source model file = " << ToUTF8String(forward_model_file) << "\n";