Serialize proto in binary rather than text (#4810)

* Serialize proto in binary rather than text

* refine

Co-authored-by: oneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>

oneflow/api/python/framework/framework.cpp

@@ -26,10 +26,11 @@ ONEFLOW_API_PYBIND11_MODULE("", m) {
   m.def("RegisterWatcherOnlyOnce", &RegisterWatcherOnlyOnce);
   m.def("LaunchJob", &LaunchJob, py::call_guard<py::gil_scoped_release>());
 
-  m.def("GetSerializedInterUserJobInfo", &GetSerializedInterUserJobInfo);
-  m.def("GetSerializedJobSet", &GetSerializedJobSet);
-  m.def("GetSerializedStructureGraph", &GetSerializedStructureGraph);
-  m.def("GetSerializedCurrentJob", &GetSerializedCurrentJob);
+  m.def("GetSerializedInterUserJobInfo",
+        []() { return py::bytes(GetSerializedInterUserJobInfo()); });
+  m.def("GetSerializedJobSet", []() { return py::bytes(GetSerializedJobSet()); });
+  m.def("GetSerializedStructureGraph", &GetSerializedStructureGraph /* a prototxt saved to file*/);
+  m.def("GetSerializedCurrentJob", []() { return py::bytes(GetSerializedCurrentJob()); });
 
   m.def("GetFunctionConfigDef", &GetFunctionConfigDef);
   m.def("GetScopeConfigDef", &GetScopeConfigDef);

oneflow/api/python/framework/framework.h

@@ -85,9 +85,7 @@ inline Maybe<std::string> GetSerializedInterUserJobInfo() {
   CHECK_OR_RETURN(GlobalProcessCtx::IsThisProcessMaster());
   CHECK_NOTNULL_OR_RETURN(Global<Oneflow>::Get());
   CHECK_NOTNULL_OR_RETURN(Global<InterUserJobInfo>::Get());
-  std::string ret;
-  google::protobuf::TextFormat::PrintToString(*Global<InterUserJobInfo>::Get(), &ret);
-  return ret;
+  return Global<InterUserJobInfo>::Get()->SerializeAsString();
 }
 
 inline Maybe<const JobSet&> GetJobSet() {
@@ -101,7 +99,7 @@ inline Maybe<const JobSet&> GetJobSet() {
   return job_ctx_mgr->job_set();
 }
 
-inline Maybe<std::string> GetSerializedJobSet() { return PbMessage2TxtString(JUST(GetJobSet())); }
+inline Maybe<std::string> GetSerializedJobSet() { return JUST(GetJobSet()).SerializeAsString(); }
 
 inline Maybe<std::string> GetSerializedCurrentJob() {
   auto* job_ctx_mgr = Global<LazyJobBuildAndInferCtxMgr>::Get();
@@ -109,7 +107,7 @@ inline Maybe<std::string> GetSerializedCurrentJob() {
   auto* job_ctx =
       JUST(job_ctx_mgr->FindJobBuildAndInferCtx(*JUST(job_ctx_mgr->GetCurrentJobName())));
   CHECK_NOTNULL_OR_RETURN(job_ctx);
-  return PbMessage2TxtString(job_ctx->job());
+  return job_ctx->job().SerializeAsString();
 }
 
 inline Maybe<std::string> GetFunctionConfigDef() {

oneflow/core/job/inter_user_job_info.proto

@@ -4,7 +4,7 @@ package oneflow;
 message InterUserJobInfo {
   map<string, string> input_or_var_op_name2push_job_name = 1;
   map<string, string> output_or_var_op_name2pull_job_name = 2;
-  required string global_model_init_job_name = 4;
-  required string global_model_load_job_name = 5;
-  required string global_model_save_job_name = 6;
+  optional string global_model_init_job_name = 4;
+  optional string global_model_load_job_name = 5;
+  optional string global_model_save_job_name = 6;
 }

oneflow/python/framework/c_api_util.py

@@ -65,7 +65,9 @@ def InitLazyGlobalSession(config_proto):
 
 def GetInterUserJobInfo():
     inter_user_job_info = oneflow._oneflow_internal.GetSerializedInterUserJobInfo()
-    return text_format.Parse(inter_user_job_info, InterUserJobInfo())
+    ret = InterUserJobInfo()
+    ret.ParseFromString(inter_user_job_info)
+    return ret
 
 
 def JobBuildAndInferCtx_Open(job_name):
@@ -252,9 +254,13 @@ def GetInterfaceOpAttributes():
 @oneflow_export("experimental.get_job_set")
 def GetJobSet():
     job_set = oneflow._oneflow_internal.GetSerializedJobSet()
-    return text_format.Parse(job_set, job_set_pb.JobSet())
+    ret = job_set_pb.JobSet()
+    ret.ParseFromString(job_set)
+    return ret
 
 
 def GetCurrentJob():
     serialized_job = oneflow._oneflow_internal.GetSerializedCurrentJob()
-    return text_format.Parse(serialized_job, job_pb.Job())
+    ret = job_pb.Job()
+    ret.ParseFromString(serialized_job)
+    return ret