Control Graph / Session / Env's python c++ object destruction

oneflow/api/python/session/session.cpp

@@ -33,7 +33,9 @@ ONEFLOW_API_PYBIND11_MODULE("", m) {
   // multi-client lazy global session context
   m.def("CreateMultiClientSessionContext", &CreateMultiClientSessionContext);
   m.def("InitMultiClientSessionContext", &InitMultiClientSessionContext);
-  m.def("TryDestroyMultiClientSessionContext", &TryDestroyMultiClientSessionContext);
+  m.def("MultiClientSessionContextAddCGraph", &MultiClientSessionContextAddCGraph);
+  m.def("TryDestroyMultiClientSessionContext", &TryDestroyMultiClientSessionContext,
+        py::call_guard<py::gil_scoped_release>());
 
   using namespace oneflow;
   m.def("NewSessionId", &NewSessionId);

oneflow/api/python/session/session.h

@@ -30,6 +30,7 @@ limitations under the License.
 #include "oneflow/core/job/resource_desc.h"
 #include "oneflow/core/framework/config_def.h"
 #include "oneflow/core/framework/multi_client_session_context.h"
+#include "oneflow/core/framework/nn_graph.h"
 #include "oneflow/core/persistence/tee_persistent_log_stream.h"
 
 namespace oneflow {
@@ -127,6 +128,12 @@ inline Maybe<void> InitMultiClientSessionContext(const std::string& config_proto
   return Maybe<void>::Ok();
 }
 
+inline Maybe<void> MultiClientSessionContextAddCGraph(
+    const std::shared_ptr<oneflow::NNGraph>& c_graph_ptr) {
+  JUST(Global<MultiClientSessionContext>::Get()->AddCGraph(c_graph_ptr));
+  return Maybe<void>::Ok();
+}
+
 inline Maybe<void> TryDestroyMultiClientSessionContext() {
   // Global<T>::Delete is not allowed to be called here
   // because glog is not constructed yet and LOG(INFO) has bad bahavior

oneflow/api/python/session/session_api.h

@@ -17,6 +17,7 @@ limitations under the License.
 #define ONEFLOW_API_PYTHON_SESSION_SESSION_API_H_
 
 #include "oneflow/api/python/session/session.h"
+#include "oneflow/core/framework/nn_graph.h"
 
 inline bool IsSessionInited() { return oneflow::IsSessionInited().GetOrThrow(); }
 
@@ -42,6 +43,11 @@ inline void InitMultiClientSessionContext(const std::string& config_proto_str) {
   return oneflow::InitMultiClientSessionContext(config_proto_str).GetOrThrow();
 }
 
+inline void MultiClientSessionContextAddCGraph(
+    const std::shared_ptr<oneflow::NNGraph>& c_graph_ptr) {
+  return oneflow::MultiClientSessionContextAddCGraph(c_graph_ptr).GetOrThrow();
+}
+
 inline void TryDestroyMultiClientSessionContext() {
   return oneflow::TryDestroyMultiClientSessionContext().GetOrThrow();
 }

oneflow/core/framework/multi_client_session_context.cpp

@@ -119,11 +119,21 @@ Maybe<void> MultiClientSessionContext::TryInit(const ConfigProto& config_proto)
   return Maybe<void>::Ok();
 }
 
+Maybe<void> MultiClientSessionContext::AddCGraph(
+    const std::shared_ptr<oneflow::NNGraph>& c_graph_ptr) {
+  graphs_.push_back(c_graph_ptr);
+  return Maybe<void>::Ok();
+}
+
 Maybe<void> MultiClientSessionContext::TryClose() {
   if (is_inited_) {
-    VLOG(2) << "Try to delete multi client session." << std::endl;
-    JUST(vm::MultiClientSync());
-    VLOG(2) << "Start to delete multi client session." << std::endl;
+    VLOG(2) << "Try to delete multi client session context." << std::endl;
+    for (auto wk_graph_ptr : graphs_) {
+      if (auto sh_graph_ptr = wk_graph_ptr.lock()) {
+        VLOG(2) << "grap name " << sh_graph_ptr->job_name() << " not closed, try to close it.";
+        JUST(sh_graph_ptr->Close());
+      }
+    }
     {
       // NOTE(chengcheng): delete runtime global objects
       Global<boxing::collective::CollectiveBoxingDeviceCtxPoller>::Delete();
@@ -148,6 +158,7 @@ Maybe<void> MultiClientSessionContext::TryClose() {
     Global<ResourceDesc, ForSession>::New(Global<ResourceDesc, ForEnv>::Get()->resource(),
                                           GlobalProcessCtx::NumOfProcessPerNode());
   }
+  VLOG(2) << "Finish delete multi client session context." << std::endl;
   return Maybe<void>::Ok();
 }
 

oneflow/core/framework/multi_client_session_context.h

@@ -19,6 +19,7 @@ limitations under the License.
 #include "oneflow/core/common/util.h"
 #include "oneflow/core/job/job_set.pb.h"
 #include "oneflow/core/common/maybe.h"
+#include "oneflow/core/framework/nn_graph.h"
 #include "oneflow/core/framework/tensor.h"
 
 namespace oneflow {
@@ -30,6 +31,7 @@ class MultiClientSessionContext {
   ~MultiClientSessionContext() {}
 
   Maybe<void> TryInit(const ConfigProto& config_proto);
+  Maybe<void> AddCGraph(const std::shared_ptr<oneflow::NNGraph>& c_graph_ptr);
   Maybe<void> TryClose();
 
   // NOTE(chengcheng): for nn.Graph catch free EagerTensor in Graph.build().
@@ -48,6 +50,7 @@ class MultiClientSessionContext {
   bool is_inited_;
   HashMap<std::string, std::vector<std::pair<std::string, std::shared_ptr<one::Tensor>>>>
       graph_name2free_eager_tensors_;
+  std::vector<std::weak_ptr<NNGraph>> graphs_;
 };
 
 }  // namespace oneflow

oneflow/core/framework/nn_graph.cpp

@@ -25,17 +25,28 @@ limitations under the License.
 #include "oneflow/core/job/job_desc.h"
 #include "oneflow/core/job/job_instance.h"
 #include "oneflow/core/job/plan_util.h"
-#include "oneflow/core/job/runtime.h"
 #include "oneflow/core/persistence/tee_persistent_log_stream.h"
 #include "oneflow/core/vm/vm_util.h"
 
 namespace oneflow {
 
 NNGraph::~NNGraph() {
-  VLOG(2) << "Try to delete c nn graph name " << name_ << "." << std::endl;
-  CloseRuntimeBuffers();
-  runtime_.reset();
-  Global<MultiClientSessionContext>::Get()->RemoveGraphFreeEagerTensors(name_);
+  if (!is_closed_) {
+    VLOG(2) << "graph destructor Try to close c nn graph name " << name_ << "." << std::endl;
+    Close();
+  }
+}
+
+Maybe<void> NNGraph::Close() {
+  if (!is_closed_) {
+    VLOG(2) << "Try to close c nn graph name " << name_ << "." << std::endl;
+    CloseRuntimeBuffers();
+    runtime_.reset();
+    Global<MultiClientSessionContext>::Get()->RemoveGraphFreeEagerTensors(name_);
+    is_closed_ = true;
+    VLOG(2) << "Finish close c nn graph name " << name_ << "." << std::endl;
+  }
+  return Maybe<void>::Ok();
 }
 
 const std::vector<std::string>& NNGraph::inputs_op_names() const { return input_op_names_; }

oneflow/core/framework/nn_graph.h

@@ -21,15 +21,16 @@ limitations under the License.
 #include "oneflow/core/framework/tensor_tuple.h"
 #include "oneflow/core/job/job.pb.h"
 #include "oneflow/core/job/plan.pb.h"
+#include "oneflow/core/job/runtime.h"
 
 namespace oneflow {
 
 class Blob;
-class Runtime;
 
 class NNGraph final : public NNGraphIf {
  public:
-  explicit NNGraph(const std::string& name) : name_(name), runtime_inited_(false) {}
+  explicit NNGraph(const std::string& name)
+      : name_(name), runtime_inited_(false), is_closed_(false) {}
   ~NNGraph();
 
   const std::string& job_name() const { return name_; }
@@ -43,6 +44,7 @@ class NNGraph final : public NNGraphIf {
       const std::vector<std::string>& variable_op_names,
       const std::vector<std::shared_ptr<one::Tensor>>& variable_tensors);
   Maybe<void> CompileAndInitRuntime();
+  Maybe<void> Close();
 
  private:
   Maybe<void> RegisterFreeEagerTensorsToVariableOpNames();
@@ -59,6 +61,7 @@ class NNGraph final : public NNGraphIf {
   // TODO(chengcheng): temp impl using runtime now, need reimplement for dynamic multi nn.Graph.
   std::unique_ptr<Runtime> runtime_;
   bool runtime_inited_;
+  bool is_closed_;
 };
 
 Maybe<void> RunLazyNNGraph(const one::TensorTuple& inputs, const one::TensorTuple& outputs,

oneflow/core/job/env_global_objects_scope.cpp

@@ -36,6 +36,7 @@ limitations under the License.
 #include "oneflow/core/transport/transport.h"
 #include "oneflow/core/device/node_device_descriptor_manager.h"
 #include "oneflow/core/vm/symbol_storage.h"
+#include "oneflow/core/framework/multi_client_session_context.h"
 #include "oneflow/core/framework/symbol_id_cache.h"
 #include "oneflow/core/operator/op_node_signature.cfg.h"
 #include "oneflow/core/operator/op_conf.cfg.h"
@@ -199,6 +200,9 @@ Maybe<void> EnvGlobalObjectsScope::Init(const EnvProto& env_proto) {
 }
 
 EnvGlobalObjectsScope::~EnvGlobalObjectsScope() {
+  auto session_ctx = Global<MultiClientSessionContext>::Get();
+  if (session_ctx != nullptr) { session_ctx->TryClose(); }
+
   if (!Global<ResourceDesc, ForSession>::Get()->enable_dry_run()) {
 #ifdef __linux__
     if (Global<ResourceDesc, ForSession>::Get()->process_ranks().size() > 1) {

python/oneflow/__init__.py

@@ -73,17 +73,16 @@ def is_deprecated(func_or_class):
 import oneflow.framework.session_context as session_ctx
 from oneflow.framework.multi_client_session import MultiClientSession
 
-if not env_util.HasAllMultiClientEnvVars():
-    env_util.SetDefaultMultiClientEnvVars()
-oneflow._oneflow_internal.SetIsMultiClient(True)
-env_util.api_env_init()
+env_util.GetEnvHolder()
+
 oneflow._oneflow_internal.InitDefaultConsistentTransportTokenScope()
 session_ctx.OpenDefaultSession(
-    MultiClientSession(oneflow._oneflow_internal.NewSessionId())
+    MultiClientSession(
+        oneflow._oneflow_internal.NewSessionId(), env_util.GetEnvHolder()
+    )
 )
 scope_util.InitScopeStack()
 oneflow._oneflow_internal.EnableEagerEnvironment(True)
-del env_util
 from oneflow.framework import python_callback, register_python_callback
 
 oneflow._oneflow_internal.RegisterGlobalForeignCallback(
@@ -93,24 +92,13 @@ def is_deprecated(func_or_class):
 del register_python_callback
 
 
-def _SyncOnMasterFn():
-    import oneflow
-
-    def Sync():
-        if not oneflow._oneflow_internal.IsEnvInited():
-            return
-        if oneflow.framework.distribute.is_multi_client():
-            oneflow._oneflow_internal.eager.multi_client.Sync()
-        elif oneflow.framework.distribute.get_rank() == 0:
-            oneflow._oneflow_internal.eager.single_client.Sync()
+def _ExitOneFlow():
+    session_ctx.TryCloseDefaultSession()
+    env_util.DelEnvHolder()
 
-    return Sync
 
+atexit.register(_ExitOneFlow)
 
-atexit.register(oneflow._oneflow_internal.SetShuttingDown)
-atexit.register(oneflow._oneflow_internal.DestroyEnv)
-atexit.register(oneflow.framework.session_context.TryCloseDefaultSession)
-atexit.register(_SyncOnMasterFn)
 del atexit
 del oneflow
 import oneflow.framework.docstr as docstr

python/oneflow/framework/env_util.py

@@ -359,12 +359,52 @@ def str2int(env_config):
     if os.getenv("GLOG_log_dir"):
         cpp_logging_conf.log_dir = os.getenv("GLOG_log_dir")
     if os.getenv("GLOG_logtostderr"):
-        cpp_logging_conf.logtostderr = os.getenv("GLOG_logtostderr")
+        cpp_logging_conf.logtostderr = int(os.getenv("GLOG_logtostderr"))
     if os.getenv("GLOG_logbuflevel"):
         cpp_logging_conf.logbuflevel = os.getenv("GLOG_logbuflevel")
     env_proto.cpp_logging_conf.CopyFrom(cpp_logging_conf)
 
 
+def _SyncOnMasterFn():
+    def Sync():
+        if not oneflow._oneflow_internal.IsEnvInited():
+            return
+        if oneflow.framework.distribute.is_multi_client():
+            oneflow._oneflow_internal.eager.multi_client.Sync()
+        elif oneflow.framework.distribute.get_rank() == 0:
+            oneflow._oneflow_internal.eager.single_client.Sync()
+
+    return Sync
+
+
+class EnvHolder(object):
+    def __init__(self):
+        if not HasAllMultiClientEnvVars():
+            SetDefaultMultiClientEnvVars()
+        oneflow._oneflow_internal.SetIsMultiClient(True)
+        api_env_init()
+
+    def __del__(self):
+        _SyncOnMasterFn()
+        oneflow._oneflow_internal.DestroyEnv()
+        oneflow._oneflow_internal.SetShuttingDown()
+
+
+def GetEnvHolder():
+    global _env_holder
+    if _env_holder is not None:
+        return _env_holder
+    else:
+        _env_holder = EnvHolder()
+        return _env_holder
+
+
+def DelEnvHolder():
+    global _env_holder
+    del _env_holder
+
+
+_env_holder = None
 device_tag2default_parallel_conf = {}
 default_env_proto = _DefaultEnvProto()
 config_master_addr = ctrl_bootstrap_pb.Address()

python/oneflow/framework/multi_client_session.py

@@ -29,7 +29,8 @@ class Status(enum.Enum):
         INITED = 2
         CLOSED = 3
 
-    def __init__(self, sess_id):
+    def __init__(self, sess_id, env_holder):
+        self._env_holder = env_holder
         self.sess_ = oneflow._oneflow_internal.RegsiterSession(sess_id)
         oneflow._oneflow_internal.CreateMultiClientSessionContext()
         self.config_proto_ = self._make_config_proto()
@@ -39,22 +40,26 @@ def __init__(self, sess_id):
         self._update_scope_attr_name2defaultVal()
         self.status_ = self.Status.CREATED
 
-    def __del__(self):
-        self.TryClose()
-
     def TryInit(self):
         self._check_status(self.Status.CREATED, self.Status.INITED)
         if self.status_ == self.Status.CREATED:
             config_proto_str = text_format.MessageToString(self.config_proto)
             oneflow._oneflow_internal.InitMultiClientSessionContext(config_proto_str)
             self.status_ = self.Status.INITED
 
-    def TryClose(self):
+    def AddCGraph(self, graph):
+        self._check_status(self.Status.INITED)
+        oneflow._oneflow_internal.MultiClientSessionContextAddCGraph(graph)
+
+    def _try_close(self):
         if self.status_ != self.Status.CLOSED:
             oneflow._oneflow_internal.TryDestroyMultiClientSessionContext()
             oneflow._oneflow_internal.ClearSessionById(self.id)
         self.status_ = self.Status.CLOSED
 
+    def __del__(self):
+        self._try_close()
+
     @property
     def status(self):
         return self.status_

python/oneflow/framework/session_context.py

@@ -43,8 +43,7 @@ def TryCloseDefaultSession():
     default_sess_id = oneflow._oneflow_internal.GetDefaultSessionId()
     assert default_sess_id in _sess_id2sess
     if default_sess_id in _sess_id2sess:
-        _sess_id2sess[default_sess_id].TryClose()
-    del _sess_id2sess[default_sess_id]
+        del _sess_id2sess[default_sess_id]
 
 
 def try_init_default_session(func):

python/oneflow/nn/graph.py

@@ -48,6 +48,7 @@ def __init__(self):
         self._args_repr = []
         self._outs_repr = []
         self._debug = False
+        self._session = None
 
     @property
     def name(self):
@@ -163,11 +164,12 @@ def _compile(self, *args):
     def _build_forward_graph(self, *args):
         self._generate_optimizer_and_variable_configs()
 
-        session = session_ctx.GetDefaultSession()
-        assert type(session) is MultiClientSession
-        session.TryInit()
+        self._session = session_ctx.GetDefaultSession()
+        assert type(self._session) is MultiClientSession
+        self._session.TryInit()
+        self._session.AddCGraph(self._c_nn_graph)
 
-        with graph_build_util.graph_build_context(self.config.proto, session):
+        with graph_build_util.graph_build_context(self.config.proto, self._session):
             # Deal with inputs
             arg_op_names, lazy_args, self._args_repr = self._build_io(
                 "input", graph_build_util.build_graph_input_arg, *args
@@ -459,6 +461,9 @@ def _shallow_repr(self):
         shallow_repr = "(GRAPH:" + self._name + ":" + self.__class__.__name__ + ")"
         return shallow_repr
 
+    def __del__(self):
+        oneflow._oneflow_internal.eager.multi_client.Sync()
+
 
 class GraphConfig(FunctionConfig):
     def __init__(self):