replace ForeignJobInstance using JobInstance (#5374)

Co-authored-by: oneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>
Oneflow-Inc · Jul 5, 2021 · 94ffe85 · 94ffe85
1 parent 5806e2b
commit 94ffe85
Show file tree

Hide file tree

Showing 12 changed files with 97 additions and 98 deletions.
diff --git a/oneflow/api/python/framework/foreign_job_instance.cpp b/oneflow/api/python/framework/foreign_job_instance.cpp
diff --git a/oneflow/api/python/framework/framework.h b/oneflow/api/python/framework/framework.h
@@ -27,7 +27,7 @@ limitations under the License.
 #include "oneflow/core/job/inter_user_job_info.pb.h"
 #include "oneflow/core/job/foreign_callback.h"
 #include "oneflow/core/job/foreign_watcher.h"
-#include "oneflow/core/job/foreign_job_instance.h"
+#include "oneflow/core/job/job_instance.h"
 #include "oneflow/core/job/oneflow.h"
 #include "oneflow/core/job/placement.pb.h"
 #include "oneflow/core/framework/config_def.h"
@@ -58,11 +58,11 @@ inline Maybe<void> RegisterWatcherOnlyOnce(const std::shared_ptr<ForeignWatcher>
   return Maybe<void>::Ok();
 }
 
-inline Maybe<void> LaunchJob(const std::shared_ptr<oneflow::ForeignJobInstance>& cb) {
+inline Maybe<void> LaunchJob(const std::shared_ptr<oneflow::JobInstance>& cb) {
   CHECK_OR_RETURN(GlobalProcessCtx::IsThisProcessMaster());
   CHECK_NOTNULL_OR_RETURN(Global<Oneflow>::Get());
   const auto& job_name = cb->job_name();
-  auto* buffer_mgr = Global<BufferMgr<std::shared_ptr<ForeignJobInstance>>>::Get();
+  auto* buffer_mgr = Global<BufferMgr<std::shared_ptr<JobInstance>>>::Get();
   int64_t job_id = Global<JobName2JobId>::Get()->at(job_name);
   if (IsPullJob(job_name, *Global<InterUserJobInfo>::Get())) {
     buffer_mgr->Get(GetForeignOutputBufferName(job_name))->Send(cb);

diff --git a/oneflow/api/python/framework/framework_api.h b/oneflow/api/python/framework/framework_api.h
@@ -28,7 +28,7 @@ inline void RegisterWatcherOnlyOnce(const std::shared_ptr<oneflow::ForeignWatche
   return oneflow::RegisterWatcherOnlyOnce(watcher).GetOrThrow();
 }
 
-inline void LaunchJob(const std::shared_ptr<oneflow::ForeignJobInstance>& cb) {
+inline void LaunchJob(const std::shared_ptr<oneflow::JobInstance>& cb) {
   return oneflow::LaunchJob(cb).GetOrThrow();
 }
 

diff --git a/oneflow/api/python/framework/job_instance.cpp b/oneflow/api/python/framework/job_instance.cpp
@@ -0,0 +1,72 @@
+/*
+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 <pybind11/pybind11.h>
+#include <string>
+#include <memory>
+#include "oneflow/api/python/of_api_registry.h"
+#include "oneflow/core/common/util.h"
+#include "oneflow/core/job/job_instance.h"
+
+namespace py = pybind11;
+
+namespace oneflow {
+
+class PyJobInstance : public JobInstance {
+ public:
+  // Inherit the constructors
+  using JobInstance::JobInstance;
+
+  // Trampoline (need one for each virtual function)
+  std::string job_name() const override {
+    PYBIND11_OVERRIDE(std::string, /* Return type */
+                      JobInstance, /* Parent class */
+                      job_name,    /* Name of function in C++ (must match Python name) */
+    );
+  }
+
+  std::string sole_input_op_name_in_user_job() const override {
+    PYBIND11_OVERRIDE(std::string, JobInstance, sole_input_op_name_in_user_job, );
+  }
+
+  std::string sole_output_op_name_in_user_job() const override {
+    PYBIND11_OVERRIDE(std::string, JobInstance, sole_output_op_name_in_user_job, );
+  }
+
+  void PushBlob(uint64_t ofblob_ptr) const override {
+    PYBIND11_OVERRIDE(void, JobInstance, PushBlob, ofblob_ptr);
+  }
+
+  void PullBlob(uint64_t ofblob_ptr) const override {
+    PYBIND11_OVERRIDE(void, JobInstance, PullBlob, ofblob_ptr);
+  }
+
+  void Finish() const override { PYBIND11_OVERRIDE(void, JobInstance, Finish, ); }
+};
+
+}  // namespace oneflow
+
+ONEFLOW_API_PYBIND11_MODULE("", m) {
+  using namespace oneflow;
+
+  py::class_<JobInstance, PyJobInstance, std::shared_ptr<JobInstance>>(m, "JobInstance")
+      .def(py::init<>())
+      .def("job_name", &JobInstance::job_name)
+      .def("sole_input_op_name_in_user_job", &JobInstance::sole_input_op_name_in_user_job)
+      .def("sole_output_op_name_in_user_job", &JobInstance::sole_output_op_name_in_user_job)
+      .def("PushBlob", &JobInstance::PushBlob)
+      .def("PullBlob", &JobInstance::PullBlob)
+      .def("Finish", &JobInstance::Finish);
+}
diff --git a/oneflow/core/job/foreign_job_instance.h → oneflow/core/job/job_instance.h b/oneflow/core/job/foreign_job_instance.h → oneflow/core/job/job_instance.h
@@ -20,11 +20,11 @@ limitations under the License.
 
 namespace oneflow {
 
-class ForeignJobInstance {
+class JobInstance {
  public:
-  ForeignJobInstance() = default;
+  JobInstance() = default;
 
-  virtual ~ForeignJobInstance() = default;
+  virtual ~JobInstance() = default;
 
   virtual std::string job_name() const { UNIMPLEMENTED(); }
   virtual std::string sole_input_op_name_in_user_job() const { UNIMPLEMENTED(); }

diff --git a/oneflow/core/job/runtime_buffer_managers_scope.cpp b/oneflow/core/job/runtime_buffer_managers_scope.cpp
@@ -15,17 +15,17 @@ limitations under the License.
 */
 #include "oneflow/core/common/buffer_manager.h"
 #include "oneflow/core/job/runtime_buffer_managers_scope.h"
-#include "oneflow/core/job/foreign_job_instance.h"
+#include "oneflow/core/job/job_instance.h"
 
 namespace oneflow {
 
 RuntimeBufferManagersScope::RuntimeBufferManagersScope() {
   Global<BufferMgr<int64_t>>::New();
-  Global<BufferMgr<std::shared_ptr<ForeignJobInstance>>>::New();
+  Global<BufferMgr<std::shared_ptr<JobInstance>>>::New();
 }
 
 RuntimeBufferManagersScope::~RuntimeBufferManagersScope() {
-  Global<BufferMgr<std::shared_ptr<ForeignJobInstance>>>::Delete();
+  Global<BufferMgr<std::shared_ptr<JobInstance>>>::Delete();
   Global<BufferMgr<int64_t>>::Delete();
 }
 

diff --git a/oneflow/core/job/runtime_buffers_scope.cpp b/oneflow/core/job/runtime_buffers_scope.cpp
@@ -16,14 +16,14 @@ limitations under the License.
 #include "oneflow/core/common/buffer_manager.h"
 #include "oneflow/core/job/runtime_buffers_scope.h"
 #include "oneflow/core/job/job_desc.h"
-#include "oneflow/core/job/foreign_job_instance.h"
+#include "oneflow/core/job/job_instance.h"
 
 namespace oneflow {
 
 RuntimeBuffersScope::RuntimeBuffersScope(const JobConfs& job_confs) {
   size_t job_size = Global<JobName2JobId>::Get()->size();
   Global<BufferMgr<int64_t>>::Get()->NewBuffer(kBufferNameGlobalWaitJobId, job_size);
-  auto* buffer_mgr = Global<BufferMgr<std::shared_ptr<ForeignJobInstance>>>::Get();
+  auto* buffer_mgr = Global<BufferMgr<std::shared_ptr<JobInstance>>>::Get();
   for (const auto& pair : job_confs.job_id2job_conf()) {
     const auto& job_name = pair.second.job_name();
     CHECK_EQ(pair.first, Global<JobName2JobId>::Get()->at(job_name));
@@ -35,7 +35,7 @@ RuntimeBuffersScope::RuntimeBuffersScope(const JobConfs& job_confs) {
 }
 
 RuntimeBuffersScope::~RuntimeBuffersScope() {
-  auto* buffer_mgr = Global<BufferMgr<std::shared_ptr<ForeignJobInstance>>>::Get();
+  auto* buffer_mgr = Global<BufferMgr<std::shared_ptr<JobInstance>>>::Get();
   for (const auto& pair : *Global<JobName2JobId>::Get()) {
     const auto& job_name = pair.first;
     buffer_mgr->Get(GetCallbackNotifierBufferName(job_name))->Close();

diff --git a/oneflow/core/job/session_global_objects_scope.cpp b/oneflow/core/job/session_global_objects_scope.cpp
@@ -21,7 +21,7 @@ limitations under the License.
 #include "oneflow/core/job/available_memory_desc.pb.h"
 #include "oneflow/core/job/id_manager.h"
 #include "oneflow/core/job/profiler.h"
-#include "oneflow/core/job/foreign_job_instance.h"
+#include "oneflow/core/job/job_instance.h"
 #include "oneflow/core/job/inter_user_job_info.pb.h"
 #include "oneflow/core/job/job_desc.h"
 #include "oneflow/core/job/critical_section_desc.h"

diff --git a/oneflow/core/kernel/callback_notify_kernel.cpp b/oneflow/core/kernel/callback_notify_kernel.cpp
@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 #include "oneflow/core/kernel/callback_notify_kernel.h"
-#include "oneflow/core/job/foreign_job_instance.h"
+#include "oneflow/core/job/job_instance.h"
 
 namespace oneflow {
 
@@ -23,8 +23,8 @@ void CallbackNotifyKernel<T>::ForwardDataContent(
     const KernelCtx& ctx, std::function<Blob*(const std::string&)> BnInOp2Blob) const {
   T job_id = *BnInOp2Blob("in")->dptr<T>();
   const auto& buffer_name = this->op_conf().callback_notify_conf().callback_buffer_name(job_id);
-  std::shared_ptr<ForeignJobInstance> foreign_job_instance;
-  BufferStatus buffer_status = Global<BufferMgr<std::shared_ptr<ForeignJobInstance>>>::Get()
+  std::shared_ptr<JobInstance> foreign_job_instance;
+  BufferStatus buffer_status = Global<BufferMgr<std::shared_ptr<JobInstance>>>::Get()
                                    ->Get(buffer_name)
                                    ->TryReceive(&foreign_job_instance);
   CHECK_NE(buffer_status, kBufferStatusEmpty);

diff --git a/oneflow/core/kernel/foreign_input_kernel.cpp b/oneflow/core/kernel/foreign_input_kernel.cpp
@@ -16,15 +16,15 @@ limitations under the License.
 #include "oneflow/core/kernel/foreign_input_kernel.h"
 #include "oneflow/core/common/buffer_manager.h"
 #include "oneflow/core/register/ofblob.h"
-#include "oneflow/core/job/foreign_job_instance.h"
+#include "oneflow/core/job/job_instance.h"
 
 namespace oneflow {
 
 void ForeignInputKernel::ForwardDataContent(
     const KernelCtx& ctx, std::function<Blob*(const std::string&)> BnInOp2Blob) const {
   const auto& buffer_name = op_conf().foreign_input_conf().ofblob_buffer_name();
-  std::shared_ptr<ForeignJobInstance> foreign_job_instance;
-  BufferStatus buffer_status = Global<BufferMgr<std::shared_ptr<ForeignJobInstance>>>::Get()
+  std::shared_ptr<JobInstance> foreign_job_instance;
+  BufferStatus buffer_status = Global<BufferMgr<std::shared_ptr<JobInstance>>>::Get()
                                    ->Get(buffer_name)
                                    ->TryReceive(&foreign_job_instance);
   CHECK_NE(buffer_status, kBufferStatusEmpty);

diff --git a/oneflow/core/kernel/foreign_output_kernel.cpp b/oneflow/core/kernel/foreign_output_kernel.cpp
@@ -16,15 +16,15 @@ limitations under the License.
 #include "oneflow/core/kernel/foreign_output_kernel.h"
 #include "oneflow/core/common/buffer_manager.h"
 #include "oneflow/core/register/ofblob.h"
-#include "oneflow/core/job/foreign_job_instance.h"
+#include "oneflow/core/job/job_instance.h"
 
 namespace oneflow {
 
 void ForeignOutputKernel::ForwardDataContent(
     const KernelCtx& ctx, std::function<Blob*(const std::string&)> BnInOp2Blob) const {
   const auto& buffer_name = op_conf().foreign_output_conf().ofblob_buffer_name();
-  std::shared_ptr<ForeignJobInstance> foreign_job_instance;
-  BufferStatus buffer_status = Global<BufferMgr<std::shared_ptr<ForeignJobInstance>>>::Get()
+  std::shared_ptr<JobInstance> foreign_job_instance;
+  BufferStatus buffer_status = Global<BufferMgr<std::shared_ptr<JobInstance>>>::Get()
                                    ->Get(buffer_name)
                                    ->TryReceive(&foreign_job_instance);
   CHECK_NE(buffer_status, kBufferStatusEmpty);

diff --git a/oneflow/python/framework/job_instance.py b/oneflow/python/framework/job_instance.py
@@ -73,7 +73,7 @@ def DereferenceJobInstance(job_instance):
     return job_instance
 
 
-class JobInstance(oneflow._oneflow_internal.ForeignJobInstance):
+class JobInstance(oneflow._oneflow_internal.JobInstance):
     def __init__(
         self,
         job_name,
@@ -83,7 +83,7 @@ def __init__(
         pull_cb=None,
         finish_cb=None,
     ):
-        oneflow._oneflow_internal.ForeignJobInstance.__init__(self)
+        oneflow._oneflow_internal.JobInstance.__init__(self)
         self.thisown = 0
         self.job_name_ = str(job_name)
         self.sole_input_op_name_in_user_job_ = str(sole_input_op_name_in_user_job)