MultiClientSessionContext

oneflow/api/python/session/session.cpp

@@ -30,6 +30,11 @@ ONEFLOW_API_PYBIND11_MODULE("", m) {
   m.def("StartLazyGlobalSession", &StartLazyGlobalSession);
   m.def("StopLazyGlobalSession", &StopLazyGlobalSession);
 
+  // multi-client lazy global session context
+  m.def("CreateMultiClientSessionContext", &CreateMultiClientSessionContext);
+  m.def("InitMultiClientSessionContext", &InitMultiClientSessionContext);
+  m.def("DestroyMultiClientSessionContext", &DestroyMultiClientSessionContext);
+
   using namespace oneflow;
   m.def("NewSessionId", &NewSessionId);
   py::class_<LogicalConfigProtoContext>(m, "LogicalConfigProtoContext")

oneflow/api/python/session/session.h

@@ -22,11 +22,13 @@ limitations under the License.
 #include "oneflow/core/control/ctrl_client.h"
 #include "oneflow/core/control/global_process_ctx.h"
 #include "oneflow/core/job/global_for.h"
+#include "oneflow/core/job/env_global_objects_scope.h"
 #include "oneflow/core/job/session_global_objects_scope.h"
 #include "oneflow/core/job/cluster_instruction.h"
 #include "oneflow/core/job/oneflow.h"
 #include "oneflow/core/job/job_build_and_infer_ctx_mgr.h"
 #include "oneflow/core/framework/config_def.h"
+#include "oneflow/core/framework/multi_client_session_context.h"
 #include "oneflow/core/persistence/tee_persistent_log_stream.h"
 
 namespace oneflow {
@@ -106,6 +108,29 @@ inline Maybe<void> StopLazyGlobalSession() {
   return Maybe<void>::Ok();
 }
 
+inline Maybe<void> CreateMultiClientSessionContext() {
+  CHECK_ISNULL_OR_RETURN(Global<MultiClientSessionContext>::Get());
+  Global<MultiClientSessionContext>::New();
+  return Maybe<void>::Ok();
+}
+
+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;
+  CHECK_OR_RETURN(TxtString2PbMessage(config_proto_str, &config_proto))
+      << "failed to parse config_proto: " << config_proto_str;
+  JUST(Global<MultiClientSessionContext>::Get()->TryInit(config_proto));
+  return Maybe<void>::Ok();
+}
+
+inline Maybe<void> DestroyMultiClientSessionContext() {
+  Global<MultiClientSessionContext>::Delete();
+  return Maybe<void>::Ok();
+}
+
 }  // namespace oneflow
 
 #endif  // ONEFLOW_API_PYTHON_SESSION_SESSION_H_

oneflow/api/python/session/session_api.h

