Export env to python

oneflow/api/python/env/env.cpp

@@ -16,19 +16,41 @@ limitations under the License.
 #include <pybind11/pybind11.h>
 #include "oneflow/api/python/env/env.h"
 #include "oneflow/api/python/of_api_registry.h"
+#include "oneflow/core/common/global.h"
+#include "oneflow/core/vm/vm_util.h"
+#include "oneflow/core/vm/virtual_machine.h"
+#include "oneflow/core/framework/shut_down_util.h"
 
 namespace py = pybind11;
 
 namespace oneflow {
 
+Maybe<void> SwitchToShuttingDownPhase(EnvGlobalObjectsScope* env, bool is_normal_exit) {
+  if (is_normal_exit) {
+    JUST(vm::ClusterSync());
+    auto* vm = JUST(GlobalMaybe<VirtualMachine>());
+    JUST(vm->CloseVMThreads());
+  }
+  JUST(env->init_is_normal_exit(is_normal_exit));
+  SetShuttingDown(true);
+  return Maybe<void>::Ok();
+}
+
 ONEFLOW_API_PYBIND11_MODULE("", m) {
   m.def("CurrentResource", &CurrentResource);
   m.def("EnvResource", &EnvResource);
   m.def("EnableEagerEnvironment", &EnableEagerEnvironment);
 
-  m.def("IsEnvInited", &IsEnvInited);
-  m.def("InitEnv", &InitEnv);
-  m.def("DestroyEnv", &DestroyEnv, py::call_guard<py::gil_scoped_release>());
+  py::class_<EnvGlobalObjectsScope, std::shared_ptr<EnvGlobalObjectsScope>>(m, "Env")
+      .def(py::init([](const std::string& env_proto_str) {
+        return CreateEnv(env_proto_str).GetPtrOrThrow();
+      }))
+      .def(
+          "SwitchToShuttingDownPhase",
+          [](EnvGlobalObjectsScope* env, bool is_normal_exit) {
+            SwitchToShuttingDownPhase(env, is_normal_exit).GetOrThrow();
+          },
+          py::call_guard<py::gil_scoped_release>());
 
   m.def("CurrentMachineId", &CurrentMachineId);
 

oneflow/api/python/env/env.h

@@ -46,25 +46,13 @@ inline Maybe<void> EnableEagerEnvironment(bool enable_eager_execution) {
   return Maybe<void>::Ok();
 }
 
-inline Maybe<bool> IsEnvInited() { return Global<EnvGlobalObjectsScope>::Get() != nullptr; }
-
-inline Maybe<void> DestroyEnv() {
-  if (Global<EnvGlobalObjectsScope>::Get() == nullptr) { return Maybe<void>::Ok(); }
-  OF_ENV_BARRIER();
-  Global<EnvGlobalObjectsScope>::Delete();
-  return Maybe<void>::Ok();
-}
-
-inline Maybe<void> InitEnv(const std::string& env_proto_str) {
+inline Maybe<EnvGlobalObjectsScope> CreateEnv(const std::string& env_proto_str) {
   EnvProto env_proto;
   CHECK_OR_RETURN(TxtString2PbMessage(env_proto_str, &env_proto))
       << "failed to parse env_proto" << env_proto_str;
-  CHECK_ISNULL_OR_RETURN(Global<EnvGlobalObjectsScope>::Get());
-  // Global<T>::New is not allowed to be called here
-  // because glog is not constructed yet and has bad bahavior.
-  Global<EnvGlobalObjectsScope>::SetAllocated(new EnvGlobalObjectsScope());
-  JUST(Global<EnvGlobalObjectsScope>::Get()->Init(env_proto));
-  return Maybe<void>::Ok();
+  auto env = std::make_shared<EnvGlobalObjectsScope>();
+  JUST(env->Init(env_proto));
+  return env;
 }
 
 inline Maybe<long long> CurrentMachineId() { return GlobalProcessCtx::Rank(); }

oneflow/api/python/init.cpp

@@ -51,12 +51,6 @@ bool Int2IntListMapContaining(const Int2IntListMap& bigger, const Int2IntListMap
 }  // namespace
 
 PYBIND11_MODULE(_oneflow_internal, m) {
-  m.def("MasterSendAbort", []() {
-    if (Global<EnvGlobalObjectsScope>::Get() != nullptr) {
-      return ClusterInstruction::MasterSendAbort();
-    }
-  });
-
   using IntList = std::vector<int64_t>;
   using Int2IntListMap = std::unordered_map<int64_t, std::shared_ptr<IntList>>;
 

oneflow/api/python/session/session.h

@@ -119,7 +119,6 @@ inline Maybe<void> CreateMultiClientSessionContext() {
 
 inline Maybe<void> InitMultiClientSessionContext(const std::string& config_proto_str) {
   CHECK_NOTNULL_OR_RETURN(Global<MultiClientSessionContext>::Get());
-  CHECK_NOTNULL_OR_RETURN(Global<EnvGlobalObjectsScope>::Get());
   CHECK_NOTNULL_OR_RETURN(Global<EnvDesc>::Get()) << "env not found";
 
   ConfigProto config_proto;

oneflow/core/framework/random_generator_impl.cpp

@@ -20,7 +20,7 @@ limitations under the License.
 #include "oneflow/core/framework/instructions_builder.h"
 #include "oneflow/core/framework/tensor_util.h"
 #include "oneflow/core/functional/functional.h"
-#include "oneflow/core/job/env_global_objects_scope.h"
+#include "oneflow/core/vm/virtual_machine.h"
 #include "oneflow/core/register/ofblob.h"
 #include "oneflow/core/vm/vm_util.h"
 #ifdef WITH_CUDA
@@ -35,7 +35,7 @@ namespace one {
 namespace {
 
 Maybe<void> CPUSynchronize() {
-  if (Global<EnvGlobalObjectsScope>::Get() != nullptr) { return vm::CurrentRankSync(); }
+  if (Global<VirtualMachine>::Get() != nullptr) { return vm::CurrentRankSync(); }
   return Maybe<void>::Ok();
 }
 

oneflow/core/job/env_global_objects_scope.cpp

@@ -229,6 +229,7 @@ Maybe<void> EnvGlobalObjectsScope::Init(const EnvProto& env_proto) {
 }
 
 EnvGlobalObjectsScope::~EnvGlobalObjectsScope() {
+  if (is_normal_exit_.has_value() && !CHECK_JUST(is_normal_exit_)) { return; }
   auto session_ctx = Global<MultiClientSessionContext>::Get();
   if (session_ctx != nullptr) {
     VLOG(1) << "Multi client session has not closed , env close it at env scope destruction.";

oneflow/core/job/env_global_objects_scope.h

@@ -17,6 +17,7 @@ limitations under the License.
 #define ONEFLOW_CORE_JOB_CLUSTER_OBJECTS_SCOPE_H_
 #include "oneflow/core/common/util.h"
 #include "oneflow/core/common/maybe.h"
+#include "oneflow/core/common/optional.h"
 #include "oneflow/core/job/env_desc.h"
 #include "oneflow/core/framework/device.h"
 
@@ -31,6 +32,15 @@ class EnvGlobalObjectsScope final {
   ~EnvGlobalObjectsScope();
 
   Maybe<void> Init(const EnvProto& env_proto);
+
+  Maybe<void> init_is_normal_exit(bool is_normal_exit) {
+    CHECK_OR_RETURN(!is_normal_exit_.has_value());
+    is_normal_exit_ = is_normal_exit;
+    return Maybe<void>::Ok();
+  }
+
+ private:
+  Optional<bool> is_normal_exit_;
 };
 
 }  // namespace oneflow