From 85215df087d1d6f8f9af19f71feb4140c72765fe Mon Sep 17 00:00:00 2001
From: tangwei12 <tangwei12@baidu.com>
Date: Mon, 18 Jun 2018 22:08:58 +0800
Subject: [PATCH] move checkpoint message to variable message

---
 paddle/fluid/operators/detail/grpc_client.cc  | 14 +++++++++++---
 paddle/fluid/operators/detail/send_recv.proto |  7 ++++---
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/paddle/fluid/operators/detail/grpc_client.cc b/paddle/fluid/operators/detail/grpc_client.cc
index 9a25ec8fdb4..17476ab513b 100644
--- a/paddle/fluid/operators/detail/grpc_client.cc
+++ b/paddle/fluid/operators/detail/grpc_client.cc
@@ -233,12 +233,20 @@ void GRPCClient::AsyncCheckpointNotify(const std::string& ep,
                                        const std::string& dir,
                                        int64_t time_out) {
   const auto ch = GetChannel(ep);
+
   CheckpointNotifyProcessor* s = new CheckpointNotifyProcessor(ch);
   s->Prepare(time_out);
+  s->response_call_back_ = nullptr;
 
-  sendrecv::CheckpointMessage req;
-  req.set_notify_type(CHECKPOINT_SAVE_MESSAGE);
-  req.set_checkpoint_dir(dir);
+  sendrecv::VariableMessage req;
+  req.set_varname(CHECKPOINT_SAVE_MESSAGE);
+  req.out_varname(dir);
+
+  auto call = s->stub_g_.PrepareUnaryCall(
+      s->context_.get(), "/sendrecv.SendRecvService/CheckpointNotify", req,
+      &cq_);
+  call->StartCall();
+  call->Finish(&s->reply_, &s->status_, reinterpret_cast<void*>(s));
 
   auto rpc = s->stub_->AsyncCheckpointNotify(s->context_.get(), req, &cq_);
   rpc->Finish(&s->reply_, &s->status_, reinterpret_cast<void*>(s));
diff --git a/paddle/fluid/operators/detail/send_recv.proto b/paddle/fluid/operators/detail/send_recv.proto
index cc6529cea74..f5800cdb7f7 100644
--- a/paddle/fluid/operators/detail/send_recv.proto
+++ b/paddle/fluid/operators/detail/send_recv.proto
@@ -26,7 +26,7 @@ service SendRecvService {
   // pre-fetch variable by given variable name and Ids
   rpc PrefetchVariable(VariableMessage) returns (VariableMessage) {}
 
-  rpc CheckpointNotify(CheckpointMessage) returns (VoidMessage) {}
+  rpc CheckpointNotify(VariableMessage) returns (VoidMessage) {}
 }
 
 // VariableMessage is serialized paddle variable message.
@@ -83,6 +83,7 @@ message VariableMessage {
 message VoidMessage {}
 
 message CheckpointMessage {
-  string notify_type = 1;
-  string checkpoint_dir = 2;
+  string varname = 1;
+  string notify_type = 2;
+  string checkpoint_dir = 3;
 }
-- 
GitLab