@@ -34,4 +34,16 @@ inline void StartLazyGlobalSession() { return oneflow::StartLazyGlobalSession().
 
 inline void StopLazyGlobalSession() { return oneflow::StopLazyGlobalSession().GetOrThrow(); }
 
+inline void CreateMultiClientSessionContext() {
+  return oneflow::CreateMultiClientSessionContext().GetOrThrow();
+}
+
+inline void InitMultiClientSessionContext(const std::string& config_proto_str) {
+  return oneflow::InitMultiClientSessionContext(config_proto_str).GetOrThrow();
+}
+
+inline void DestroyMultiClientSessionContext() {
+  return oneflow::DestroyMultiClientSessionContext().GetOrThrow();
+}
+
 #endif  // ONEFLOW_API_PYTHON_SESSION_SESSION_API_H_

oneflow/core/framework/multi_client_session_context.cpp

@@ -0,0 +1,103 @@
+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+#include "oneflow/core/framework/multi_client_session_context.h"
+#include "oneflow/core/framework/load_library.h"
+#include "oneflow/core/job/version.h"
+#include "oneflow/core/job/global_for.h"
+#include "oneflow/core/job/id_manager.h"
+#include "oneflow/core/job/job_instance.h"
+#include "oneflow/core/job/job_build_and_infer_ctx_mgr.h"
+#include "oneflow/core/common/buffer_manager.h"
+#include "oneflow/core/rpc/include/global_process_ctx.h"
+#ifdef WITH_CUDA
+#include <cuda.h>
+#endif  // WITH_CUDA
+
+namespace oneflow {
+
+namespace {
+
+int32_t GetGpuDeviceNum() {
+#ifndef WITH_CUDA
+  return 0;
+#else
+  int device_count = 0;
+  cudaGetDeviceCount(&device_count);
+  return device_count;
+#endif
+}
+
+}  // namespace
+
+MultiClientSessionContext::~MultiClientSessionContext() {
+  if (is_inited_) {
+    {
+      // NOTE(chengcheng): delete runtime global objects
+      Global<BufferMgr<std::shared_ptr<JobInstance>>>::Delete();
+    }
+
+    Global<LazyJobBuildAndInferCtxMgr>::Delete();
+    Global<IDMgr>::Delete();
+
+    // TODO(chengcheng): remove ForEnv
+    Global<ResourceDesc, ForSession>::Delete();
+    Global<ResourceDesc, ForSession>::New(Global<ResourceDesc, ForEnv>::Get()->resource(),
+                                          GlobalProcessCtx::NumOfProcessPerNode());
+  }
+}
+
+Maybe<void> MultiClientSessionContext::TryInit(const ConfigProto& config_proto) {
+  if (!is_inited_) {
+    CHECK_OR_RETURN(GlobalProcessCtx::IsMultiClient());
+    DumpVersionInfo();
+
+    Resource resource = config_proto.resource();
+
+    {
+      // TODO(chengcheng): remove this hack
+      //   env config for multi-client
+      resource.set_machine_num(GlobalProcessCtx::NodeSize());
+      const int32_t gpu_device_num = GetGpuDeviceNum();
+      resource.set_gpu_device_num(gpu_device_num);
+      if (gpu_device_num == 0) {
+        resource.set_cpu_device_num(GlobalProcessCtx::NumOfProcessPerNode());
+      } else {
+        resource.set_cpu_device_num(gpu_device_num);
+      }
+    }
+
+    Global<ResourceDesc, ForSession>::Delete();
+    Global<ResourceDesc, ForSession>::New(resource, GlobalProcessCtx::NumOfProcessPerNode());
+    Global<IDMgr>::New();
+    // TODO(chengcheng): refactor JobBuildAndInferCtxMgr
+    Global<LazyJobBuildAndInferCtxMgr>::New();
+
+    for (const std::string& lib_path : config_proto.load_lib_path()) {
+      JUST(LoadLibrary(lib_path));
+    }
+
+    {
+      // NOTE(chengcheng): init runtime global objects
+      Global<BufferMgr<std::shared_ptr<JobInstance>>>::New();
+    }
+
+    is_inited_ = true;
+  }
+  return Maybe<void>::Ok();
+}
+
+}  // namespace oneflow

oneflow/core/framework/multi_client_session_context.h

@@ -0,0 +1,39 @@
+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#ifndef ONEFLOW_CORE_FRAMEWORK_MULTI_CLIENT_SESSION_CONTEXT_H_
+#define ONEFLOW_CORE_FRAMEWORK_MULTI_CLIENT_SESSION_CONTEXT_H_
+
+#include "oneflow/core/common/util.h"
+#include "oneflow/core/job/job_set.pb.h"
+#include "oneflow/core/common/maybe.h"
+
+namespace oneflow {
+
+class MultiClientSessionContext {
+ public:
+  OF_DISALLOW_COPY_AND_MOVE(MultiClientSessionContext);
+  MultiClientSessionContext() : is_inited_(false) {}
+  ~MultiClientSessionContext();
+
+  Maybe<void> TryInit(const ConfigProto& config_proto);
+
+ private:
+  bool is_inited_;
+};
+
+}  // namespace oneflow
+
+#endif  // ONEFLOW_CORE_FRAMEWORK_MULTI_CLIENT_SESSION_CONTEXT_H_